From 09b6a1604a3f7f32c8592cd6ce115c4b89610e92 Mon Sep 17 00:00:00 2001 From: m0g3r <87276771+m0g3r@users.noreply.github.com> Date: Sat, 15 Aug 2026 15:00:27 +0200 Subject: [PATCH] [Bugfix] Restore per-stage runtime env during launch (#6214) Signed-off-by: m0g3r <87276771+m0g3r@users.noreply.github.com> --- .../test_async_omni_engine_stage_init.py | 88 +++++++++++++++++++ vllm_omni/engine/stage_runtime.py | 8 +- 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 27aee8706..9ee59f153 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -161,6 +161,17 @@ def _make_diffusion_plan( ) +def _make_stage_runtime() -> StageRuntime: + return StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) + + def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): """Regression test for forward references in make_async_mp_client.""" import vllm_omni.engine.stage_engine_core_client as client_mod @@ -233,6 +244,48 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): assert cleanup_clients == [shared, extra] +def test_initialize_local_diffusion_replica_scopes_runtime_env(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.platforms import current_omni_platform + + runtime = _make_stage_runtime() + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] + + runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV" + device_env_var = current_omni_platform.device_control_env_var + runtime._init_visible_devices_baseline = "0,1" + plan.metadata.runtime_cfg = { + "devices": "0", + "env": {runtime_env_var: "stage-value"}, + } + monkeypatch.delenv(runtime_env_var, raising=False) + monkeypatch.setenv(device_env_var, "0,1") + + captured: dict[str, str | None] = {} + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) + + def _capture_launch_diffusion_stage_replica(**_kwargs): + captured["runtime_env"] = os.environ.get(runtime_env_var) + captured["device_env"] = os.environ.get(device_env_var) + raise RuntimeError("stop after capturing launch environment") + + monkeypatch.setattr( + runtime_mod, + "launch_diffusion_stage_replica", + _capture_launch_diffusion_stage_replica, + ) + + with pytest.raises(RuntimeError, match="stop after capturing launch environment"): + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) + + assert captured == { + "runtime_env": "stage-value", + "device_env": "0", + } + assert runtime_env_var not in os.environ + assert os.environ[device_env_var] == "0,1" + + def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.engine.stage_engine_startup import StageReplicaResources @@ -311,6 +364,41 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline } +def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + plan = _make_llm_plan(0, stage_id=0, vllm_config=object()).replicas[0] + plan.engine_args_dict = {} + + runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV" + plan.metadata.runtime_cfg = { + "devices": "0", + "env": {runtime_env_var: "stage-value"}, + } + monkeypatch.setenv(runtime_env_var, "parent-value") + + captured: dict[str, str | None] = {} + + @contextlib.contextmanager + def _capture_launch_stage_replica(*, stage_visible_devices, **_kwargs): + captured["runtime_env"] = os.environ.get(runtime_env_var) + captured["stage_visible_devices"] = stage_visible_devices + yield None + + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica) + + with pytest.raises(RuntimeError, match="launcher returned no resources"): + runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) + + assert captured == { + "runtime_env": "stage-value", + "stage_visible_devices": "0", + } + assert os.environ[runtime_env_var] == "parent-value" + + def test_stage_runtime_initializes_stage_pools(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index bf84f7793..c0aafc86a 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -58,6 +58,7 @@ from vllm_omni.engine.stage_init_utils import ( load_omni_transfer_config_for_model, prepare_engine_environment, release_device_locks, + stage_runtime_env, ) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.entrypoints.stage_utils import resolve_stage_physical_devices @@ -576,7 +577,7 @@ class StageRuntime: ) # Serialize engine-core spawning across all LLM replicas to avoid # ZMQ port-allocation races and simultaneous CUDA context init. - with self._replica_launch_lock: + with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): with launch_stage_replica( vllm_config=vllm_config, executor_class=executor_class, @@ -645,7 +646,10 @@ class StageRuntime: client = None resources = None try: - with self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg): + with ( + stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg), + self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg), + ): omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector if omni_conn_cfg: inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to)