[Test] Fix flaky parallel test

ghstack-source-id: aecea30
Pull-Request: #3204

Author: vmoens

test/test_env.py

@@ -13,6 +13,7 @@
 import pickle
 import random
 import re
+import time
 from collections import defaultdict
 from functools import partial
 from sys import platform
@@ -3715,26 +3716,39 @@ def test_batched_nondynamic(self, penv):
             use_buffers=True,
             mp_start_method=mp_ctx if penv is ParallelEnv else None,
         )
-        env_buffers.set_seed(0)
-        torch.manual_seed(0)
-        rollout_buffers = env_buffers.rollout(
-            20, return_contiguous=True, break_when_any_done=False
-        )
-        del env_buffers
+        try:
+            env_buffers.set_seed(0)
+            torch.manual_seed(0)
+            rollout_buffers = env_buffers.rollout(
+                20, return_contiguous=True, break_when_any_done=False
+            )
+        finally:
+            env_buffers.close(raise_if_closed=False)
+            del env_buffers
         gc.collect()
+        # Add a small delay to allow multiprocessing resource_sharer threads
+        # to fully clean up before creating the next environment. This prevents
+        # a race condition where the old resource_sharer service thread is still
+        # active when the new environment starts, causing a deadlock.
+        # See: https://bugs.python.org/issue30289
+        if penv is ParallelEnv:
+            time.sleep(0.1)
 
         env_no_buffers = penv(
             3,
             lambda: GymEnv(CARTPOLE_VERSIONED(), device=None),
             use_buffers=False,
             mp_start_method=mp_ctx if penv is ParallelEnv else None,
         )
-        env_no_buffers.set_seed(0)
-        torch.manual_seed(0)
-        rollout_no_buffers = env_no_buffers.rollout(
-            20, return_contiguous=True, break_when_any_done=False
-        )
-        del env_no_buffers
+        try:
+            env_no_buffers.set_seed(0)
+            torch.manual_seed(0)
+            rollout_no_buffers = env_no_buffers.rollout(
+                20, return_contiguous=True, break_when_any_done=False
+            )
+        finally:
+            env_no_buffers.close(raise_if_closed=False)
+            del env_no_buffers
         gc.collect()
         assert_allclose_td(rollout_buffers, rollout_no_buffers)
 

torchrl/data/datasets/d4rl.py

@@ -279,6 +279,7 @@ def _get_dataset_direct(self, name, env_kwargs):
         # so we need to ensure we're using the gym backend
         with set_gym_backend("gym"):
             import gym
+
             env = GymWrapper(gym.make(name))
         with tempfile.TemporaryDirectory() as tmpdir:
             os.environ["D4RL_DATASET_DIR"] = tmpdir
@@ -358,6 +359,7 @@ def _get_dataset_from_env(self, name, env_kwargs):
         # so we need to ensure we're using the gym backend
         with set_gym_backend("gym"), tempfile.TemporaryDirectory() as tmpdir:
             import gym
+
             os.environ["D4RL_DATASET_DIR"] = tmpdir
             env = GymWrapper(gym.make(name))
             dataset = make_tensordict(