From fdc668a006d98f869aba20697c21f16aefab4a2a Mon Sep 17 00:00:00 2001 From: vmoens Date: Thu, 16 Oct 2025 19:38:40 +0100 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- test/test_env.py | 38 ++++++++++++++++++++++++++------------ 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/test/test_env.py b/test/test_env.py index b092be80e1c..7cd3c61ae0c 100644 --- a/test/test_env.py +++ b/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,13 +3716,23 @@ 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, @@ -3729,12 +3740,15 @@ def test_batched_nondynamic(self, penv): 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) From 361e2c87df0b9ce26f4ddc90a7d22183e595663c Mon Sep 17 00:00:00 2001 From: vmoens Date: Thu, 16 Oct 2025 22:22:27 +0100 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- torchrl/data/datasets/d4rl.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/torchrl/data/datasets/d4rl.py b/torchrl/data/datasets/d4rl.py index ceb9e55e5d9..6c3719d216a 100644 --- a/torchrl/data/datasets/d4rl.py +++ b/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(