From abbcc7aefade50d57503d442c1e693d282fb7168 Mon Sep 17 00:00:00 2001
From: wozeparrot <wozeparrot@gmail.com>
Date: Tue, 21 Nov 2023 01:03:43 -0500
Subject: [PATCH] missed cleanup from cache_id removal (#2376)

---
 extra/dist/world.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/extra/dist/world.py b/extra/dist/world.py
index 26cfb8bed0..70c1728499 100644
--- a/extra/dist/world.py
+++ b/extra/dist/world.py
@@ -60,7 +60,7 @@ def _send_rb(x:RawBuffer, target_rank:int):
 
     # jit support
     CacheCollector.add(__send_rb, [x, target_rank, y], {})
-setattr(_send_rb, "shared_memory_cache", {})
+  if DEBUG >= 2: print(f"****  rank {getenv('RANK')} sent {x} to rank {target_rank}")
 
 # receive a rawbuffer from the target rank
 def _recv_rb(x:RawBuffer, target_rank:int):
@@ -74,7 +74,6 @@ def _recv_rb(x:RawBuffer, target_rank:int):
     y = RawHIPBuffer(x.size, x.dtype, device=str(y_device), buf=ptr, allocator=None)
     x._transfer(y)
 
-    if DEBUG >= 2: print(f"****  rank {getenv('RANK')} got {x} from rank {target_rank}")
     CacheCollector.add(__recv_rb, [x, target_rank, y], {})
   else:
     shm_name = dist.OOB.recv(target_rank)
@@ -83,10 +82,10 @@ def _recv_rb(x:RawBuffer, target_rank:int):
     # fast path when we can directly copyin
     if isinstance(x, RawBufferCopyIn): x._copyin(y.toCPU())
     else: x.fromCPU(y.toCPU())
-    if DEBUG >= 2: print(f"****  rank {getenv('RANK')} got {x} from rank {target_rank}")
 
     # jit support
     CacheCollector.add(__recv_rb, [x, target_rank, y], {})
+  if DEBUG >= 2: print(f"****  rank {getenv('RANK')} got {x} from rank {target_rank}")
 
 # sends a lazybuffer from our rank to the target rank
 def _send_lb(x:LazyBuffer, target_rank:int) -> None: