Rebase multi-API serving on latest main

Signed-off-by: Sy03 <1370724210@qq.com>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -42,7 +42,10 @@ from __future__ import annotations
 
 import asyncio
 import json
+import signal
+from argparse import Namespace
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from types import SimpleNamespace
 
 import pytest
@@ -59,6 +62,66 @@ from vllm_omni.entrypoints.serve.utils import errors as serve_errors
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+
+@dataclass
+class _FakeParallelConfig:
+    _api_process_rank: int
+
+
+@dataclass
+class _FakeVllmConfig:
+    parallel_config: _FakeParallelConfig
+
+
+def test_omni_api_worker_sets_parent_death_signal_before_validation(mocker):
+    args = _minimal_args(_omni_stage_client_configs=[{"stage_addresses": {}}])
+    death_signal = mocker.patch.object(api_server, "set_death_signal")
+
+    with pytest.raises(RuntimeError, match="client configuration"):
+        api_server.run_omni_api_server_worker_proc(
+            "127.0.0.1:8000",
+            mocker.Mock(),
+            args,
+            {"client_count": 2, "client_index": -1},
+        )
+
+    death_signal.assert_called_once_with(signal.SIGTERM)
+
+
+@pytest.mark.asyncio
+async def test_omni_run_server_worker_forwards_client_config(monkeypatch):
+    captured: dict[str, object] = {}
+
+    @asynccontextmanager
+    async def fake_build_async_omni(*_args, **kwargs):
+        captured.update(kwargs)
+        yield _FakeEngineClient(vllm_config=_FakeVllmConfig(_FakeParallelConfig(_api_process_rank=1)))
+
+    monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+    monkeypatch.setattr(api_server, "build_openai_app", lambda *_args: FastAPI())
+    monkeypatch.setattr(api_server, "omni_init_app_state", lambda *_args: asyncio.sleep(0))
+    monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_args: None)
+
+    async def fake_storage_start() -> None:
+        return None
+
+    monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start)
+
+    async def fake_serve_http(app, **_kwargs):
+        assert app.state.api_server_count == 2
+        task = asyncio.create_task(asyncio.sleep(0))
+        await asyncio.sleep(0)
+        return task
+
+    monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+    args = _minimal_args(model="dummy")
+    config = {"client_count": 2, "client_index": 1, "stage_addresses": {}}
+
+    await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), args, config)
+    assert captured["client_config"] is config
+
+
 # FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
 # would invent ("HEAD", path) entries for every GET and fail the current census.
 # TODO: extend this set and the expected census if @router.head / @router.options appear.
@@ -708,6 +771,38 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None:
     assert exc_info.value.detail == "The model does not support Speech API"
 
 
+@pytest.mark.asyncio
+async def test_multi_api_rejects_runtime_voice_upload() -> None:
+    app = FastAPI()
+    app.state.api_server_count = 2
+    raw_request = _request_for(app, method="POST", path="/v1/audio/voices")
+
+    with pytest.raises(HTTPException) as exc_info:
+        await api_server.upload_voice(
+            raw_request,
+            audio_sample=None,
+            speaker_embedding=None,
+            consent="consent-id",
+            name="probe",
+        )
+
+    assert exc_info.value.status_code == 409
+    assert "process-local frontend state" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_multi_api_rejects_sleep_route() -> None:
+    app = FastAPI()
+    app.state.api_server_count = 2
+    raw_request = _request_for(app, method="POST", path="/v1/omni/sleep")
+
+    with pytest.raises(HTTPException) as exc_info:
+        await api_server.omni_sleep(api_server.OmniSleepRequest(stage_ids=[0]), raw_request)
+
+    assert exc_info.value.status_code == 409
+    assert "process-local frontend state" in exc_info.value.detail
+
+
 def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
     """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
 
@@ -931,3 +1026,57 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
         must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
         must_be_none=_MULTISTAGE_MUST_BE_NONE,
     )
+
+
+@pytest.mark.parametrize("count", [None, 0, -1, "2", True])
+def test_process_local_guard_rejects_unknown_topology(count):
+    app = FastAPI()
+    if count is not None:
+        app.state.api_server_count = count
+    # CLI args cannot substitute for missing/invalid actual runtime state.
+    app.state.args = Namespace(api_server_count=1)
+    with pytest.raises(HTTPException) as exc:
+        api_server._reject_process_local_state_with_multiple_api_workers(_request_for(app), "Sleep")
+    assert exc.value.status_code == 503
+
+
+@pytest.mark.parametrize(
+    "requested,config,expected",
+    [(None, None, 1), (1, None, 1), (2, {"client_count": 2}, 2), (None, {"client_count": 2}, 2)],
+)
+def test_api_topology_resolved_from_launch(requested, config, expected):
+    assert api_server._resolve_api_server_count(Namespace(api_server_count=requested), config) == expected
+
+
+@pytest.mark.parametrize("requested,config", [(2, None), (2, {}), (2, {"client_count": 1}), (2, {"client_count": "2"})])
+def test_api_topology_rejects_inconsistent_launch(requested, config):
+    with pytest.raises(ValueError):
+        api_server._resolve_api_server_count(Namespace(api_server_count=requested), config)
+
+
+@pytest.mark.parametrize("count,status", [(2, 409), (None, 503)])
+@pytest.mark.parametrize(
+    "method,path",
+    [
+        ("POST", "/v1/videos"),
+        ("GET", "/v1/videos"),
+        ("GET", "/v1/videos/job"),
+        ("GET", "/v1/videos/job/content"),
+        ("DELETE", "/v1/videos/job"),
+    ],
+)
+def test_video_lifecycle_guard_precedes_processing(monkeypatch, count, status, method, path):
+    from fastapi.testclient import TestClient
+
+    app = FastAPI()
+    if count is not None:
+        app.state.api_server_count = count
+    app.include_router(api_server.router)
+
+    async def forbidden_form():
+        pytest.fail("video form processing must not run before topology rejection")
+
+    app.dependency_overrides[api_server._parse_video_form] = forbidden_form
+    with TestClient(app) as client:
+        response = client.request(method, path)
+    assert response.status_code == status

tests/entrypoints/openai_api/test_omni_sleep_wakeup.py

@@ -20,6 +20,7 @@ def _make_app(engine_client):
     from vllm_omni.entrypoints.openai.api_server import router
 
     app = FastAPI()
+    app.state.api_server_count = 1
     app.include_router(router)
     app.state.engine_client = engine_client
     app.state.sleeping_stages = set()
@@ -253,6 +254,7 @@ def pure_diffusion_app(pure_diffusion_engine, mocker):
     )
 
     app = FastAPI()
+    app.state.api_server_count = 1
     app.include_router(router)
 
     mocker.patch(

tests/entrypoints/openai_api/test_serving_speech.py

@@ -3793,6 +3793,7 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt
 def _make_api_server_request(handler, *, method: str = "POST", path: str = "/v1/audio/voices") -> Request:
     app = FastAPI()
     app.state.openai_serving_speech = handler
+    app.state.api_server_count = 1
     scope = {
         "type": "http",
         "app": app,

tests/entrypoints/openai_api/test_video_server.py

@@ -345,6 +345,7 @@ async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch)
 @pytest.fixture
 def test_client():
     app = FastAPI()
+    app.state.api_server_count = 1
     app.include_router(router)
     app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
         diffusion_engine=FakeAsyncOmni(),
@@ -1629,6 +1630,7 @@ def test_action_extraction_accepts_multimodal_actions_payload():
 
 def test_missing_handler_returns_503():
     app = FastAPI()
+    app.state.api_server_count = 1
     app.include_router(router)
     app.state.openai_serving_video = None
     client = TestClient(app)
@@ -2541,6 +2543,7 @@ def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, tes
 
 def test_sync_missing_handler_returns_503():
     app = FastAPI()
+    app.state.api_server_count = 1
     app.include_router(router)
     app.state.openai_serving_video = None
     client = TestClient(app)

tests/entrypoints/test_serve.py

@@ -6,12 +6,17 @@
 from __future__ import annotations
 
 import argparse
+from dataclasses import dataclass
+from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
 
 import pytest
 from pytest_mock import MockerFixture
+from vllm.v1.engine.utils import EngineZmqAddresses
 
 from vllm_omni.config.resolver import OmniConfigResolution
+from vllm_omni.engine.stage_engine_startup import StageReplicaResources
+from vllm_omni.engine.stage_runtime import StageEngineLaunch
 from vllm_omni.entrypoints.cli.serve import (
     OmniServeCommand,
     _parse_stage_overrides,
@@ -23,6 +28,31 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@dataclass(frozen=True)
+class _FakeProcess:
+    sentinel: str
+    exitcode: int | None
+    name: str
+    pid: int
+
+
+@dataclass
+class _FakeProcessOwner:
+    processes: list[_FakeProcess]
+
+
+@dataclass
+class _FakeStageEngineArgs:
+    async_chunk: bool
+    enable_sleep_mode: bool
+
+
+@dataclass
+class _FakeStageConfig:
+    stage_id: int
+    engine_args: _FakeStageEngineArgs
+
+
 def _resolved(*stages: SimpleNamespace) -> OmniConfigResolution:
     return OmniConfigResolution(config_path="/fake/stages.yaml", stage_configs=tuple(stages))
 
@@ -42,6 +72,8 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
     explicit = args.get_explicit_kwargs_dict()
     assert args.get_explicit_kwargs_dict()
     assert not explicit["async_chunk"]
+    assert (args.api_server_count or 1) == 1
+    assert "api_server_count" not in explicit
 
 
 def _parse_serve_args(argv: list[str]) -> TrackingNamespace:
@@ -167,6 +199,214 @@ def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None
         parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value])
 
 
+def test_tracking_namespace_is_picklable_for_spawned_api_workers() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2"])
+    args._omni_stage_client_configs = [{"stage_addresses": {0: {0: {"input_address": "ipc://input"}}}}]
+
+    restored = ForkingPickler.loads(ForkingPickler.dumps(args))
+
+    assert restored.api_server_count == 2
+    assert restored.get_explicit_kwargs_dict()["api_server_count"] == 2
+    assert restored._omni_stage_client_configs == args._omni_stage_client_configs
+
+
+def test_serve_validate_rejects_multiple_api_servers_for_diffusion(mocker: MockerFixture) -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    cmd = OmniServeCommand()
+    cmd.subparser_init(subparsers)
+    args = parser.parse_args(["serve", "fake-diffusion-model", "--omni", "--api-server-count", "2"])
+
+    mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=True)
+    validate = mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
+
+    with pytest.raises(ValueError, match="not supported for diffusion"):
+        cmd.validate(args)
+
+    validate.assert_not_called()
+
+
+def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: MockerFixture) -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    cmd = OmniServeCommand()
+    cmd.subparser_init(subparsers)
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--api-server-count", "2", "--enable-sleep-mode"])
+
+    mocker.patch("vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", return_value=False)
+    mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args")
+
+    with pytest.raises(ValueError, match="enable-sleep-mode"):
+        cmd.validate(args)
+
+
+def test_build_multi_api_stage_runtime_matches_current_constructor(mocker: MockerFixture) -> None:
+    from vllm_omni.entrypoints.cli import serve as serve_module
+
+    args = TrackingNamespace(
+        argparse.Namespace(
+            model="dummy-model",
+            model_tag=None,
+            stage_init_timeout=1,
+            tokenizer=None,
+            disable_log_stats=False,
+        ),
+        frozenset({"model"}),
+    )
+    stage_config = _FakeStageConfig(
+        stage_id=0,
+        engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=False),
+    )
+    mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
+    mocker.patch(
+        "vllm_omni.config.config_factory.with_trust_remote_code_override",
+        side_effect=lambda kwargs, _trust: kwargs,
+    )
+    mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
+    mocker.patch(
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)),
+    )
+
+    runtime = serve_module._build_multi_api_stage_runtime(args, 2)
+
+    assert runtime._client_count == 1
+    assert runtime._client_index == 0
+
+
+def test_build_multi_api_stage_runtime_rejects_sleep_enabled_in_stage_config(mocker: MockerFixture) -> None:
+    from vllm_omni.entrypoints.cli import serve as serve_module
+
+    args = TrackingNamespace(
+        argparse.Namespace(
+            model="dummy-model",
+            model_tag=None,
+            stage_init_timeout=1,
+            tokenizer=None,
+            disable_log_stats=False,
+        ),
+        frozenset({"model"}),
+    )
+    stage_config = _FakeStageConfig(
+        stage_id=3,
+        engine_args=_FakeStageEngineArgs(async_chunk=False, enable_sleep_mode=True),
+    )
+    mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model")
+    mocker.patch(
+        "vllm_omni.config.config_factory.with_trust_remote_code_override",
+        side_effect=lambda kwargs, _trust: kwargs,
+    )
+    mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={})
+    mocker.patch(
+        "vllm_omni.config.resolver.resolve_omni_config",
+        return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=(stage_config,)),
+    )
+
+    with pytest.raises(ValueError, match=r"sleep mode.*stage\(s\) \[3\]"):
+        serve_module._build_multi_api_stage_runtime(args, 2)
+
+
+def test_run_multi_api_server_omni_starts_workers_after_shared_engine_launch(mocker: MockerFixture) -> None:
+    from vllm_omni.entrypoints.cli import serve as serve_module
+
+    args = TrackingNamespace(
+        argparse.Namespace(api_server_count=2, shutdown_timeout=1),
+        frozenset({"api_server_count", "shutdown_timeout"}),
+    )
+    socket = mocker.Mock()
+    primary_addresses = EngineZmqAddresses(
+        inputs=["ipc://input-0", "ipc://input-1"],
+        outputs=["ipc://output-0", "ipc://output-1"],
+    )
+    engine_launch = StageEngineLaunch(
+        client_configs=[{"client_count": 2, "client_index": index, "stage_addresses": {}} for index in range(2)],
+        resources=[StageReplicaResources(addresses=primary_addresses)],
+    )
+    runtime = mocker.MagicMock()
+    runtime.launch_stage_engines.return_value.__enter__.return_value = engine_launch
+    runtime.launch_stage_engines.return_value.__exit__.return_value = False
+
+    manager = mocker.Mock()
+    manager.processes = [mocker.Mock()]
+    manager.gather_actual_addresses.return_value = (primary_addresses.inputs, primary_addresses.outputs)
+
+    mocker.patch.object(serve_module, "_build_multi_api_stage_runtime", return_value=runtime)
+    mocker.patch.object(serve_module, "_wait_for_multi_api_server_completion")
+    mocker.patch("signal.signal")
+    mocker.patch("vllm.entrypoints.openai.api_server.setup_server", return_value=("127.0.0.1:8000", socket))
+    mocker.patch("vllm.v1.metrics.prometheus.setup_multiprocess_prometheus")
+    start_manager = mocker.patch.object(serve_module, "_start_api_server_process_manager", return_value=manager)
+
+    serve_module.run_multi_api_server_omni(args)
+
+    assert args._omni_stage_client_configs is engine_launch.client_configs
+    start_manager.assert_called_once()
+    manager_kwargs = start_manager.call_args.kwargs
+    assert manager_kwargs["cleanup_timeout"] == 1
+    assert manager_kwargs["num_servers"] == 2
+    assert manager_kwargs["input_addresses"] == ["ipc://input-0", "ipc://input-1"]
+    assert manager_kwargs["target_server_fn"] is serve_module.run_omni_api_server_worker_proc
+    assert engine_launch.watched_frontend_processes == manager.processes
+    runtime.shutdown.assert_called_once_with()
+    socket.close.assert_called_once_with()
+
+
+def test_start_api_server_process_manager_cleans_up_partial_start(mocker: MockerFixture) -> None:
+    from vllm_omni.entrypoints.cli import serve as serve_module
+
+    first_process = mocker.Mock(name="first_process")
+    second_process = mocker.Mock(name="second_process")
+    first_process.pid = 1234
+    second_process.pid = None
+    second_process.start.side_effect = RuntimeError("spawn failed")
+    parent_pipes = [mocker.Mock(name="parent_pipe_0"), mocker.Mock(name="parent_pipe_1")]
+    child_pipes = [mocker.Mock(name="child_pipe_0"), mocker.Mock(name="child_pipe_1")]
+    spawn_context = mocker.Mock()
+    spawn_context.Process.side_effect = [first_process, second_process]
+    spawn_context.Pipe.side_effect = list(zip(parent_pipes, child_pipes))
+    mocker.patch("vllm.v1.utils.multiprocessing.get_context", return_value=spawn_context)
+    shutdown = mocker.patch("vllm.v1.utils.shutdown")
+
+    with pytest.raises(RuntimeError, match="spawn failed"):
+        serve_module._start_api_server_process_manager(
+            cleanup_timeout=1.5,
+            listen_address="127.0.0.1:8000",
+            sock=mocker.Mock(),
+            args=argparse.Namespace(),
+            num_servers=2,
+            input_addresses=["ipc://input-0", "ipc://input-1"],
+            output_addresses=["ipc://output-0", "ipc://output-1"],
+            target_server_fn=mocker.Mock(),
+        )
+
+    first_process.start.assert_called_once_with()
+    second_process.start.assert_called_once_with()
+    for pipe in parent_pipes:
+        pipe.close.assert_called_once_with()
+    shutdown.assert_called_once_with([first_process], timeout=1.5)
+
+
+def test_wait_for_multi_api_server_completion_rejects_engine_failure(
+    mocker: MockerFixture,
+) -> None:
+    from vllm_omni.entrypoints.cli import serve as serve_module
+
+    api_process = _FakeProcess(sentinel="api-0", exitcode=None, name="api-0", pid=10)
+    engine_process = _FakeProcess(sentinel="engine-0", exitcode=4, name="engine-0", pid=20)
+    manager = _FakeProcessOwner(processes=[api_process])
+    engine_launch = StageEngineLaunch(
+        client_configs=[],
+        resources=[StageReplicaResources(manager=_FakeProcessOwner(processes=[engine_process]))],
+    )
+    mocker.patch("multiprocessing.connection.wait", return_value=["engine-0"])
+
+    with pytest.raises(RuntimeError, match=r"Shared stage engine process engine-0 .* exited with code 4"):
+        serve_module._wait_for_multi_api_server_completion(manager, engine_launch)
+
+
 def test_serve_parser_rejects_stage_configs_path() -> None:
     parser = TrackingArgumentParser()
     subparsers = parser.add_subparsers(dest="subcommand")
@@ -754,3 +994,22 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
 
     with pytest.raises(RuntimeError, match=r"exited with code 137"):
         run_headless(_make_headless_args(stage_id=1))
+
+
+@pytest.mark.parametrize("name", ["unfiltered_ns", "explicit_keys", "missing", "__dict__"])
+def test_tracking_namespace_uninitialized_access_raises_attribute_error(name):
+    namespace = object.__new__(TrackingNamespace)
+    with pytest.raises(AttributeError):
+        getattr(namespace, name)
+
+
+@pytest.mark.parametrize("deep", [False, True])
+def test_tracking_namespace_copy_preserves_tracking(deep):
+    import copy
+
+    namespace = TrackingNamespace(
+        argparse.Namespace(model="example", api_server_count=2), frozenset({"api_server_count"})
+    )
+    restored = copy.deepcopy(namespace) if deep else copy.copy(namespace)
+    assert restored.get_explicit_kwargs_dict() == {"api_server_count": 2}
+    assert restored.model == "example"

tests/entrypoints/test_utils.py

@@ -529,3 +529,21 @@ class TestCumulativeStreamingCoercion:
         assert all([isinstance(r, SamplingParams) for r in result])
         assert result[0].output_kind == RequestOutputKind.DELTA
         assert result[1].output_kind == RequestOutputKind.DELTA
+
+
+@pytest.mark.parametrize("api_server_count", [1, 2])
+def test_prepare_stage_config_inputs_filters_frontend_options(api_server_count):
+    from vllm_omni.entrypoints.utils import prepare_stage_config_inputs
+
+    kwargs = {
+        "api_server_count": api_server_count,
+        "disable_log_stats": True,
+        "model_tag": "example",
+        "max_num_seqs": 8,
+        "stage_overrides": '{"0": {"max_num_seqs": 4}}',
+    }
+    inputs = prepare_stage_config_inputs("example", kwargs, trust_remote_code=None)
+
+    assert inputs.kwargs == {"max_num_seqs": 8}
+    assert inputs.stage_overrides == {"0": {"max_num_seqs": 4}}
+    assert kwargs["api_server_count"] == api_server_count