Update

[ghstack-poisoned]

Author: vmoens

test/test_env.py

@@ -3836,6 +3836,8 @@ def test_parallel(self, bwad, use_buffers, maybe_fork_ParallelEnv):
         finally:
             env.close(raise_if_closed=False)
             del env
+            time.sleep(0.1)
+            gc.collect()
 
     class AddString(Transform):
         def __init__(self):

torchrl/envs/batched_envs.py

@@ -2489,14 +2489,15 @@ def look_for_cuda(tensor, has_cuda=has_cuda):
             # Make sure the root is updated
             root_shared_tensordict.update_(env._step_mdp(input))
 
+            # Set event before sending non-tensor data so parent knows worker is done
+            # The recv() call itself will provide synchronization for the pipe
+            mp_event.set()
+            
             if _non_tensor_keys:
                 child_pipe.send(
                     ("non_tensor", next_td.select(*_non_tensor_keys, strict=False))
                 )
 
-            # Set event only after non-tensor data is sent to avoid race condition
-            mp_event.set()
-
             del next_td
 
         elif cmd == "step_and_maybe_reset":
@@ -2530,14 +2531,15 @@ def look_for_cuda(tensor, has_cuda=has_cuda):
                 event.record()
                 event.synchronize()
 
+            # Set event before sending non-tensor data so parent knows worker is done
+            # The recv() call itself will provide synchronization for the pipe
+            mp_event.set()
+            
             if _non_tensor_keys:
                 ntd = root_next_td.select(*_non_tensor_keys)
                 ntd.set("next", td_next.select(*_non_tensor_keys))
                 child_pipe.send(("non_tensor", ntd))
 
-            # Set event only after non-tensor data is sent to avoid race condition
-            mp_event.set()
-
             del td, root_next_td
 
         elif cmd == "close":