201 lines
8.4 KiB
Python
201 lines
8.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.config.omni_config import OmniStageRuntimeConfig
|
|
from vllm_omni.engine import cuda_mps
|
|
from vllm_omni.engine.stage_runtime import StageRuntime
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.fixture
|
|
def controls(monkeypatch):
|
|
monkeypatch.delenv("CUDA_MPS_PIPE_DIRECTORY", raising=False)
|
|
monkeypatch.setattr(cuda_mps.shutil, "which", lambda name: "/bin/mps-control")
|
|
run = Mock(return_value=SimpleNamespace(stdout=""))
|
|
monkeypatch.setattr(cuda_mps.subprocess, "run", run)
|
|
return run
|
|
|
|
|
|
def test_private_server_has_scoped_environment_and_idempotent_close(controls):
|
|
original = dict(os.environ)
|
|
server = cuda_mps.CudaMPSServer("GPU-example")
|
|
pipe = Path(server.env["CUDA_MPS_PIPE_DIRECTORY"])
|
|
assert pipe.is_dir()
|
|
assert controls.call_args_list[0].args[0] == ["/bin/mps-control", "-d"]
|
|
assert controls.call_args_list[0].kwargs["env"]["CUDA_VISIBLE_DEVICES"] == "GPU-example"
|
|
assert dict(os.environ) == original
|
|
server.close()
|
|
count = controls.call_count
|
|
server.close()
|
|
assert controls.call_count == count
|
|
assert controls.call_args.kwargs["input"] == "quit\n"
|
|
assert not pipe.exists()
|
|
|
|
|
|
def test_operator_daemon_is_checked_but_never_stopped(controls, monkeypatch):
|
|
monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/mps")
|
|
server = cuda_mps.CudaMPSServer("GPU-example")
|
|
server.close()
|
|
assert controls.call_count == 1
|
|
assert controls.call_args.kwargs["input"] == "get_server_list\n"
|
|
|
|
|
|
def test_explicit_pipe_overrides_parent_without_taking_ownership(controls, monkeypatch):
|
|
monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
|
|
server = cuda_mps.CudaMPSServer("GPU-example", pipe_directory="/operator/stage")
|
|
|
|
assert controls.call_args.kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage"
|
|
assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
|
|
server.close()
|
|
assert controls.call_count == 1
|
|
|
|
|
|
def test_explicit_empty_pipe_selects_private_daemon(controls, monkeypatch):
|
|
monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
|
|
server = cuda_mps.CudaMPSServer("GPU-example", pipe_directory="")
|
|
pipe = server.env["CUDA_MPS_PIPE_DIRECTORY"]
|
|
try:
|
|
assert pipe != "/operator/parent"
|
|
assert controls.call_args_list[0].args[0] == ["/bin/mps-control", "-d"]
|
|
assert controls.call_args_list[0].kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == pipe
|
|
assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
|
|
finally:
|
|
server.close()
|
|
assert controls.call_args.kwargs["input"] == "quit\n"
|
|
assert controls.call_args.kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == pipe
|
|
assert not Path(pipe).exists()
|
|
|
|
|
|
@pytest.fixture
|
|
def mps_runtime(monkeypatch):
|
|
from vllm_omni.engine import stage_runtime
|
|
|
|
runtime = StageRuntime.__new__(StageRuntime)
|
|
runtime._mps_servers = {}
|
|
monkeypatch.setattr(stage_runtime, "physical_gpu_uuid", lambda device: "GPU-example")
|
|
monkeypatch.setattr(stage_runtime.current_omni_platform, "is_cuda", lambda: True)
|
|
yield runtime
|
|
runtime._close_mps_servers()
|
|
|
|
|
|
@pytest.mark.parametrize("pipe_directory", ["/operator/stage", ""])
|
|
def test_runtime_shares_first_stage_pipe_policy(controls, monkeypatch, mps_runtime, pipe_directory):
|
|
monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
|
|
config = OmniStageRuntimeConfig(cuda_mps=True, env={"CUDA_MPS_PIPE_DIRECTORY": pipe_directory})
|
|
first_env = mps_runtime._mps_environment("0", config)
|
|
selected_pipe = first_env["CUDA_MPS_PIPE_DIRECTORY"]
|
|
if pipe_directory:
|
|
assert selected_pipe == pipe_directory
|
|
else:
|
|
assert selected_pipe != "/operator/parent"
|
|
assert Path(selected_pipe).is_dir()
|
|
call_count = controls.call_count
|
|
|
|
assert mps_runtime._mps_environment("0", {"cuda_mps": True}) is first_env
|
|
assert mps_runtime._mps_environment("0", config) is first_env
|
|
assert controls.call_count == call_count
|
|
assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"first_pipe,conflicting_pipe",
|
|
[("/operator/first", "/operator/other"), ("/operator/first", ""), ("", "/operator/other")],
|
|
)
|
|
def test_runtime_rejects_conflicting_explicit_stage_pipes(controls, mps_runtime, first_pipe, conflicting_pipe):
|
|
mps_runtime._mps_environment("0", {"cuda_mps": True, "env": {"CUDA_MPS_PIPE_DIRECTORY": first_pipe}})
|
|
call_count = controls.call_count
|
|
|
|
with pytest.raises(ValueError, match="Conflicting CUDA_MPS_PIPE_DIRECTORY"):
|
|
mps_runtime._mps_environment("0", {"cuda_mps": True, "env": {"CUDA_MPS_PIPE_DIRECTORY": conflicting_pipe}})
|
|
|
|
assert controls.call_count == call_count
|
|
|
|
|
|
def test_startup_failure_closes_only_the_private_daemon(controls):
|
|
controls.side_effect = [SimpleNamespace(stdout=""), RuntimeError("health failed"), SimpleNamespace(stdout="")]
|
|
with pytest.raises(RuntimeError, match="health failed"):
|
|
cuda_mps.CudaMPSServer("GPU-example")
|
|
assert controls.call_args.kwargs["input"] == "quit\n"
|
|
|
|
|
|
def test_runtime_opt_in_and_shared_gpu_ownership(monkeypatch):
|
|
from vllm_omni.engine import stage_runtime
|
|
|
|
runtime = StageRuntime.__new__(StageRuntime)
|
|
runtime._mps_servers = {}
|
|
create = Mock(return_value=SimpleNamespace(env={"CUDA_VISIBLE_DEVICES": "GPU-example"}, close=Mock()))
|
|
monkeypatch.setattr(stage_runtime, "CudaMPSServer", create)
|
|
monkeypatch.setattr(stage_runtime, "physical_gpu_uuid", lambda device: "GPU-example")
|
|
monkeypatch.setattr(stage_runtime, "current_omni_platform", SimpleNamespace(is_cuda=lambda: True))
|
|
assert runtime._mps_environment("0", {}) == {}
|
|
assert create.call_count == 0
|
|
for _ in range(2):
|
|
assert runtime._mps_environment("0", {"cuda_mps": True}) == {"CUDA_VISIBLE_DEVICES": "GPU-example"}
|
|
assert create.call_count == 1
|
|
runtime._close_mps_servers()
|
|
runtime._close_mps_servers()
|
|
create.return_value.close.assert_called_once()
|
|
|
|
|
|
@pytest.mark.parametrize("devices", [None, "0,1"])
|
|
def test_runtime_rejects_ambiguous_placement(devices, monkeypatch):
|
|
from vllm_omni.engine import stage_runtime
|
|
|
|
runtime = StageRuntime.__new__(StageRuntime)
|
|
monkeypatch.setattr(stage_runtime, "current_omni_platform", SimpleNamespace(is_cuda=lambda: True))
|
|
with pytest.raises(ValueError, match="exactly one explicit GPU"):
|
|
runtime._mps_environment(devices, {"cuda_mps": True})
|
|
|
|
|
|
def test_failed_shutdown_preserves_socket_for_retry(controls):
|
|
server = cuda_mps.CudaMPSServer("GPU-example")
|
|
pipe = Path(server.env["CUDA_MPS_PIPE_DIRECTORY"])
|
|
controls.side_effect = RuntimeError("quit failed")
|
|
with pytest.raises(RuntimeError, match="quit failed"):
|
|
server.close()
|
|
assert pipe.is_dir()
|
|
controls.side_effect = None
|
|
server.close()
|
|
assert not pipe.exists()
|
|
|
|
|
|
@pytest.mark.parametrize("launch_mode,stage_type", [("remote", "llm"), ("local", "diffusion")])
|
|
def test_mps_rejects_unsupported_stages(launch_mode, stage_type):
|
|
runtime = StageRuntime.__new__(StageRuntime)
|
|
metadata = SimpleNamespace(runtime_cfg={"cuda_mps": True}, stage_type=stage_type)
|
|
stage_cfg = SimpleNamespace(runtime=metadata.runtime_cfg)
|
|
plans = [
|
|
SimpleNamespace(replicas=[SimpleNamespace(metadata=metadata, stage_cfg=stage_cfg, launch_mode=launch_mode)])
|
|
]
|
|
with pytest.raises(ValueError, match="local EngineCore"):
|
|
runtime._validate_mps_topology(plans)
|
|
metadata.runtime_cfg["cuda_mps"] = False
|
|
runtime._validate_mps_topology(plans)
|
|
|
|
|
|
def test_mps_rejects_parallel_initialization():
|
|
runtime = StageRuntime.__new__(StageRuntime)
|
|
runtime._parallel_stage_init = True
|
|
replica = SimpleNamespace(
|
|
metadata=SimpleNamespace(stage_type="llm"),
|
|
stage_cfg=SimpleNamespace(runtime={"cuda_mps": True}),
|
|
launch_mode="local",
|
|
)
|
|
with pytest.raises(ValueError, match="parallel_stage_init=false"):
|
|
runtime._validate_mps_topology([SimpleNamespace(replicas=[replica])])
|
|
|
|
|
|
def test_failed_launch_command_still_checks_private_daemon_cleanup(controls):
|
|
controls.side_effect = [RuntimeError("launch failed after fork"), SimpleNamespace(stdout="")]
|
|
with pytest.raises(RuntimeError, match="launch failed after fork"):
|
|
cuda_mps.CudaMPSServer("GPU-example")
|
|
assert controls.call_args.kwargs["input"] == "quit\n"
|