# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """CPU tests of Ray placement, rank assignment, RPC output, and cleanup.""" import threading from types import SimpleNamespace from unittest.mock import Mock import pytest import torch from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.executor import ray_executor as ray_module pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] def test_cross_node_workers_get_global_rank_and_explicit_rendezvous(monkeypatch): monkeypatch.setenv("NCCL_SOCKET_IFNAME", "eth-test") monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,4,5") actors = [] for ip in ["10.0.0.2", "10.0.0.1", "10.0.0.2"]: actor = Mock() actor.get_node_ip.remote.return_value = ip actor.get_open_port.remote.return_value = 23456 actors.append(actor) factory = Mock() factory.remote.side_effect = actors fake_ray = Mock() fake_ray.remote.return_value = lambda cls: factory fake_ray.get.side_effect = lambda values, **kwargs: values monkeypatch.setattr(ray_module, "ray", fake_ray) monkeypatch.setattr(ray_module, "get_ip", lambda: "10.0.0.1") monkeypatch.setattr(ray_module, "PlacementGroupSchedulingStrategy", Mock()) executor = object.__new__(ray_module.RayDiffusionExecutor) executor.od_config = SimpleNamespace(num_gpus=3, ray_worker_env={"CUSTOM_PLUGIN_SETTING": "enabled"}) executor.workers = [] executor._init_workers(object()) assert fake_ray.remote.call_args.kwargs["max_concurrency"] == 1 assert fake_ray.remote.call_args.kwargs["concurrency_groups"] == {"health": 1, "monitor": 1} worker_env = fake_ray.remote.call_args.kwargs["runtime_env"]["env_vars"] assert worker_env["NCCL_SOCKET_IFNAME"] == "eth-test" assert worker_env["CUSTOM_PLUGIN_SETTING"] == "enabled" assert "CUDA_VISIBLE_DEVICES" not in worker_env assert len(fake_ray.get.call_args_list) == 3 assert all(call.kwargs["timeout"] == ray_module._WORKER_INIT_TIMEOUT_S for call in fake_ray.get.call_args_list) assert [worker.ip for worker in executor.workers] == ["10.0.0.1", "10.0.0.2", "10.0.0.2"] for rank, metadata in enumerate(executor.workers): metadata.worker.init_worker.remote.assert_called_once_with(executor.od_config, rank, "tcp://10.0.0.1:23456") def test_worker_env_forwards_stage_overrides_without_driver_identity(monkeypatch): from unittest.mock import patch driver_env = { "HF_TOKEN": "test-token", "VLLM_PLUGINS": "test-plugin", "NCCL_DEBUG": "WARN", "DIFFUSION_ATTENTION_BACKEND": "FLASH_ATTN", "PYTHONPATH": "/test/modules", "UNRELATED_SECRET": "do-not-copy", "HOME": "/driver/home", } protected = { "CUDA_VISIBLE_DEVICES": "7", "VLLM_HOST_IP": "10.0.0.1", "VLLM_HOST_PORT": "1234", "VLLM_NIXL_SIDE_CHANNEL_HOST": "10.0.0.1", "LOCAL_RANK": "7", "RANK": "7", "WORLD_SIZE": "8", "MASTER_ADDR": "driver-host", "MASTER_PORT": "1234", "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", } config = SimpleNamespace(ray_worker_env={**protected, "NCCL_DEBUG": "INFO", "CUSTOM_PLUGIN_SETTING": "stage"}) with patch.dict(ray_module.os.environ, {**driver_env, **protected}, clear=True): env = ray_module._worker_env(config) assert env == {key: value for key, value in driver_env.items() if key not in {"UNRELATED_SECRET", "HOME"}} | { "NCCL_DEBUG": "INFO", "CUSTOM_PLUGIN_SETTING": "stage", } assert ray_module.os.environ["NCCL_DEBUG"] == "WARN" @pytest.mark.parametrize("rank", [0, 1]) def test_collective_executes_all_ranks_but_only_output_rank_returns_tensors(rank): actor = ray_module.RayDiffusionWorkerWrapper(rank) value = torch.ones(2, requires_grad=True) actor.worker = Mock() actor.worker.execute_method.return_value = DiffusionOutput(output=value) result = actor.execute_rpc("execute_model", output_rank=0, exec_all_ranks=True) actor.worker.execute_method.assert_called_once() assert result["replied"] == (rank == 0) if rank == 0: assert not result["result"].output.requires_grad torch.testing.assert_close(result["result"].output, value) else: assert result["result"] is None def test_finalizer_kills_workers_without_waiting_and_releases_owned_group(monkeypatch): events = [] fake_ray = Mock() fake_ray.get.side_effect = lambda futures, **kw: events.append("wait") fake_ray.kill.side_effect = lambda worker, **kw: events.append("kill") fake_ray.util.remove_placement_group.side_effect = lambda group: events.append("remove") monkeypatch.setattr(ray_module, "ray", fake_ray) actor = Mock() actor.shutdown.remote.side_effect = lambda: events.append("shutdown") resources = ray_module._RayExecutorResources([ray_module._RayWorkerMetadata(actor)], object(), True) resources() assert events == ["kill", "remove"] fake_ray.kill.assert_called_once_with(actor, no_restart=True) actor.shutdown.remote.assert_not_called() fake_ray.get.assert_not_called() resources() assert events == ["kill", "remove"] assert resources.workers == [] def test_wrapper_shutdown_releases_worker_once(): actor = ray_module.RayDiffusionWorkerWrapper(0) worker = Mock() actor.worker = worker actor.shutdown() actor.shutdown() worker.shutdown.assert_called_once_with() assert not actor.check_alive() def _make_executor(fake_ray, actors, group=None, owns_group=True): executor = object.__new__(ray_module.RayDiffusionExecutor) executor._closed = False executor._is_failed = False executor._state_lock = threading.Lock() executor._failure_callbacks = [] executor.workers = [ray_module._RayWorkerMetadata(actor, rank=rank) for rank, actor in enumerate(actors)] executor._resources = ray_module._RayExecutorResources(executor.workers, group, owns_group) executor._finalizer = executor._resources return executor @pytest.mark.parametrize("stuck_rank", [None, 1]) def test_shutdown_drains_workers_before_killing_them(monkeypatch, stuck_rank): events = [] actors = [Mock(), Mock()] refs = [] for rank, actor in enumerate(actors): ref = f"shutdown-{rank}" refs.append(ref) actor.shutdown.remote.side_effect = lambda ref=ref: events.append(ref) or ref fake_ray = Mock() def wait(pending, num_returns, timeout): events.append("wait") assert num_returns == len(pending) assert timeout == ray_module._WORKER_SHUTDOWN_TIMEOUT_S return [ref for ref in pending if ref != f"shutdown-{stuck_rank}"], [ ref for ref in pending if ref == f"shutdown-{stuck_rank}" ] fake_ray.wait.side_effect = wait fake_ray.kill.side_effect = lambda worker, **kw: events.append("kill") fake_ray.util.remove_placement_group.side_effect = lambda group: events.append("remove") monkeypatch.setattr(ray_module, "ray", fake_ray) executor = _make_executor(fake_ray, actors, group=object()) callback = Mock() executor._failure_callbacks = [callback] executor.shutdown() assert events == [*refs, "wait", "kill", "kill", "remove"] assert [call.args[0] for call in fake_ray.get.call_args_list] == [ ref for ref in refs if ref != f"shutdown-{stuck_rank}" ] assert executor._closed assert not executor._is_failed callback.assert_not_called() executor.shutdown() assert fake_ray.kill.call_count == len(actors) def test_shutdown_kills_workers_when_drain_fails(monkeypatch): fake_ray = Mock() fake_ray.wait.side_effect = RuntimeError("Ray is shutting down") monkeypatch.setattr(ray_module, "ray", fake_ray) actors = [Mock(), Mock()] executor = _make_executor(fake_ray, actors) executor.shutdown() assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors def test_shutdown_preserves_borrowed_placement_group(monkeypatch): fake_ray = Mock() monkeypatch.setattr(ray_module, "ray", fake_ray) ray_module._RayExecutorResources([], object(), False)() fake_ray.util.remove_placement_group.assert_not_called() @pytest.mark.parametrize("failure_kind", ["GetTimeoutError", "RayActorError", "RayTaskError"]) @pytest.mark.parametrize("output_rank", [None, 0]) def test_rpc_failure_kills_peers_and_rejects_subsequent_work(monkeypatch, failure_kind, output_rank): fake_ray = Mock() fake_ray.exceptions = SimpleNamespace( **{name: type(name, (Exception,), {}) for name in ("GetTimeoutError", "RayActorError", "RayTaskError")} ) failure = getattr(fake_ray.exceptions, failure_kind)("rank failed while peers are still running") fake_ray.get.side_effect = failure monkeypatch.setattr(ray_module, "ray", fake_ray) actors = [Mock(), Mock()] group = object() executor = _make_executor(fake_ray, actors, group=group) callback = Mock() executor._failure_callbacks = [callback] expected_error = TimeoutError if failure_kind == "GetTimeoutError" else ray_module.EngineDeadError with pytest.raises(expected_error) as raised: executor.collective_rpc("execute_model", unique_reply_rank=output_rank, exec_all_ranks=True) assert raised.value.__cause__ is failure assert executor.is_dead assert executor._is_failed assert executor._closed callback.assert_called_once_with() assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors fake_ray.util.remove_placement_group.assert_called_once_with(group) # Peers may be stuck in a collective, so failure skips the graceful drain. fake_ray.wait.assert_not_called() for actor in actors: actor.shutdown.remote.assert_not_called() with pytest.raises(RuntimeError, match="closed"): executor.collective_rpc("execute_model", unique_reply_rank=output_rank, exec_all_ranks=True) for actor in actors: actor.execute_rpc.remote.assert_called_once() fake_ray.get.assert_called_once() executor.shutdown() assert fake_ray.kill.call_count == len(actors) callback.assert_called_once_with() def test_actor_death_during_shutdown_does_not_report_failure(): executor = object.__new__(ray_module.RayDiffusionExecutor) executor._closed = True executor._is_failed = False executor._state_lock = threading.Lock() executor._finalizer = Mock() callback = Mock() executor._failure_callbacks = [callback] executor._fail(RuntimeError("actor terminated")) assert not executor._is_failed callback.assert_not_called() executor._finalizer.assert_not_called() class _InlineThread: """Run the monitor synchronously so tests observe its effects.""" def __init__(self, target, **kwargs): self._target = target def start(self): self._target() @pytest.mark.parametrize("closed_first", [False, True]) def test_worker_monitor_fails_executor_when_idle_actor_dies(monkeypatch, closed_first): fake_ray = Mock() death = RuntimeError("actor died") actors = [Mock(), Mock()] sentinels = [actor.wait_for_exit.options.return_value.remote.return_value for actor in actors] def wait(refs, num_returns): assert refs == sentinels assert num_returns == 1 if closed_first: executor._closed = True return [sentinels[1]], [sentinels[0]] fake_ray.wait.side_effect = wait fake_ray.get.side_effect = death monkeypatch.setattr(ray_module, "ray", fake_ray) monkeypatch.setattr(ray_module.threading, "Thread", _InlineThread) group = object() executor = _make_executor(fake_ray, actors, group=group) callback = Mock() executor.register_failure_callback(callback) executor._start_worker_monitor() for actor in actors: actor.wait_for_exit.options.assert_called_once_with(concurrency_group="monitor") if closed_first: assert not executor._is_failed callback.assert_not_called() fake_ray.kill.assert_not_called() return fake_ray.get.assert_called_once_with(sentinels[1]) assert executor._is_failed assert executor._closed callback.assert_called_once_with() assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors fake_ray.util.remove_placement_group.assert_called_once_with(group) fake_ray.wait.assert_called_once() with pytest.raises(RuntimeError, match="closed"): executor.check_health() def test_worker_monitor_does_not_keep_executor_alive(monkeypatch): fake_ray = Mock() monkeypatch.setattr(ray_module, "ray", fake_ray) started = [] monkeypatch.setattr( ray_module.threading, "Thread", lambda target, **kwargs: started.append(target) or SimpleNamespace(start=lambda: None), ) executor = _make_executor(fake_ray, [Mock()]) executor._start_worker_monitor() (monitor,) = started assert all(executor is not cell.cell_contents for cell in monitor.__closure__ or ()) @pytest.mark.parametrize( "bundles, required, error", [ ([{"GPU": 1}] * 4, 4, None), ([{"GPU": 1}] * 8, 4, None), ([{"CPU": 1}, {"GPU": 0.5}, {"GPU": 1}], 1, None), ([{"GPU": 1}] * 2, 4, "requires 4 full GPUs.*only 2"), ([{"GPU": 0.5}] * 4, 4, "requires 4 full GPUs.*only 0"), ([{"CPU": 4}], 1, "requires 1 full GPUs.*only 0"), ([{"GPU": 4}], 4, None), ([{"GPU": 2}, {"GPU": 2}], 4, None), ([{"GPU": 1.5}, {"GPU": 1.5}], 3, "requires 3 full GPUs.*only 2"), ], ) def test_reused_placement_group_requires_full_gpu_capacity(monkeypatch, bundles, required, error): group = SimpleNamespace(bundle_specs=bundles) fake_ray = Mock() fake_ray.util.get_current_placement_group.return_value = group monkeypatch.setattr(ray_module, "ray", fake_ray) executor = object.__new__(ray_module.RayDiffusionExecutor) executor.od_config = SimpleNamespace(num_gpus=required) if error is not None: with pytest.raises(ValueError, match=error): executor._get_or_create_placement_group() else: actual_group, owned = executor._get_or_create_placement_group() assert actual_group is group assert not owned fake_ray.util.placement_group.assert_not_called() fake_ray.util.remove_placement_group.assert_not_called() @pytest.mark.parametrize("failure_phase", ["creation", "discovery", "port", "initialization"]) def test_worker_startup_failure_cleans_up_registered_actors(monkeypatch, failure_phase): actors = [Mock(), Mock()] for actor in actors: actor.get_node_ip.remote.return_value = "10.0.0.1" actor.get_open_port.remote.return_value = 23456 factory = Mock() error = RuntimeError("startup failed") factory.remote.side_effect = [actors[0], error] if failure_phase == "creation" else actors fake_ray = Mock() fake_ray.is_initialized.return_value = True fake_ray.remote.return_value = lambda cls: factory call_count = 0 def get(values, **kwargs): nonlocal call_count call_count += 1 assert kwargs["timeout"] == ray_module._WORKER_INIT_TIMEOUT_S if call_count == {"discovery": 1, "port": 2, "initialization": 3}.get(failure_phase): raise error return values fake_ray.get.side_effect = get monkeypatch.setattr(ray_module, "ray", fake_ray) monkeypatch.setattr(ray_module, "get_ip", lambda: "10.0.0.1") monkeypatch.setattr(ray_module, "PlacementGroupSchedulingStrategy", Mock()) executor = object.__new__(ray_module.RayDiffusionExecutor) executor.od_config = SimpleNamespace(num_gpus=2) group = object() executor._get_or_create_placement_group = Mock(return_value=(group, True)) with pytest.raises(RuntimeError, match="startup failed"): executor._init_executor() expected = actors[:1] if failure_phase == "creation" else actors assert [call.args[0] for call in fake_ray.kill.call_args_list] == expected for actor in expected: actor.shutdown.remote.assert_not_called() actor.wait_for_exit.options.assert_not_called() fake_ray.util.remove_placement_group.assert_called_once_with(group) assert executor.workers == [] assert executor._closed @pytest.mark.parametrize("failure_type", [TimeoutError, RuntimeError, KeyboardInterrupt]) @pytest.mark.parametrize("cleanup_fails", [False, True]) def test_placement_group_wait_failure_cleans_up_and_preserves_error(monkeypatch, failure_type, cleanup_fails): fake_ray = Mock() fake_ray.exceptions.GetTimeoutError = TimeoutError fake_ray.util.get_current_placement_group.return_value = None group = fake_ray.util.placement_group.return_value original_error = failure_type("readiness failed") fake_ray.get.side_effect = original_error if cleanup_fails: fake_ray.util.remove_placement_group.side_effect = RuntimeError("cleanup failed") monkeypatch.setattr(ray_module, "ray", fake_ray) executor = object.__new__(ray_module.RayDiffusionExecutor) executor.od_config = SimpleNamespace(num_gpus=2) expected_type = ValueError if failure_type is TimeoutError else failure_type with pytest.raises(expected_type) as raised: executor._get_or_create_placement_group() if failure_type is TimeoutError: assert raised.value.__cause__ is original_error assert "Cannot reserve 2 Ray GPU bundle(s)" in str(raised.value) else: assert raised.value is original_error fake_ray.util.remove_placement_group.assert_called_once_with(group) def test_health_probes_use_separate_concurrency_group(monkeypatch): fake_ray = Mock() fake_ray.get.return_value = [True, True] monkeypatch.setattr(ray_module, "ray", fake_ray) actors = [Mock(), Mock()] executor = _make_executor(fake_ray, actors) executor.check_health() for actor in actors: actor.check_alive.options.assert_called_once_with(concurrency_group="health") actor.check_alive.options.return_value.remote.assert_called_once_with() actor.check_alive.remote.assert_not_called() fake_ray.get.assert_called_once_with( [actor.check_alive.options.return_value.remote.return_value for actor in actors], timeout=ray_module._HEALTH_CHECK_TIMEOUT_S, ) assert not executor._is_failed @pytest.mark.parametrize("probe_result", ["error", "not_alive"]) def test_health_failure_releases_actors_and_notifies(monkeypatch, probe_result): fake_ray = Mock() if probe_result == "error": fake_ray.get.side_effect = RuntimeError("actor died") else: fake_ray.get.return_value = [True, False] monkeypatch.setattr(ray_module, "ray", fake_ray) actors = [Mock(), Mock()] group = object() executor = _make_executor(fake_ray, actors, group=group) callback = Mock() executor.register_failure_callback(callback) with pytest.raises(ray_module.EngineDeadError): executor.check_health() assert executor._is_failed assert executor._closed callback.assert_called_once_with() assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors fake_ray.util.remove_placement_group.assert_called_once_with(group) def _make_request_executor(monkeypatch, prompts, component="text_encoder", allgather=True): executor = object.__new__(ray_module.RayDiffusionExecutor) executor._closed = False executor._is_failed = False executor._state_lock = threading.Lock() executor._failure_callbacks = [] executor._finalizer = Mock() executor.od_config = SimpleNamespace( parallel_config=SimpleNamespace(data_parallel_size=len(prompts)), diffusion_offload_config={ "mode": "layer", "components": [component], "layer_options": {component: {"weight_transfer": "allgather" if allgather else "rank-local"}}, }, ) monkeypatch.setattr(ray_module, "build_request_batch_sampling_params_key", lambda req: ()) requests = [] for index, prompt in enumerate(prompts): request_id = str(index) req = SimpleNamespace(request_id=request_id, prompt=prompt, sampling_params=SimpleNamespace(extra_args=None)) requests.append(SimpleNamespace(request_id=request_id, req=req, diffusion_kv_metadata=None)) return executor, SimpleNamespace(scheduled_new_reqs=requests, kv_prefetch_job=None) @pytest.mark.parametrize("field", ["prompt_embeds", "negative_prompt_embeds"]) def test_text_encoder_allgather_rejects_mismatched_embedding_presence(monkeypatch, field): executor, scheduled = _make_request_executor(monkeypatch, ["first", {"prompt": "second", field: object()}]) executor.collective_rpc = Mock() with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"): executor.execute_request(scheduled) executor.collective_rpc.assert_not_called() def test_text_encoder_allgather_accepts_matching_precomputed_embedding_paths(monkeypatch): executor, scheduled = _make_request_executor( monkeypatch, [{"prompt_embeds": object()}, {"prompt_embeds": object()}] ) results = [DiffusionOutput(output="first"), DiffusionOutput(output="second")] executor.collective_rpc = Mock(return_value=[{"dp_rank": rank, "output": out} for rank, out in enumerate(results)]) output = executor.execute_request(scheduled) assert [item.result for item in output.runner_outputs] == results executor.collective_rpc.assert_called_once() @pytest.mark.parametrize("allgather", [False, True]) def test_single_request_allgather_timeout_option(monkeypatch, allgather): executor, scheduled = _make_request_executor(monkeypatch, ["prompt"], allgather=allgather) executor.collective_rpc = Mock(return_value=DiffusionOutput(output="done")) executor.execute_request(scheduled) kwargs = executor.collective_rpc.call_args.kwargs if allgather: assert kwargs["timeout"] == ray_module._DLO_DP_WAVE_TIMEOUT_S else: assert "timeout" not in kwargs def test_single_request_allgather_timeout_fails_executor_and_cleans_up(monkeypatch): executor, scheduled = _make_request_executor(monkeypatch, ["prompt"]) callback = Mock() executor.register_failure_callback(callback) executor.collective_rpc = Mock(side_effect=TimeoutError("AllGather timed out")) output = executor.execute_request(scheduled) assert executor._is_failed assert executor._closed executor._finalizer.assert_called_once() callback.assert_called_once() assert "AllGather timed out" in output.runner_outputs[0].result.error @pytest.mark.parametrize("wrapper", ["media", "output", "batch"]) def test_move_to_cpu_handles_typed_video_media(wrapper): from vllm_omni.diffusion.media import ( DiffusionMediaOutput, VideoMediaOutput, VideoTensorEncoding, VideoTensorLayout, VideoTensorSpec, VideoValueRange, ) from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput tensor = torch.ones((1, 3, 2, 4, 4), requires_grad=True) media = DiffusionMediaOutput( video=VideoMediaOutput( tensor=tensor, spec=VideoTensorSpec( layout=VideoTensorLayout.BCTHW, encoding=VideoTensorEncoding.NORMALIZED_FLOAT, value_range=VideoValueRange.ZERO_TO_ONE, ), ), prepared_for_transport=True, ) value = media if wrapper != "media": value = DiffusionOutput(media=media, to_cpu=False) if wrapper == "batch": value = BatchRunnerOutput.from_list([RunnerOutput(request_id="video", finished=True, result=value)]) moved = ray_module._move_to_cpu(value) if wrapper == "batch": moved = moved.runner_outputs[0].result if wrapper != "media": moved = moved.media moved.validate() assert moved.video.tensor.device.type == "cpu" assert not moved.video.tensor.requires_grad assert moved.video.spec == media.video.spec assert moved.prepared_for_transport torch.testing.assert_close(moved.video.tensor, tensor.detach()) assert media.video.tensor is tensor assert tensor.requires_grad