From 41c2d21dc4839912f59c42fb1696a47905668ef3 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Thu, 16 Jul 2026 15:36:10 +0100 Subject: [PATCH 1/2] [Performance] Avoid parent ParallelEnv metadata construction --- benchmarks/benchmark_parallel_env_startup.py | 158 +++++++++ docs/source/reference/envs_vectorized.rst | 22 ++ sota-implementations/vla_grpo/test_openvla.py | 122 +++++++ sota-implementations/vla_grpo/utils.py | 55 +-- test/envs/test_parallel.py | 128 +++++++ test/envs/test_special.py | 19 ++ torchrl/envs/batched_envs.py | 321 +++++++++++++++--- torchrl/envs/env_creator.py | 9 +- 8 files changed, 752 insertions(+), 82 deletions(-) create mode 100644 benchmarks/benchmark_parallel_env_startup.py diff --git a/benchmarks/benchmark_parallel_env_startup.py b/benchmarks/benchmark_parallel_env_startup.py new file mode 100644 index 00000000000..5336051f17d --- /dev/null +++ b/benchmarks/benchmark_parallel_env_startup.py @@ -0,0 +1,158 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +"""Benchmark ``ParallelEnv`` startup metadata strategies. + +The benchmark records constructor calls in separate marker files so parent-side +temporary environments and long-lived worker environments can be counted +independently. For example:: + + python benchmarks/benchmark_parallel_env_startup.py --workers 80 --delay 0.1 + +Use ``--mode`` to run one strategy, and ``--start-method forkserver`` to compare +the no-shadow-environment path with the default ``spawn`` worker startup. +""" + +from __future__ import annotations + +import argparse +import json +import os +import tempfile +import time +from functools import partial +from pathlib import Path +from typing import Literal + +from torchrl import timeit +from torchrl._utils import logger as torchrl_logger +from torchrl.envs import ParallelEnv +from torchrl.testing.mocking_classes import CountingEnv + + +class _StartupEnv(CountingEnv): + def __init__( + self, + worker_idx: int, + *, + marker_dir: str, + parent_pid: int, + delay: float, + ) -> None: + if delay: + time.sleep(delay) + role = "parent" if os.getpid() == parent_pid else "worker" + self._marker_path = Path(marker_dir) / ( + f"constructed-{role}-{worker_idx}-{os.getpid()}" + ) + self._marker_path.touch() + super().__init__() + + def close(self, *, raise_if_closed: bool = True) -> None: + self._marker_path.with_name(f"closed-{self._marker_path.name}").touch() + super().close(raise_if_closed=raise_if_closed) + + +def _run_mode( + mode: Literal["legacy", "homogeneous", "workers"], + *, + num_workers: int, + delay: float, + start_method: Literal["spawn", "forkserver", "fork"], + marker_dir: Path, +) -> dict[str, float | int | str]: + parent_pid = os.getpid() + common_factory = partial( + _StartupEnv, + marker_dir=str(marker_dir), + parent_pid=parent_pid, + delay=delay, + ) + create_env_kwargs = [ + {"worker_idx": worker_idx} for worker_idx in range(num_workers) + ] + if mode == "legacy": + create_env_fn = [ + partial(common_factory, worker_idx=worker_idx) + for worker_idx in range(num_workers) + ] + create_env_kwargs = None + else: + create_env_fn = common_factory + + with timeit(f"{mode}/construct") as construct_timer: + env = ParallelEnv( + num_workers, + create_env_fn, + create_env_kwargs=create_env_kwargs, + metadata_from_workers=mode == "workers", + use_buffers=False, + mp_start_method=start_method, + ) + construct_seconds = construct_timer.elapsed() + try: + with timeit(f"{mode}/reset") as reset_timer: + tensordict = env.reset() + reset_seconds = reset_timer.elapsed() + with timeit(f"{mode}/first_step") as step_timer: + env.rand_step(tensordict) + first_step_seconds = step_timer.elapsed() + finally: + env.close(raise_if_closed=False) + + parent_constructions = len(list(marker_dir.glob("constructed-parent-*"))) + worker_constructions = len(list(marker_dir.glob("constructed-worker-*"))) + closed_constructions = len(list(marker_dir.glob("closed-constructed-*"))) + return { + "mode": mode, + "start_method": start_method, + "num_workers": num_workers, + "parent_constructions": parent_constructions, + "worker_constructions": worker_constructions, + "closed_constructions": closed_constructions, + "construct_seconds": construct_seconds, + "reset_seconds": reset_seconds, + "first_step_seconds": first_step_seconds, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--workers", type=int, default=80) + parser.add_argument("--delay", type=float, default=0.05) + parser.add_argument( + "--mode", + choices=("all", "legacy", "homogeneous", "workers"), + default="all", + ) + parser.add_argument( + "--start-method", + choices=("spawn", "forkserver", "fork"), + default="spawn", + ) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + modes = ("legacy", "homogeneous", "workers") if args.mode == "all" else (args.mode,) + results = [] + for mode in modes: + with tempfile.TemporaryDirectory() as marker_dir: + result = _run_mode( + mode, + num_workers=args.workers, + delay=args.delay, + start_method=args.start_method, + marker_dir=Path(marker_dir), + ) + results.append(result) + torchrl_logger.info("ParallelEnv startup benchmark: %s", result) + + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(results, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/docs/source/reference/envs_vectorized.rst b/docs/source/reference/envs_vectorized.rst index 38f5b9e4e8c..cf44f9bf221 100644 --- a/docs/source/reference/envs_vectorized.rst +++ b/docs/source/reference/envs_vectorized.rst @@ -63,6 +63,28 @@ needed. parallel env can be a bottleneck. This is why, for instance, TorchRL tests are so slow. Once the environment is launched, a great speedup should be observed. + Environments with especially expensive constructors can avoid a second, + parent-side construction pass by setting ``metadata_from_workers=True``. + In this opt-in mode, the real worker environments report their metadata before + normal initialization. The workers start eagerly, their tensor schemas are + validated for compatibility, and no temporary environment is created in the + parent. This mode currently requires pipe-based communication with + ``use_buffers=False``. Use one common factory with ``create_env_kwargs`` for + worker-specific arguments: + + .. code-block:: python + + from functools import partial + + make_env = partial(GymEnv, "Pendulum-v1") + env = ParallelEnv( + 4, + make_env, + create_env_kwargs=[{"g": 9.0 + worker_idx} for worker_idx in range(4)], + metadata_from_workers=True, + use_buffers=False, + ) + .. note:: *TorchRL requires precise specs*: Another thing to take in consideration is diff --git a/sota-implementations/vla_grpo/test_openvla.py b/sota-implementations/vla_grpo/test_openvla.py index 5c9e34912ce..4aa41d57fc7 100644 --- a/sota-implementations/vla_grpo/test_openvla.py +++ b/sota-implementations/vla_grpo/test_openvla.py @@ -818,6 +818,128 @@ def test_chunk_transform_openvla_tokens_decodes_then_postprocesses(self): class TestLiberoWorkers: + def test_standalone_env_factory_uses_worker_metadata(self, monkeypatch): + captured = {} + + def fake_parallel_env(num_envs, env_factory, **kwargs): + captured["num_envs"] = num_envs + captured["env_factory"] = env_factory + captured["kwargs"] = kwargs + return object() + + cfg = SimpleNamespace( + env=SimpleNamespace( + backend="libero", + task_ids=[0, 1], + num_envs=2, + parallel_group_repeats=False, + ), + collector=SimpleNamespace(group_size=8, candidate_group_size=None), + ) + monkeypatch.setattr(utils, "ParallelEnv", fake_parallel_env) + monkeypatch.setattr(utils, "_chunk_transform", lambda *args, **kwargs: object()) + monkeypatch.setattr(utils, "TransformedEnv", lambda base, transform: base) + + result = utils.make_env(cfg, tokenizer=None) + + assert result is not None + assert captured["num_envs"] == 2 + assert captured["kwargs"]["create_env_kwargs"] == [ + {"worker_idx": 0}, + {"worker_idx": 1}, + ] + assert captured["kwargs"]["metadata_from_workers"] + assert captured["kwargs"]["use_buffers"] is False + assert captured["env_factory"].func is utils._make_libero_worker + + def test_collector_factory_uses_common_env_factory_and_worker_kwargs( + self, monkeypatch + ): + captured = {} + + def fake_make_env_worker(*args, **kwargs): + return args, kwargs + + def fake_parallel_env(num_envs, env_factory, **kwargs): + captured["num_envs"] = num_envs + captured["env_factory"] = env_factory + captured["kwargs"] = kwargs + return object() + + cfg = SimpleNamespace(env=SimpleNamespace(backend="libero")) + tokenizer = object() + monkeypatch.setattr(utils, "_make_env_worker", fake_make_env_worker) + monkeypatch.setattr(utils, "ParallelEnv", fake_parallel_env) + + result = utils._make_collector_env( + cfg, + tokenizer=tokenizer, + num_envs=3, + group_repeats=8, + seed=100, + device=torch.device("cpu"), + worker_idx_offset=16, + render_gpu_device_id=2, + ) + + assert result is not None + assert captured["num_envs"] == 3 + assert captured["kwargs"]["create_env_kwargs"] == [ + {"worker_idx": 0}, + {"worker_idx": 1}, + {"worker_idx": 2}, + ] + assert captured["kwargs"]["metadata_from_workers"] + assert captured["kwargs"]["use_buffers"] is False + assert captured["kwargs"]["mp_start_method"] == "spawn" + args, kwargs = captured["env_factory"](worker_idx=2) + assert args == (cfg, tokenizer) + assert kwargs == { + "worker_idx": 2, + "group_repeats": 8, + "seed": 100, + "device": None, + "worker_idx_offset": 16, + "render_gpu_device_id": 2, + } + + def test_collector_factory_preserves_group_offset_and_seed(self, monkeypatch): + captured = {} + + class FakeLiberoEnv: + def set_seed(self, seed): + captured["seed"] = seed + + def fake_make_libero_worker(cfg, worker_idx, **kwargs): + captured["worker_idx"] = worker_idx + captured["worker_kwargs"] = kwargs + return FakeLiberoEnv() + + monkeypatch.setattr(utils, "_make_libero_worker", fake_make_libero_worker) + monkeypatch.setattr(utils, "_chunk_transform", lambda *args, **kwargs: object()) + monkeypatch.setattr(utils, "TransformedEnv", lambda base, transform: base) + cfg = SimpleNamespace(env=SimpleNamespace(backend="libero")) + + utils._make_env_worker( + cfg, + tokenizer=None, + worker_idx=3, + group_repeats=8, + seed=100, + worker_idx_offset=16, + render_gpu_device_id=2, + ) + + assert captured["worker_idx"] == 3 + assert captured["worker_kwargs"] == { + "group_repeats": 8, + "eval_mode": False, + "from_pixels": False, + "worker_idx_offset": 16, + "render_gpu_device_id": 2, + } + assert captured["seed"] == 119 + def test_libero_worker_assignment_serial_and_parallel_groups(self): class _EnvCfg(dict): def __getattr__(self, name): diff --git a/sota-implementations/vla_grpo/utils.py b/sota-implementations/vla_grpo/utils.py index 37ef2b55f34..88ab92a056f 100644 --- a/sota-implementations/vla_grpo/utils.py +++ b/sota-implementations/vla_grpo/utils.py @@ -623,22 +623,24 @@ def make_env( eval_mode=eval_mode, override=override, ) + env_factory = partial( + _make_libero_worker, + cfg, + group_repeats=group_repeats, + eval_mode=eval_mode, + from_pixels=from_pixels, + worker_idx_offset=worker_idx_offset, + render_gpu_device_id=render_gpu_device_id, + ) base = ParallelEnv( num_envs, - [ - partial( - _make_libero_worker, - cfg, - worker_idx, - group_repeats=group_repeats, - eval_mode=eval_mode, - from_pixels=from_pixels, - worker_idx_offset=worker_idx_offset, - render_gpu_device_id=render_gpu_device_id, - ) - for worker_idx in range(num_envs) + env_factory, + create_env_kwargs=[ + {"worker_idx": worker_idx} for worker_idx in range(num_envs) ], mp_start_method="spawn", + metadata_from_workers=True, + use_buffers=False, # MuJoCo runs on CPU; pin the env device so the collector/rollout # cast the GPU policy's action back to CPU before stepping (else a # cuda action reaches the CPU transforms -> mixed-device error) @@ -669,24 +671,27 @@ def _make_collector_env( worker_idx_offset: int, render_gpu_device_id: int | None, ) -> ParallelEnv: + metadata_from_workers = cfg.env.backend == "libero" + env_factory = partial( + _make_env_worker, + cfg, + tokenizer, + group_repeats=group_repeats, + seed=seed, + device=device if cfg.env.backend == "toy" else None, + worker_idx_offset=worker_idx_offset, + render_gpu_device_id=render_gpu_device_id, + ) return ParallelEnv( num_envs, - [ - partial( - _make_env_worker, - cfg, - tokenizer, - worker_idx, - group_repeats=group_repeats, - seed=seed, - device=device if cfg.env.backend == "toy" else None, - worker_idx_offset=worker_idx_offset, - render_gpu_device_id=render_gpu_device_id, - ) - for worker_idx in range(num_envs) + env_factory, + create_env_kwargs=[ + {"worker_idx": worker_idx} for worker_idx in range(num_envs) ], mp_start_method="spawn", device=device, + metadata_from_workers=metadata_from_workers, + use_buffers=False if metadata_from_workers else None, ) diff --git a/test/envs/test_parallel.py b/test/envs/test_parallel.py index c3ef83c5976..c34eaf8b825 100644 --- a/test/envs/test_parallel.py +++ b/test/envs/test_parallel.py @@ -6,6 +6,8 @@ import gc import os +import time +from pathlib import Path import numpy as np import pytest @@ -31,6 +33,7 @@ EnvCreator, ParallelEnv, SerialEnv, + ToyVLAEnv, TransformedEnv, ) from torchrl.envs.batched_envs import _stackable @@ -56,6 +59,42 @@ ) +class _WorkerMetadataToyEnv(ToyVLAEnv): + def __init__( + self, + worker_idx: int, + marker_dir: str, + *, + state_dim: int = 4, + fail: bool = False, + ) -> None: + self._worker_idx = worker_idx + self._marker_dir = Path(marker_dir) + self._marker_dir.joinpath(f"constructed-{worker_idx}-{os.getpid()}").touch() + if fail: + raise RuntimeError("intentional worker construction failure") + super().__init__( + action_dim=2, + state_dim=state_dim, + instruction=f"instruction-{worker_idx}", + seed=worker_idx, + ) + + def close(self, *, raise_if_closed: bool = True) -> None: + self._marker_dir.joinpath(f"closed-{self._worker_idx}-{os.getpid()}").touch() + super().close(raise_if_closed=raise_if_closed) + + +class _SlowCloseCountingEnv(CountingEnv): + def __init__(self, close_delay: float) -> None: + self.close_delay = close_delay + super().__init__() + + def close(self, *, raise_if_closed: bool = True) -> None: + time.sleep(self.close_delay) + super().close(raise_if_closed=raise_if_closed) + + class TestParallel: @pytest.fixture(autouse=True, scope="class") def disable_autowrap(self): @@ -104,6 +143,95 @@ def make_env(): 4, make_env, create_env_kwargs=[{"seed": 0}, {"seed": 1}] ) + def test_metadata_from_workers_uses_live_envs(self, tmp_path): + env = ParallelEnv( + 2, + _WorkerMetadataToyEnv, + create_env_kwargs=[ + {"worker_idx": i, "marker_dir": str(tmp_path)} for i in range(2) + ], + metadata_from_workers=True, + use_buffers=False, + mp_start_method="spawn", + ) + try: + assert not env.is_closed + assert env._use_buffers is False + constructed = list(tmp_path.glob("constructed-*")) + assert len(constructed) == 2 + assert all(f"-{os.getpid()}" not in path.name for path in constructed) + td = env.reset() + assert td["language_instruction"][0] == "instruction-0" + assert td["language_instruction"][1] == "instruction-1" + finally: + env.close(raise_if_closed=False) + assert len(list(tmp_path.glob("closed-*"))) == 2 + + def test_metadata_from_workers_rejects_buffers(self): + with pytest.raises(RuntimeError, match="requires use_buffers=False"): + ParallelEnv( + 2, + CountingEnv, + metadata_from_workers=True, + use_buffers=True, + ) + + def test_metadata_from_workers_rejects_incompatible_schemas(self, tmp_path): + with pytest.raises(RuntimeError, match="metadata are incompatible"): + ParallelEnv( + 2, + _WorkerMetadataToyEnv, + create_env_kwargs=[ + { + "worker_idx": 0, + "marker_dir": str(tmp_path), + "state_dim": 4, + }, + { + "worker_idx": 1, + "marker_dir": str(tmp_path), + "state_dim": 5, + }, + ], + metadata_from_workers=True, + use_buffers=False, + mp_start_method="spawn", + ) + assert len(list(tmp_path.glob("closed-*"))) == 2 + + def test_metadata_from_workers_reports_construction_failure(self, tmp_path): + with pytest.raises( + RuntimeError, match="intentional worker construction failure" + ): + ParallelEnv( + 2, + _WorkerMetadataToyEnv, + create_env_kwargs=[ + {"worker_idx": 0, "marker_dir": str(tmp_path)}, + {"worker_idx": 1, "marker_dir": str(tmp_path), "fail": True}, + ], + metadata_from_workers=True, + use_buffers=False, + mp_start_method="spawn", + ) + assert len(list(tmp_path.glob("closed-0-*"))) == 1 + + def test_metadata_from_workers_shutdown_is_bounded(self): + env = ParallelEnv( + 1, + _SlowCloseCountingEnv, + create_env_kwargs={"close_delay": 10.0}, + metadata_from_workers=True, + use_buffers=False, + mp_start_method="spawn", + ) + workers = list(env._workers) + env._timeout = 0.1 + + env.close() + + assert all(not worker.is_alive() for worker in workers) + def test_compact_collector_skips_next_observation_copy( self, maybe_fork_ParallelEnv ): diff --git a/test/envs/test_special.py b/test/envs/test_special.py index e746a02156b..af0447f7d22 100644 --- a/test/envs/test_special.py +++ b/test/envs/test_special.py @@ -61,6 +61,25 @@ ) +def test_callable_metadata_env_closes_when_extraction_fails(monkeypatch): + env = ContinuousActionVecMockEnv() + closed = False + + def fake_tensordict(): + raise RuntimeError("metadata extraction failed") + + def close(): + nonlocal closed + closed = True + + monkeypatch.setattr(env, "fake_tensordict", fake_tensordict) + monkeypatch.setattr(env, "close", close) + + with pytest.raises(RuntimeError, match="metadata extraction failed"): + get_env_metadata(lambda: env) + assert closed + + @pytest.mark.parametrize( "envclass", [ diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py index b40e458028d..1d02ca363aa 100644 --- a/torchrl/envs/batched_envs.py +++ b/torchrl/envs/batched_envs.py @@ -264,6 +264,11 @@ def __call__(cls, *args, **kwargs): # multiprocessing with the spawn start method. Lambda functions cannot # be serialized with standard pickle, but EnvCreator uses cloudpickle. auto_wrap_envs = kwargs.pop("auto_wrap_envs", True) + if kwargs.get("metadata_from_workers", False): + # The worker path wraps callables with CloudpickleWrapper directly. + # EnvCreator would eagerly construct a shadow env in the parent, + # defeating worker-originated metadata. + auto_wrap_envs = False def _warn_lambda(): if rl_warnings(): @@ -378,6 +383,11 @@ class BatchedEnvBase(EnvBase): defaults to ``use_buffers=False`` (with a warning) and the data exchanged between processes is staged on CPU. Passing ``use_buffers=True`` in that case raises a ``RuntimeError``. + metadata_from_workers (bool, optional): if ``True``, each worker constructs + its environment and sends its metadata to the parent during startup. This + avoids constructing temporary environments in the parent process. The mode + is only supported by :class:`~torchrl.envs.ParallelEnv`, starts its workers + eagerly, and currently requires ``use_buffers=False``. Defaults to ``False``. daemon (bool, optional): whether the processes should be daemonized. This is only applicable to parallel environments such as :class:`~torchrl.envs.ParallelEnv`. Defaults to ``False``. @@ -509,6 +519,7 @@ def __init__( non_blocking: bool = False, mp_start_method: str | None = None, use_buffers: bool | None = None, + metadata_from_workers: bool = False, consolidate: bool = True, daemon: bool = False, ): @@ -522,9 +533,22 @@ def __init__( self.num_threads = num_threads self._cache_in_keys = None self._use_buffers = use_buffers + self._metadata_from_workers = metadata_from_workers self.consolidate = consolidate self.daemon = daemon + if metadata_from_workers: + if not isinstance(self, ParallelEnv): + raise TypeError( + "metadata_from_workers=True is only supported by ParallelEnv." + ) + if use_buffers: + raise RuntimeError( + "metadata_from_workers=True currently requires use_buffers=False " + "because shared buffers must be allocated before workers start." + ) + self._use_buffers = False + self._single_task = callable(create_env_fn) or (len(set(create_env_fn)) == 1) if callable(create_env_fn): create_env_fn = [create_env_fn for _ in range(num_workers)] @@ -573,15 +597,17 @@ def __init__( self._seeds = None self.__dict__["_input_spec"] = None self.__dict__["_output_spec"] = None - # self._prepare_dummy_env(create_env_fn, create_env_kwargs) self._properties_set = False - self._get_metadata(create_env_fn, create_env_kwargs) self._non_blocking = non_blocking if mp_start_method is not None and not isinstance(self, ParallelEnv): raise TypeError( f"Cannot use mp_start_method={mp_start_method} with envs of type {type(self)}." ) self._mp_start_method = mp_start_method + if metadata_from_workers: + self._start_workers() + else: + self._get_metadata(create_env_fn, create_env_kwargs) is_spec_locked = EnvBase.is_spec_locked @@ -910,6 +936,21 @@ def _get_metadata( create_env_kwargs[0], env_validator=self._validate_worker_env, ) + else: + meta_data = [ + get_env_metadata( + create_env_fn[i], + create_env_kwargs[i], + env_validator=self._validate_worker_env, + ).clone() + for i in range(len(create_env_fn)) + ] + self._set_metadata(meta_data) + + def _set_metadata(self, meta_data: EnvMetaData | list[EnvMetaData]) -> None: + if self._single_task: + if isinstance(meta_data, list): + raise TypeError("Expected one metadata object for a homogeneous env.") self.meta_data = meta_data.expand( *(self.num_workers, *meta_data.batch_size) ) @@ -928,16 +969,9 @@ def _get_metadata( if self.share_individual_td is None: self.share_individual_td = False else: - n_tasks = len(create_env_fn) - self.meta_data: list[EnvMetaData] = [] - for i in range(n_tasks): - self.meta_data.append( - get_env_metadata( - create_env_fn[i], - create_env_kwargs[i], - env_validator=self._validate_worker_env, - ).clone() - ) + if not isinstance(meta_data, list): + raise TypeError("Expected a metadata list for heterogeneous envs.") + self.meta_data = meta_data if self.share_individual_td is not True: share_individual_td = not _stackable( *[meta_data.tensordict for meta_data in self.meta_data] @@ -967,6 +1001,64 @@ def _get_metadata( self._set_properties() + @staticmethod + def _worker_metadata_schema(meta_data: EnvMetaData) -> tuple: + """Return metadata fields that determine safe batched communication. + + Non-tensor payload values are intentionally excluded: workers may expose + different strings (for example, language instructions) while retaining an + identical tensor schema. + """ + spec_schema = tuple( + ( + key, + type(spec), + torch.Size(spec.shape), + getattr(spec, "dtype", None), + getattr(spec, "device", None), + ) + for key, spec in meta_data.specs.items(True, True) + ) + tensordict_schema = [] + for key, value in meta_data.tensordict.items(True, True): + if isinstance(value, torch.Tensor): + value_schema = ( + type(value), + torch.Size(value.shape), + value.dtype, + value.device, + ) + else: + value_schema = (type(value), getattr(value, "shape", None)) + tensordict_schema.append((key, value_schema)) + return ( + meta_data.batch_size, + meta_data.device, + meta_data.batch_locked, + meta_data.supports_set_state, + tuple(meta_data.device_map.items()), + spec_schema, + tuple(tensordict_schema), + ) + + def _set_worker_metadata(self, metadata: list[EnvMetaData]) -> None: + if len(metadata) != self.num_workers: + raise RuntimeError( + "ParallelEnv received metadata from " + f"{len(metadata)} workers, expected {self.num_workers}." + ) + reference_schema = self._worker_metadata_schema(metadata[0]) + for worker_idx, worker_metadata in enumerate(metadata[1:], 1): + if self._worker_metadata_schema(worker_metadata) != reference_schema: + raise RuntimeError( + "ParallelEnv worker metadata are incompatible: " + f"worker {worker_idx} has a different tensor schema from worker 0." + ) + if self._single_task: + self._set_metadata(metadata[0]) + else: + self._set_metadata(metadata) + def update_kwargs(self, kwargs: dict | list[dict]) -> None: """Updates the kwargs of each environment given a dictionary or a list of dictionaries. @@ -1954,6 +2046,7 @@ class ParallelEnv(BatchedEnvBase, metaclass=_PEnvMeta): def _start_workers(self) -> None: import torchrl + initial_num_threads = torch.get_num_threads() self._timeout = 10.0 self.BATCHED_PIPE_TIMEOUT = torchrl._utils.BATCHED_PIPE_TIMEOUT @@ -2031,7 +2124,7 @@ def look_for_cuda(tensor, has_cuda=has_cuda): "child_pipe": child_pipe, "env_fun": env_fun, "env_fun_kwargs": self.create_env_kwargs[idx], - "has_lazy_inputs": self.has_lazy_inputs, + "has_lazy_inputs": self.__dict__.get("has_lazy_inputs", False), "num_threads": num_sub_threads, "non_blocking": self.non_blocking, "filter_warnings": self._filter_warnings_subprocess(), @@ -2051,6 +2144,7 @@ def look_for_cuda(tensor, has_cuda=has_cuda): kwargs[idx].update( { "consolidate": self.consolidate, + "metadata_from_worker": self._metadata_from_workers, } ) process = proc_fun(target=func, kwargs=kwargs[idx]) @@ -2060,16 +2154,103 @@ def look_for_cuda(tensor, has_cuda=has_cuda): self.parent_channels.append(parent_pipe) self._workers.append(process) - for parent_pipe in self.parent_channels: - # use msg as sync point - parent_pipe.recv() + try: + if self._metadata_from_workers: + metadata = self._receive_worker_metadata() + self._set_worker_metadata(metadata) + else: + for parent_pipe in self.parent_channels: + # use msg as sync point + parent_pipe.recv() + + for channel in self.parent_channels: + channel.send(("init", None)) + except Exception: + self._cleanup_failed_worker_startup() + torch.set_num_threads(initial_num_threads) + raise - # send shared tensordict to workers - for channel in self.parent_channels: - channel.send(("init", None)) self.is_closed = False self.set_spec_lock_() + def _receive_worker_metadata(self) -> list[EnvMetaData]: + metadata = [None] * self.num_workers + pending = { + channel: worker_idx + for worker_idx, channel in enumerate(self.parent_channels) + } + deadline = time.monotonic() + self.BATCHED_PIPE_TIMEOUT + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "ParallelEnv timed out while waiting for worker metadata. " + "Increase BATCHED_PIPE_TIMEOUT if environment construction " + "legitimately takes longer." + ) + ready = connection_wait(list(pending), timeout=remaining) + if not ready: + continue + for channel in ready: + worker_idx = pending.pop(channel) + try: + message, payload = channel.recv() + except EOFError as err: + raise RuntimeError( + f"ParallelEnv worker {worker_idx} exited before sending metadata." + ) from err + if message == "metadata_error": + raise RuntimeError( + f"ParallelEnv worker {worker_idx} failed during metadata " + f"construction:\n{payload}" + ) + if message != "metadata" or not isinstance(payload, EnvMetaData): + raise RuntimeError( + f"ParallelEnv worker {worker_idx} sent an invalid metadata " + f"message: {message!r}." + ) + metadata[worker_idx] = payload + return metadata + + def _cleanup_failed_worker_startup(self) -> None: + for channel, process in zip(self.parent_channels, self._workers): + if process.is_alive(): + try: + channel.send(("close", None)) + except (EOFError, OSError): + pass + deadline = time.monotonic() + self._timeout + for event, process in zip(self._events, self._workers): + if process.is_alive(): + event.wait(max(0.0, deadline - time.monotonic())) + join_deadline = time.monotonic() + min(self._timeout, 2.0) + for process in self._workers: + process.join(timeout=max(0.0, join_deadline - time.monotonic())) + self._terminate_and_join_workers(self._workers, timeout=min(self._timeout, 2.0)) + for channel in self.parent_channels: + channel.close() + self.parent_channels = [] + self._workers = [] + self._events = None + self.event = None + self.is_closed = True + + @staticmethod + def _terminate_and_join_workers(workers: list[mp.Process], timeout: float) -> None: + for process in workers: + if process.is_alive(): + process.terminate() + deadline = time.monotonic() + timeout + for process in workers: + process.join(timeout=max(0.0, deadline - time.monotonic())) + for process in workers: + if process.is_alive(): + process.kill() + deadline = time.monotonic() + 1.0 + for process in workers: + if process.is_alive(): + process.join(timeout=max(0.0, deadline - time.monotonic())) + def _filter_warnings_subprocess(self) -> bool: from torchrl import filter_warnings_subprocess @@ -2950,34 +3131,46 @@ def _shutdown_workers(self) -> None: raise RuntimeError( "calling {self.__class__.__name__}._shutdown_workers only allowed when env.is_closed = False" ) - for i, channel in enumerate(self.parent_channels): - if self._verbose: - torchrl_logger.info(f"closing {i}") - channel.send(("close", None)) - # Wait on mp.Event (not _wait_for_workers) because the "close" - # handler doesn't send data on the pipe — it just closes it — and - # connection_wait does not reliably detect socketpair closure - # on all platforms (macOS with forked workers). - for i in range(self.num_workers): - self._events[i].wait(self._timeout) - self._events[i].clear() + if self._metadata_from_workers: + # Expensive native environments may retain renderer resources + # until process teardown. Reap each process before asking the + # next one to close to avoid a second resource spike. + worker_groups = [ + (worker_idx,) for worker_idx in range(self.num_workers) + ] + else: + worker_groups = [range(self.num_workers)] + for worker_indices in worker_groups: + for worker_idx in worker_indices: + if self._verbose: + torchrl_logger.info(f"closing {worker_idx}") + self.parent_channels[worker_idx].send(("close", None)) + # Wait on mp.Event (not _wait_for_workers) because the "close" + # handler doesn't send data on the pipe — it just closes it — and + # connection_wait does not reliably detect socketpair closure + # on all platforms (macOS with forked workers). + for worker_idx in worker_indices: + self._events[worker_idx].wait(self._timeout) + self._events[worker_idx].clear() + self.parent_channels[worker_idx].close() + processes = [self._workers[worker_idx] for worker_idx in worker_indices] + exit_deadline = time.monotonic() + self._timeout + while ( + any(process.is_alive() for process in processes) + and time.monotonic() < exit_deadline + ): + time.sleep(0.01) + self._terminate_and_join_workers( + processes, timeout=min(self._timeout, 2.0) + ) if self._use_buffers: del self.shared_tensordicts, self.shared_tensordict_parent - + finally: for channel in self.parent_channels: channel.close() - start_time = time.time() - while ( - any(proc.is_alive() for proc in self._workers) - and (time.time() - start_time) < self._timeout - ): - time.sleep(0.01) - for proc in self._workers: - proc.join() - finally: - for proc in self._workers: - if proc.is_alive(): - proc.terminate() + self._terminate_and_join_workers( + self._workers, timeout=min(self._timeout, 2.0) + ) del self._workers del self.parent_channels self._cuda_events = None @@ -3449,6 +3642,7 @@ def _run_worker_pipe_direct( num_threads: int | None = None, # for fork start method consolidate: bool = True, filter_warnings: bool = False, + metadata_from_worker: bool = False, ) -> None: # Handle warning filtering (moved from _ProcessNoWarn) if filter_warnings: @@ -3458,14 +3652,34 @@ def _run_worker_pipe_direct( parent_pipe.close() pid = os.getpid() - if not isinstance(env_fun, EnvBase): - env = env_fun(**env_fun_kwargs) - else: - if env_fun_kwargs: - raise RuntimeError( - "env_fun_kwargs must be empty if an environment is passed to a process." - ) - env = env_fun + env = None + try: + if not isinstance(env_fun, EnvBase): + env = env_fun(**env_fun_kwargs) + else: + if env_fun_kwargs: + raise RuntimeError( + "env_fun_kwargs must be empty if an environment is passed to a process." + ) + env = env_fun + if metadata_from_worker: + BatchedEnvBase._validate_worker_env(env) + child_pipe.send(("metadata", EnvMetaData.metadata_from_env(env))) + except Exception as err: + if not metadata_from_worker: + raise + try: + child_pipe.send(("metadata_error", f"{type(err).__name__}: {err}")) + except (EOFError, OSError): + pass + if env is not None: + try: + env.close() + except Exception: + pass + mp_event.set() + child_pipe.close() + return del env_fun for spec in env.output_spec.values(True, True): if spec.device is not None and spec.device.type == "cuda": @@ -3504,7 +3718,8 @@ def _run_worker_pipe_direct( initialized = False - child_pipe.send("started") + if not metadata_from_worker: + child_pipe.send("started") while True: try: if child_pipe.poll(_timeout): @@ -3653,8 +3868,6 @@ def _run_worker_pipe_direct( mp_event.set() elif cmd == "close": - if not initialized: - raise RuntimeError("call 'init' before closing") env.close() mp_event.set() child_pipe.close() diff --git a/torchrl/envs/env_creator.py b/torchrl/envs/env_creator.py index 2078e8cb305..e52efcaba00 100644 --- a/torchrl/envs/env_creator.py +++ b/torchrl/envs/env_creator.py @@ -260,9 +260,12 @@ def get_env_metadata( if kwargs is None: kwargs = {} env = env_or_creator(**kwargs) - if env_validator is not None: - env_validator(env) - return EnvMetaData.metadata_from_env(env) + try: + if env_validator is not None: + env_validator(env) + return EnvMetaData.metadata_from_env(env) + finally: + env.close() elif isinstance(env_or_creator, EnvCreator): if not ( kwargs == env_or_creator.create_env_kwargs From efac773da245491bb09a4b7673ef71a4a3aca958 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Fri, 17 Jul 2026 09:58:49 +0100 Subject: [PATCH 2/2] Address review: serial_for_single fallback, shutdown_timeout, worker tracebacks - Drop ParallelEnv-only kwargs (mp_start_method, metadata_from_workers) when serial_for_single falls back to SerialEnv instead of raising a TypeError, and retain mp_start_method when the env stays parallel - Expose the worker shutdown grace period as a shutdown_timeout constructor argument (also settable via configure_parallel) instead of a hardcoded 10s - Send the full worker traceback in metadata_error messages so construction failures are debuggable from the parent - Close transient metadata envs with raise_if_closed=False so a close failure cannot mask the metadata result or the original error - Document the identical-schema requirement and sequential shutdown of metadata_from_workers in the docstring and the vectorized-envs docs Co-Authored-By: Claude Fable 5 --- docs/source/reference/envs_vectorized.rst | 8 ++++- test/envs/test_parallel.py | 19 +++++++++++- test/envs/test_special.py | 2 +- torchrl/envs/batched_envs.py | 36 ++++++++++++++++++----- torchrl/envs/env_creator.py | 4 ++- 5 files changed, 58 insertions(+), 11 deletions(-) diff --git a/docs/source/reference/envs_vectorized.rst b/docs/source/reference/envs_vectorized.rst index cf44f9bf221..605eabd07d3 100644 --- a/docs/source/reference/envs_vectorized.rst +++ b/docs/source/reference/envs_vectorized.rst @@ -69,7 +69,13 @@ needed. normal initialization. The workers start eagerly, their tensor schemas are validated for compatibility, and no temporary environment is created in the parent. This mode currently requires pipe-based communication with - ``use_buffers=False``. Use one common factory with ``create_env_kwargs`` for + ``use_buffers=False``. All workers must expose the same tensor schema: specs + and example tensors may only differ in non-tensor payload values (such as + language instructions). Environments with genuinely heterogeneous specs + should keep the default metadata path. At shutdown, workers started in this + mode are closed one at a time to bound teardown resource spikes; the + per-worker grace period is controlled by the ``shutdown_timeout`` argument. + Use one common factory with ``create_env_kwargs`` for worker-specific arguments: .. code-block:: python diff --git a/test/envs/test_parallel.py b/test/envs/test_parallel.py index c34eaf8b825..36df3181d41 100644 --- a/test/envs/test_parallel.py +++ b/test/envs/test_parallel.py @@ -224,14 +224,28 @@ def test_metadata_from_workers_shutdown_is_bounded(self): metadata_from_workers=True, use_buffers=False, mp_start_method="spawn", + shutdown_timeout=0.1, ) workers = list(env._workers) - env._timeout = 0.1 env.close() assert all(not worker.is_alive() for worker in workers) + def test_metadata_from_workers_serial_for_single_fallback(self): + env = ParallelEnv( + 1, + CountingEnv, + serial_for_single=True, + metadata_from_workers=True, + use_buffers=False, + ) + try: + assert isinstance(env, SerialEnv) + env.reset() + finally: + env.close(raise_if_closed=False) + def test_compact_collector_skips_next_observation_copy( self, maybe_fork_ParallelEnv ): @@ -320,6 +334,9 @@ def test_serial_for_single(self, maybe_fork_ParallelEnv, start_method): mp_start_method=start_method, ) assert isinstance(env, ParallelEnv) + # serial_for_single must not swallow the start method when the + # env stays parallel + assert env._mp_start_method == start_method finally: env.close(raise_if_closed=False) diff --git a/test/envs/test_special.py b/test/envs/test_special.py index af0447f7d22..8a049fa3326 100644 --- a/test/envs/test_special.py +++ b/test/envs/test_special.py @@ -68,7 +68,7 @@ def test_callable_metadata_env_closes_when_extraction_fails(monkeypatch): def fake_tensordict(): raise RuntimeError("metadata extraction failed") - def close(): + def close(*, raise_if_closed: bool = True): nonlocal closed closed = True diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py index 1d02ca363aa..806ab6cc7dd 100644 --- a/torchrl/envs/batched_envs.py +++ b/torchrl/envs/batched_envs.py @@ -9,6 +9,7 @@ import gc import os import time +import traceback import warnings import weakref from collections import OrderedDict @@ -252,12 +253,14 @@ def __call__(cls, *args, **kwargs): serial_for_single = kwargs.pop("serial_for_single", False) if serial_for_single: num_workers = kwargs.get("num_workers") - # Remove start method from kwargs - kwargs.pop("mp_start_method", None) if num_workers is None: num_workers = args[0] if num_workers == 1: - # We still use a serial to keep the shape unchanged + # We still use a serial to keep the shape unchanged. + # SerialEnv constructs its envs in-process, so ParallelEnv-only + # kwargs are dropped rather than forwarded (they would raise). + kwargs.pop("mp_start_method", None) + kwargs.pop("metadata_from_workers", None) return SerialEnv(*args, **kwargs) # Wrap lambda functions with EnvCreator so they can be pickled for @@ -387,10 +390,21 @@ class BatchedEnvBase(EnvBase): its environment and sends its metadata to the parent during startup. This avoids constructing temporary environments in the parent process. The mode is only supported by :class:`~torchrl.envs.ParallelEnv`, starts its workers - eagerly, and currently requires ``use_buffers=False``. Defaults to ``False``. + eagerly, and currently requires ``use_buffers=False``. All workers must + report the same tensor schema: specs and example tensors may only differ + in non-tensor payload values (such as language instructions). In this + mode workers are also closed one at a time at shutdown to bound teardown + resource spikes. Defaults to ``False``. daemon (bool, optional): whether the processes should be daemonized. This is only applicable to parallel environments such as :class:`~torchrl.envs.ParallelEnv`. Defaults to ``False``. + shutdown_timeout (float, optional): grace period in seconds granted to + workers to exit cleanly when the environment is closed. Workers that + are still alive after this window are terminated, then killed. + Increase it for environments whose ``close()`` legitimately takes + long (e.g. flushing recorders or tearing down native renderers). + This is only applicable to parallel environments such as + :class:`~torchrl.envs.ParallelEnv`. Defaults to ``10.0``. auto_wrap_envs (bool, optional): if ``True`` (default), lambda functions passed as ``create_env_fn`` will be automatically wrapped in an :class:`~torchrl.envs.EnvCreator` to enable pickling for multiprocessing with the ``spawn`` start method. @@ -522,6 +536,7 @@ def __init__( metadata_from_workers: bool = False, consolidate: bool = True, daemon: bool = False, + shutdown_timeout: float = 10.0, ): super().__init__(device=device) self.serial_for_single = serial_for_single @@ -536,6 +551,7 @@ def __init__( self._metadata_from_workers = metadata_from_workers self.consolidate = consolidate self.daemon = daemon + self.shutdown_timeout = float(shutdown_timeout) if metadata_from_workers: if not isinstance(self, ParallelEnv): @@ -622,6 +638,7 @@ def configure_parallel( num_sub_threads: int | None = None, non_blocking: bool | None = None, daemon: bool | None = None, + shutdown_timeout: float | None = None, ) -> BatchedEnvBase: """Configure parallel execution parameters before the environment starts. @@ -642,6 +659,8 @@ def configure_parallel( non_blocking (bool, optional): if ``True``, device moves will be done using the ``non_blocking=True`` option. daemon (bool, optional): whether the processes should be daemonized. + shutdown_timeout (float, optional): grace period in seconds granted to + workers to exit cleanly at shutdown before they are terminated. Returns: self: Returns self for method chaining. @@ -677,6 +696,8 @@ def configure_parallel( self._non_blocking = non_blocking if daemon is not None: self.daemon = daemon + if shutdown_timeout is not None: + self.shutdown_timeout = float(shutdown_timeout) return self def select_and_clone(self, name, tensor, selected_keys=None): @@ -2047,7 +2068,8 @@ def _start_workers(self) -> None: import torchrl initial_num_threads = torch.get_num_threads() - self._timeout = 10.0 + # getattr: envs unpickled from older versions lack the attribute + self._timeout = getattr(self, "shutdown_timeout", 10.0) self.BATCHED_PIPE_TIMEOUT = torchrl._utils.BATCHED_PIPE_TIMEOUT num_threads = max( @@ -3665,11 +3687,11 @@ def _run_worker_pipe_direct( if metadata_from_worker: BatchedEnvBase._validate_worker_env(env) child_pipe.send(("metadata", EnvMetaData.metadata_from_env(env))) - except Exception as err: + except Exception: if not metadata_from_worker: raise try: - child_pipe.send(("metadata_error", f"{type(err).__name__}: {err}")) + child_pipe.send(("metadata_error", traceback.format_exc())) except (EOFError, OSError): pass if env is not None: diff --git a/torchrl/envs/env_creator.py b/torchrl/envs/env_creator.py index e52efcaba00..618d398c485 100644 --- a/torchrl/envs/env_creator.py +++ b/torchrl/envs/env_creator.py @@ -265,7 +265,9 @@ def get_env_metadata( env_validator(env) return EnvMetaData.metadata_from_env(env) finally: - env.close() + # raise_if_closed=False: a close failure must not mask the + # metadata result or the original validator error. + env.close(raise_if_closed=False) elif isinstance(env_or_creator, EnvCreator): if not ( kwargs == env_or_creator.create_env_kwargs