diff --git a/benchmarks/test_envs_benchmark.py b/benchmarks/test_envs_benchmark.py index 56ac2487b70..bf5c17e6029 100644 --- a/benchmarks/test_envs_benchmark.py +++ b/benchmarks/test_envs_benchmark.py @@ -3,6 +3,8 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. import argparse +import time +from functools import partial import pytest import torch @@ -12,6 +14,30 @@ from torchrl.envs.libs.dm_control import DMControlEnv from torchrl.envs.libs.libero import _has_libero, LiberoEnv from torchrl.envs.transforms.functional import cat_frames +from torchrl.testing.mocking_classes import CountingEnv + + +class _SlowStartEnv(CountingEnv): + def __init__(self, delay: float = 0.1): + time.sleep(delay) + super().__init__() + + +def _make_slow_start_env(delay: float) -> _SlowStartEnv: + return _SlowStartEnv(delay) + + +def _run_parallel_cold_start(metadata_from_workers: bool) -> None: + env = ParallelEnv( + 4, + partial(_make_slow_start_env, 0.1), + mp_start_method="spawn", + metadata_from_workers=metadata_from_workers, + ) + try: + env.reset() + finally: + env.close(raise_if_closed=False) def make_simple_env(): @@ -111,6 +137,16 @@ def test_parallel(benchmark): benchmark(execute_env, c) +@pytest.mark.parametrize("metadata_from_workers", [False, True]) +def test_parallel_cold_start(benchmark, metadata_from_workers): + benchmark.pedantic( + _run_parallel_cold_start, + args=(metadata_from_workers,), + rounds=3, + iterations=1, + ) + + @pytest.mark.skipif(not _has_libero, reason="libero not found") def test_libero(benchmark): # raw simulation + render throughput of the LIBERO adapter (steps/s) diff --git a/docs/source/reference/envs_vectorized.rst b/docs/source/reference/envs_vectorized.rst index 38f5b9e4e8c..1c985fb4585 100644 --- a/docs/source/reference/envs_vectorized.rst +++ b/docs/source/reference/envs_vectorized.rst @@ -24,6 +24,28 @@ It is important that your environment specs match the input and output that it s :class:`ParallelEnv` will create buffers from these specs to communicate with the spawn processes. Check the :func:`~torchrl.envs.check_env_specs` method for a sanity check. +By default, :class:`ParallelEnv` obtains those specs by creating temporary +environments in the parent process. For environments that are expensive or +unsafe to construct in the parent, ``metadata_from_workers=True`` obtains the +metadata from the real, long-lived worker environments instead. The workers +start during :class:`ParallelEnv` construction, their schemas are validated, +and no parent-side environment is created. This opt-in mode currently requires +``use_buffers=False`` (which is selected automatically when the argument is +unset); passing ``use_buffers=True`` raises an error. A common factory can still +receive per-worker settings through ``create_env_kwargs``: + +.. code-block:: python + + from functools import partial + + make_env = partial(GymEnv, "Pendulum-v1") + env = ParallelEnv( + 4, + make_env, + create_env_kwargs=[{"frame_skip": i + 1} for i in range(4)], + metadata_from_workers=True, + ) + .. code-block:: :caption: Parallel environment diff --git a/sota-implementations/vla_grpo/README.md b/sota-implementations/vla_grpo/README.md index cd4981acae2..87513941047 100644 --- a/sota-implementations/vla_grpo/README.md +++ b/sota-implementations/vla_grpo/README.md @@ -243,6 +243,30 @@ workers plus the evaluator share one process policy server. python sota-implementations/vla_grpo/vla-grpo.py --config-name vla_grpo_libero ``` +The LIBERO recipe enables worker-originated environment metadata. This keeps +MuJoCo/EGL construction inside the long-lived rollout workers and avoids a +temporary parent-side environment wave. The nested startup path can be measured +without loading the VLA policy: + +```bash +python sota-implementations/vla_grpo/bench_libero_startup.py \ + --mode legacy-parent --inner-start-method spawn \ + --output-dir /root/artifacts/libero-startup/legacy-spawn +python sota-implementations/vla_grpo/bench_libero_startup.py \ + --mode worker-metadata --inner-start-method spawn \ + --output-dir /root/artifacts/libero-startup/worker-spawn +python sota-implementations/vla_grpo/bench_libero_startup.py \ + --mode worker-metadata --inner-start-method forkserver \ + --output-dir /root/artifacts/libero-startup/worker-forkserver +``` + +Outer subcollector processes always use ``spawn``. The script permits +``forkserver`` or diagnostic ``fork`` only in worker-metadata mode, after the +subcollector parent has been kept free of EGL contexts. Each output directory +contains the exact command, per-construction marker files, and a JSON summary +covering nested environment readiness, first-step latency, peak process count, +and shutdown cleanup. + Rollout clients request random sampling; eval and video clients request deterministic decoding from the same server, so rollout and eval are synced by one explicit TensorDict weight update after each optimizer step. The replay diff --git a/sota-implementations/vla_grpo/bench_libero_startup.py b/sota-implementations/vla_grpo/bench_libero_startup.py new file mode 100644 index 00000000000..8ab16a8ee3f --- /dev/null +++ b/sota-implementations/vla_grpo/bench_libero_startup.py @@ -0,0 +1,459 @@ +# 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 nested LIBERO ``ParallelEnv`` startup without loading a policy. + +The benchmark mirrors the production five-by-sixty-four collector topology, +records every environment construction, executes one random outer step, and +verifies that shutdown leaves no environment processes behind. + +Example: + python sota-implementations/vla_grpo/bench_libero_startup.py \ + --mode worker-metadata --inner-start-method spawn \ + --output-dir /root/artifacts/libero-startup/worker-spawn +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +import traceback +from functools import partial +from pathlib import Path +from queue import Empty + +import torch +from omegaconf import OmegaConf +from torch import multiprocessing as mp +from torchrl._utils import logger as torchrl_logger, timeit +from torchrl.envs import ParallelEnv + +_VLA_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(_VLA_DIR)) + +import utils as vla_utils + + +def _instrumented_make_env_worker( + cfg, + tokenizer, + worker_idx: int, + *, + marker_dir: str, + subcollector_idx: int, + **kwargs, +): + parent_pid = int(os.environ["TORCHRL_LIBERO_BENCH_PARENT_PID"]) + role = "parent_metadata" if os.getpid() == parent_pid else "worker" + marker_path = Path(marker_dir) + marker_path.mkdir(parents=True, exist_ok=True) + marker = marker_path / ( + f"construct-{role}-sub{subcollector_idx}-pid{os.getpid()}-" + f"worker{worker_idx}" + ) + marker.touch() + return vla_utils._make_env_worker( + cfg, + tokenizer, + worker_idx, + **kwargs, + ) + + +def _make_nested_env( + cfg, + *, + mode: str, + inner_start_method: str, + num_envs: int, + subcollector_idx: int, + marker_dir: str, +) -> ParallelEnv: + worker_idx_offset = subcollector_idx * num_envs + render_gpu_device_id = vla_utils._render_gpu_for_subcollector(cfg, subcollector_idx) + factory = partial( + _instrumented_make_env_worker, + cfg, + None, + marker_dir=marker_dir, + subcollector_idx=subcollector_idx, + group_repeats=int(cfg.collector.group_size), + seed=int(cfg.env.seed), + device=None, + worker_idx_offset=worker_idx_offset, + render_gpu_device_id=render_gpu_device_id, + ) + common_kwargs = { + "mp_start_method": inner_start_method, + "device": torch.device("cpu"), + } + if mode == "legacy-parent": + return ParallelEnv( + num_envs, + [partial(factory, worker_idx=worker_idx) for worker_idx in range(num_envs)], + **common_kwargs, + ) + return ParallelEnv( + num_envs, + factory, + create_env_kwargs=[ + {"worker_idx": worker_idx} for worker_idx in range(num_envs) + ], + metadata_from_workers=mode == "worker-metadata", + **common_kwargs, + ) + + +def _subcollector_main( + cfg_path: str, + *, + mode: str, + inner_start_method: str, + num_envs: int, + subcollector_idx: int, + marker_dir: str, + status_queue, + step_event, +) -> None: + env = None + os.environ["TORCHRL_LIBERO_BENCH_PARENT_PID"] = str(os.getpid()) + try: + cfg = OmegaConf.load(cfg_path) + construct_timer = timeit( + f"libero_startup/subcollector_{subcollector_idx}/construct" + ).start() + env = _make_nested_env( + cfg, + mode=mode, + inner_start_method=inner_start_method, + num_envs=num_envs, + subcollector_idx=subcollector_idx, + marker_dir=marker_dir, + ) + if env.is_closed: + env.start() + status_queue.put( + { + "event": "ready", + "subcollector": subcollector_idx, + "construct_s": construct_timer.elapsed(), + "worker_pids": [process.pid for process in env._workers], + } + ) + if not step_event.wait(timeout=1800): + raise TimeoutError("timed out waiting for the first-step barrier") + + step_timer = timeit( + f"libero_startup/subcollector_{subcollector_idx}/first_step" + ).start() + reset = env.reset() + step = env.rand_step(reset) + instructions = reset.get("language_instruction") + group_ids = reset.get("group_id") + instruction_values = [str(instruction) for instruction in instructions] + status_queue.put( + { + "event": "stepped", + "subcollector": subcollector_idx, + "first_step_s": step_timer.elapsed(), + "reset_batch_size": list(reset.batch_size), + "step_batch_size": list(step.batch_size), + "instruction_count": len(instruction_values), + "unique_instruction_count": len(set(instruction_values)), + "group_ids": torch.as_tensor(group_ids).reshape(-1).tolist(), + } + ) + except Exception as err: + status_queue.put( + { + "event": "error", + "subcollector": subcollector_idx, + "error": repr(err), + "traceback": traceback.format_exc(), + } + ) + finally: + if env is not None: + env.close(raise_if_closed=False) + status_queue.put( + { + "event": "closed", + "subcollector": subcollector_idx, + } + ) + + +def _descendant_pids(pid: int) -> set[int]: + descendants = set() + pending = [pid] + while pending: + parent = pending.pop() + children_path = Path(f"/proc/{parent}/task/{parent}/children") + try: + children = [int(child) for child in children_path.read_text().split()] + except (FileNotFoundError, ProcessLookupError): + continue + for child in children: + if child not in descendants: + descendants.add(child) + pending.append(child) + return descendants + + +def _process_cmdline(pid: int) -> str: + try: + return Path(f"/proc/{pid}/cmdline").read_bytes().replace(b"\0", b" ").decode() + except (FileNotFoundError, ProcessLookupError, UnicodeDecodeError): + return "" + + +def _collect_until( + status_queue, + *, + event: str, + count: int, + processes, + timeout_s: float, + total_timer, + messages: list[dict], +) -> tuple[list[dict], int]: + selected = [] + selected_subcollectors = set() + peak_processes = 0 + deadline_s = total_timer.elapsed() + timeout_s + while len(selected) < count: + if total_timer.elapsed() > deadline_s: + raise TimeoutError( + f"timed out waiting for {event!r} messages: " + f"received {len(selected)}/{count}" + ) + peak_processes = max(peak_processes, len(_descendant_pids(os.getpid()))) + try: + message = status_queue.get(timeout=1.0) + except Empty: + dead = [ + process.pid + for subcollector_idx, process in enumerate(processes) + if subcollector_idx not in selected_subcollectors + and not process.is_alive() + ] + if dead: + raise RuntimeError( + f"outer subcollector processes exited before {event!r}: {dead}" + ) + continue + peak_processes = max(peak_processes, len(_descendant_pids(os.getpid()))) + messages.append(message) + if message["event"] == "error": + raise RuntimeError( + f"subcollector {message['subcollector']} failed: " + f"{message['error']}\n{message['traceback']}" + ) + if message["event"] == event: + selected.append(message) + selected_subcollectors.add(message["subcollector"]) + return selected, peak_processes + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--mode", + choices=("legacy-parent", "homogeneous-parent", "worker-metadata"), + required=True, + ) + parser.add_argument( + "--inner-start-method", + choices=("spawn", "forkserver", "fork"), + default="spawn", + ) + parser.add_argument("--num-collectors", type=int, default=5) + parser.add_argument("--envs-per-collector", type=int, default=64) + parser.add_argument( + "--timeout-s", + type=float, + default=1800, + help="Maximum wait for each startup, step, and shutdown phase.", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--config", + type=Path, + default=_VLA_DIR / "config" / "vla_grpo_libero.yaml", + ) + return parser.parse_args() + + +def _main() -> None: + args = _parse_args() + if args.inner_start_method != "spawn" and args.mode != "worker-metadata": + raise ValueError( + "forkserver/fork benchmarking is only allowed with worker-metadata, " + "which guarantees that the subcollector parent has not created an " + "EGL environment." + ) + + output_dir = args.output_dir.resolve() + marker_dir = output_dir / "construction_markers" + output_dir.mkdir(parents=True, exist_ok=True) + marker_dir.mkdir(parents=True, exist_ok=True) + command = " ".join(sys.argv) + (output_dir / "command.txt").write_text(command + "\n") + + ctx = mp.get_context("spawn") + status_queue = ctx.Queue() + step_event = ctx.Event() + messages = [] + total_timer = timeit("libero_startup/total").start() + processes = [ + ctx.Process( + target=_subcollector_main, + kwargs={ + "cfg_path": str(args.config.resolve()), + "mode": args.mode, + "inner_start_method": args.inner_start_method, + "num_envs": args.envs_per_collector, + "subcollector_idx": subcollector_idx, + "marker_dir": str(marker_dir), + "status_queue": status_queue, + "step_event": step_event, + }, + ) + for subcollector_idx in range(args.num_collectors) + ] + for process in processes: + process.start() + + error = None + peak_processes = 0 + ready = [] + stepped = [] + all_ready_s = None + first_batch_s = None + try: + ready, peak = _collect_until( + status_queue, + event="ready", + count=args.num_collectors, + processes=processes, + timeout_s=args.timeout_s, + total_timer=total_timer, + messages=messages, + ) + peak_processes = max(peak_processes, peak) + all_ready_s = total_timer.elapsed() + step_event.set() + stepped, peak = _collect_until( + status_queue, + event="stepped", + count=args.num_collectors, + processes=processes, + timeout_s=args.timeout_s, + total_timer=total_timer, + messages=messages, + ) + peak_processes = max(peak_processes, peak) + first_batch_s = total_timer.elapsed() + closed, peak = _collect_until( + status_queue, + event="closed", + count=args.num_collectors, + processes=processes, + timeout_s=args.timeout_s, + total_timer=total_timer, + messages=messages, + ) + del closed + peak_processes = max(peak_processes, peak) + except Exception: + error = traceback.format_exc() + step_event.set() + finally: + for process in processes: + process.join(timeout=30) + for process in processes: + if process.is_alive(): + process.terminate() + for process in processes: + process.join() + + time.sleep(1) + marker_names = [path.name for path in marker_dir.glob("construct-*")] + parent_constructions = sum("parent_metadata" in name for name in marker_names) + worker_constructions = sum("-worker-" in name for name in marker_names) + remaining_descendants = sorted(_descendant_pids(os.getpid())) + remaining_processes = {pid: _process_cmdline(pid) for pid in remaining_descendants} + remaining_environment_processes = { + pid: command + for pid, command in remaining_processes.items() + if "multiprocessing.resource_tracker" not in command + } + worker_pids = [pid for item in ready for pid in item["worker_pids"]] + lingering_worker_processes = { + pid: _process_cmdline(pid) + for pid in worker_pids + if Path(f"/proc/{pid}").exists() + } + summary = { + "command": command, + "mode": args.mode, + "inner_start_method": args.inner_start_method, + "num_collectors": args.num_collectors, + "envs_per_collector": args.envs_per_collector, + "total_envs": args.num_collectors * args.envs_per_collector, + "parent_metadata_constructions": parent_constructions, + "worker_constructions": worker_constructions, + "subcollector_construct_s": [ + item["construct_s"] + for item in sorted(ready, key=lambda item: item["subcollector"]) + ], + "all_ready_s": all_ready_s, + "subcollector_first_step_s": [ + item["first_step_s"] + for item in sorted(stepped, key=lambda item: item["subcollector"]) + ], + "first_batch_s": first_batch_s, + "peak_descendant_processes": peak_processes, + "outer_exit_codes": [process.exitcode for process in processes], + "remaining_descendant_processes": remaining_processes, + "remaining_environment_processes": remaining_environment_processes, + "lingering_worker_processes": lingering_worker_processes, + "messages": messages, + "error": error, + } + (output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + torchrl_logger.info(json.dumps(summary, indent=2)) + if error is not None: + raise RuntimeError(error) + if parent_constructions != ( + args.num_collectors * args.envs_per_collector + if args.mode == "legacy-parent" + else args.num_collectors + if args.mode == "homogeneous-parent" + else 0 + ): + raise RuntimeError( + f"unexpected parent construction count: {parent_constructions}" + ) + if worker_constructions != args.num_collectors * args.envs_per_collector: + raise RuntimeError( + f"unexpected worker construction count: {worker_constructions}" + ) + if any(process.exitcode for process in processes): + raise RuntimeError(f"outer process failures: {summary['outer_exit_codes']}") + if remaining_environment_processes: + raise RuntimeError( + f"orphaned environment processes: {remaining_environment_processes}" + ) + if lingering_worker_processes: + raise RuntimeError( + f"worker processes survived shutdown: {lingering_worker_processes}" + ) + + +if __name__ == "__main__": + _main() diff --git a/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml b/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml index 6e7cbd1ebbf..71766189a33 100644 --- a/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml +++ b/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml @@ -36,6 +36,7 @@ env: render_gpu_ids: [2, 3, 4, 5] # H100 fast mode render GPUs for rollout workers eval_render_gpu_ids: [7] # reserve GPU 7 for eval/video rendering render_gpu_device_zero_fallback: true + metadata_from_workers: true # construct LIBERO/EGL only in long-lived worker processes env_kwargs: null chunk_size: 8 # NUM_ACTIONS_CHUNK of the checkpoint max_env_steps: 512 # base env steps per episode (the paper's cap) diff --git a/sota-implementations/vla_grpo/test_openvla.py b/sota-implementations/vla_grpo/test_openvla.py index 9563e547ba3..979ba3bfa92 100644 --- a/sota-implementations/vla_grpo/test_openvla.py +++ b/sota-implementations/vla_grpo/test_openvla.py @@ -464,6 +464,7 @@ def _complete_toy_env_cfg(**env_overrides): "render_gpu_ids": [2, 3], "eval_render_gpu_ids": None, "render_gpu_device_zero_fallback": True, + "metadata_from_workers": False, "env_kwargs": None, } env.update(env_overrides) @@ -633,6 +634,116 @@ def client(self): assert isinstance(collector, _FakeMultiCollector) assert len(captured["multi_args"][0]) == 4 + def test_make_collector_env_uses_common_factory_with_worker_kwargs( + self, monkeypatch + ): + captured = {} + + def _fake_make_env_worker(cfg, tokenizer, worker_idx, **kwargs): + worker_idx_with_offset = worker_idx + kwargs["worker_idx_offset"] + logical_worker = worker_idx_with_offset // kwargs["group_repeats"] + task_id = cfg.env.task_ids[logical_worker % len(cfg.env.task_ids)] + return { + "task_id": task_id, + "language_instruction": f"instruction-{task_id}", + "group_id": logical_worker, + "worker_idx": worker_idx, + "worker_idx_with_offset": worker_idx_with_offset, + "seed": kwargs["seed"] + worker_idx_with_offset, + "init_state_slot": logical_worker, + } + + class _FakeParallelEnv: + def __init__( + self, + num_envs, + create_env_fn, + *, + create_env_kwargs, + **kwargs, + ): + captured["num_envs"] = num_envs + captured["create_env_fn"] = create_env_fn + captured["create_env_kwargs"] = create_env_kwargs + captured["kwargs"] = kwargs + captured["workers"] = [ + create_env_fn(**worker_kwargs) + for worker_kwargs in create_env_kwargs + ] + + monkeypatch.setattr(utils, "_make_env_worker", _fake_make_env_worker) + monkeypatch.setattr(utils, "ParallelEnv", _FakeParallelEnv) + cfg = SimpleNamespace( + collector=_complete_collector_cfg(group_size=2), + env=_complete_toy_env_cfg( + backend="libero", + metadata_from_workers=True, + task_ids=[10, 11], + ), + ) + + env = utils._make_collector_env( + cfg, + tokenizer=None, + num_envs=4, + group_repeats=2, + seed=5, + device=torch.device("cpu"), + worker_idx_offset=4, + render_gpu_device_id=3, + ) + + assert isinstance(env, _FakeParallelEnv) + assert captured["num_envs"] == 4 + assert callable(captured["create_env_fn"]) + assert captured["create_env_kwargs"] == [ + {"worker_idx": 0}, + {"worker_idx": 1}, + {"worker_idx": 2}, + {"worker_idx": 3}, + ] + assert captured["kwargs"]["mp_start_method"] == "spawn" + assert captured["kwargs"]["device"] == torch.device("cpu") + assert captured["kwargs"]["metadata_from_workers"] + assert captured["workers"] == [ + { + "task_id": 10, + "language_instruction": "instruction-10", + "group_id": 2, + "worker_idx": 0, + "worker_idx_with_offset": 4, + "seed": 9, + "init_state_slot": 2, + }, + { + "task_id": 10, + "language_instruction": "instruction-10", + "group_id": 2, + "worker_idx": 1, + "worker_idx_with_offset": 5, + "seed": 10, + "init_state_slot": 2, + }, + { + "task_id": 11, + "language_instruction": "instruction-11", + "group_id": 3, + "worker_idx": 2, + "worker_idx_with_offset": 6, + "seed": 11, + "init_state_slot": 3, + }, + { + "task_id": 11, + "language_instruction": "instruction-11", + "group_id": 3, + "worker_idx": 3, + "worker_idx_with_offset": 7, + "seed": 12, + "init_state_slot": 3, + }, + ] + class TestEvaluatorFactory: def test_eval_env_records_root_pixels_from_scalar_env(self, monkeypatch): diff --git a/sota-implementations/vla_grpo/utils.py b/sota-implementations/vla_grpo/utils.py index 329789ce337..c662889c7db 100644 --- a/sota-implementations/vla_grpo/utils.py +++ b/sota-implementations/vla_grpo/utils.py @@ -623,26 +623,29 @@ 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", # 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) device="cpu", + metadata_from_workers=bool( + _cfg_get(cfg.env, "metadata_from_workers", False) + ), ) if seed is not None: base.set_seed(seed) @@ -669,24 +672,25 @@ def _make_collector_env( worker_idx_offset: int, render_gpu_device_id: int | None, ) -> ParallelEnv: + 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=bool(_cfg_get(cfg.env, "metadata_from_workers", False)), ) diff --git a/test/envs/test_parallel.py b/test/envs/test_parallel.py index c3ef83c5976..5e2da31f8e0 100644 --- a/test/envs/test_parallel.py +++ b/test/envs/test_parallel.py @@ -6,6 +6,8 @@ import gc import os +from functools import partial +from pathlib import Path import numpy as np import pytest @@ -31,9 +33,11 @@ EnvCreator, ParallelEnv, SerialEnv, + ToyVLAEnv, TransformedEnv, ) from torchrl.envs.batched_envs import _stackable +from torchrl.envs.env_creator import get_env_metadata from torchrl.envs.libs.dm_control import _has_dmc, DMControlEnv from torchrl.envs.libs.gym import GymEnv from torchrl.envs.transforms import Compose, StepCounter @@ -56,6 +60,59 @@ ) +def _make_parent_guarded_env(parent_pid: int) -> ContinuousActionVecMockEnv: + if os.getpid() == parent_pid: + raise RuntimeError("environment factory was called in the parent process") + return ContinuousActionVecMockEnv() + + +def _make_instruction_env(instruction: str) -> ToyVLAEnv: + return ToyVLAEnv(action_dim=2, state_dim=2, instruction=instruction) + + +class _StartupTrackingEnv(CountingEnv): + def __init__( + self, + worker_idx: int, + marker_dir: str, + *, + batch_size: list[int] | None = None, + fail_metadata: bool = False, + ): + self.worker_idx = worker_idx + self.marker_dir = Path(marker_dir) + self.fail_metadata = fail_metadata + super().__init__(batch_size=batch_size) + (self.marker_dir / f"started-{worker_idx}").touch() + + def fake_tensordict(self, *args, **kwargs): + if self.fail_metadata: + raise RuntimeError("metadata extraction failed") + return super().fake_tensordict(*args, **kwargs) + + def close(self, *args, **kwargs): + (self.marker_dir / f"closed-{self.worker_idx}").touch() + return super().close(*args, **kwargs) + + +def _make_startup_tracking_env( + worker_idx: int, + marker_dir: str, + *, + fail_worker: int | None = None, + metadata_failure_worker: int | None = None, + incompatible: bool = False, +) -> _StartupTrackingEnv: + if worker_idx == fail_worker: + raise RuntimeError("worker construction failed") + return _StartupTrackingEnv( + worker_idx, + marker_dir, + batch_size=[worker_idx + 1] if incompatible else None, + fail_metadata=worker_idx == metadata_failure_worker, + ) + + class TestParallel: @pytest.fixture(autouse=True, scope="class") def disable_autowrap(self): @@ -104,6 +161,101 @@ def make_env(): 4, make_env, create_env_kwargs=[{"seed": 0}, {"seed": 1}] ) + def test_get_env_metadata_closes_only_callable_created_env(self): + closed = [] + + class CloseCountingEnv(CountingEnv): + def __init__(self, *, fail_metadata=False): + self.fail_metadata = fail_metadata + super().__init__() + + def fake_tensordict(self, *args, **kwargs): + if self.fail_metadata: + raise RuntimeError("metadata extraction failed") + return super().fake_tensordict(*args, **kwargs) + + def close(self, *args, **kwargs): + closed.append(True) + return super().close(*args, **kwargs) + + get_env_metadata(CloseCountingEnv) + assert closed == [True] + + with pytest.raises(RuntimeError, match="metadata extraction failed"): + get_env_metadata(partial(CloseCountingEnv, fail_metadata=True)) + assert closed == [True, True] + + env = CloseCountingEnv() + get_env_metadata(env) + assert closed == [True, True] + env.close() + assert closed == [True, True, True] + + creator = EnvCreator(CloseCountingEnv) + assert closed == [True, True, True, True] + get_env_metadata(creator) + assert closed == [True, True, True, True] + + def test_metadata_from_workers(self, maybe_fork_ParallelEnv): + parent_pid = os.getpid() + env = maybe_fork_ParallelEnv( + 2, + partial(_make_parent_guarded_env, parent_pid), + metadata_from_workers=True, + ) + try: + assert env._use_buffers is False + assert not env.is_closed + rollout = env.rollout(3) + assert rollout.shape[:2] == torch.Size([2, 3]) + finally: + env.close(raise_if_closed=False) + + def test_metadata_from_workers_rejects_buffers(self, maybe_fork_ParallelEnv): + with pytest.raises( + RuntimeError, + match="metadata_from_workers=True is incompatible with use_buffers=True", + ): + maybe_fork_ParallelEnv( + 2, + ContinuousActionVecMockEnv, + metadata_from_workers=True, + use_buffers=True, + ) + + @pytest.mark.parametrize( + ("factory_kwargs", "error"), + [ + ({"incompatible": True}, "metadata is incompatible"), + ({"fail_worker": 1}, "worker construction failed"), + ({"metadata_failure_worker": 1}, "metadata extraction failed"), + ], + ) + def test_metadata_from_workers_startup_failure_cleans_workers( + self, maybe_fork_ParallelEnv, tmp_path, factory_kwargs, error + ): + factory = partial( + _make_startup_tracking_env, + marker_dir=str(tmp_path), + **factory_kwargs, + ) + with pytest.raises(RuntimeError, match=error): + maybe_fork_ParallelEnv( + 3, + factory, + create_env_kwargs=[{"worker_idx": idx} for idx in range(3)], + metadata_from_workers=True, + ) + + started = { + path.name.removeprefix("started-") for path in tmp_path.glob("started-*") + } + closed = { + path.name.removeprefix("closed-") for path in tmp_path.glob("closed-*") + } + assert started + assert closed == started + def test_compact_collector_skips_next_observation_copy( self, maybe_fork_ParallelEnv ): @@ -1475,8 +1627,6 @@ def test_heterogeneous_non_tensor_workers(cls, maybe_fork_ParallelEnv): # instructions) stack their specs into a Stacked spec: the batched env # must still register the entry as non-tensor and route it through the # non-tensor channel instead of the shared buffers - from torchrl.envs import ToyVLAEnv - if cls is ParallelEnv: cls = maybe_fork_ParallelEnv @@ -1494,6 +1644,26 @@ def make(instruction): env.close(raise_if_closed=False) +@set_list_to_stack(True) +def test_worker_metadata_allows_different_non_tensor_values(maybe_fork_ParallelEnv): + env = maybe_fork_ParallelEnv( + 2, + _make_instruction_env, + create_env_kwargs=[ + {"instruction": "pick up the apple"}, + {"instruction": "pick up the pear"}, + ], + metadata_from_workers=True, + ) + try: + rollout = env.rollout(3) + instructions = rollout["language_instruction"] + assert instructions[0][0] == "pick up the apple" + assert instructions[1][0] == "pick up the pear" + finally: + env.close(raise_if_closed=False) + + def test_stackable(): # Tests the _stackable util stack = [TensorDict({"a": 0}, []), TensorDict({"b": 1}, [])] diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py index b40e458028d..7fefd1ea241 100644 --- a/torchrl/envs/batched_envs.py +++ b/torchrl/envs/batched_envs.py @@ -264,6 +264,7 @@ 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) + metadata_from_workers = kwargs.get("metadata_from_workers", False) def _warn_lambda(): if rl_warnings(): @@ -302,7 +303,7 @@ def _wrap_lambdas(create_env_fn): return result return create_env_fn - if auto_wrap_envs: + if auto_wrap_envs and not metadata_from_workers: if "create_env_fn" in kwargs: kwargs["create_env_fn"] = _wrap_lambdas(kwargs["create_env_fn"]) elif len(args) >= 2: @@ -381,6 +382,12 @@ class BatchedEnvBase(EnvBase): 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``. + metadata_from_workers (bool, optional): if ``True``, :class:`~torchrl.envs.ParallelEnv` + starts its workers during construction and obtains metadata from the + real worker environments instead of creating temporary environments + in the parent process. Worker schemas are validated before normal + initialization. This mode uses direct pipe communication and is + incompatible with ``use_buffers=True``. Defaults to ``False``. 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. @@ -511,6 +518,7 @@ def __init__( use_buffers: bool | None = None, consolidate: bool = True, daemon: bool = False, + metadata_from_workers: bool = False, ): super().__init__(device=device) self.serial_for_single = serial_for_single @@ -524,6 +532,7 @@ def __init__( self._use_buffers = use_buffers self.consolidate = consolidate self.daemon = daemon + self._metadata_from_workers = metadata_from_workers self._single_task = callable(create_env_fn) or (len(set(create_env_fn)) == 1) if callable(create_env_fn): @@ -575,13 +584,38 @@ def __init__( 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 + self.has_lazy_inputs = False + if self._metadata_from_workers: + if not isinstance(self, ParallelEnv): + raise TypeError( + "metadata_from_workers=True is only supported by ParallelEnv." + ) + if self._use_buffers is True: + raise RuntimeError( + "metadata_from_workers=True is incompatible with use_buffers=True. " + "Worker metadata must be collected before shared buffers can be " + "allocated; pass use_buffers=False or leave use_buffers unset." + ) + self._use_buffers = False + try: + worker_metadata = self._start_workers(gather_metadata=True) + self._get_metadata( + create_env_fn, + create_env_kwargs, + worker_metadata=worker_metadata, + ) + self._initialize_started_workers() + except Exception: + self._abort_worker_startup() + raise + else: + self._get_metadata(create_env_fn, create_env_kwargs) is_spec_locked = EnvBase.is_spec_locked @@ -901,15 +935,26 @@ def _check_mps_use_buffers(self) -> None: self._use_buffers = False def _get_metadata( - self, create_env_fn: list[Callable], create_env_kwargs: list[dict] + self, + create_env_fn: list[Callable], + create_env_kwargs: list[dict], + *, + worker_metadata: list[EnvMetaData] | None = None, ): if self._single_task: - # if EnvCreator, the metadata are already there - meta_data: EnvMetaData = get_env_metadata( - create_env_fn[0], - create_env_kwargs[0], - env_validator=self._validate_worker_env, - ) + if worker_metadata is None: + # if EnvCreator, the metadata are already there + meta_data: EnvMetaData = get_env_metadata( + create_env_fn[0], + create_env_kwargs[0], + env_validator=self._validate_worker_env, + ) + else: + self._validate_worker_metadata(worker_metadata) + meta_data = worker_metadata[0].clone() + meta_data.supports_set_state = all( + metadata.supports_set_state for metadata in worker_metadata + ) self.meta_data = meta_data.expand( *(self.num_workers, *meta_data.batch_size) ) @@ -931,13 +976,15 @@ def _get_metadata( n_tasks = len(create_env_fn) self.meta_data: list[EnvMetaData] = [] for i in range(n_tasks): - self.meta_data.append( - get_env_metadata( + if worker_metadata is None: + meta_data = get_env_metadata( create_env_fn[i], create_env_kwargs[i], env_validator=self._validate_worker_env, - ).clone() - ) + ) + else: + meta_data = worker_metadata[i] + self.meta_data.append(meta_data.clone()) 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 +1014,55 @@ def _get_metadata( self._set_properties() + @staticmethod + def _validate_worker_metadata(worker_metadata: list[EnvMetaData]) -> None: + """Validate schemas before reusing one common factory's metadata.""" + if not worker_metadata: + raise RuntimeError("No worker metadata was received.") + reference = worker_metadata[0] + reference_specs = reference.specs + reference_spec_keys = set(reference_specs.keys(True, True)) + for worker_idx, metadata in enumerate(worker_metadata[1:], 1): + mismatch = None + if metadata.batch_size != reference.batch_size: + mismatch = ( + f"batch size {metadata.batch_size} does not match " + f"worker 0 batch size {reference.batch_size}" + ) + elif metadata.device != reference.device: + mismatch = ( + f"device {metadata.device} does not match worker 0 device " + f"{reference.device}" + ) + elif metadata.batch_locked != reference.batch_locked: + mismatch = "batch_locked differs from worker 0" + elif metadata.device_map != reference.device_map: + mismatch = "tensor devices differ from worker 0" + else: + specs = metadata.specs + spec_keys = set(specs.keys(True, True)) + if spec_keys != reference_spec_keys: + mismatch = "input or output spec keys differ from worker 0" + else: + for key in reference_spec_keys: + reference_spec = reference_specs[key] + spec = specs[key] + if isinstance(reference_spec, NonTensor) and isinstance( + spec, NonTensor + ): + continue + if spec != reference_spec: + mismatch = ( + f"input or output tensor spec {key!r} differs " + "from worker 0" + ) + break + if mismatch is not None: + raise RuntimeError( + f"Worker {worker_idx} metadata is incompatible with worker 0: " + f"{mismatch}." + ) + def update_kwargs(self, kwargs: dict | list[dict]) -> None: """Updates the kwargs of each environment given a dictionary or a list of dictionaries. @@ -1951,15 +2047,22 @@ class ParallelEnv(BatchedEnvBase, metaclass=_PEnvMeta): """ - def _start_workers(self) -> None: + def _start_workers( + self, *, gather_metadata: bool = False + ) -> list[EnvMetaData] | None: import torchrl + if gather_metadata and self._use_buffers: + raise RuntimeError( + "Worker metadata can only be gathered with use_buffers=False." + ) self._timeout = 10.0 self.BATCHED_PIPE_TIMEOUT = torchrl._utils.BATCHED_PIPE_TIMEOUT num_threads = max( 1, torch.get_num_threads() - self.num_workers ) # 1 more thread for this proc + self._num_threads_before_start = torch.get_num_threads() if self.num_threads is None: self.num_threads = num_threads @@ -2015,61 +2118,181 @@ def look_for_cuda(tensor, has_cuda=has_cuda): "shm_done_flags": self._shm_done_flags, } ) - with clear_mpi_env_vars(): - for idx in range(_num_workers): - if self._verbose: - torchrl_logger.info(f"initiating worker {idx}") - # No certainty which module multiprocessing_context is - parent_pipe, child_pipe = ctx.Pipe() - env_fun = self.create_env_fn[idx] - if not isinstance(env_fun, (EnvCreator, CloudpickleWrapper)): - env_fun = CloudpickleWrapper(env_fun) - - kwargs[idx].update( - { - "parent_pipe": parent_pipe, - "child_pipe": child_pipe, - "env_fun": env_fun, - "env_fun_kwargs": self.create_env_kwargs[idx], - "has_lazy_inputs": self.has_lazy_inputs, - "num_threads": num_sub_threads, - "non_blocking": self.non_blocking, - "filter_warnings": self._filter_warnings_subprocess(), - } - ) - if self._use_buffers: + try: + with clear_mpi_env_vars(): + for idx in range(_num_workers): + if self._verbose: + torchrl_logger.info(f"initiating worker {idx}") + # No certainty which module multiprocessing_context is + parent_pipe, child_pipe = ctx.Pipe() + env_fun = self.create_env_fn[idx] + if not isinstance(env_fun, (EnvCreator, CloudpickleWrapper)): + env_fun = CloudpickleWrapper(env_fun) + kwargs[idx].update( { - "shared_tensordict": self.shared_tensordicts[idx], - "_selected_input_keys": self._selected_input_keys, - "_selected_reset_keys": self._selected_reset_keys, - "_selected_step_keys": self._selected_step_keys, - "_non_tensor_keys": self._non_tensor_keys, + "parent_pipe": parent_pipe, + "child_pipe": child_pipe, + "env_fun": env_fun, + "env_fun_kwargs": self.create_env_kwargs[idx], + "has_lazy_inputs": self.has_lazy_inputs, + "num_threads": num_sub_threads, + "non_blocking": self.non_blocking, + "filter_warnings": self._filter_warnings_subprocess(), } ) - else: - kwargs[idx].update( - { - "consolidate": self.consolidate, - } + if self._use_buffers: + kwargs[idx].update( + { + "shared_tensordict": self.shared_tensordicts[idx], + "_selected_input_keys": self._selected_input_keys, + "_selected_reset_keys": self._selected_reset_keys, + "_selected_step_keys": self._selected_step_keys, + "_non_tensor_keys": self._non_tensor_keys, + } + ) + else: + kwargs[idx].update( + { + "consolidate": self.consolidate, + "send_metadata": gather_metadata, + } + ) + process = proc_fun(target=func, kwargs=kwargs[idx]) + process.daemon = self.daemon + try: + process.start() + except Exception: + child_pipe.close() + parent_pipe.close() + raise + child_pipe.close() + self.parent_channels.append(parent_pipe) + self._workers.append(process) + + metadata = self._receive_worker_startup(gather_metadata=gather_metadata) + if gather_metadata: + return metadata + self._initialize_started_workers() + except Exception: + self._abort_worker_startup() + raise + return None + + def _receive_worker_startup( + self, *, gather_metadata: bool + ) -> list[EnvMetaData] | None: + metadata: list[EnvMetaData | None] | None = ( + [None] * self.num_workers if gather_metadata else None + ) + pipes_pending = { + channel: worker_idx + for worker_idx, channel in enumerate(self.parent_channels) + } + startup_timer = timeit(f"parallel_env/{id(self)}/worker_startup").start() + while pipes_pending: + remaining = self.BATCHED_PIPE_TIMEOUT - startup_timer.elapsed() + if remaining <= 0: + pending = sorted(pipes_pending.values()) + raise RuntimeError( + "Failed to start ParallelEnv workers " + f"{pending} within the {self.BATCHED_PIPE_TIMEOUT} sec time " + "limit. This threshold can be increased via the " + "BATCHED_PIPE_TIMEOUT environment variable." + ) + ready = connection_wait(list(pipes_pending), timeout=remaining) + if not ready: + dead = [ + worker_idx + for worker_idx in pipes_pending.values() + if not self._workers[worker_idx].is_alive() + ] + if dead: + raise RuntimeError( + f"ParallelEnv workers {dead} exited during startup." + ) + continue + for pipe in ready: + worker_idx = pipes_pending.pop(pipe) + try: + message = pipe.recv() + except EOFError as err: + raise RuntimeError( + f"ParallelEnv worker {worker_idx} exited during startup." + ) from err + if gather_metadata: + if ( + isinstance(message, tuple) + and len(message) == 2 + and message[0] == "startup_error" + ): + raise RuntimeError( + f"ParallelEnv worker {worker_idx} failed during " + f"metadata collection: {message[1]}" + ) + if not ( + isinstance(message, tuple) + and len(message) == 2 + and message[0] == "metadata" + and isinstance(message[1], EnvMetaData) + ): + raise RuntimeError( + f"Expected metadata from ParallelEnv worker " + f"{worker_idx}, got {message!r}." + ) + metadata[worker_idx] = message[1] + elif message != "started": + raise RuntimeError( + f"Expected startup confirmation from ParallelEnv worker " + f"{worker_idx}, got {message!r}." ) - process = proc_fun(target=func, kwargs=kwargs[idx]) - process.daemon = self.daemon - process.start() - child_pipe.close() - 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() - - # send shared tensordict to workers + if metadata is None: + return None + if any(item is None for item in metadata): + raise RuntimeError("Metadata was not received from every worker.") + return [item for item in metadata if item is not None] + + def _initialize_started_workers(self) -> None: for channel in self.parent_channels: channel.send(("init", None)) self.is_closed = False self.set_spec_lock_() + def _abort_worker_startup(self) -> None: + workers = getattr(self, "_workers", ()) + channels = getattr(self, "parent_channels", ()) + if not workers and not channels: + if hasattr(self, "_num_threads_before_start"): + torch.set_num_threads(self._num_threads_before_start) + return + + for channel, process in zip(channels, workers): + if process.is_alive(): + try: + channel.send(("close", None)) + except (EOFError, OSError): + pass + + timeout = getattr(self, "_timeout", 10.0) + shutdown_timer = timeit(f"parallel_env/{id(self)}/abort_startup").start() + for process in workers: + remaining = max(0.0, timeout - shutdown_timer.elapsed()) + process.join(timeout=remaining) + for process in workers: + if process.is_alive(): + process.terminate() + for process in workers: + process.join() + for channel in channels: + channel.close() + + self._workers = [] + self.parent_channels = [] + self._events = None + self.event = None + self.is_closed = True + torch.set_num_threads(self._num_threads_before_start) + def _filter_warnings_subprocess(self) -> bool: from torchrl import filter_warnings_subprocess @@ -3449,6 +3672,7 @@ def _run_worker_pipe_direct( num_threads: int | None = None, # for fork start method consolidate: bool = True, filter_warnings: bool = False, + send_metadata: bool = False, ) -> None: # Handle warning filtering (moved from _ProcessNoWarn) if filter_warnings: @@ -3458,14 +3682,32 @@ 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 send_metadata: + BatchedEnvBase._validate_worker_env(env) + metadata = EnvMetaData.metadata_from_env(env) + except Exception as err: + if not send_metadata: + raise + try: + child_pipe.send(("startup_error", repr(err))) + finally: + if env is not None: + try: + env.close() + except Exception: + pass + 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 +3746,11 @@ def _run_worker_pipe_direct( initialized = False - child_pipe.send("started") + if send_metadata: + child_pipe.send(("metadata", metadata)) + del metadata + else: + child_pipe.send("started") while True: try: if child_pipe.poll(_timeout): @@ -3653,8 +3899,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