diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index db204bdd552..e113c9d7c81 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -85,7 +85,7 @@ depends on the installed kernels and model path. | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | | `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental | -| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental | +| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; defaults to on for Qwen3-TTS and off for other pipelines | Pipeline default computed from `pipeline_config.model_type` at engine initialization; env override resolved at `Orchestrator` construction and separately when the serving-side final-output drain starts | An explicit env value wins; otherwise both consumers use the engine's pipeline default. Values are stripped and case-normalized; any unrecognized value selects the legacy poll loop. Set before server startup. | Experimental | | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | | `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `/vllm_omni_orch_monitor_.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic | | `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 62fff29c558..4e61d2531a1 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -22,7 +22,7 @@ profiler_config: Supported fields: | Field | Description | -|---|---| +| --- | --- | | `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. | | `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. | | `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. | @@ -287,6 +287,52 @@ Recommended viewers: For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/). +## 6. Profiling Omni AR/TTS stages + +For Omni AR and TTS stages, put `profiler_config` under each stage that should +be captured. The serving API then controls a narrow window with +`POST /start_profile` and `POST /stop_profile`; model initialization and +warmup stay outside the trace. + +Example stage override for a host-stack torch trace: + +```bash +PROFILE='{"profiler":"torch","torch_profiler_dir":"/tmp/qwen3_tts_profile",'\ +'"torch_profiler_with_stack":true,"torch_profiler_record_shapes":false,'\ +'"torch_profiler_use_gzip":true,"torch_profiler_dump_cuda_time_total":true}' + +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency.yaml \ + --stage-overrides "{\"0\":{\"profiler_config\":$PROFILE},\"1\":{\"profiler_config\":$PROFILE}}" +``` + +After one or more warmup requests, capture one short measured request: + +```bash +curl -sS -X POST http://127.0.0.1:8000/start_profile \ + -H 'Content-Type: application/json' -d '{"stages":[0,1]}' +# send one fixed /v1/audio/speech or /v1/chat/completions request here +curl -sS -X POST http://127.0.0.1:8000/stop_profile \ + -H 'Content-Type: application/json' -d '{"stages":[0,1]}' +``` + +For an Nsight Systems CUDA timeline, use the same stage override with +`{"profiler":"cuda"}` and wrap the server: + +```bash +nsys profile --trace=cuda,nvtx,osrt --sample=process-tree \ + --cpuctxsw=process-tree --trace-fork-before-exec=true \ + --capture-range=cudaProfilerApi --capture-range-end=stop-shutdown \ + --cuda-graph-trace=node -o /tmp/qwen3_tts_nsys \ + vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni ... +``` + +For Omni AR workers, `profiler="cuda"` requires the native CUDA profiler +wrapper and uses `cudaProfilerStart/Stop`; the report is only valid if the +`/start_profile` and `/stop_profile` calls both succeed. Use `nsys stats` for +CUDA API, GPU kernel, OS runtime, and NVTX summaries. Treat profiler latency as +diagnostic only; use a separate no-profiler run for performance numbers. + ## 7. Orchestrator Monitor > **Warning:** Diagnostic only. Adds lightweight counters on the orchestrator poll loop and per-replica queue depth. Disable in production unless you are actively debugging orchestrator saturation. @@ -314,9 +360,11 @@ If unset, the monitor writes `./vllm_omni_orch_monitor_.json` at shutd Each 1-second window records: | Series | Description | -|---|---| +| --- | --- | | `windows.duration_s` | Wall time covered by the window | | `windows.loop_idle` / `windows.loop_active` | Orchestrator poll-loop iterations with no work vs. active forwarding | +| `windows.dispatch_queue_size` | Ready-queue size sampled at each monitor window boundary | +| `windows.dispatch_queue_high_water` | Maximum Event Driven ready-queue size observed in each window | | `replicas..outputs_queue_size` | MP client `outputs_queue` backlog for that replica | | `replicas..inflight` | Requests currently bound/routed to the replica | @@ -331,12 +379,16 @@ orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still records one iteration per routed output. Window counts and the `loop_active_pct` summary are therefore not comparable across the two modes. +For Event Driven runs, use `dispatch_queue_high_water` rather than only the +window-boundary `dispatch_queue_size`: a serial dispatcher can drain a burst +before the window closes, hiding a transient ready-queue backlog in the latter. + ### Relationship to other diagnostics This monitor is intentionally separate from the existing profiling tools: | Tool | Scope | Output | -|---|---|---| +| --- | --- | --- | | `--enable-diffusion-pipeline-profiler` | Diffusion stage functions (`vae.decode`, `diffuse`, …) | Per-request `stage_durations` / logs | | `--enable-ar-profiler` | AR stage generation time | Per-request `stage_durations` | | `profiler_config` (`torch` / `cuda`) | GPU/CPU kernels inside a stage worker | `trace.json`, Nsight, operator tables | diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index df22f3ce1f7..c4dd4bb3699 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -79,7 +79,7 @@ the boundary affected by the in-flight stage client/process refactor in responsibilities proposed only by that PR are not current contracts. The orchestration loop also has an opt-in event-driven mode -(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in +(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off except for Qwen3-TTS) proposed in [#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll cadence only: the routing, ordering, and terminal-state contracts below hold identically on both loops. diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index c0b8fe5e213..d973b51a565 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -957,22 +957,24 @@ Use `/v1/audio/voices` to list available voices for the loaded model. ## Orchestration Loop (experimental) Multi-stage omni deployments route stage outputs through a single orchestrator -loop. By default that loop polls every stage replica on a 1 ms cadence. An -opt-in event-driven mode replaces the poll with one reader task per live stage +loop. The legacy loop polls every stage replica on a 1 ms cadence. The +event-driven mode replaces the poll with one reader task per live stage replica awaiting its client directly, and switches the serving-side -final-output drain to a condition-variable wakeup at the same time. +final-output drain to a condition-variable wakeup at the same time. Qwen3-TTS +uses this mode by default; other pipelines keep the legacy poll unless enabled. **Configuration (environment variables):** | Variable | Default | Description | | --- | --- | --- | -| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. | +| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | On for Qwen3-TTS; off for other pipelines | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. An explicit value wins; otherwise the pipeline default computed at engine initialization is used. The override is resolved at orchestrator construction and separately when the final-output drain starts. `1`, `true`, `yes`, or `on` enables it, ignoring case and surrounding whitespace; other values select the legacy poll loop. | -Set it on the process that runs the orchestrator (stage 0 of an omni -deployment) before starting the server: +Set it on the process that runs the orchestrator (stage 0 of an omni deployment) +before starting the server when overriding the pipeline default. For example, +explicitly disable event-driven orchestration for Qwen3-TTS: ```bash -export VLLM_OMNI_EVENT_DRIVEN_ORCH=1 +VLLM_OMNI_EVENT_DRIVEN_ORCH=0 \ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --omni \ --port 8091 @@ -982,8 +984,9 @@ The server logs the selected loop mode and its reader/poller counts once at startup, so you can confirm which loop is live. Routing, output ordering, and terminal-state behavior are identical on both -loops; only the poll cadence changes. Leaving the variable unset keeps the -legacy poll loop, which is the supported default. +loops; only the poll cadence changes. Leaving the variable unset selects the +pipeline default: event-driven for Qwen3-TTS, and legacy polling for other +pipelines, including Qwen3-Omni. **Known limitations:** diff --git a/tests/engine/test_orchestrator_event_driven.py b/tests/engine/test_orchestrator_event_driven.py index 284dde55d1c..10972144688 100644 --- a/tests/engine/test_orchestrator_event_driven.py +++ b/tests/engine/test_orchestrator_event_driven.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Event-driven orchestration loop (``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``) tests. Parity suite: re-runs the legacy orchestration scenarios from @@ -19,7 +19,10 @@ from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import ShutdownRequestMessage -from vllm_omni.engine.orchestrator import _event_driven_orch_enabled +from vllm_omni.engine.orchestrator import ( + _event_driven_orch_default_for_pipeline, + _event_driven_orch_enabled, +) from . import test_orchestrator as legacy from . import test_orchestrator_error_handling as legacy_errors @@ -125,6 +128,19 @@ def test_flag_parsing(monkeypatch) -> None: assert _event_driven_orch_enabled() is False +def test_only_qwen3_tts_has_a_pipeline_default() -> None: + assert _event_driven_orch_default_for_pipeline("qwen3_tts") is True + for model_type in (None, "qwen3_omni_moe", "minicpmo_4_5", "moss_tts_delay"): + assert _event_driven_orch_default_for_pipeline(model_type) is False + + +def test_explicit_event_driven_value_overrides_pipeline_default(monkeypatch) -> None: + monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False) + assert _event_driven_orch_enabled(default=True) is True + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", "0") + assert _event_driven_orch_enabled(default=True) is False + + def test_default_is_legacy_loop(monkeypatch) -> None: """Without the env flag, the harness runs the legacy poll loop.""" monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False) diff --git a/tests/engine/test_orchestrator_monitor.py b/tests/engine/test_orchestrator_monitor.py index ee9126cf470..7ac3f395880 100644 --- a/tests/engine/test_orchestrator_monitor.py +++ b/tests/engine/test_orchestrator_monitor.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import json import time @@ -45,6 +48,8 @@ def sampler() -> dict[str, tuple[int, int]]: monitor.register_replica(0, 0) monitor.note_loop(idle=False) + monitor.set_dispatch_queue_size(9) + monitor.set_dispatch_queue_size(7) monitor.note_loop(idle=True) now[0] = 1.5 @@ -54,10 +59,37 @@ def sampler() -> dict[str, tuple[int, int]]: assert payload["configured_window_s"] == monitor_mod._WINDOW_S assert payload["windows"]["loop_idle"] == [1] assert payload["windows"]["loop_active"] == [1] + assert payload["windows"]["dispatch_queue_size"] == [7] + assert payload["windows"]["dispatch_queue_high_water"] == [9] assert payload["replicas"]["stage=0,replica=0"]["outputs_queue_size"] == [5] assert payload["replicas"]["stage=0,replica=0"]["inflight"] == [10] +@pytest.mark.parametrize("drain", [False, True]) +def test_dispatch_backlog_carries_across_windows(tmp_path, monkeypatch, drain): + out_path = tmp_path / "orch_monitor.json" + monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path)) + now = [0.0] + monkeypatch.setattr(time, "monotonic", lambda: now[0]) + monitor = OrchestratorMonitor(replica_sampler=lambda: {}) + monitor.set_dispatch_queue_size(3) + monitor.set_dispatch_queue_size(1) + + now[0] = monitor_mod._WINDOW_S + monitor.note_loop(idle=False) + if drain: + monitor.set_dispatch_queue_size(0) + now[0] += monitor_mod._WINDOW_S + monitor.note_loop(idle=False) + now[0] += monitor_mod._WINDOW_S + monitor.flush() + + windows = json.loads(out_path.read_text())["windows"] + remaining = 0 if drain else 1 + assert windows["dispatch_queue_size"] == [1, remaining, remaining] + assert windows["dispatch_queue_high_water"] == [3, 1, remaining] + + def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch): out_path = tmp_path / "orch_monitor.json" monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path)) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 2e01e910b11..0d5e47960d8 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -17,7 +17,12 @@ from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.omni_config import VllmOmniARStageConfig, VllmOmniDiffusionStageConfig -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, StageExecutionType, StagePipelineConfig +from vllm_omni.config.stage_config import ( + DuplexSessionRuntimeConfig, + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase from vllm_omni.engine.stage_engine_startup import ( @@ -393,6 +398,7 @@ def _make_engine_no_thread( stage_cfgs: list[Any] | None = None, resolved_config_path: str = "/fake/path", patch_deploy_config: bool = True, + resolved_pipeline: PipelineConfig | None = None, **kwargs: Any, ) -> AsyncOmniEngine: mock_stage_configs = stage_cfgs or [_make_stage_cfg(0)] @@ -407,10 +413,19 @@ def _make_engine_no_thread( "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), ) + + def resolve_stage_configs(engine, *args, **kwargs): + if resolved_pipeline is not None: + engine._config_resolution = SimpleNamespace( + pipeline_config=resolved_pipeline, config_path=resolved_config_path + ) + return resolved_config_path, mock_stage_configs + mocker.patch.object( AsyncOmniEngine, "_resolve_stage_configs", - return_value=(resolved_config_path, mock_stage_configs), + autospec=True, + side_effect=resolve_stage_configs, ) mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator") mock_thread_cls = mocker.patch("threading.Thread") @@ -435,6 +450,19 @@ def test_explicit_single_stage_mode_true(self, mocker: MockerFixture): ) assert engine.single_stage_mode is True + @pytest.mark.parametrize("model_type, expected", [("qwen3_tts", True), ("qwen3_omni_moe", False)]) + def test_event_driven_default_uses_resolved_pipeline(self, mocker, model_type, expected): + pipeline = PipelineConfig( + model_type=model_type, + stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),), + ) + engine = self._make_engine_no_thread(mocker, resolved_pipeline=pipeline) + + StageConfigFactory.get_pipeline_config.assert_called_once_with( + model="fake-model", trust_remote_code=False, deploy_config_path=None + ) + assert engine._event_driven_orch_default is expected + def test_stage_id_kwarg_promotes_to_single_stage_mode(self, mocker: MockerFixture): engine = self._make_engine_no_thread( mocker, diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py index a04cc35c2b1..a2779f93771 100644 --- a/tests/worker/test_worker_base.py +++ b/tests/worker/test_worker_base.py @@ -239,3 +239,25 @@ def wake_up(self, tags): assert OmniGPUWorkerBase.wake_up(worker, tags=["weights"]) is True assert calls["tags"] == ["weights"] + + +def test_cuda_profiler_is_wired_for_omni_worker(monkeypatch): + import vllm.profiler.wrapper as profiler_wrapper + + def fake_gpu_worker_init(worker, *args, **kwargs): + worker.vllm_config = SimpleNamespace(profiler_config=SimpleNamespace(profiler="cuda")) + worker.local_rank = 0 + worker.rank = 0 + worker.profiler = None + + class FakeCudaProfiler: + def __init__(self, profiler_config): + self.profiler_config = profiler_config + + monkeypatch.setattr(base.GPUWorker, "__init__", fake_gpu_worker_init) + monkeypatch.setattr(profiler_wrapper, "CudaProfilerWrapper", FakeCudaProfiler) + + worker = OmniGPUWorkerBase() + + assert isinstance(worker.profiler, FakeCudaProfiler) + assert worker.profiler.profiler_config.profiler == "cuda" diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index e22bfb241b9..0a89f41411e 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -55,7 +55,7 @@ ErrorMessage, OutputMessage, ) -from vllm_omni.engine.orchestrator import OrchestratorBase +from vllm_omni.engine.orchestrator import OrchestratorBase, _event_driven_orch_default_for_pipeline from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient from vllm_omni.engine.stage_client import StageClient from vllm_omni.engine.stage_init_utils import build_stage0_input_processor @@ -459,6 +459,7 @@ async def _run_orchestrator() -> None: prom_metrics=self._prom_metrics, log_stats=self._log_stats, enable_orch_monitor=self._enable_orch_monitor, + event_driven_orch_default=self._event_driven_orch_default, ) if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) @@ -879,6 +880,9 @@ def _set_pipeline_runtime_config( ) -> None: """Initialize engine-wide settings resolved from pipeline metadata.""" self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () + self._event_driven_orch_default = _event_driven_orch_default_for_pipeline( + pipeline_config.model_type if pipeline_config is not None else None + ) # No duplex_runtime_extension / duplex_serving_adapter / duplex_control # here: the pre-framework wiring they named is gone, and design rule 3 # in docs/design/fullduplex.md keeps duplex vocabulary out of the base. diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index e69ebc7aea0..bb360a8315c 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -79,7 +79,7 @@ def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: # 1 ms poll cadence to event-driven wakeups: one reader task per live LLM stage # replica awaits `client.get_output_async()` directly — the same pattern vLLM's # own AsyncLLM output handler uses — and feeds a single serial dispatch queue. -# Default off; the legacy poll loop remains the fallback. +# Default is off except for pipelines with an explicit validated default. _EVENT_DRIVEN_ORCH_ENV = "VLLM_OMNI_EVENT_DRIVEN_ORCH" # How often the event-driven loop reconciles its reader-task set against @@ -87,8 +87,16 @@ def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: _ORCH_READER_RECONCILE_INTERVAL_S = 0.5 -def _event_driven_orch_enabled() -> bool: - return os.environ.get(_EVENT_DRIVEN_ORCH_ENV, "0").strip().lower() in ("1", "true", "yes", "on") +def _event_driven_orch_enabled(*, default: bool = False) -> bool: + value = os.environ.get(_EVENT_DRIVEN_ORCH_ENV) + if value is None: + return default + return value.strip().lower() in ("1", "true", "yes", "on") + + +def _event_driven_orch_default_for_pipeline(pipeline_model_type: str | None) -> bool: + """Return whether a pipeline has a validated event-driven default.""" + return pipeline_model_type == "qwen3_tts" def _build_terminal_empty_output( @@ -298,6 +306,7 @@ def __init__( prom_metrics: Any = None, log_stats: bool = False, enable_orch_monitor: bool = False, + event_driven_orch_default: bool = False, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -340,7 +349,7 @@ def __init__( self._shutdown_event = asyncio.Event() self._stages_shutdown = False - self._event_driven_orch = _event_driven_orch_enabled() + self._event_driven_orch = _event_driven_orch_enabled(default=event_driven_orch_default) # Distributed membership (optional, injected by DistStageRuntime) self._membership = membership_controller @@ -890,7 +899,8 @@ async def _orchestration_loop(self) -> None: async def _orchestration_loop_event_driven(self) -> None: """Event-driven variant of ``_orchestration_loop``. - Selected by ``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``. One reader task per + Selected by the explicit ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` value or the + pipeline default computed at engine initialization. One reader task per available LLM stage replica awaits ``client.get_output_async()`` directly — the same pattern vLLM's own ``AsyncLLM`` output handler uses — and feeds a single dispatch queue. This coroutine consumes that queue @@ -923,6 +933,7 @@ async def _llm_replica_reader(stage_id: int, replica_id: int, client: Any) -> No await asyncio.sleep(0.001) continue await ready_q.put(("llm", stage_id, replica_id, raw_outputs)) + self._orch_monitor.set_dispatch_queue_size(ready_q.qsize()) except asyncio.CancelledError: raise except BaseException as e: # noqa: BLE001 - routed to the dispatcher @@ -941,6 +952,7 @@ async def _diffusion_poller(stage_id: int, pool: StagePool) -> None: if output is None: continue await ready_q.put(("diffusion", stage_id, replica_id, output)) + self._orch_monitor.set_dispatch_queue_size(ready_q.qsize()) got = True await asyncio.sleep(0 if got else 0.001) except asyncio.CancelledError: @@ -1059,6 +1071,9 @@ def _reconcile_readers() -> None: continue kind, stage_id, replica_id, payload = pending_get.result() pending_get = None + # Record the queue after dequeue so the gauge reflects work + # still waiting for dispatch rather than producer-side depth. + self._orch_monitor.set_dispatch_queue_size(ready_q.qsize()) if kind == "error": # replica_id < 0 means the failure was raised by a poller diff --git a/vllm_omni/engine/orchestrator_monitor.py b/vllm_omni/engine/orchestrator_monitor.py index 8be903a88dd..e91e8219617 100644 --- a/vllm_omni/engine/orchestrator_monitor.py +++ b/vllm_omni/engine/orchestrator_monitor.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Optional orchestrator window monitor (diagnostic only). Enable with: @@ -56,6 +59,8 @@ def register_replica(self, stage_id: int, replica_id: int) -> None: ... def note_loop(self, *, idle: bool) -> None: ... + def set_dispatch_queue_size(self, size: int) -> None: ... + def flush(self) -> None: ... @@ -66,6 +71,9 @@ def register_replica(self, stage_id: int, replica_id: int) -> None: def note_loop(self, *, idle: bool) -> None: return + def set_dispatch_queue_size(self, size: int) -> None: + return + def flush(self) -> None: return @@ -94,6 +102,10 @@ def __init__(self, *, replica_sampler: ReplicaSampler) -> None: self._duration_s: list[float] = [] self._loop_idle_buf: list[int] = [] self._loop_active_buf: list[int] = [] + self._dispatch_queue_size = 0 + self._dispatch_queue_size_buf: list[int] = [] + self._dispatch_queue_high_water = 0 + self._dispatch_queue_high_water_buf: list[int] = [] self._replicas: dict[str, _ReplicaSeries] = {} def register_replica(self, stage_id: int, replica_id: int) -> None: @@ -107,6 +119,11 @@ def note_loop(self, *, idle: bool) -> None: else: self._loop_active += 1 + def set_dispatch_queue_size(self, size: int) -> None: + size = max(int(size), 0) + self._dispatch_queue_size = size + self._dispatch_queue_high_water = max(self._dispatch_queue_high_water, size) + def flush(self) -> None: if self._flushed: return @@ -118,6 +135,8 @@ def flush(self) -> None: "duration_s": self._duration_s, "loop_idle": self._loop_idle_buf, "loop_active": self._loop_active_buf, + "dispatch_queue_size": self._dispatch_queue_size_buf, + "dispatch_queue_high_water": self._dispatch_queue_high_water_buf, }, "replicas": dict(sorted(self._replicas.items())), } @@ -154,6 +173,8 @@ def _roll_window(self, now_mono: float) -> None: self._duration_s.append(max(now_mono - self._window_start_mono, 1e-9)) self._loop_idle_buf.append(self._loop_idle) self._loop_active_buf.append(self._loop_active) + self._dispatch_queue_size_buf.append(self._dispatch_queue_size) + self._dispatch_queue_high_water_buf.append(self._dispatch_queue_high_water) for key, series in self._replicas.items(): outputs_qsize, inflight = sampled.get(key, (0, 0)) series["outputs_queue_size"].append(outputs_qsize) @@ -161,6 +182,8 @@ def _roll_window(self, now_mono: float) -> None: self._loop_idle = 0 self._loop_active = 0 + # Queued work carries into the next window even without new arrivals. + self._dispatch_queue_high_water = self._dispatch_queue_size self._window_start_mono = now_mono def _log_summary(self, output_path: str) -> None: diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index e9bf2c77ebe..6fa98f8cd8a 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -35,7 +35,7 @@ logger = init_logger(__name__) _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 # Blocking-wait interval for the event-driven final-output drain -# (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): a message wakes the drain immediately via +# (explicit env value or the engine pipeline default): a message wakes the drain immediately via # the janus queue's condition variable; this timeout only bounds how often the # orchestrator liveness check runs while the pipeline is idle. _FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 @@ -313,13 +313,15 @@ def _final_output_handler(self) -> None: engine = self.engine - # Event-driven drain (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): block on the + # Event-driven drain (explicit env value or the engine pipeline default): block on the # queue's condition variable in a dedicated thread instead of the # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side # event-driven loop (vllm_omni/engine/orchestrator.py). from vllm_omni.engine.orchestrator import _event_driven_orch_enabled - event_driven_drain = _event_driven_orch_enabled() and hasattr(engine, "get_output_blocking_async") + event_driven_drain = _event_driven_orch_enabled( + default=bool(getattr(engine, "_event_driven_orch_default", False)) + ) and hasattr(engine, "get_output_blocking_async") async def _final_output_loop(): """Background coroutine that dispatches final outputs to request queues.""" diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 1d0fcd5dcff..6ee27c71958 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -61,6 +61,13 @@ def __init__(self, *args, **kwargs): worker_name=worker_name, local_rank=self.local_rank, ) + elif profiler_config and profiler_config.profiler == "cuda": + # OmniGPUWorkerBase replaces vLLM's torch profiler above; keep the + # native CUDA profiler wrapper available for Nsight Systems capture + # through the /start_profile and /stop_profile endpoints as well. + from vllm.profiler.wrapper import CudaProfilerWrapper + + self.profiler = CudaProfilerWrapper(profiler_config) def profile(self, is_start: bool = True, profile_prefix: str | None = None): """Override to set trace filename before starting the profiler.