From 723c55020983b1c31d8d449cbf888be90c38e1b3 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 22 Sep 2026 13:12:33 +0800 Subject: [PATCH] Integrate portable edge graph backends with Omni runtime Extract versioned stage contracts, model shared CPU/iGPU/NPU memory constraints, and add explicit graph routes, reservations, output acknowledgements, and worker retirement. Fix native Windows shared-buffer lifetime and host paths. Document qualified hardware paths and remaining rollout gates. --- .gitignore | 3 + benchmarks/tts/stream_latency_bench.py | 5 + docs/design/omni_edge_backends.md | 235 ++++++++ examples/edge/README.md | 18 + examples/edge/run_graph_stage.py | 145 +++++ omni_stage_contracts/__init__.py | 27 + packages/omni-stage-contracts/MANIFEST.in | 1 + packages/omni-stage-contracts/README.md | 9 + .../omni-stage-contracts/conformance/v1.json | 24 + .../omni_stage_contracts/__init__.py | 23 + .../omni_stage_contracts/legacy.py | 87 +++ .../omni_stage_contracts/types.py | 194 +++++++ .../omni_stage_contracts/wire.py | 233 ++++++++ packages/omni-stage-contracts/pyproject.toml | 16 + pyproject.toml | 5 +- tests/edge/local/external/fake_worker.py | 25 +- tests/edge/local/external/test_protocol.py | 34 +- tests/engine/test_graph_backend.py | 526 ++++++++++++++++++ tests/windows/test_shared_buffers.py | 75 +++ tests/windows/test_windows_compat.py | 19 +- tests/worker/test_forward_kwargs_filtering.py | 16 +- vllm_omni/config/__init__.py | 2 + vllm_omni/config/omni_config.py | 44 +- vllm_omni/config/stage_config.py | 5 + .../connectors/shm_connector.py | 40 +- vllm_omni/edge/local/external/client.py | 133 ++++- vllm_omni/edge/local/external/launch.py | 51 +- vllm_omni/edge/local/external/protocol.py | 192 +------ vllm_omni/edge/local/external/worker_ort.py | 33 +- vllm_omni/edge/local/session.py | 95 +--- vllm_omni/edge/topology.py | 54 ++ vllm_omni/engine/async_omni_engine.py | 40 +- vllm_omni/engine/backends/__init__.py | 10 + vllm_omni/engine/backends/graph.py | 219 ++++++++ vllm_omni/engine/orchestrator.py | 39 +- vllm_omni/engine/resource_ledger.py | 101 ++++ vllm_omni/engine/stage_client.py | 8 + vllm_omni/engine/stage_init_utils.py | 37 +- vllm_omni/engine/stage_phase_lock.py | 3 +- vllm_omni/engine/stage_pool.py | 16 + vllm_omni/engine/stage_runtime.py | 97 +++- vllm_omni/entrypoints/async_omni.py | 5 +- vllm_omni/entrypoints/omni.py | 7 +- vllm_omni/entrypoints/stage_utils.py | 8 + vllm_omni/host/__init__.py | 2 + vllm_omni/host/paths.py | 12 + vllm_omni/host/process.py | 58 ++ vllm_omni/host/routes.py | 32 ++ vllm_omni/host/shared_memory.py | 109 ++++ vllm_omni/host/windows_process.py | 70 +++ vllm_omni/host/worker_bootstrap.py | 20 + vllm_omni/outputs/__init__.py | 8 + vllm_omni/windows/paths.py | 2 +- vllm_omni/windows/shims.py | 5 +- vllm_omni/worker/gpu_model_runner.py | 6 +- 55 files changed, 2900 insertions(+), 383 deletions(-) create mode 100644 docs/design/omni_edge_backends.md create mode 100644 examples/edge/README.md create mode 100644 examples/edge/run_graph_stage.py create mode 100644 omni_stage_contracts/__init__.py create mode 100644 packages/omni-stage-contracts/MANIFEST.in create mode 100644 packages/omni-stage-contracts/README.md create mode 100644 packages/omni-stage-contracts/conformance/v1.json create mode 100644 packages/omni-stage-contracts/omni_stage_contracts/__init__.py create mode 100644 packages/omni-stage-contracts/omni_stage_contracts/legacy.py create mode 100644 packages/omni-stage-contracts/omni_stage_contracts/types.py create mode 100644 packages/omni-stage-contracts/omni_stage_contracts/wire.py create mode 100644 packages/omni-stage-contracts/pyproject.toml create mode 100644 tests/engine/test_graph_backend.py create mode 100644 tests/windows/test_shared_buffers.py create mode 100644 vllm_omni/edge/topology.py create mode 100644 vllm_omni/engine/backends/__init__.py create mode 100644 vllm_omni/engine/backends/graph.py create mode 100644 vllm_omni/engine/resource_ledger.py create mode 100644 vllm_omni/host/__init__.py create mode 100644 vllm_omni/host/paths.py create mode 100644 vllm_omni/host/process.py create mode 100644 vllm_omni/host/routes.py create mode 100644 vllm_omni/host/shared_memory.py create mode 100644 vllm_omni/host/windows_process.py create mode 100644 vllm_omni/host/worker_bootstrap.py diff --git a/.gitignore b/.gitignore index 94dd0eb61..183ed522f 100644 --- a/.gitignore +++ b/.gitignore @@ -295,3 +295,6 @@ benchmarks/edge_harness/results/**/*.npz benchmarks/edge_harness/results/**/*.onnx benchmarks/edge_harness/results/**/*.npy benchmarks/edge_harness/results/**/*.wav + +# Portable wire fixtures are source inputs, not generated measurements. +!packages/omni-stage-contracts/conformance/*.json diff --git a/benchmarks/tts/stream_latency_bench.py b/benchmarks/tts/stream_latency_bench.py index 1cca70d34..52ce32f38 100644 --- a/benchmarks/tts/stream_latency_bench.py +++ b/benchmarks/tts/stream_latency_bench.py @@ -408,4 +408,9 @@ async def main() -> int: if __name__ == "__main__": + # Install the Windows selector policy before creating the caller's loop. + # Importing AsyncOmni inside main() is too late for that existing loop. + from vllm_omni.windows.aio import install_selector_policy + + install_selector_policy() raise SystemExit(asyncio.run(main())) diff --git a/docs/design/omni_edge_backends.md b/docs/design/omni_edge_backends.md new file mode 100644 index 000000000..43ecd9254 --- /dev/null +++ b/docs/design/omni_edge_backends.md @@ -0,0 +1,235 @@ +# Local edge backends in Omni + +The first backend is `external.graph.v1`: one complete, stateless graph operation +per request, with a fixed input bucket and copied host tensors. It uses the normal +`Omni` / `AsyncOmni` APIs, `PipelineConfig`, `StageRuntime`, `StagePool` and +orchestrator. Native vLLM stages retain their model execution, KV cache, sampling +and batching. No accelerator is mandatory. + +## Boundaries + +| Owner | Responsibility | +|---|---| +| `omni_stage_contracts` | Version/features, buffer descriptors, request/event identities, opaque state, physical device descriptors, artifact manifests. Root import uses only the standard library. | +| `vllm_omni.host` | Explicit interpreter/OS routes, process bootstrap, process identity and bounded tree retirement. | +| `vllm_omni.engine.backends` | Adapt backend operations to `StageClientBase`; observe placement; fence epochs and worker generations. | +| `StageRuntime` | Resolve all replicas, reserve the whole plan before loading, create clients, unwind failures. | +| `ResourceLedger` | Atomically charge every affected memory constraint; keep uncertain allocations quarantined. | +| `StagePool` / orchestrator | Existing stage dispatch, handoff, polling, metrics, cancellation and request cleanup. | +| Model adapter | Tokenization, media preparation, stage transforms, sampling semantics, artifact export and numerical/task validation. | + +Native Windows device-init locks use the host temporary directory, independent +of a local or UNC checkout. This coordinates same-user native processes. It does +not establish a cross-user or Windows/WSL global lock authority; cross-domain +workers in one graph plan are coordinated by that controller's reservations. + +Native Windows SHM transport also uses a host-local directory for locks and +acknowledgements. A producer retains each named mapping until its single consumer +copies the payload and acknowledges its generation, or the producer cancels/closes +it. A length header handles Windows page-rounded mapping sizes. Retained mappings +are limited to 64 MiB and 256 buffers per producer process; exceeding either limit +fails the write explicitly. This host transport fix does not enable cross-OS or +accelerator zero-copy. POSIX keeps its existing unlink-based lifetime. + +`edge.local.session` and `edge.local.external.protocol` retain compatibility imports. +The old direct `ExternalStage` API remains available; new pipelines should use +the runtime-owned graph backend. New state handles default to non-replayable and +non-migratable. Legacy M0 defaults are explicitly preserved in `contracts.legacy`. + +## Hardware composition + +The abstraction supports these single-machine compositions. This table separates +topology/admission support from successful execution of a particular model; it +does not promise that every listed accelerator can run every model. + +| Composition | Resource accounting | Current execution evidence | +|---|---|---| +| CPU only | Host RAM; WSL quota when applicable | ORT CPU graphs through public Omni on Windows and WSL. Native Windows CPU text remains unqualified. | +| CPU + integrated GPU | Shared host RAM | Radeon 890M DirectML graph through Omni. This extraction's probe is not a complete vision/language pipeline. | +| CPU + integrated NPU | Shared host RAM | AMD NPU VitisAI graph through Omni; permitted CPU partitions are explicit in the placement report. | +| CPU + integrated GPU + integrated NPU | Shared RAM plus optional observed package/power/bandwidth relationships | Topology and admission contracts tested; individual laptop routes execute. Simultaneous full-model overlap is not qualified. | +| CPU + integrated GPU + integrated NPU + discrete GPU | Shared RAM and separate discrete VRAM | Current Ryzen AI 9 HX370 / Radeon 890M / AMD NPU / RTX 5090 Laptop machine. Spark CUDA text and native two-stage TTS run on Windows and WSL; using all accelerators together is not qualified. | +| CPU + discrete GPU, with integrated devices unused or absent | Host RAM plus discrete VRAM | Native CUDA text/TTS execution uses this subset. NVIDIA is optional in the abstraction. | + +Mobile/embedded SoCs fit the CPU + integrated GPU + NPU composition, but their +local controller and model paths still require device-specific qualification. +Galaxy S25 AI Hub inference/profiling passed for one Spark attention component; +this is component evidence, not full local mobile generation or streaming. +AI Hub is a test facility, never a deployment dependency. Intel, Apple and other +vendor combinations are not certified by the AMD/NVIDIA/Qualcomm results. + +`DeviceDescriptor` separates physical device identity, execution domain and +integration. A CPU, integrated GPU and NPU can all reference `machine:ram`. +An additional discrete GPU references its own `machine:vram:`. +A WSL allocation can consume both `machine:ram` and `wsl:quota` constraints; +those are two limits on the same allocation, not two amounts of physical RAM. +Use the same physical GPU ID when Windows and WSL expose the same adapter. + +Package IDs, bandwidth groups and power domains are optional observed relationships. +Missing values mean unknown. Fitting memory does not authorize overlap on a shared +package: selection and overlap still require model quality and complete workload +measurements. Synthetic topology tests cover CPU-only and all four integrated / +discrete combinations; they are not hardware certification. + +## Deploy a graph + +Register a `PipelineConfig` with `StageExecutionType.GRAPH`, `model_stage="graph"` +and an appropriate final output type. Set `async_chunk=False`. A graph following +another stage must have a model-owned `custom_process_input_func` that returns +`{"tensors": {name: numpy_array}}`. Names, dtypes and shapes must match the +manifest's placement examples. Graph stages do not tokenize prompts or apply +model-specific transformations. + +Graph-containing v1 pipelines must be a linear chain with exactly one terminal +output at the last stage. Fan-out, joins and multiple final outputs are refused +before loading; those require reference-counted consumer leases. + +Each stage in a graph-containing plan declares the same capacities and its own +demands, including native stages: + +```yaml +pipeline: my_registered_pipeline +async_chunk: false +stages: + - stage_id: 0 + backend: + name: external.graph.v1 + manifest: /absolute/path/to/bundle/manifest.json + route: + name: local-cpu + interpreter: /absolute/path/to/worker/python + worker: /absolute/path/to/worker_ort.py + ep: cpu + os_domain: posix + max_io_bytes: 8388608 + min_fraction_on_target: 1.0 + allowed_providers: [CPUExecutionProvider] + resource_budget: + capacities: {machine:ram: 2147483648, wsl:quota: 1073741824} + demands: {machine:ram: 536870912, wsl:quota: 536870912} +``` + +Quote pool keys or use block mappings if your YAML producer treats colon-containing +keys specially. On native Windows use Windows paths and `os_domain: windows`. +From WSL use controller-visible paths; host services translate worker paths. Named +legacy routes still honor `VLLM_OMNI_EXTERNAL_PYTHON_`, but explicit routes +have no developer-machine defaults. + +Capacities are **pre-load controller ceilings**, below observed free host memory / +guest limits with safety margin. They are not continuously sampled available RAM. +Demands must include loading peaks, weights, state, activations, workspace, +transport copies, retained output and headroom. The ledger does not claim to be +an OS memory limiter or to attribute shared vendor allocations from RSS. +Driver/cache memory outside the worker remains part of the deployment estimate. + +The graph worker environment needs NumPy, ONNX Runtime and ONNX (for checking +external weight references against the manifest). Vendor providers need their +own qualified environment. The controller need not install their runtimes. + +## Artifact and placement gate + +```json +{ + "schema_version": 1, + "component": "model-owned-component-name", + "files": {"graph.onnx": "", "weights.bin": "", "inputs.npz": ""}, + "metadata": { + "graph_file": "graph.onnx", "example_inputs_file": "inputs.npz", + "checkpoint_revision": "", "precision": "", + "layout": "", "exporter": "", + "calibration": "", "validation": "" + } +} +``` + +All payload paths are relative to the manifest. Hashes are checked before startup; +ONNX external weight references must also belong to that verified set. Keep bundles +immutable during execution. Checkpoint, exporter, precision and validation metadata +are model-owned: supplying a manifest does not establish model quality. + +Warmup records actual provider node assignments. Unverified placement, absent +providers and undeclared CPU partitions refuse startup. The default target fraction +is 1.0. A deliberately partitioned NPU graph can lower it and explicitly allow CPU, +but reports always retain the real split. `execution_plan` records the manifest +identity, worker generation, route, reservations and placement. It labels this B +(backend execution), not complete-model P. + +## Flow, cancellation and ownership + +* Graph pipelines admit one request before the engine ingress queue. Extra work + receives `ResourceUnavailable`; await consumption or cancellation and retry. + The current API does not queue a batch of graph requests. +* Each graph client has one in-flight call and one retained result. Delivery does + not free the slot. Public synchronous/asynchronous generators acknowledge when + iteration advances or closes. Direct engine consumers call + `OmniRequestOutput.release_stage_buffers()` in `finally`. +* After acknowledgement, retained result arrays belong to the application. Keeping + arbitrarily many consumed results is application memory, outside engine admission. +* Intermediate graph handoffs copy their next-stage inputs before acknowledgement. + Buffer descriptors carry owner, generation, exact byte size, dtype and shape. + Framing rejects unknown required features and malformed lengths before payload allocation. +* Cancellation fences the request epoch immediately. An active non-preemptible graph + call retires its worker. It does not replay work or restart transparently; recreate + the pipeline to use that backend again. Completed calls can cancel without killing + an otherwise reusable worker. +* Windows workers announce PID plus process creation time. Cross-OS termination + verifies that identity and retires the native process tree before stopping the WSL + launcher. An absent/unverified tree keeps its reservation quarantined. Python + call references stay charged until the call unwinds. + +Native APIs and their admission logic remain compatible. This change adds a shared +explicit reservation transaction for graph-containing plans; it does not yet replace +every native memory observation/ledger or qualify mixed-model fan-out and streaming. + +## Packaging, validation and scope + +Build `packages/omni-stage-contracts` with `uv build`. Its wheel imports without +torch/vLLM; the optional `wire` extra installs NumPy. The main Omni wheel bundles +the same source. Python/native implementers can use `conformance/v1.json` from +the source distribution. A native mobile implementation has not yet run these fixtures. + +`examples/edge/run_graph_stage.py` drives the public API, writes a deployment profile, +records the resolved plan and raw request samples, and saves outputs for numerical +comparison. It supports explicit CPU, DirectML and VitisAI worker routes. + +The implementation has tests for real runtime startup, two-stage handoff on both +orchestrator modes, public APIs, ingress/consumer backpressure, malformed framing, +placement refusal, shared-pool overcommit, stale acknowledgements, cancellation, +worker failure and artifact membership. Native Spark acceptance and small actual +Windows/WSL graph probes are recorded separately in the workspace experiment report. + +Validation snapshot (2026-09-22, base `00510ce38bbb96ef2e4ce462618d06ab142fae7c`): + +* WSL focused regression: 415 passed, 1 skipped, 20 deselected. Windows focused + regression: 78 passed, 2 skipped; the controller lacks ONNX for those two tests, + while actual graph probes use their separate worker environments. +* Spark 4B BF16 CUDA M0: 11 tests passed on each host. These runs check the + specified text/state/cancellation workload, not every model or device. +* CPU and DirectML Add graph outputs are exact. AMD NPU versus the same A16W8 + CPU graph has maximum absolute error `4.458427429199219e-05`, normalized + maximum `3.111236521857185e-05` against a `0.002` bound. This is backend + validation, not original-checkpoint model-quality validation. +* Real Qwen3-TTS 0.6B two-stage streaming completes on both hosts with checkpoint + revision `85e237c12c027371202489a0ec509ded67b5e4b5`. One measured utterance per + host still has playback stalls; these diagnostic samples do not satisfy the + full streaming, interruption, quality or sustained-thermal acceptance gates. +* The built Omni wheel runs 20 public-API CPU graph requests outside the checkout + with exact outputs and zero retained reservation after shutdown. The standalone + contracts source archive imports in a dependency-free environment. +* Fresh Galaxy S25 [inference](https://workbench.aihub.qualcomm.com/jobs/jp0mvkj6g/) + and [NPU profiling](https://workbench.aihub.qualcomm.com/jobs/jpyo7nx75/) passed + for the existing W8A16 attention component. Three outputs match the same + artifact/input reference exactly; the `full1024` artifact label denotes one + attention component's cache bucket, not the complete Spark model. + +Raw records, machine/package inventories, failed attempts and reproduction commands +are preserved in the surrounding edge-infer workspace under +`analysis/experiments/omni_abstraction_20260922/`. They are not bundled as runtime +dependencies. Tests of reduced-device configurations are not measurements on a +different physical hardware SKU. + +Still outside the completed v1 scope: arbitrary device/model certification, mobile +C++ control integration, persistent external AR sessions, shared-device zero-copy, +state migration/replay, mixed-model streaming/fan-out qualification, joint thermal / +bandwidth policy, native Windows CPU text qualification, and the full M1–M4 device +matrix. These remain explicit rollout gates, not implied supported features. diff --git a/examples/edge/README.md b/examples/edge/README.md new file mode 100644 index 000000000..7eed32230 --- /dev/null +++ b/examples/edge/README.md @@ -0,0 +1,18 @@ +# Local graph example + +See [the backend contract and deployment guide](../../docs/design/omni_edge_backends.md). + +From a configured Omni environment: + +```bash +python examples/edge/run_graph_stage.py /path/to/manifest.json \ + --python /path/to/ort-venv/bin/python --report /path/to/run.json +``` + +The worker environment needs `numpy`, `onnx` and `onnxruntime`. For Windows, +use `--windows-worker` and the controller-visible path to `python.exe`. +`--ep vitisai` or `--ep dml` selects a vendor provider; declare every permitted +provider using `--allowed-provider`. No provider substitution is hidden. + +The example reads fixed-bucket inputs from the bundle. It reports backend +execution only; use your own model adapter and quality fixtures for a model claim. diff --git a/examples/edge/run_graph_stage.py b/examples/edge/run_graph_stage.py new file mode 100644 index 000000000..7e956559e --- /dev/null +++ b/examples/edge/run_graph_stage.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run a verified graph bundle through the public Omni API on one machine. + +The bundle manifest identifies graph_file and example_inputs_file. This example +uses those fixed-bucket examples as inputs; real pipelines supply model-owned +input/handoff functions. All memory arguments are controller ceilings/estimates. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import tempfile +import time +from pathlib import Path + +import numpy as np +import yaml + +from omni_stage_contracts import ArtifactManifest +from vllm_omni.config.pipeline_registry import register_pipeline +from vllm_omni.config.stage_config import PipelineConfig, StageExecutionType, StagePipelineConfig +from vllm_omni.entrypoints.omni import Omni + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("manifest", type=Path) + parser.add_argument("--route", default="ort-cpu") + parser.add_argument("--python", dest="interpreter") + parser.add_argument("--windows-worker", action="store_true") + parser.add_argument("--ep", default="cpu", choices=["cpu", "dml", "vitisai"]) + parser.add_argument("--allowed-provider", action="append", default=[]) + parser.add_argument("--min-fraction", type=float, default=1.0) + parser.add_argument("--capacity-mib", type=int, default=1024) + parser.add_argument("--stage-mib", type=int, default=512) + parser.add_argument("--max-io-mib", type=int, default=8) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--report", type=Path, required=True) + args = parser.parse_args() + if args.repeats < 1: + parser.error("--repeats must be positive") + manifest_path = args.manifest.resolve() + manifest = ArtifactManifest.read(manifest_path) + with np.load(manifest_path.parent / manifest.metadata["example_inputs_file"], allow_pickle=False) as samples: + inputs = {name: samples[name] for name in samples.files} + register_pipeline( + PipelineConfig( + model_type="omni_external_graph", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="graph", + execution_type=StageExecutionType.GRAPH, + final_output=True, + final_output_type="latent", + ), + ), + ) + ) + route = args.route + if args.interpreter: + import vllm_omni.edge.local.external.worker_ort as worker_module + + route = { + "name": args.route, + "interpreter": args.interpreter, + "worker": worker_module.__file__, + "os_domain": "windows" if args.windows_worker else "posix", + "ep": args.ep, + } + backend = { + "name": "external.graph.v1", + "route": route, + "manifest": str(manifest_path), + "min_fraction_on_target": args.min_fraction, + "max_io_bytes": args.max_io_mib << 20, + "profile_prefix": str(args.report.resolve().with_suffix(".placement")), + } + if args.allowed_provider: + backend["allowed_providers"] = args.allowed_provider + result = { + "schema_version": 1, + "evidence": "B", + "platform": platform.platform(), + "python": sys.executable, + "workload": {"repeats": args.repeats, "concurrency": 1}, + "artifact": str(manifest_path), + "samples_s": [], + "started_unix": time.time(), + } + engine = None + try: + with tempfile.TemporaryDirectory(prefix="omni-graph-") as tmp: + deploy = Path(tmp) / "deploy.yaml" + deploy.write_text( + yaml.safe_dump( + { + "pipeline": "omni_external_graph", + "async_chunk": False, + "stages": [ + { + "stage_id": 0, + "backend": backend, + "resource_budget": { + "capacities": {"machine:ram": args.capacity_mib << 20}, + "demands": {"machine:ram": args.stage_mib << 20}, + }, + } + ], + } + ) + ) + started = time.perf_counter() + engine = Omni(model=str(manifest_path.parent), deploy_config=str(deploy)) + result["startup_s"] = time.perf_counter() - started + client = engine.engine.stage_pools[0].stage_client + result["execution_plan"] = client.execution_plan + outputs = None + for _ in range(args.repeats): + started = time.perf_counter() + (output,) = engine.generate({"tensors": inputs}, use_tqdm=False) + result["samples_s"].append(time.perf_counter() - started) + outputs = output.custom_output["tensors"] + result["worker_stats"] = client._worker.stats() + result["finite"] = all(np.isfinite(v).all().item() for v in outputs.values()) + args.report.parent.mkdir(parents=True, exist_ok=True) + np.savez(args.report.with_suffix(".npz"), **outputs) + result["status"] = "passed" + except Exception as exc: + result.update(status="failed", error=f"{type(exc).__name__}: {exc}") + raise + finally: + if engine is not None: + engine.shutdown() + result["ledger_after_shutdown"] = engine.engine._runtime.resource_ledger.snapshot() + result["ended_unix"] = time.time() + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/omni_stage_contracts/__init__.py b/omni_stage_contracts/__init__.py new file mode 100644 index 000000000..3e9e046e9 --- /dev/null +++ b/omni_stage_contracts/__init__.py @@ -0,0 +1,27 @@ +"""Source-checkout bridge; wheels use the canonical package directly.""" + +from pathlib import Path + +__path__ = [str(Path(__file__).resolve().parents[1] / "packages/omni-stage-contracts/omni_stage_contracts")] + +from .types import ( + PROTOCOL_VERSION, + ArtifactManifest, + BufferRef, + DeviceDescriptor, + StageEvent, + StageRequest, + StateHandle, + negotiate, +) + +__all__ = [ + "PROTOCOL_VERSION", + "ArtifactManifest", + "BufferRef", + "DeviceDescriptor", + "StageEvent", + "StageRequest", + "StateHandle", + "negotiate", +] diff --git a/packages/omni-stage-contracts/MANIFEST.in b/packages/omni-stage-contracts/MANIFEST.in new file mode 100644 index 000000000..d1c5b37eb --- /dev/null +++ b/packages/omni-stage-contracts/MANIFEST.in @@ -0,0 +1 @@ +recursive-include conformance *.json diff --git a/packages/omni-stage-contracts/README.md b/packages/omni-stage-contracts/README.md new file mode 100644 index 000000000..a9109b35a --- /dev/null +++ b/packages/omni-stage-contracts/README.md @@ -0,0 +1,9 @@ +# Omni stage contracts + +Versioned, dependency-free data contracts for local Omni stage adapters. +The optional `wire` extra provides the NumPy copied-host tensor transport. +Importing the contracts never imports torch, vLLM or vLLM-Omni. + +Build from this directory with `uv build`; both the wheel and source archive +contain the canonical package. The parent Omni distribution bundles the same +source for compatibility. Install matching releases when using both distributions. diff --git a/packages/omni-stage-contracts/conformance/v1.json b/packages/omni-stage-contracts/conformance/v1.json new file mode 100644 index 000000000..2a458cd53 --- /dev/null +++ b/packages/omni-stage-contracts/conformance/v1.json @@ -0,0 +1,24 @@ +{ + "protocol_version": 1, + "required_features": ["host-copy", "request-epochs"], + "buffer": { + "object_id": "output-1", "owner": "stage-0", "generation": "worker-A", + "dtype": "float32", "shape": [2, 3], "nbytes": 24, + "memory_domain": "host-copy", "layout": "C", "ready": true + }, + "request": { + "request_id": "request-1", "stage_id": 0, "epoch": 2, + "worker_generation": "worker-A", "operation": "run", "inputs": [] + }, + "event": { + "request_id": "request-1", "stage_id": 0, "epoch": 2, "seq": 1, + "kind": "tensor", "worker_generation": "worker-A", "buffers": [], "terminal": true + }, + "state": { + "session_id": "session-1", "backend": "backend-1", "artifact_id": "sha256:fixture", + "backend_instance_id": "backend-instance-A", "worker_generation": "worker-A", + "layout_version": 1, "epoch": 2, "replayable": false, "migratable": false + }, + "reject_buffer_overrides": [{"nbytes": 23}, {"shape": [-1, 3]}, {"memory_domain": "cuda-ipc"}], + "stale_state_overrides": [{"epoch": 1}, {"worker_generation": "worker-B"}, {"backend_instance_id": "backend-instance-B"}] +} diff --git a/packages/omni-stage-contracts/omni_stage_contracts/__init__.py b/packages/omni-stage-contracts/omni_stage_contracts/__init__.py new file mode 100644 index 000000000..7037d65e2 --- /dev/null +++ b/packages/omni-stage-contracts/omni_stage_contracts/__init__.py @@ -0,0 +1,23 @@ +"""Portable Omni data contracts. Importing this module needs only the stdlib.""" + +from .types import ( + PROTOCOL_VERSION, + ArtifactManifest, + BufferRef, + DeviceDescriptor, + StageEvent, + StageRequest, + StateHandle, + negotiate, +) + +__all__ = [ + "PROTOCOL_VERSION", + "ArtifactManifest", + "BufferRef", + "DeviceDescriptor", + "StageEvent", + "StageRequest", + "StateHandle", + "negotiate", +] diff --git a/packages/omni-stage-contracts/omni_stage_contracts/legacy.py b/packages/omni-stage-contracts/omni_stage_contracts/legacy.py new file mode 100644 index 000000000..db75ef3e9 --- /dev/null +++ b/packages/omni-stage-contracts/omni_stage_contracts/legacy.py @@ -0,0 +1,87 @@ +"""Compatible M0 contracts; new backend state uses types.StateHandle.""" + +from __future__ import annotations + +import time +from dataclasses import asdict, dataclass, field +from typing import Any + +STATE_LAYOUT_VERSION = 1 + + +@dataclass(frozen=True) +class StateHandle: + """A reference to session state the backend owns. + + The engine never reads the state. It reads this, to decide whether a + request may continue a session or must start a new one. + """ + + session_id: str + backend: str + artifact_id: str + """Ties the state to the checkpoint that produced it. The same shapes from + a different quantization are not the same state.""" + layout_version: int = STATE_LAYOUT_VERSION + epoch: int = 0 + """Incremented on every cancel. Events from a lower epoch are stale.""" + replayable: bool = True + """A text session can be rebuilt by re-submitting its committed prompt, so + recovery after a backend crash is possible at an output boundary.""" + migratable: bool = False + """Never true in M0. Moving KV between backends needs a measured layout + conversion for the specific pair; assuming it is the failure the proposal + warns about.""" + created_unix: float = field(default_factory=time.time) + + def next_epoch(self) -> StateHandle: + from dataclasses import replace + + return replace(self, epoch=self.epoch + 1) + + backend_instance_id: str = "" + worker_generation: str = "" + + def accepts(self, other: StateHandle) -> bool: + """Whether ``other`` may continue this session's state.""" + return ( + other.backend == self.backend + and other.backend_instance_id == self.backend_instance_id + and other.worker_generation == self.worker_generation + and other.session_id == self.session_id + and other.artifact_id == self.artifact_id + and other.layout_version == self.layout_version + and other.epoch == self.epoch + ) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ChunkEvent: + """One unit of output, with everything needed to order and retire it.""" + + request_id: str + stage_id: int + seq: int + """Monotonic within (request, epoch). Gaps mean loss; repeats mean a bug.""" + epoch: int + kind: str + """``token`` | ``done`` | ``error`` | ``cancelled``.""" + payload: Any + """For ``token``: the incremental text and token ids. Never the whole + output so far -- a cumulative payload makes a dropped chunk invisible.""" + started_unix: float + emitted_unix: float + final: bool = False + error: str | None = None + input_watermark: int = 0 + """How much of the input the producer had consumed when this was emitted. + One number in M0 (prompt tokens); it is what a duplex turn will align on.""" + release_token: str = "" + """Returned by the consumer to give credit back. Empty when the event + carries no resource the producer is holding on its behalf.""" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) diff --git a/packages/omni-stage-contracts/omni_stage_contracts/types.py b/packages/omni-stage-contracts/omni_stage_contracts/types.py new file mode 100644 index 000000000..595f7f98e --- /dev/null +++ b/packages/omni-stage-contracts/omni_stage_contracts/types.py @@ -0,0 +1,194 @@ +"""Versioned host-copy contracts shared by Python and native stage adapters.""" + +from __future__ import annotations + +import hashlib +import json +import math +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +PROTOCOL_VERSION = 1 +FEATURES = frozenset({"host-copy", "request-epochs"}) +_ITEM_SIZES = { + "bool": 1, + "uint8": 1, + "int8": 1, + "int16": 2, + "uint16": 2, + "int32": 4, + "uint32": 4, + "int64": 8, + "uint64": 8, + "float16": 2, + "float32": 4, + "float64": 8, +} + + +def negotiate(version: int, required: list[str] | tuple[str, ...] = ()) -> None: + if type(version) is not int or version != PROTOCOL_VERSION: + raise ValueError(f"unsupported stage protocol version: {version!r}") + missing = set(required) - FEATURES + if missing: + raise ValueError(f"unsupported required stage features: {sorted(missing)}") + + +@dataclass(frozen=True) +class BufferRef: + object_id: str + owner: str + generation: str + dtype: str + shape: tuple[int, ...] + nbytes: int + memory_domain: str = "host-copy" + layout: str = "C" + ready: bool = True + + def __post_init__(self) -> None: + if not self.object_id or not self.owner or not self.generation: + raise ValueError("buffer identity/owner/generation are required") + if self.memory_domain != "host-copy" or self.layout != "C" or not self.ready: + raise ValueError("v1 only supports ready, contiguous copied host buffers") + if self.dtype not in _ITEM_SIZES: + raise ValueError(f"unsupported dtype: {self.dtype}") + if len(self.shape) > 32 or any(type(x) is not int or x < 0 for x in self.shape): + raise ValueError("invalid tensor dimensions") + expected = math.prod(self.shape) * _ITEM_SIZES[self.dtype] + if type(self.nbytes) is not int or self.nbytes != expected or self.nbytes > 2 << 30: + raise ValueError("tensor shape/dtype/byte length mismatch or payload exceeds limit") + + +@dataclass(frozen=True) +class StateHandle: + session_id: str + backend: str + artifact_id: str + layout_version: int = 1 + epoch: int = 0 + replayable: bool = False + migratable: bool = False + backend_instance_id: str = "" + worker_generation: str = "" + + def accepts(self, other: StateHandle) -> bool: + return all( + getattr(self, k) == getattr(other, k) + for k in ( + "session_id", + "backend", + "artifact_id", + "layout_version", + "epoch", + "backend_instance_id", + "worker_generation", + ) + ) + + def next_epoch(self) -> StateHandle: + from dataclasses import replace + + return replace(self, epoch=self.epoch + 1) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class StageRequest: + request_id: str + stage_id: int + epoch: int = 0 + worker_generation: str = "" + operation: str = "run" + inputs: tuple[BufferRef, ...] = () + + def __post_init__(self): + if not self.request_id or not self.worker_generation or not self.operation: + raise ValueError("request identity, worker generation and operation are required") + if any(type(n) is not int or n < 0 for n in (self.stage_id, self.epoch)): + raise ValueError("stage_id and epoch must be nonnegative integers") + + +@dataclass(frozen=True) +class StageEvent: + request_id: str + stage_id: int + epoch: int + seq: int + kind: str + worker_generation: str + buffers: tuple[BufferRef, ...] = () + terminal: bool = False + error: str | None = None + + def __post_init__(self): + if not self.request_id or not self.worker_generation or not self.kind: + raise ValueError("event identity, worker generation and kind are required") + if any(type(n) is not int or n < 0 for n in (self.stage_id, self.epoch, self.seq)): + raise ValueError("stage_id, epoch and seq must be nonnegative integers") + + +@dataclass(frozen=True) +class DeviceDescriptor: + """Physical identity is independent from per-domain execution routes. + + Unknown package/power/bandwidth relationships are empty, never inferred + from an integrated/discrete label. Pool IDs identify backing constraints. + """ + + device_id: str + kind: str + domain_id: str + memory_pool_ids: tuple[str, ...] + integration: str = "unknown" + package_id: str | None = None + bandwidth_group_ids: tuple[str, ...] = () + power_domain_ids: tuple[str, ...] = () + + def __post_init__(self) -> None: + if self.kind not in {"cpu", "gpu", "npu"}: + raise ValueError(f"unknown device kind: {self.kind}") + if self.integration not in {"integrated", "discrete", "unknown"}: + raise ValueError("unknown integration class") + if not self.device_id or not self.domain_id or not self.memory_pool_ids: + raise ValueError("device, domain and physical memory constraints are required") + if len(set(self.memory_pool_ids)) != len(self.memory_pool_ids): + raise ValueError("duplicate physical memory constraint") + + +@dataclass(frozen=True) +class ArtifactManifest: + """A complete artifact payload set; paths are relative to the manifest.""" + + schema_version: int + component: str + files: dict[str, str] + metadata: dict[str, Any] = field(default_factory=dict) + + @classmethod + def read(cls, path: str | Path) -> ArtifactManifest: + path = Path(path) + manifest = cls(**json.loads(path.read_text(encoding="utf-8"))) + if manifest.schema_version != 1 or not manifest.files: + raise ValueError("unsupported or empty artifact manifest") + root = path.parent.resolve() + for name, digest in manifest.files.items(): + if not isinstance(name, str) or "\\" in name or ":" in name: + raise ValueError("artifact payloads must use portable relative POSIX paths") + target = (root / name).resolve() + if not target.is_relative_to(root): + raise ValueError(f"artifact path leaves bundle: {name}") + if not target.is_file() or file_digest(target) != digest: + raise ValueError(f"artifact payload missing or digest mismatch: {name}") + return manifest + + +def file_digest(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1 << 20), b""): + h.update(block) + return h.hexdigest() diff --git a/packages/omni-stage-contracts/omni_stage_contracts/wire.py b/packages/omni-stage-contracts/omni_stage_contracts/wire.py new file mode 100644 index 000000000..efcfe274b --- /dev/null +++ b/packages/omni-stage-contracts/omni_stage_contracts/wire.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The wire between the engine and a stage that runs in another process. + +Deliberately tiny, and deliberately **stdlib + numpy only**. The worker on the +other end of this socket may be a native-Windows interpreter that has no torch, +no vllm and no vllm_omni on its path -- ``C:\\Users\\zhout\\npu-ep`` is exactly +that: Python 3.12, onnxruntime, onnx, numpy, nothing else. Anything imported +here has to be installable there, so this module imports neither. + +**Why a socket at all.** Two devices on this laptop cannot be reached from the +engine's own interpreter, and for different reasons: + +* the **XDNA2 NPU** is an MCDM device. WSL2's GPU-PV forwards WDDM display + adapters only, so there is no ``/dev/accel`` inside WSL and no amount of + installing changes that. VitisAI is a Windows DLL. +* the **Radeon 890M** is reachable from WSL through DirectML, but + ``torch-directml`` pins torch 2.4.1 and Omni runs 2.13. + +So out-of-process is forced, not preferred. One protocol covers both, plus the +native-Windows DirectML route, which needs its own venv anyway because +``onnxruntime-directml`` is pinned at 1.24.4 and cannot load the VitisAI EP +(that needs ORT >= 1.25) -- the two AMD devices cannot share an interpreter. + +**Framing.** ``[4-byte big-endian header length][JSON header][raw tensor bytes]``. +Tensors are described in the header (name, dtype, shape, nbytes) and their +buffers are concatenated after it in the order listed, so a receive is two +reads and one ``frombuffer`` per tensor -- no per-element work and no pickle. +Measured on this machine, WSL <-> native Windows: 11 us half round trip, +330 MB/s. That is ample for a stage boundary (a 2 s Code2Wav chunk is ~192 KB, +about 0.6 ms) and nowhere near enough for per-layer traffic inside a decode +loop, which is the arithmetic behind "coarse stages only". + +``pickle`` is not used anywhere in here on purpose: the header is JSON, the +payload is raw buffers, so a hostile or merely broken peer cannot execute code +by replying. +""" + +from __future__ import annotations + +import json +import math +import os +import socket +import struct +from typing import Any + +import numpy as np + +from .types import negotiate + +# Ops the client sends. +OP_LOAD = "load" +"""Open a session on a graph. Replies with the placement report.""" +OP_RUN = "run" +"""One forward. Inputs in, outputs out.""" +OP_STATS = "stats" +"""Counters and RSS. The worker's memory is invisible to the engine's sampler +when the worker is a Windows process, so it has to report its own.""" +OP_CLOSE = "close" + +# Ops the worker sends. +OP_HELLO = "hello" +"""First message after connect-back: interpreter, versions, available EPs.""" +OP_OK = "ok" +OP_ERR = "err" + +MAX_HEADER_BYTES = 1 << 20 +"""A header is metadata; a megabyte of it is a bug or an attack, not a graph.""" + +MAX_PAYLOAD_BYTES = 2 << 30 +"""Refuse rather than allocate on a corrupt length. Two gibibytes is far above +any legitimate stage payload here -- the largest is a vision tower's patch +tensor, single-digit MB.""" + + +class ProtocolError(RuntimeError): + """The peer sent something this protocol cannot represent.""" + + +class WorkerError(RuntimeError): + """The peer ran and failed. Carries the remote traceback when there is one.""" + + def __init__(self, message: str, *, remote_traceback: str = "", code: str = "") -> None: + # The remote traceback goes in the message, not only in an attribute: + # a bare "AssertionError:" surfacing here with the actual stack one + # attribute away is a debugging loop nobody should have to discover. + detail = f"{message}\n--- worker traceback ---\n{remote_traceback}" if remote_traceback else message + super().__init__(detail) + self.message = message + self.remote_traceback = remote_traceback + self.code = code + + +def _recv_exactly(sock: socket.socket, n: int) -> bytes: + """Read exactly ``n`` bytes or raise. + + ``recv`` is allowed to return short reads on a stream socket and does, once + payloads get past a page or two; a partial read here would be silently + misparsed as the next tensor. + """ + if n == 0: + return b"" + chunks: list[bytes] = [] + remaining = n + while remaining: + chunk = sock.recv(min(remaining, 1 << 20)) + if not chunk: + raise ProtocolError(f"peer closed with {remaining} of {n} bytes outstanding") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) if len(chunks) > 1 else chunks[0] + + +def send_message( + sock: socket.socket, + op: str, + body: dict[str, Any] | None = None, + tensors: dict[str, np.ndarray] | None = None, +) -> None: + """Send one message. ``tensors`` are sent as raw buffers after the header.""" + tensors = tensors or {} + described: list[dict[str, Any]] = [] + buffers: list[memoryview] = [] + for name, array in tensors.items(): + array = np.asarray(array) + if ( + not isinstance(name, str) + or not name + or array.dtype.kind not in "biuf" + or array.dtype.itemsize not in (1, 2, 4, 8) + ): + raise ProtocolError("unsupported tensor name or dtype") + if sum(b.nbytes for b in buffers) + array.nbytes > MAX_PAYLOAD_BYTES: + raise ProtocolError("outgoing payload exceeds limit") + # Record the shape *before* making it contiguous: ``ascontiguousarray`` + # promotes a 0-d array to shape (1,), which would silently turn a + # scalar graph input into a length-1 vector on the far side. + shape = list(array.shape) + contiguous = np.ascontiguousarray(array) + described.append( + { + "name": name, + # ``str(dtype)`` loses endianness for multibyte types; ``.str`` + # keeps the '<'/'>' so a big-endian peer would still decode. + "dtype": contiguous.dtype.str, + "shape": shape, + "nbytes": int(contiguous.nbytes), + } + ) + # ``reshape(-1)`` rather than a direct cast: ``memoryview.cast`` refuses + # a 0-dimensional source. + buffers.append(memoryview(contiguous.reshape(-1)).cast("B")) + + if op == OP_HELLO and os.environ.get("VLLM_OMNI_WORKER_IDENTITY"): + body = {**(body or {}), "process_identity": json.loads(os.environ["VLLM_OMNI_WORKER_IDENTITY"])} + header = json.dumps( + {"version": 1, "required_features": ["host-copy"], "op": op, "body": body or {}, "tensors": described} + ).encode("utf-8") + if len(header) > MAX_HEADER_BYTES: + raise ProtocolError(f"header is {len(header)} bytes, over the {MAX_HEADER_BYTES} limit") + sock.sendall(struct.pack(">I", len(header))) + sock.sendall(header) + for buffer in buffers: + sock.sendall(buffer) + + +def recv_message( + sock: socket.socket, *, max_payload_bytes: int = MAX_PAYLOAD_BYTES +) -> tuple[str, dict[str, Any], dict[str, np.ndarray]]: + """Receive one message as ``(op, body, tensors)``. + + Raises :class:`WorkerError` on an ``err`` reply so callers do not each have + to remember to check; every other op is returned as-is. + """ + (header_len,) = struct.unpack(">I", _recv_exactly(sock, 4)) + if header_len > MAX_HEADER_BYTES: + raise ProtocolError(f"peer announced a {header_len}-byte header, over the limit") + try: + header = json.loads(_recv_exactly(sock, header_len).decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ProtocolError(f"header is not JSON: {exc}") from exc + + if not isinstance(header, dict): + raise ProtocolError("invalid message header") + try: + required = header.get("required_features", []) + if not isinstance(required, list) or any(not isinstance(f, str) for f in required): + raise ValueError("required_features must be a list of names") + # Missing version/features is the explicit legacy-v1 conversion. + negotiate(header.get("version", 1), required) + except ValueError as exc: + raise ProtocolError(str(exc)) from exc + if not isinstance(header.get("body", {}), dict) or not isinstance(header.get("tensors", []), list): + raise ProtocolError("invalid message schema") + total = 0 + names = set() + try: + for desc in header.get("tensors", ()): + dtype = np.dtype(desc["dtype"]) + shape, size, name = desc["shape"], desc["nbytes"], desc["name"] + if dtype.hasobject or dtype.kind not in "biuf" or dtype.itemsize not in (1, 2, 4, 8): + raise ValueError("unsupported tensor dtype") + if not isinstance(shape, list) or len(shape) > 32 or any(type(x) is not int or x < 0 for x in shape): + raise ValueError("invalid tensor dimensions") + if type(size) is not int or size != math.prod(shape) * dtype.itemsize: + raise ValueError("tensor byte length disagrees with dtype/shape") + if not isinstance(name, str) or not name or name in names: + raise ValueError("invalid or duplicate tensor name") + names.add(name) + total += size + except (KeyError, TypeError, ValueError) as exc: + raise ProtocolError(f"invalid tensor descriptor: {exc}") from exc + if total > min(max_payload_bytes, MAX_PAYLOAD_BYTES): + raise ProtocolError(f"peer announced a {total}-byte payload, over the limit") + + tensors: dict[str, np.ndarray] = {} + for spec in header.get("tensors", ()): + raw = _recv_exactly(sock, int(spec["nbytes"])) + array = np.frombuffer(raw, dtype=np.dtype(spec["dtype"])) + # ``frombuffer`` is read-only and aliases ``raw``; a consumer that + # writes in place (ORT's IOBinding does) needs its own buffer. + tensors[spec["name"]] = array.reshape(spec["shape"]).copy() + + op = str(header.get("op", "")) + body = header.get("body") or {} + if op == OP_ERR: + raise WorkerError( + str(body.get("message", "worker failed with no message")), + remote_traceback=str(body.get("traceback", "")), + code=str(body.get("code", "")), + ) + return op, body, tensors diff --git a/packages/omni-stage-contracts/pyproject.toml b/packages/omni-stage-contracts/pyproject.toml new file mode 100644 index 000000000..04066a97f --- /dev/null +++ b/packages/omni-stage-contracts/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "omni-stage-contracts" +version = "0.1.0" +description = "Dependency-free stage contracts from the vLLM-Omni repository" +requires-python = ">=3.10" +readme = "README.md" + +[project.optional-dependencies] +wire = ["numpy>=1.24"] + +[tool.setuptools] +packages = ["omni_stage_contracts"] diff --git a/pyproject.toml b/pyproject.toml index 922ac04ac..f0b71a7d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,9 +168,12 @@ vllm_omni_windows = "vllm_omni.windows:activate" vllm_omni_windows = "vllm_omni.windows:activate" +[tool.setuptools.package-dir] +omni_stage_contracts = "packages/omni-stage-contracts/omni_stage_contracts" + [tool.setuptools.packages.find] where = ["."] -include = ["vllm_omni*"] +include = ["vllm_omni*", "omni_stage_contracts*"] [tool.setuptools.package-data] "vllm_omni" = ["_version.py", "py.typed", "transformers_utils/chat_templates/*.jinja"] diff --git a/tests/edge/local/external/fake_worker.py b/tests/edge/local/external/fake_worker.py index 7d422ac1f..65dba636c 100644 --- a/tests/edge/local/external/fake_worker.py +++ b/tests/edge/local/external/fake_worker.py @@ -25,9 +25,7 @@ from pathlib import Path import numpy as np _HERE = Path(__file__).resolve() -_PROTOCOL = ( - _HERE.parents[4] / "vllm_omni" / "edge" / "local" / "external" / "protocol.py" -) +_PROTOCOL = _HERE.parents[4] / "vllm_omni" / "edge" / "local" / "external" / "protocol.py" spec = importlib.util.spec_from_file_location("_fake_protocol", _PROTOCOL) proto = importlib.util.module_from_spec(spec) @@ -78,13 +76,15 @@ def main() -> int: "session_create_s": 0.01, "warmup_s": 0.001, "rss_bytes": 111 * 2**20, - "inputs": [{"name": n, "type": "tensor(float)", "shape": list(t.shape)} - for n, t in tensors.items()], + "inputs": [{"name": n, "type": "tensor(float)", "shape": list(t.shape)} for n, t in tensors.items()], "outputs": [{"name": "y", "type": "tensor(float)", "shape": []}], } reply.update(config.get("load", {})) proto.send_message(sock, proto.OP_OK, reply) elif op == proto.OP_RUN: + import time + + time.sleep(float(config.get("run", {}).get("sleep_s", 0))) scale = float(config.get("run", {}).get("scale", 2.0)) proto.send_message( sock, @@ -94,16 +94,21 @@ def main() -> int: ) elif op == proto.OP_STATS: proto.send_message( - sock, proto.OP_OK, - {"runs": 1, "total_run_s": 0.001, "rss_bytes": 111 * 2**20, - "peak_rss_bytes": 122 * 2**20, "report": {}}, + sock, + proto.OP_OK, + { + "runs": 1, + "total_run_s": 0.001, + "rss_bytes": 111 * 2**20, + "peak_rss_bytes": 122 * 2**20, + "report": {}, + }, ) elif op == proto.OP_CLOSE: proto.send_message(sock, proto.OP_OK, {}) return 0 else: - proto.send_message(sock, proto.OP_ERR, - {"message": f"unknown op {op}", "code": "ValueError"}) + proto.send_message(sock, proto.OP_ERR, {"message": f"unknown op {op}", "code": "ValueError"}) if __name__ == "__main__": diff --git a/tests/edge/local/external/test_protocol.py b/tests/edge/local/external/test_protocol.py index 6dbe37445..e39600a4c 100644 --- a/tests/edge/local/external/test_protocol.py +++ b/tests/edge/local/external/test_protocol.py @@ -16,9 +16,7 @@ from vllm_omni.edge.local.external import protocol def _round_trip(op, body=None, tensors=None): near, far = socket.socketpair() try: - threading.Thread( - target=lambda: protocol.send_message(near, op, body, tensors), daemon=True - ).start() + threading.Thread(target=lambda: protocol.send_message(near, op, body, tensors), daemon=True).start() return protocol.recv_message(far) finally: near.close() @@ -66,8 +64,7 @@ def test_empty_message_and_no_tensors(): def test_error_reply_raises_with_remote_detail(): with pytest.raises(protocol.WorkerError) as excinfo: - _round_trip(protocol.OP_ERR, {"message": "boom", "code": "ValueError", - "traceback": "far side traceback"}) + _round_trip(protocol.OP_ERR, {"message": "boom", "code": "ValueError", "traceback": "far side traceback"}) assert excinfo.value.code == "ValueError" assert "far side traceback" in excinfo.value.remote_traceback @@ -92,3 +89,30 @@ def test_oversized_header_is_refused_before_allocation(): finally: near.close() far.close() + + +@pytest.mark.parametrize( + "override", + [ + {"version": 2}, + {"version": True}, + {"required_features": ["unknown-required-feature"]}, + {"tensors": [{"name": "x", "dtype": "I", len(header)) + header) + with pytest.raises(protocol.ProtocolError): + protocol.recv_message(far) + finally: + near.close() + far.close() diff --git a/tests/engine/test_graph_backend.py b/tests/engine/test_graph_backend.py new file mode 100644 index 000000000..d9cfc6c01 --- /dev/null +++ b/tests/engine/test_graph_backend.py @@ -0,0 +1,526 @@ +"""The real StageRuntime/StagePool with a subprocess graph worker.""" + +import asyncio +import hashlib +import json +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from omni_stage_contracts import BufferRef, StateHandle, negotiate +from vllm_omni.config.stage_config import ( + DeployConfig, + PipelineConfig, + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, + merge_pipeline_deploy, +) +from vllm_omni.engine.resource_ledger import ResourceLedger, ResourceUnavailable +from vllm_omni.engine.stage_runtime import StageRuntime + +ROOT = Path(__file__).resolve().parents[2] + + +def artifact(tmp_path): + (tmp_path / "graph.onnx").write_bytes(b"synthetic graph, not model evidence") + np.savez(tmp_path / "inputs.npz", x=np.ones((2, 3), dtype=np.float32)) + files = {n: hashlib.sha256((tmp_path / n).read_bytes()).hexdigest() for n in ("graph.onnx", "inputs.npz")} + path = tmp_path / "manifest.json" + path.write_text( + json.dumps( + { + "schema_version": 1, + "component": "test", + "files": files, + "metadata": {"graph_file": "graph.onnx", "example_inputs_file": "inputs.npz"}, + } + ) + ) + return path + + +@pytest.fixture +def graph_runtime(tmp_path, monkeypatch): + config_path = tmp_path / "worker.json" + config_path.write_text("{}") + monkeypatch.setenv("VLLM_OMNI_FAKE_WORKER_JSON", str(config_path)) + monkeypatch.setenv("VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU", str(ROOT / "tests/edge/local/external/fake_worker.py")) + monkeypatch.setenv("VLLM_OMNI_EXTERNAL_PYTHON_ORT_CPU", sys.executable) + runtimes = [] + + def create( + *, stages=1, capacity=1024 << 20, worker_config=None, start=True, chain=False, real_ort=False, entrypoint=None + ): + config_path.write_text(json.dumps(worker_config or {})) + manifest = artifact(tmp_path) + if real_ort: + onnx = pytest.importorskip("onnx") + pytest.importorskip("onnxruntime") + from onnx import TensorProto, helper + + graph = helper.make_graph( + [helper.make_node("Add", ["x", "x"], ["x_out"])], + "double", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("x_out", TensorProto.FLOAT, [2, 3])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], ir_version=9) + onnx.save(model, tmp_path / "graph.onnx") + data = json.loads(manifest.read_text()) + data["files"]["graph.onnx"] = hashlib.sha256((tmp_path / "graph.onnx").read_bytes()).hexdigest() + manifest.write_text(json.dumps(data)) + monkeypatch.delenv("VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU", raising=False) + pipeline = PipelineConfig( + model_type="test_graph", + stages=tuple( + StagePipelineConfig( + stage_id=i, + model_stage="graph", + execution_type=StageExecutionType.GRAPH, + final_output=(not chain or i == stages - 1), + final_output_type="latent", + input_sources=(i - 1,) if chain and i else (), + custom_process_input_func=( + "tests.engine.test_graph_backend.graph_handoff" if chain and i else None + ), + ) + for i in range(stages) + ), + ) + deploy = DeployConfig( + async_chunk=False, + stages=[ + StageDeployConfig( + stage_id=i, + backend={ + "name": "external.graph.v1", + "route": "ort-cpu", + "manifest": str(manifest), + "start_timeout_s": 5, + }, + resource_budget={"capacities": {"host_ram": capacity}, "demands": {"host_ram": 512 << 20}}, + ) + for i in range(stages) + ], + ) + if entrypoint is not None: + import yaml + + from vllm_omni.config.pipeline_registry import OMNI_PIPELINES + + monkeypatch.setitem(OMNI_PIPELINES, pipeline.model_type, pipeline) + (tmp_path / "config.json").write_text('{"model_type":"bert"}') + deploy_path = tmp_path / "deploy.yaml" + deploy_path.write_text( + yaml.safe_dump( + { + "pipeline": pipeline.model_type, + "async_chunk": False, + "stages": [ + {"stage_id": s.stage_id, "backend": s.backend, "resource_budget": s.resource_budget} + for s in deploy.stages + ], + } + ) + ) + engine = entrypoint( + model=str(tmp_path), deploy_config=str(deploy_path), stage_init_timeout=10, init_timeout=30 + ) + runtimes.append(engine) + return engine + configs = [s.to_omegaconf() for s in merge_pipeline_deploy(pipeline, deploy)] + runtime = StageRuntime(configs, "local-graph", "", stage_init_timeout=5, async_chunk=False) + runtimes.append(runtime) + if start: + runtime.initialize() + return runtime + + yield create + for runtime in runtimes: + runtime.shutdown() + + +def test_contract_import_does_not_import_engine(): + code = "import omni_stage_contracts; import sys; assert not ({'torch','vllm','vllm_omni'} & set(sys.modules))" + subprocess.run([sys.executable, "-S", "-c", code], cwd=ROOT, check=True) + + +def test_device_init_lock_is_host_local_and_preserves_exclusion(tmp_path, monkeypatch): + import os + + from vllm_omni import host + from vllm_omni._filelock_compat import flock_exclusive_nb, funlock + from vllm_omni.engine.stage_init_utils import device_init_lock_path, open_device_lock_file + from vllm_omni.host import paths + + assert host is not None + monkeypatch.setattr(paths, "device_lock_directory", lambda: str(tmp_path)) + lock_path = device_init_lock_path(0) + assert Path(lock_path).parent == tmp_path + first, writable = open_device_lock_file(lock_path) + second = None + try: + assert writable + flock_exclusive_nb(first) + second, _ = open_device_lock_file(lock_path) + with pytest.raises(BlockingIOError): + flock_exclusive_nb(second) + funlock(first) + flock_exclusive_nb(second) + funlock(second) + finally: + os.close(first) + if second is not None: + os.close(second) + + +def test_graph_metadata_matches_structured_and_legacy_config(): + from vllm_omni.config import VllmOmniConfig, VllmOmniGraphStageConfig + from vllm_omni.engine.stage_init_utils import ( + extract_legacy_stage_metadata, + extract_stage_metadata_from_omni_stage_config, + ) + + pipeline = PipelineConfig( + model_type="test_graph", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="graph", + execution_type=StageExecutionType.GRAPH, + final_output=True, + final_output_type="latent", + ), + ), + ) + deploy = DeployConfig( + async_chunk=False, + stages=[ + StageDeployConfig( + stage_id=0, + backend={"name": "external.graph.v1"}, + resource_budget={"capacities": {"ram": 1024}, "demands": {"ram": 512}}, + ) + ], + ) + config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0) + assert isinstance(config, VllmOmniGraphStageConfig) + structured = extract_stage_metadata_from_omni_stage_config(config) + legacy = extract_legacy_stage_metadata(merge_pipeline_deploy(pipeline, deploy)[0].to_omegaconf()) + assert structured.stage_type == legacy.stage_type == "graph" + assert type(structured.default_sampling_params) is type(legacy.default_sampling_params) + + +def test_language_independent_contract_fixtures(): + from omni_stage_contracts import StageEvent, StageRequest + + fixture = json.loads((ROOT / "packages/omni-stage-contracts/conformance/v1.json").read_text()) + negotiate(fixture["protocol_version"], fixture["required_features"]) + BufferRef(**fixture["buffer"]) + StageRequest(**fixture["request"]) + StageEvent(**fixture["event"]) + state = StateHandle(**fixture["state"]) + for override in fixture["reject_buffer_overrides"]: + with pytest.raises(ValueError): + BufferRef(**{**fixture["buffer"], **override}) + for override in fixture["stale_state_overrides"]: + assert not state.accepts(StateHandle(**{**fixture["state"], **override})) + + +def test_external_weight_payload_must_be_in_manifest(tmp_path): + onnx = pytest.importorskip("onnx") + from onnx import helper, numpy_helper + + from vllm_omni.edge.local.external.worker_ort import _verify_artifact_members + + weight = numpy_helper.from_array(np.ones((32, 32), dtype=np.float32), "w") + graph = helper.make_graph([helper.make_node("Add", ["x", "w"], ["y"])], "external", [], [], [weight]) + path = tmp_path / "model.onnx" + onnx.save_model( + helper.make_model(graph), + path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="weights.bin", + size_threshold=0, + ) + with pytest.raises(ValueError, match="absent from verified manifest"): + _verify_artifact_members(str(path), [str(path)]) + _verify_artifact_members(str(path), [str(path), str(tmp_path / "weights.bin")]) + + +def test_contract_rejects_shape_and_generation_mismatch(): + with pytest.raises(ValueError): + BufferRef("x", "a", "1", "float32", (2, 3), 23) + with pytest.raises(ValueError): + negotiate(2) + with pytest.raises(ValueError): + negotiate(1, ["zero-copy"]) + a = StateHandle("s", "ort", "sha", worker_generation="old") + b = StateHandle("s", "ort", "sha", worker_generation="new") + assert not a.accepts(b) + + +def test_atomic_shared_constraints_and_quarantine(): + ledger = ResourceLedger({"ram": 100, "wsl": 60}) + first = ledger.reserve("wsl", {"ram": 50, "wsl": 50}) + with pytest.raises(ResourceUnavailable): + ledger.reserve("igpu+npu", {"ram": 60}) + with pytest.raises(ResourceUnavailable): + ledger.reserve("second-wsl", {"ram": 10, "wsl": 20}) + assert ledger.snapshot()["reserved"] == {"ram": 50, "wsl": 50} + assert not ledger.release(first, drained=False) + assert ledger.snapshot()["quarantined"] == ["wsl"] + assert ledger.release(first, drained=True) + second = ledger.reserve("wsl", {"ram": 5}) + assert not ledger.release(first, drained=True) + assert ledger.snapshot()["reserved"]["ram"] == 5 + ledger.release(second, drained=True) + + +@pytest.mark.parametrize( + "kinds", + [ + ("cpu",), + ("cpu", "gpu_integrated"), + ("cpu", "npu"), + ("cpu", "gpu_integrated", "npu"), + ("cpu", "gpu_integrated", "npu", "gpu_discrete"), + ], +) +def test_integrated_topology_has_one_physical_ram(kinds): + from vllm_omni.edge.topology import describe_topology + + devices = [ + SimpleNamespace(device_id=k, kind=k, extra={}, memory_pool="vram" if k == "gpu_discrete" else "host_ram") + for k in kinds + ] + descriptions = describe_topology(devices, machine_id="pc", domain_id="windows") + for descriptor in descriptions: + if descriptor.integration != "discrete": + assert descriptor.memory_pool_ids == ("pc:ram",) + assert descriptor.power_domain_ids == () # unknown is not invented + + +async def wait_output(pool): + async def poll(): + while (out := pool.poll_graph_output(0)) is None: + await asyncio.sleep(0.005) + return out + + return await asyncio.wait_for(poll(), timeout=10) + + +@pytest.mark.asyncio +async def test_runtime_load_submit_lease_and_shutdown(graph_runtime): + runtime = graph_runtime() + pool = runtime.stage_pools[0] + state = SimpleNamespace(sampling_params_list=[None]) + prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + await pool.submit_initial("a", state, prompt) + output = await wait_output(pool) + np.testing.assert_array_equal(output.custom_output["tensors"]["x_out"], 2) + with pytest.raises(ResourceUnavailable): + await pool.submit_initial("b", state, prompt) + output.release_stage_buffers() + await asyncio.sleep(0) + await pool.submit_initial("b", state, prompt) + (await wait_output(pool)).release_stage_buffers() + runtime.shutdown() + assert runtime.resource_ledger.snapshot()["reserved"] == {"host_ram": 0} + + +def test_joint_admission_before_any_worker(graph_runtime): + runtime = graph_runtime(stages=2, chain=True, capacity=600 << 20, start=False) + with pytest.raises(ResourceUnavailable): + runtime.initialize() + assert runtime.resource_ledger.snapshot()["owners"] == [] + assert runtime.stage_pools == [] + + +@pytest.mark.parametrize("fanout", [False, True]) +def test_unimplemented_graph_topologies_refused_before_loading(graph_runtime, fanout): + runtime = graph_runtime(stages=3, chain=fanout, start=False) + if fanout: + runtime._stage_configs[2].engine_input_source = [0] + with pytest.raises(ValueError, match="linear chain with exactly one final output"): + runtime.initialize() + assert runtime.resource_ledger is None + assert runtime.stage_pools == [] + + +def test_placement_failure_releases_load(graph_runtime): + runtime = graph_runtime(worker_config={"load": {"fraction_on_target": 0.0}}, start=False) + with pytest.raises(RuntimeError, match="REFUSE_EP_PLACEMENT"): + runtime.initialize() + assert runtime.resource_ledger.snapshot()["owners"] == [] + + +@pytest.mark.asyncio +async def test_worker_failure_produces_error_not_success(graph_runtime): + runtime = graph_runtime(worker_config={"die_on": "run"}) + pool = runtime.stage_pools[0] + await pool.submit_initial( + "a", SimpleNamespace(sampling_params_list=[None]), {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + ) + output = await wait_output(pool) + assert output.finished and output.error + + +@pytest.mark.asyncio +async def test_cancel_fences_inflight_graph(graph_runtime): + runtime = graph_runtime(worker_config={"run": {"sleep_s": 30}}) + pool = runtime.stage_pools[0] + await pool.submit_initial( + "a", SimpleNamespace(sampling_params_list=[None]), {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + ) + await asyncio.sleep(0.05) + await asyncio.wait_for(pool.abort_requests(["a"]), timeout=8) + assert pool.poll_graph_output(0) is None + assert runtime.resource_ledger.snapshot()["reserved"]["host_ram"] == 0 + + +def graph_handoff(outputs, prompt, requires_multimodal_data): + return {"tensors": {"x": outputs[0].custom_output["tensors"]["x_out"]}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_driven", [False, True]) +async def test_two_graphs_through_existing_orchestrator(graph_runtime, event_driven): + from vllm_omni.engine.messages import StageSubmissionMessage + from vllm_omni.engine.orchestrator import Orchestrator + + runtime = graph_runtime(stages=2, chain=True) + request_queue, output_queue, rpc_queue = asyncio.Queue(), asyncio.Queue(), asyncio.Queue() + orch = Orchestrator(request_queue, output_queue, rpc_queue, runtime.stage_pools) + loop = asyncio.create_task(orch._orchestration_loop_event_driven() if event_driven else orch._orchestration_loop()) + prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + try: + await orch._handle_add_request( + StageSubmissionMessage( + type="add_request", + request_id="pipeline", + prompt=prompt, + original_prompt=prompt, + output_prompt_text=None, + sampling_params_list=[None, None], + final_stage_id=1, + final_output_stage_ids=[1], + preprocess_ms=0, + request_timestamp=0, + enqueue_ts=0, + ) + ) + from vllm_omni.engine.messages import StageMetricsMessage + + while True: + result = await asyncio.wait_for(output_queue.get(), timeout=10) + if not isinstance(result, StageMetricsMessage): + break + assert result.finished + np.testing.assert_array_equal(result.engine_outputs.custom_output["tensors"]["x_out"], 4) + result.engine_outputs.release_stage_buffers() + await asyncio.sleep(0) + assert not orch.request_states + assert all(pool.stage_client._active is None for pool in runtime.stage_pools) + finally: + orch._shutdown_event.set() + await asyncio.wait_for(loop, timeout=5) + + +@pytest.mark.asyncio +async def test_stale_ack_cannot_release_reused_request_id(graph_runtime): + runtime = graph_runtime() + pool = runtime.stage_pools[0] + state = SimpleNamespace(sampling_params_list=[None]) + prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + await pool.submit_initial("same", state, prompt) + output = await wait_output(pool) + stale_ack = output._stage_release + output.release_stage_buffers() + await asyncio.sleep(0) + await pool.submit_initial("same", state, prompt) + next_output = await wait_output(pool) + stale_ack() + await asyncio.sleep(0) + assert pool.stage_client._active == "same" + next_output.release_stage_buffers() + + +@pytest.mark.asyncio +async def test_real_ort_cpu_graph_through_stage_runtime(graph_runtime): + runtime = graph_runtime(real_ort=True) + pool = runtime.stage_pools[0] + await pool.submit_initial( + "ort", SimpleNamespace(sampling_params_list=[None]), {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + ) + output = await wait_output(pool) + assert not output.error + np.testing.assert_array_equal(output.custom_output["tensors"]["x_out"], 2) + assert pool.stage_client.load_report.target_nodes > 0 + output.release_stage_buffers() + + +@pytest.mark.asyncio +async def test_public_async_omni_graph_requests(graph_runtime): + from vllm_omni.entrypoints.async_omni import AsyncOmni + + engine = graph_runtime(entrypoint=AsyncOmni) + for request_id in ("first", "second"): + outputs = [ + out + async for out in engine.generate( + {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}, request_id=request_id + ) + ] + assert len(outputs) == 1 + assert outputs[0].request_id == request_id + assert outputs[0].custom_output["tensors"]["x_out"].shape == (2, 3) + + +def test_public_sync_omni_graph_requests(graph_runtime): + from vllm_omni.entrypoints.omni import Omni + + engine = graph_runtime(entrypoint=Omni) + for _ in range(2): + outputs = engine.generate({"tensors": {"x": np.ones((2, 3), dtype=np.float32)}}, use_tqdm=False) + assert len(outputs) == 1 + assert outputs[0].custom_output["tensors"]["x_out"].shape == (2, 3) + + +@pytest.mark.asyncio +async def test_graph_ingress_stays_bounded_until_consumer_ack(graph_runtime): + from vllm_omni.engine.messages import OutputMessage + from vllm_omni.entrypoints.async_omni import AsyncOmni + + frontend = graph_runtime(entrypoint=AsyncOmni) + engine = frontend.engine + prompt = {"tensors": {"x": np.ones((2, 3), dtype=np.float32)}} + engine.add_request("slow", prompt) + for i in range(100): + with pytest.raises(ResourceUnavailable, match="capacity is one"): + engine.add_request(f"blocked-{i}", prompt) + + async def take(): + while True: + msg = await engine.try_get_output_async() + if isinstance(msg, OutputMessage): + return msg.engine_outputs + await asyncio.sleep(0.005) + + output = await asyncio.wait_for(take(), 5) + with pytest.raises(ResourceUnavailable): + engine.add_request("still-blocked", prompt) + stale_release = output._stage_release + output.release_stage_buffers() + await asyncio.sleep(0.01) + engine.add_request("next", prompt) + stale_release() + with pytest.raises(ResourceUnavailable): + engine.add_request("late-ack-must-not-release-next", prompt) + (await asyncio.wait_for(take(), 5)).release_stage_buffers() diff --git a/tests/windows/test_shared_buffers.py b/tests/windows/test_shared_buffers.py new file mode 100644 index 000000000..75718107f --- /dev/null +++ b/tests/windows/test_shared_buffers.py @@ -0,0 +1,75 @@ +"""Ownership/framing tests also run on POSIX; native connector tests run on Windows.""" + +import subprocess +import sys +import uuid + +import pytest + +from vllm_omni.host.shared_memory import SharedBufferOwner, read_buffer + + +def test_producer_retains_until_a_different_process_reads(): + owner = SharedBufferOwner() + meta = owner.write(b"payload", f"omni-test-{uuid.uuid4().hex}") + try: + assert owner.retained_bytes > 0 + code = ( + "from vllm_omni.host.shared_memory import read_buffer; assert read_buffer(" + repr(meta) + ") == b'payload'" + ) + subprocess.run([sys.executable, "-c", code], check=True, timeout=60) + owner.collect() + assert owner.retained_bytes == 0 + finally: + owner.release(meta["name"]) + + +def test_bounds_stale_generation_and_name_collision(): + owner = SharedBufferOwner(max_bytes=128, max_buffers=1) + meta = owner.write(b"one", f"omni-test-{uuid.uuid4().hex}") + try: + with pytest.raises(FileExistsError): + SharedBufferOwner().write(b"collision", meta["name"]) + with pytest.raises(ValueError, match="stale"): + read_buffer({**meta, "generation": "wrong"}) + with pytest.raises(BufferError): + owner.write(b"two") + assert owner.retained_bytes > 0 + assert read_buffer(meta) == b"one" + owner.collect() + assert owner.retained_bytes == 0 + with pytest.raises(BufferError): + owner.write(b"x" * 129) + finally: + owner.release(meta["name"]) + + +@pytest.mark.skipif(sys.platform != "win32", reason="native Windows path and mapping lifetime") +def test_native_connector_roundtrip_and_cleanup(): + from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector + from vllm_omni.host.shared_memory import owner + from vllm_omni.windows.paths import shm_dir + + assert not shm_dir().startswith("/dev/shm") + producer, consumer = SharedMemoryConnector({}), SharedMemoryConnector({}) + key = f"omni-test-{uuid.uuid4().hex}" + try: + ok, size, _ = producer.put("0", "1", key, {"data": [1, 2, 3]}) + assert ok and owner.retained_bytes > 0 + assert consumer.get("0", "1", key) == ({"data": [1, 2, 3]}, size) + owner.collect() + assert owner.retained_bytes == 0 + assert producer.put("0", "1", key, {"unconsumed": True})[0] + assert owner.retained_bytes > 0 + producer.cleanup(key) + assert owner.retained_bytes == 0 + assert consumer.get("0", "1", key) is None + for index in range(8): + next_key = f"{key}_{index}" + assert producer.put("0", "1", next_key, {"index": index})[0] + assert len(producer._pending_keys) <= 1 + assert consumer.get("0", "1", next_key)[0] == {"index": index} + owner.collect() + finally: + producer.close() + consumer.close() diff --git a/tests/windows/test_windows_compat.py b/tests/windows/test_windows_compat.py index e32d0f461..e2024ff1f 100644 --- a/tests/windows/test_windows_compat.py +++ b/tests/windows/test_windows_compat.py @@ -7,6 +7,7 @@ that decide whether a vLLM function is replaced or left alone.""" from __future__ import annotations import asyncio +import errno import os import sys import tempfile @@ -34,7 +35,7 @@ class FakeLocking: self.calls.append((fd, mode)) if mode == self.NBLCK: if fd in self.held: - raise OSError(36, "Resource deadlock avoided") + raise OSError(errno.EACCES, "Lock held") self.held.add(fd) elif mode == self.UNLCK: self.held.discard(fd) @@ -86,6 +87,20 @@ def test_flock_blocking_waits_until_released(lock_fd): assert len(sleeps) == 2 and lock_fd in locking.held +@pytest.mark.parametrize("operation", [shims.LOCK_EX, shims.LOCK_EX | shims.LOCK_NB]) +def test_flock_permanent_failure_does_not_retry(lock_fd, operation): + sleeps = [] + + def invalid_lock(*args): + raise OSError(errno.EINVAL, "Unsupported lock path") + + shim = shims.FcntlShim(invalid_lock, 2, 0, sleep=sleeps.append) + with pytest.raises(OSError) as failure: + shim.flock(lock_fd, operation) + assert failure.value.errno == errno.EINVAL + assert sleeps == [] + + def test_shared_lock_is_served_as_exclusive_and_counted(lock_fd): locking, sleeps = FakeLocking(), [] s = _shim(locking, sleeps) @@ -237,7 +252,7 @@ def test_shm_dir_exists_and_path_joins(): d = shm_dir() assert os.path.isdir(d) assert shm_path("x.lock").startswith(d) - if os.path.isdir("/dev/shm"): + if sys.platform != "win32" and os.path.isdir("/dev/shm"): assert d == "/dev/shm" diff --git a/tests/worker/test_forward_kwargs_filtering.py b/tests/worker/test_forward_kwargs_filtering.py index e4db8483f..929d9c95c 100644 --- a/tests/worker/test_forward_kwargs_filtering.py +++ b/tests/worker/test_forward_kwargs_filtering.py @@ -43,9 +43,7 @@ def test_a_model_taking_kwargs_is_left_alone(): def test_a_vllm_native_model_reports_exactly_its_parameters(): accepted = _Probe(_VLLMNativeModel())._accepted_forward_kwargs() - assert accepted == frozenset( - {"input_ids", "positions", "intermediate_tensors", "inputs_embeds"} - ) + assert accepted == frozenset({"input_ids", "positions", "intermediate_tensors", "inputs_embeds"}) for omni_only in ("sampling_metadata", "logits_index", "sampler"): assert omni_only not in accepted @@ -67,3 +65,15 @@ def test_an_uninspectable_forward_falls_back_to_passing_everything(): accepted = _Probe(Weird())._accepted_forward_kwargs() assert accepted is None or "sampling_metadata" not in accepted + + +def test_cuda_graph_wrapper_preserves_underlying_forward_signature(): + from vllm.compilation.cuda_graph import CUDAGraphWrapper + + # No CUDA context is needed for attribute delegation. Use the actual vLLM + # wrapper rather than a fake that could conceal API changes. + wrapper = object.__new__(CUDAGraphWrapper) + wrapper.runnable = _VLLMNativeModel() + assert _Probe(wrapper)._accepted_forward_kwargs() == _Probe(wrapper.runnable)._accepted_forward_kwargs() + wrapper.runnable = _OmniNativeModel() + assert _Probe(wrapper)._accepted_forward_kwargs() is None diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index 073cf4f59..0909d0c58 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -19,6 +19,7 @@ from vllm_omni.config.omni_config import ( VllmOmniConfig, VllmOmniDiffusionStageConfig, VllmOmniGenerationStageConfig, + VllmOmniGraphStageConfig, VllmOmniOrchestratorConfig, ) from vllm_omni.config.stage_config import ( @@ -60,6 +61,7 @@ __all__ = [ "BaseVllmOmniStageConfig", "VllmOmniARStageConfig", "VllmOmniGenerationStageConfig", + "VllmOmniGraphStageConfig", "VllmOmniDiffusionStageConfig", "StageConfigType", "OmniConfigResolution", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 81eb76f0b..36bd6880e 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -56,6 +56,7 @@ _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | StageExecutionType.LLM_AR: (StageType.LLM, "ar"), StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"), StageExecutionType.DIFFUSION: (StageType.DIFFUSION, None), + StageExecutionType.GRAPH: (StageType.GRAPH, None), } _PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS @@ -197,6 +198,8 @@ class _SchedulerEngineOverrides(TypedDict, total=False): class _RuntimeEngineOverrides(TypedDict, total=False): + backend: dict[str, Any] + resource_budget: dict[str, Any] additional_config: dict[str, Any] distributed_executor_backend: Any worker_cls: str @@ -528,6 +531,8 @@ class OmniStageRuntimeConfig: # LLM backend extensions; diffusion owns these in its config projection. additional_config: dict[str, Any] | None = None + backend: dict[str, Any] | None = None + resource_budget: dict[str, Any] | None = None distributed_executor_backend: Any = None worker_cls: str | None = None devices: str | None = None @@ -1142,19 +1147,22 @@ _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = ( # validation must use the execution-type-specific sets below. _STAGE_ENGINE_FIELDS = _LLM_STAGE_ENGINE_FIELDS | _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE = { - StageExecutionType.LLM_AR: _LLM_STAGE_ENGINE_FIELDS, - StageExecutionType.LLM_GENERATION: _LLM_STAGE_ENGINE_FIELDS, - StageExecutionType.DIFFUSION: _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS, + StageExecutionType.LLM_AR: _LLM_STAGE_ENGINE_FIELDS - {"backend"}, + StageExecutionType.LLM_GENERATION: _LLM_STAGE_ENGINE_FIELDS - {"backend"}, + StageExecutionType.DIFFUSION: _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS - {"backend"}, + StageExecutionType.GRAPH: frozenset({"model", "backend", "resource_budget", "log_level", "log_stats"}), } _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = { StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS, StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS, StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS, + StageExecutionType.GRAPH: frozenset(), } _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = { StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS, StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS, StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP), + StageExecutionType.GRAPH: frozenset(), } _LOAD_STAGE_ENGINE_FIELD_MAP = { @@ -1254,6 +1262,10 @@ def _stage_engine_values( stage_cli_overrides: Mapping[str, Any] | None = None, ) -> _StageEngineValues: engine = _stage_engine_overrides(stage_deploy) + if topology.execution_type == StageExecutionType.GRAPH and engine.get("silence_ban_frames") == 0: + # StageDeployConfig carries this speech-specific default even when no + # speech option was supplied. It is not an external-graph argument. + engine.pop("silence_ban_frames") # Preserve legacy ordering: topology-owned KV roles override deploy # extras, while an explicit CLI override remains highest priority. if topology.omni_kv_config: @@ -1467,7 +1479,14 @@ class VllmOmniDiffusionStageConfig(BaseVllmOmniStageConfig): diffusion_config: _DiffusionConfigProjection = field(default_factory=_DiffusionConfigProjection) -StageConfigType: TypeAlias = VllmOmniARStageConfig | VllmOmniGenerationStageConfig | VllmOmniDiffusionStageConfig +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniGraphStageConfig(BaseVllmOmniStageConfig): + """A whole external graph with an explicit backend and resource budget.""" + + +StageConfigType: TypeAlias = ( + VllmOmniARStageConfig | VllmOmniGenerationStageConfig | VllmOmniDiffusionStageConfig | VllmOmniGraphStageConfig +) def _build_common_stage_config_kwargs( @@ -1628,10 +1647,27 @@ def _build_diffusion_stage_config( ) +def _build_graph_stage_config(pipeline, deploy, topology, stage_deploy, engine, *, model): + if deploy.async_chunk or deploy.session_mode == "duplex": + raise ValueError("external graph v1 supports complete, stateless requests with async_chunk=False") + if stage_deploy is None or not engine.runtime.get("backend") or not engine.runtime.get("resource_budget"): + raise ValueError("graph stage requires backend and resource_budget") + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + pipeline, + deploy, + topology, + stage_deploy, + engine, + model=model, + ) + return _with_resolved_processors(VllmOmniGraphStageConfig(**common_kwargs), input_proc, next_stage_proc) + + _STAGE_CONFIG_BUILDERS = { StageExecutionType.LLM_AR: _build_ar_stage_config, StageExecutionType.LLM_GENERATION: _build_generation_stage_config, StageExecutionType.DIFFUSION: _build_diffusion_stage_config, + StageExecutionType.GRAPH: _build_graph_stage_config, } diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2ac1f99dd..0bfb99c19 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -189,6 +189,7 @@ class StageType(str, Enum): # TODO(@lishunyang12): remove once all models migrate to StageExecutionType LLM = "llm" DIFFUSION = "diffusion" + GRAPH = "graph" class StageExecutionType(str, Enum): @@ -197,6 +198,7 @@ class StageExecutionType(str, Enum): LLM_AR = "llm_ar" LLM_GENERATION = "llm_generation" DIFFUSION = "diffusion" + GRAPH = "graph" def _resolve_scheduler( @@ -386,6 +388,8 @@ class StageDeployConfig: devices: str | None = None num_replicas: int = 1 env: dict[str, Any] | None = None + backend: dict[str, Any] | None = None + resource_budget: dict[str, Any] | None = None # Inter-stage connector wiring and request defaults. output_connectors: dict[str, str] | None = None @@ -867,6 +871,7 @@ _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | StageExecutionType.LLM_AR: (StageType.LLM, "ar"), StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"), StageExecutionType.DIFFUSION: (StageType.DIFFUSION, None), + StageExecutionType.GRAPH: (StageType.GRAPH, None), } diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 10999f2b0..02f817001 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -3,9 +3,8 @@ # [edge-infer W1] This module is in the import closure of AsyncOmni, so a # module-scope `import fcntl` makes the whole engine unimportable on Windows -# -- before anyone has chosen a connector. The transport itself is genuinely -# POSIX (it hardcodes /dev/shm lock paths), so the honest Windows behaviour -# is to stay importable and refuse at construction, not to fake a port. +# -- before anyone has chosen a connector. The Windows plugin supplies locks; +# host helpers keep mappings alive until the single consumer acknowledges them. try: import fcntl except ImportError: # Windows @@ -15,6 +14,7 @@ from multiprocessing import shared_memory as shm_pkg from typing import Any from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes +from vllm_omni.windows.paths import shm_path from ..utils.logging import get_connector_logger from .base import OmniConnectorBase @@ -25,7 +25,7 @@ logger = get_connector_logger(__name__) class SharedMemoryConnector(OmniConnectorBase): """Key-addressed local shared-memory connector. - SHM is a local-only transport: it reads/writes POSIX shared memory + SHM is a local-only transport: it reads/writes host shared memory segments identified purely by *key*. It does **not** understand remote-transport metadata such as ``source_host`` / ``source_port`` (that is the RDMA connector's job). When such metadata is passed in, @@ -35,9 +35,8 @@ class SharedMemoryConnector(OmniConnectorBase): def __init__(self, config: dict[str, Any]): if fcntl is None: raise RuntimeError( - "SharedMemoryConnector requires POSIX shared memory and fcntl " - "locking (it addresses lock files under /dev/shm); neither exists " - "on Windows. Select a different connector for this stage." + "SharedMemoryConnector requires host file-lock support. " + "On Windows, activate the vllm_omni.windows compatibility layer before importing connectors." ) self.config = config self.stage_id = config.get("stage_id", -1) @@ -59,7 +58,7 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) - lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock" + lock_file = shm_path(f"shm_{put_key}_lockfile.lock") with open(lock_file, "wb+") as lockf: fcntl.flock(lockf, fcntl.LOCK_EX) meta = shm_write_bytes(payload, name=put_key) @@ -67,6 +66,12 @@ class SharedMemoryConnector(OmniConnectorBase): # meta contains {'name': ..., 'size': ...} metadata = {"shm": meta, "size": size} + if os.name == "nt": + from vllm_omni.host.shared_memory import owner + + # ACKs arrive in another process. Do not retain an unbounded + # history of keys after their producer handles are released. + self._pending_keys.intersection_update(owner.retained_names) self._pending_keys.add(put_key) self._metrics["puts"] += 1 @@ -86,7 +91,7 @@ class SharedMemoryConnector(OmniConnectorBase): data_bytes = shm_read_bytes(shm_handle) fcntl.flock(lockf, fcntl.LOCK_UN) obj = self.deserialize_obj(data_bytes) - result = (obj, int(shm_handle.get("size", 0))) + result = (obj, len(data_bytes)) deserialized = True return result except Exception as e: @@ -106,7 +111,7 @@ class SharedMemoryConnector(OmniConnectorBase): shm = shm_pkg.SharedMemory(name=get_key) if shm is None or shm.size == 0: return None - lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock" + lock_file = shm_path(f"shm_{get_key}_lockfile.lock") shm_handle = {"name": get_key, "size": shm.size} result = self._get_data_with_lock(lock_file, shm_handle) if result is not None: @@ -142,7 +147,7 @@ class SharedMemoryConnector(OmniConnectorBase): if isinstance(metadata, dict) and "shm" in metadata: shm_handle = metadata["shm"] - lock_file = f"/dev/shm/shm_{shm_handle['name']}_lockfile.lock" + lock_file = shm_path(f"shm_{shm_handle['name']}_lockfile.lock") result = self._get_data_with_lock(lock_file, shm_handle) if result is not None: self._pending_keys.discard(get_key) @@ -171,6 +176,7 @@ class SharedMemoryConnector(OmniConnectorBase): ] for key in stale: self._pending_keys.discard(key) + self._release_owned_buffer(key) try: seg = shm_pkg.SharedMemory(name=key) seg.close() @@ -180,7 +186,7 @@ class SharedMemoryConnector(OmniConnectorBase): pass except Exception as e: logger.debug("cleanup: failed to unlink SHM segment %s: %s", key, e) - lock_file = f"/dev/shm/shm_{key}_lockfile.lock" + lock_file = shm_path(f"shm_{key}_lockfile.lock") if os.path.exists(lock_file): try: os.remove(lock_file) @@ -190,13 +196,14 @@ class SharedMemoryConnector(OmniConnectorBase): def close(self) -> None: """Unlink all remaining tracked SHM segments.""" for key in list(self._pending_keys): + self._release_owned_buffer(key) try: seg = shm_pkg.SharedMemory(name=key) seg.close() seg.unlink() except Exception: pass - lock_file = f"/dev/shm/shm_{key}_lockfile.lock" + lock_file = shm_path(f"shm_{key}_lockfile.lock") if os.path.exists(lock_file): try: os.remove(lock_file) @@ -204,5 +211,12 @@ class SharedMemoryConnector(OmniConnectorBase): pass self._pending_keys.clear() + @staticmethod + def _release_owned_buffer(key: str) -> None: + if os.name == "nt": + from vllm_omni.host.shared_memory import owner + + owner.release(key) + def health(self) -> dict[str, Any]: return {"status": "healthy", **self._metrics} diff --git a/vllm_omni/edge/local/external/client.py b/vllm_omni/edge/local/external/client.py index 148c615c8..1c4afc6e6 100644 --- a/vllm_omni/edge/local/external/client.py +++ b/vllm_omni/edge/local/external/client.py @@ -28,10 +28,12 @@ ledger, not a rounding error. So the worker reports its own RSS and from __future__ import annotations +import json import secrets import socket import subprocess -import sys +import tempfile +import threading import time from dataclasses import asdict, dataclass, field from pathlib import Path @@ -42,6 +44,7 @@ import numpy as np from vllm_omni.edge.local.external import launch as _launch from vllm_omni.edge.local.external import protocol as _proto from vllm_omni.edge.local.external.protocol import ProtocolError, WorkerError +from vllm_omni.host.process import spawn_worker, terminate_worker DEFAULT_START_TIMEOUT_S = 120.0 """Generous on purpose: the first VitisAI session on a cold machine spends @@ -212,12 +215,17 @@ class ExternalWorker: *, start_timeout_s: float = DEFAULT_START_TIMEOUT_S, call_timeout_s: float = DEFAULT_CALL_TIMEOUT_S, + max_payload_bytes: int = _proto.MAX_PAYLOAD_BYTES, ) -> None: if not route.available: raise WorkerStartError(f"route {route.name} is not available: {route.reason}") self.route = route self.start_timeout_s = start_timeout_s self.call_timeout_s = call_timeout_s + self.max_payload_bytes = max_payload_bytes + self._log = None + self._lifecycle_lock = threading.RLock() + self._drained = False self._process: subprocess.Popen[bytes] | None = None self._sock: socket.socket | None = None self._listener: socket.socket | None = None @@ -228,6 +236,8 @@ class ExternalWorker: def start(self) -> dict[str, Any]: """Launch the worker and complete the handshake. Returns its ``hello``.""" + if self._process is not None: + raise WorkerStartError("worker instances cannot be started twice") token = secrets.token_hex(16) host = _bind_address(is_windows=self.route.is_windows) listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -240,19 +250,29 @@ class ExternalWorker: argv = [ str(self.route.interpreter), + _launch.to_worker_path( + Path(__file__).parents[3] / "host/worker_bootstrap.py", is_windows=self.route.is_windows + ), _launch.to_worker_path(self.route.worker, is_windows=self.route.is_windows), - "--host", host, - "--port", str(port), - "--token", token, + "--host", + host, + "--port", + str(port), + "--token", + token, ] - self._process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + self._log = tempfile.TemporaryFile() + try: + self._process = spawn_worker(argv, stdout=self._log, stderr=self._log) + except BaseException: + self.close() + raise try: conn, _ = listener.accept() - except socket.timeout as exc: + except TimeoutError as exc: raise WorkerStartError( - f"worker did not connect back within {self.start_timeout_s:.0f}s: " - f"{self._drain_process()}" + f"worker did not connect back within {self.start_timeout_s:.0f}s: {self._drain_process()}" ) from exc finally: listener.close() @@ -262,7 +282,11 @@ class ExternalWorker: conn.settimeout(self.call_timeout_s) self._sock = conn - op, body, _ = _proto.recv_message(conn) + try: + op, body, _ = _proto.recv_message(conn, max_payload_bytes=0) + except BaseException: + self.terminate() + raise if op != _proto.OP_HELLO: self.close() raise WorkerStartError(f"first message was {op!r}, expected {_proto.OP_HELLO!r}") @@ -275,20 +299,74 @@ class ExternalWorker: return body def _drain_process(self) -> str: + with self._lifecycle_lock: + return self._drain_process_locked() + + def _drain_process_locked(self) -> str: process = self._process if process is None: return "worker was never started" - if process.poll() is None: - process.kill() - try: - out, err = process.communicate(timeout=5) - except subprocess.TimeoutExpired: # pragma: no cover - kill just ran + if not terminate_worker(process): return f"worker pid {process.pid} did not exit" - text = (err or b"").decode("utf-8", "replace").strip() - if not text: - text = (out or b"").decode("utf-8", "replace").strip() + text = "" + if self._log is not None: + self._log.seek(0, 2) + self._log.seek(max(0, self._log.tell() - 2000)) + text = self._log.read().decode("utf-8", "replace").strip() return f"exit={process.returncode} output={text[:2000]!r}" if text else f"exit={process.returncode}" + def terminate(self) -> bool: + """Interrupt a blocking native call without sending on its data socket. + + The caller retains reservations if the process tree did not drain. + This method may run concurrently with the one in-flight `_call`. + """ + with self._lifecycle_lock: + return self._terminate_locked() + + def _terminate_locked(self) -> bool: + if self._drained: + return True + native_drained = True + if self._process is not None and self.route.is_windows and _launch.is_wsl(): + identity = self.hello.get("process_identity") or {} + native_drained = False + if type(identity.get("pid")) is int and type(identity.get("created_filetime")) is int: + try: + helper = Path(__file__).parents[3] / "host/windows_process.py" + result = subprocess.run( + [ + str(self.route.interpreter), + _launch.to_worker_path(helper, is_windows=True), + str(identity["pid"]), + str(identity["created_filetime"]), + ], + capture_output=True, + timeout=20, + text=True, + ) + native_drained = result.returncode == 0 and json.loads(result.stdout).get("drained") is True + except (OSError, subprocess.SubprocessError, ValueError): + pass + # Keep the data connection open until native tree retirement completes: + # closing it first lets an idle worker exit before we can verify its tree. + sock, self._sock = self._sock, None + if sock is not None: + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + if self._listener is not None: + self._listener.close() + self._listener = None + drained = (self._process is None or terminate_worker(self._process)) and native_drained + self._drained = drained + if drained and self._log is not None: + self._log.close() + self._log = None + return drained + def close(self, timeout_s: float = DEFAULT_CLOSE_TIMEOUT_S) -> None: """Ask the worker to exit; kill it if it will not. Safe to call twice.""" sock, self._sock = self._sock, None @@ -304,17 +382,20 @@ class ExternalWorker: if self._listener is not None: self._listener.close() self._listener = None - process, self._process = self._process, None + process = self._process if process is not None and process.poll() is None: try: process.wait(timeout=timeout_s) except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=5) + if not self.terminate(): + raise WorkerStartError("worker tree did not drain during close") if process is not None: for stream in (process.stdout, process.stderr): if stream is not None: stream.close() + if self._log is not None: + self._log.close() + self._log = None def __enter__(self) -> ExternalWorker: self.start() @@ -336,12 +417,13 @@ class ExternalWorker: raise WorkerStartError("worker is not running; call start() first") try: _proto.send_message(sock, op, body, tensors) - reply_op, reply_body, reply_tensors = _proto.recv_message(sock) + reply_op, reply_body, reply_tensors = _proto.recv_message(sock, max_payload_bytes=self.max_payload_bytes) except (OSError, ProtocolError) as exc: # The socket is the only liveness signal we have; a dead worker # shows up as a broken pipe or a short read, and the useful part of # the diagnosis is on its stderr. detail = self._drain_process() + sock.close() self._sock = None raise WorkerStartError(f"worker connection failed ({exc}); {detail}") from exc if reply_op != _proto.OP_OK: @@ -357,6 +439,7 @@ class ExternalWorker: profile_prefix: str | Path | None = None, ep_dir: str | Path | None = None, intra_op_num_threads: int | None = None, + artifact_files: list[Path] | None = None, ) -> LoadReport: """Open a graph and measure where its nodes actually ran. @@ -371,13 +454,15 @@ class ExternalWorker: "profile": bool(profile), } if profile_prefix is not None: - body["profile_prefix"] = _launch.to_worker_path( - profile_prefix, is_windows=self.route.is_windows - ) + body["profile_prefix"] = _launch.to_worker_path(profile_prefix, is_windows=self.route.is_windows) if ep_dir is not None: body["ep_dir"] = str(ep_dir) if intra_op_num_threads is not None: body["intra_op_num_threads"] = int(intra_op_num_threads) + if artifact_files is not None: + body["artifact_files"] = [ + _launch.to_worker_path(path, is_windows=self.route.is_windows) for path in artifact_files + ] reply, _ = self._call(_proto.OP_LOAD, body, example_inputs or {}) self.report = LoadReport.from_body(reply, ep=self.route.ep, graph_path=str(graph_path)) diff --git a/vllm_omni/edge/local/external/launch.py b/vllm_omni/edge/local/external/launch.py index 001a1a7b8..818d6f12c 100644 --- a/vllm_omni/edge/local/external/launch.py +++ b/vllm_omni/edge/local/external/launch.py @@ -29,6 +29,7 @@ different work, and the planner has to be able to tell the caller which it hit. from __future__ import annotations +import getpass import os import shutil import subprocess @@ -63,20 +64,25 @@ covered on a machine with no AMD hardware and no ORT at all.""" # Where the routes live on this machine. Defaults, not requirements: the env # override above is the supported way to point at a different install, and the # reason string names this list when nothing is found. +_WINDOWS_HOME = Path.home() if os.name == "nt" else Path("/mnt/c/Users") / getpass.getuser() _DEFAULT_CANDIDATES: dict[str, tuple[str, ...]] = { ROUTE_VITISAI: ( - r"/mnt/c/Users/zhout/npu-ep/Scripts/python.exe", - r"/mnt/c/Users/zhout/ryzenai/Scripts/python.exe", + str(_WINDOWS_HOME / "npu-ep/Scripts/python.exe"), + str(_WINDOWS_HOME / "ryzenai/Scripts/python.exe"), ), ROUTE_DML: ( - r"/mnt/c/Users/zhout/dml-ep/Scripts/python.exe", - r"/mnt/c/Users/zhout/directml/Scripts/python.exe", + str(_WINDOWS_HOME / "dml-ep/Scripts/python.exe"), + str(_WINDOWS_HOME / "directml/Scripts/python.exe"), + ), + ROUTE_TORCH_DML: () if os.name == "nt" else (".venvs/dml/bin/python",), + ROUTE_CPU: ( + (".venvs/omni-cpu/Scripts/python.exe",) + if os.name == "nt" + else (".venvs/omni-cuda/bin/python", ".venvs/omni-cpu/bin/python") ), - ROUTE_TORCH_DML: (".venvs/dml/bin/python",), - ROUTE_CPU: (".venvs/omni-cuda/bin/python", ".venvs/omni-cpu/bin/python"), ROUTE_WIN_CUDA: ( - r"/mnt/c/Users/zhout/w0/w0venv/Scripts/python.exe", - r"/mnt/c/Users/zhout/vllm-win/Scripts/python.exe", + str(_WINDOWS_HOME / "w0/w0venv/Scripts/python.exe"), + str(_WINDOWS_HOME / "vllm-win/Scripts/python.exe"), ), } @@ -175,9 +181,10 @@ def is_wsl() -> bool: """ if os.name == "nt": return False - return Path("/proc/sys/fs/binfmt_misc/WSLInterop").exists() or Path( - "/proc/sys/fs/binfmt_misc/WSLInterop-late" - ).exists() + return ( + Path("/proc/sys/fs/binfmt_misc/WSLInterop").exists() + or Path("/proc/sys/fs/binfmt_misc/WSLInterop-late").exists() + ) def to_worker_path(path: str | Path, *, is_windows: bool) -> str: @@ -194,9 +201,7 @@ def to_worker_path(path: str | Path, *, is_windows: bool) -> str: if not shutil.which("wslpath"): return text try: - out = subprocess.run( - ["wslpath", "-w", text], capture_output=True, text=True, timeout=20, check=True - ) + out = subprocess.run(["wslpath", "-w", text], capture_output=True, text=True, timeout=20, check=True) except (subprocess.SubprocessError, OSError): return text return out.stdout.strip() or text @@ -206,16 +211,22 @@ def _resolve_interpreter(route: str, root: Path) -> tuple[str | None, str]: env_key = f"{ENV_PREFIX}{route.upper().replace('-', '_')}" override = os.environ.get(env_key) if override: - if Path(override).exists(): - return override, "" + try: + if Path(override).is_file(): + return override, "" + except OSError as exc: + return None, f"cannot access configured interpreter: {exc}" return None, f"{env_key}={override!r} does not exist" tried: list[str] = [] for candidate in _DEFAULT_CANDIDATES.get(route, ()): path = Path(candidate) if Path(candidate).is_absolute() else root / candidate tried.append(str(path)) - if path.exists(): - return str(path), "" + try: + if path.is_file(): + return str(path), "" + except OSError: + continue return None, f"no interpreter found; tried {tried} and {env_key} is unset" @@ -226,9 +237,7 @@ def resolve(route: str, *, root: Path | None = None) -> Route: root = root or repo_root() is_windows = route in (ROUTE_VITISAI, ROUTE_DML, ROUTE_WIN_CUDA) worker_key = f"{WORKER_ENV_PREFIX}{route.upper().replace('-', '_')}" - worker = os.environ.get(worker_key) or str( - Path(__file__).resolve().parent / _WORKER_FOR_ROUTE[route] - ) + worker = os.environ.get(worker_key) or str(Path(__file__).resolve().parent / _WORKER_FOR_ROUTE[route]) if is_windows and not is_wsl() and os.name != "nt": return Route( diff --git a/vllm_omni/edge/local/external/protocol.py b/vllm_omni/edge/local/external/protocol.py index ea890c809..2917e2479 100644 --- a/vllm_omni/edge/local/external/protocol.py +++ b/vllm_omni/edge/local/external/protocol.py @@ -1,185 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""The wire between the engine and a stage that runs in another process. +"""Compatibility import for the extracted host-copy wire implementation. -Deliberately tiny, and deliberately **stdlib + numpy only**. The worker on the -other end of this socket may be a native-Windows interpreter that has no torch, -no vllm and no vllm_omni on its path -- ``C:\\Users\\zhout\\npu-ep`` is exactly -that: Python 3.12, onnxruntime, onnx, numpy, nothing else. Anything imported -here has to be installable there, so this module imports neither. - -**Why a socket at all.** Two devices on this laptop cannot be reached from the -engine's own interpreter, and for different reasons: - -* the **XDNA2 NPU** is an MCDM device. WSL2's GPU-PV forwards WDDM display - adapters only, so there is no ``/dev/accel`` inside WSL and no amount of - installing changes that. VitisAI is a Windows DLL. -* the **Radeon 890M** is reachable from WSL through DirectML, but - ``torch-directml`` pins torch 2.4.1 and Omni runs 2.13. - -So out-of-process is forced, not preferred. One protocol covers both, plus the -native-Windows DirectML route, which needs its own venv anyway because -``onnxruntime-directml`` is pinned at 1.24.4 and cannot load the VitisAI EP -(that needs ORT >= 1.25) -- the two AMD devices cannot share an interpreter. - -**Framing.** ``[4-byte big-endian header length][JSON header][raw tensor bytes]``. -Tensors are described in the header (name, dtype, shape, nbytes) and their -buffers are concatenated after it in the order listed, so a receive is two -reads and one ``frombuffer`` per tensor -- no per-element work and no pickle. -Measured on this machine, WSL <-> native Windows: 11 us half round trip, -330 MB/s. That is ample for a stage boundary (a 2 s Code2Wav chunk is ~192 KB, -about 0.6 ms) and nowhere near enough for per-layer traffic inside a decode -loop, which is the arithmetic behind "coarse stages only". - -``pickle`` is not used anywhere in here on purpose: the header is JSON, the -payload is raw buffers, so a hostile or merely broken peer cannot execute code -by replying. +Standalone legacy worker scripts load this file directly. The source-tree +fallback loads only the portable package, never the vllm_omni initializer. """ -from __future__ import annotations +import sys +from pathlib import Path -import json -import socket -import struct -from typing import Any - -import numpy as np - -# Ops the client sends. -OP_LOAD = "load" -"""Open a session on a graph. Replies with the placement report.""" -OP_RUN = "run" -"""One forward. Inputs in, outputs out.""" -OP_STATS = "stats" -"""Counters and RSS. The worker's memory is invisible to the engine's sampler -when the worker is a Windows process, so it has to report its own.""" -OP_CLOSE = "close" - -# Ops the worker sends. -OP_HELLO = "hello" -"""First message after connect-back: interpreter, versions, available EPs.""" -OP_OK = "ok" -OP_ERR = "err" - -MAX_HEADER_BYTES = 1 << 20 -"""A header is metadata; a megabyte of it is a bug or an attack, not a graph.""" - -MAX_PAYLOAD_BYTES = 2 << 30 -"""Refuse rather than allocate on a corrupt length. Two gibibytes is far above -any legitimate stage payload here -- the largest is a vision tower's patch -tensor, single-digit MB.""" - - -class ProtocolError(RuntimeError): - """The peer sent something this protocol cannot represent.""" - - -class WorkerError(RuntimeError): - """The peer ran and failed. Carries the remote traceback when there is one.""" - - def __init__(self, message: str, *, remote_traceback: str = "", code: str = "") -> None: - # The remote traceback goes in the message, not only in an attribute: - # a bare "AssertionError:" surfacing here with the actual stack one - # attribute away is a debugging loop nobody should have to discover. - detail = f"{message}\n--- worker traceback ---\n{remote_traceback}" if remote_traceback else message - super().__init__(detail) - self.message = message - self.remote_traceback = remote_traceback - self.code = code - - -def _recv_exactly(sock: socket.socket, n: int) -> bytes: - """Read exactly ``n`` bytes or raise. - - ``recv`` is allowed to return short reads on a stream socket and does, once - payloads get past a page or two; a partial read here would be silently - misparsed as the next tensor. - """ - if n == 0: - return b"" - chunks: list[bytes] = [] - remaining = n - while remaining: - chunk = sock.recv(min(remaining, 1 << 20)) - if not chunk: - raise ProtocolError(f"peer closed with {remaining} of {n} bytes outstanding") - chunks.append(chunk) - remaining -= len(chunk) - return b"".join(chunks) if len(chunks) > 1 else chunks[0] - - -def send_message( - sock: socket.socket, - op: str, - body: dict[str, Any] | None = None, - tensors: dict[str, np.ndarray] | None = None, -) -> None: - """Send one message. ``tensors`` are sent as raw buffers after the header.""" - tensors = tensors or {} - described: list[dict[str, Any]] = [] - buffers: list[memoryview] = [] - for name, array in tensors.items(): - array = np.asarray(array) - # Record the shape *before* making it contiguous: ``ascontiguousarray`` - # promotes a 0-d array to shape (1,), which would silently turn a - # scalar graph input into a length-1 vector on the far side. - shape = list(array.shape) - contiguous = np.ascontiguousarray(array) - described.append( - { - "name": name, - # ``str(dtype)`` loses endianness for multibyte types; ``.str`` - # keeps the '<'/'>' so a big-endian peer would still decode. - "dtype": contiguous.dtype.str, - "shape": shape, - "nbytes": int(contiguous.nbytes), - } - ) - # ``reshape(-1)`` rather than a direct cast: ``memoryview.cast`` refuses - # a 0-dimensional source. - buffers.append(memoryview(contiguous.reshape(-1)).cast("B")) - - header = json.dumps({"op": op, "body": body or {}, "tensors": described}).encode("utf-8") - if len(header) > MAX_HEADER_BYTES: - raise ProtocolError(f"header is {len(header)} bytes, over the {MAX_HEADER_BYTES} limit") - sock.sendall(struct.pack(">I", len(header))) - sock.sendall(header) - for buffer in buffers: - sock.sendall(buffer) - - -def recv_message(sock: socket.socket) -> tuple[str, dict[str, Any], dict[str, np.ndarray]]: - """Receive one message as ``(op, body, tensors)``. - - Raises :class:`WorkerError` on an ``err`` reply so callers do not each have - to remember to check; every other op is returned as-is. - """ - (header_len,) = struct.unpack(">I", _recv_exactly(sock, 4)) - if header_len > MAX_HEADER_BYTES: - raise ProtocolError(f"peer announced a {header_len}-byte header, over the limit") - try: - header = json.loads(_recv_exactly(sock, header_len).decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise ProtocolError(f"header is not JSON: {exc}") from exc - - total = sum(int(t["nbytes"]) for t in header.get("tensors", ())) - if total > MAX_PAYLOAD_BYTES: - raise ProtocolError(f"peer announced a {total}-byte payload, over the limit") - - tensors: dict[str, np.ndarray] = {} - for spec in header.get("tensors", ()): - raw = _recv_exactly(sock, int(spec["nbytes"])) - array = np.frombuffer(raw, dtype=np.dtype(spec["dtype"])) - # ``frombuffer`` is read-only and aliases ``raw``; a consumer that - # writes in place (ORT's IOBinding does) needs its own buffer. - tensors[spec["name"]] = array.reshape(spec["shape"]).copy() - - op = str(header.get("op", "")) - body = header.get("body") or {} - if op == OP_ERR: - raise WorkerError( - str(body.get("message", "worker failed with no message")), - remote_traceback=str(body.get("traceback", "")), - code=str(body.get("code", "")), - ) - return op, body, tensors +try: + from omni_stage_contracts.wire import * # noqa: F403 +except ModuleNotFoundError as exc: + if exc.name != "omni_stage_contracts": + raise + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from omni_stage_contracts.wire import * # noqa: F403 diff --git a/vllm_omni/edge/local/external/worker_ort.py b/vllm_omni/edge/local/external/worker_ort.py index 3bbf6a3d0..c873a66d8 100644 --- a/vllm_omni/edge/local/external/worker_ort.py +++ b/vllm_omni/edge/local/external/worker_ort.py @@ -164,12 +164,37 @@ def _provider_counts(profile_path: str) -> dict[str, int]: return counts +def _verify_artifact_members(graph: str, files: list[str]) -> None: + """Reject ONNX external data omitted from the controller's verified bundle.""" + import onnx + + allowed = {Path(name).resolve() for name in files} + graph_path = Path(graph).resolve() + if graph_path not in allowed: + raise ValueError("graph is not a verified artifact member") + pending = [onnx.load_model(graph, load_external_data=False)] + while pending: + message = pending.pop() + if message.DESCRIPTOR.full_name == "onnx.TensorProto": + locations = [entry.value for entry in message.external_data if entry.key == "location"] + if message.data_location == onnx.TensorProto.EXTERNAL and not locations: + raise ValueError("external tensor has no payload location") + for location in locations: + if (graph_path.parent / location).resolve() not in allowed: + raise ValueError(f"external tensor payload is absent from verified manifest: {location}") + for descriptor, value in message.ListFields(): + if descriptor.type == descriptor.TYPE_MESSAGE: + pending.extend(value if descriptor.is_repeated else [value]) + + def _make_session(body: dict[str, Any]) -> tuple[Any, dict[str, Any]]: """Build the ORT session for the requested execution provider.""" import onnxruntime as ort ep = str(body.get("ep", "cpu")) graph = str(body["graph_path"]) + if "artifact_files" in body: + _verify_artifact_members(graph, body["artifact_files"]) info: dict[str, Any] = {"ep": ep, "onnxruntime": ort.__version__} options = ort.SessionOptions() @@ -241,12 +266,8 @@ def _load(state: _Session, body: dict[str, Any], tensors: dict[str, np.ndarray]) state.session = session state.output_names = [o.name for o in session.get_outputs()] - info["inputs"] = [ - {"name": i.name, "type": i.type, "shape": list(i.shape)} for i in session.get_inputs() - ] - info["outputs"] = [ - {"name": o.name, "type": o.type, "shape": list(o.shape)} for o in session.get_outputs() - ] + info["inputs"] = [{"name": i.name, "type": i.type, "shape": list(i.shape)} for i in session.get_inputs()] + info["outputs"] = [{"name": o.name, "type": o.type, "shape": list(o.shape)} for o in session.get_outputs()] # The placement evidence. One profiled run, because ORT only writes node # assignments once nodes have actually executed -- a session that has never diff --git a/vllm_omni/edge/local/session.py b/vllm_omni/edge/local/session.py index 5d563e913..69750bd1b 100644 --- a/vllm_omni/edge/local/session.py +++ b/vllm_omni/edge/local/session.py @@ -33,10 +33,7 @@ nothing that would let a caller reach inside it. from __future__ import annotations import asyncio -import time import uuid -from dataclasses import asdict, dataclass, field -from typing import Any STATE_LAYOUT_VERSION = 1 """Bumped when the meaning of a handle's state changes. A handle from a @@ -45,79 +42,10 @@ to treat state as non-migratable until a specific backend pair has been measured, and a version that silently matches would hide that.""" -@dataclass(frozen=True) -class StateHandle: - """A reference to session state the backend owns. - - The engine never reads the state. It reads this, to decide whether a - request may continue a session or must start a new one. - """ - - session_id: str - backend: str - artifact_id: str - """Ties the state to the checkpoint that produced it. The same shapes from - a different quantization are not the same state.""" - layout_version: int = STATE_LAYOUT_VERSION - epoch: int = 0 - """Incremented on every cancel. Events from a lower epoch are stale.""" - replayable: bool = True - """A text session can be rebuilt by re-submitting its committed prompt, so - recovery after a backend crash is possible at an output boundary.""" - migratable: bool = False - """Never true in M0. Moving KV between backends needs a measured layout - conversion for the specific pair; assuming it is the failure the proposal - warns about.""" - created_unix: float = field(default_factory=time.time) - - def next_epoch(self) -> StateHandle: - from dataclasses import replace - - return replace(self, epoch=self.epoch + 1) - - def accepts(self, other: StateHandle) -> bool: - """Whether ``other`` may continue this session's state.""" - return ( - other.session_id == self.session_id - and other.artifact_id == self.artifact_id - and other.layout_version == self.layout_version - and other.epoch == self.epoch - ) - - def to_dict(self) -> dict[str, Any]: - return asdict(self) +from omni_stage_contracts.legacy import ChunkEvent, StateHandle # noqa: E402, F401 -@dataclass(frozen=True) -class ChunkEvent: - """One unit of output, with everything needed to order and retire it.""" - - request_id: str - stage_id: int - seq: int - """Monotonic within (request, epoch). Gaps mean loss; repeats mean a bug.""" - epoch: int - kind: str - """``token`` | ``done`` | ``error`` | ``cancelled``.""" - payload: Any - """For ``token``: the incremental text and token ids. Never the whole - output so far -- a cumulative payload makes a dropped chunk invisible.""" - started_unix: float - emitted_unix: float - final: bool = False - error: str | None = None - input_watermark: int = 0 - """How much of the input the producer had consumed when this was emitted. - One number in M0 (prompt tokens); it is what a duplex turn will align on.""" - release_token: str = "" - """Returned by the consumer to give credit back. Empty when the event - carries no resource the producer is holding on its behalf.""" - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - -class StreamClosed(RuntimeError): +class StreamClosed(RuntimeError): # noqa: N818 - retained public API """Raised on a stream that has been closed while a consumer waited.""" @@ -157,7 +85,12 @@ class BoundedEventStream: return len(payload.encode("utf-8")) if isinstance(payload, dict): return sum( - len(str(k)) + (len(v.encode("utf-8")) if isinstance(v, str) else 8 * (len(v) if isinstance(v, (list, tuple)) else 1)) + len(str(k)) + + ( + len(v.encode("utf-8")) + if isinstance(v, str) + else 8 * (len(v) if isinstance(v, (list, tuple)) else 1) + ) for k, v in payload.items() ) return 64 @@ -172,11 +105,13 @@ class BoundedEventStream: size = self._sizeof(event) async with self._credit: await self._credit.wait_for( - lambda: self._closed - or (self._chunks < self.max_chunks and self._bytes + size <= self.max_bytes) - # A single event larger than the whole byte bound must still go - # through once the queue is empty, or it deadlocks forever. - or (self._chunks == 0 and size > self.max_bytes) + lambda: ( + self._closed + or (self._chunks < self.max_chunks and self._bytes + size <= self.max_bytes) + # A single event larger than the whole byte bound must still go + # through once the queue is empty, or it deadlocks forever. + or (self._chunks == 0 and size > self.max_bytes) + ) ) if self._closed: raise StreamClosed("stream closed while producing") diff --git a/vllm_omni/edge/topology.py b/vllm_omni/edge/topology.py new file mode 100644 index 000000000..8d6c76f74 --- /dev/null +++ b/vllm_omni/edge/topology.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Map discovery facts into physical constraints, without guessing performance.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import asdict + +from omni_stage_contracts import DeviceDescriptor + + +def describe_topology( + devices, + *, + machine_id: str, + domain_id: str, + physical_ids: Mapping[str, str] | None = None, + guest_pool: str | None = None, +) -> list[DeviceDescriptor]: + """Normalize legacy discovery; caller reconciles cross-domain identities. + + Physical IDs must be shared across routes when discovery proves they refer + to one GPU. Unknown bandwidth/power relationships stay empty. WSL quota is + an additional constraint on host allocations, never another RAM capacity. + """ + if not machine_id or not domain_id: + raise ValueError("machine and domain IDs must be explicit") + physical_ids = physical_ids or {} + result = [] + for device in devices: + extra = device.extra or {} + physical = physical_ids.get(device.device_id, device.device_id) + kind = {"gpu_integrated": "gpu", "gpu_discrete": "gpu"}.get(device.kind, device.kind) + integrated = device.kind == "gpu_integrated" + integration = "integrated" if integrated else "discrete" if device.kind == "gpu_discrete" else "unknown" + pool = f"{machine_id}:ram" if device.memory_pool == "host_ram" else f"{machine_id}:vram:{physical}" + pools = (pool,) + ((guest_pool,) if guest_pool and device.memory_pool == "host_ram" else ()) + result.append( + DeviceDescriptor( + device_id=f"{machine_id}:{physical}", + kind=kind, + domain_id=domain_id, + memory_pool_ids=pools, + integration=extra.get("integration", integration), + package_id=extra.get("package_id"), + bandwidth_group_ids=tuple(extra.get("bandwidth_group_ids", ())), + power_domain_ids=tuple(extra.get("power_domain_ids", ())), + ) + ) + return result + + +def topology_report(devices, **kwargs) -> list[dict]: + return [asdict(device) for device in describe_topology(devices, **kwargs)] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 6b31761b2..aa02d4f5f 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -244,6 +244,11 @@ class AsyncOmniEngine: ) self.num_stages = len(self.stage_configs) + from vllm_omni.engine.resource_ledger import GraphRequestGate + + self._graph_request_gate = ( + GraphRequestGate() if any(c.stage_type == "graph" for c in self.stage_configs) else None + ) stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) self.stage_pools: list[StagePool] = [] @@ -454,6 +459,7 @@ class AsyncOmniEngine: enable_duplex_control=self._duplex_control_enabled, duplex_session_config=self.duplex_session_config, ) + orchestrator.graph_request_gate = self._graph_request_gate if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) await orchestrator.run() @@ -757,7 +763,7 @@ class AsyncOmniEngine: stage_type = self.stage_metadata[0].stage_type output_prompt_text: Any = None _preprocess_ms = 0.0 - if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest): + if stage_type not in ("diffusion", "graph") and not isinstance(prompt, EngineCoreRequest): # Stage transforms and downstream stages must share the same # request identity, including when the transform replaces the # prompt object. @@ -1145,6 +1151,8 @@ class AsyncOmniEngine: a queue + coroutine-switch round-trip. The Orchestrator receives a ready-to-submit OmniEngineCoreRequest. """ + graph_gate = getattr(self, "_graph_request_gate", None) + graph_ticket = graph_gate.acquire(request_id) if graph_gate is not None else None try: msg = self._build_add_request_message( request_id=request_id, @@ -1163,6 +1171,8 @@ class AsyncOmniEngine: resumable=resumable, ) except BaseException: + if graph_gate is not None: + graph_gate.release(graph_ticket) if isinstance(prompt, dict): for artifact_dir in prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) or (): if isinstance(artifact_dir, str): @@ -1186,6 +1196,8 @@ class AsyncOmniEngine: self.request_queue.sync_q.put(msg) except BaseException: + if graph_gate is not None: + graph_gate.release(graph_ticket) for artifact_dir in msg.request_artifact_dirs or (): shutil.rmtree(artifact_dir, ignore_errors=True) raise @@ -1554,7 +1566,7 @@ class AsyncOmniEngine: def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: """Read one output message from the Orchestrator output queue.""" try: - return self.output_queue.sync_q.get(timeout=timeout) + return self._attach_graph_release(self.output_queue.sync_q.get(timeout=timeout)) except queue.Empty: if not self.is_alive(): raise RuntimeError("Orchestrator died unexpectedly. See logs above.") @@ -1563,7 +1575,7 @@ class AsyncOmniEngine: async def try_get_output_async(self) -> EngineQueueMessage | None: """Async read from the Orchestrator output queue.""" try: - return self.output_queue.sync_q.get_nowait() + return self._attach_graph_release(self.output_queue.sync_q.get_nowait()) except queue.Empty: if not self.is_alive(): raise RuntimeError("Orchestrator died unexpectedly. See logs above.") @@ -1603,6 +1615,24 @@ class AsyncOmniEngine: msg = await loop.run_in_executor(executor, _drain_get) if msg is None and not self.is_alive(): raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return self._attach_graph_release(msg) + + def _attach_graph_release(self, msg): + from vllm_omni.outputs import OmniRequestOutput + + gate = getattr(self, "_graph_request_gate", None) + if gate is not None and isinstance(msg, OutputMessage) and msg.finished: + ticket = gate.current(msg.request_id) + output = msg.engine_outputs + if isinstance(output, OmniRequestOutput): + previous = output._stage_release + + def release(): + if previous is not None: + previous() + gate.release(ticket) + + output._stage_release = release return msg def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: @@ -1673,6 +1703,10 @@ class AsyncOmniEngine: raise if not result_msg.success: raise RuntimeError(result_msg.error or "abort failed") + graph_gate = getattr(self, "_graph_request_gate", None) + if graph_gate is not None: + for request_id in request_ids: + graph_gate.release(graph_gate.current(request_id)) return list(result_msg.abort_outputs or []) def submit_interaction( diff --git a/vllm_omni/engine/backends/__init__.py b/vllm_omni/engine/backends/__init__.py new file mode 100644 index 000000000..e183ad924 --- /dev/null +++ b/vllm_omni/engine/backends/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Factories for non-vLLM stages owned by the existing StageRuntime.""" + + +def create_graph_client(metadata, config, ledger, reservation): + from .graph import GraphStageClient + + if config.get("backend", {}).get("name") != "external.graph.v1": + raise ValueError("graph stage requires backend.name=external.graph.v1") + return GraphStageClient(metadata, config["backend"], ledger, reservation) diff --git a/vllm_omni/engine/backends/graph.py b/vllm_omni/engine/backends/graph.py new file mode 100644 index 000000000..286ba5257 --- /dev/null +++ b/vllm_omni/engine/backends/graph.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +"""One bounded, non-preemptible graph stage using the existing worker. + +The runtime owns startup and capacity. This adapter owns one in-flight call, +including its output until acknowledged by the consumer. No token scheduler, +provider fallback, or automatic worker replay is introduced here. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import math +import uuid +from pathlib import Path +from typing import Any + +import numpy as np + +from omni_stage_contracts import ArtifactManifest, BufferRef, StageEvent, StageRequest +from vllm_omni.edge.local.external.client import ExternalWorker +from vllm_omni.engine.resource_ledger import ResourceUnavailable +from vllm_omni.engine.stage_client import StageClientBase +from vllm_omni.host.routes import resolve_route +from vllm_omni.outputs import OmniRequestOutput + + +class GraphStageClient(StageClientBase): + def __init__(self, metadata, config, ledger, reservation) -> None: + for name, value in vars(metadata).items(): + setattr(self, name, value) + self.stage_type = "graph" + self._ledger, self._reservation = ledger, reservation + self._generation = uuid.uuid4().hex + self._config = dict(config) + self._worker = None + self._closed = False + self._active: str | None = None + self._task: asyncio.Task | None = None + self._output: OmniRequestOutput | None = None + self._epoch = 0 + self._max_io_bytes = int(config.get("max_io_bytes", 8 << 20)) + if self._max_io_bytes <= 0 or self._max_io_bytes > min(reservation.demands.values()): + raise ValueError("graph I/O limit must be positive and fit the stage reservation") + try: + manifest_path = Path(config["manifest"]).resolve() + manifest = ArtifactManifest.read(manifest_path) + graph_name = manifest.metadata["graph_file"] + examples_name = manifest.metadata["example_inputs_file"] + if graph_name not in manifest.files or examples_name not in manifest.files: + raise ValueError("graph and placement examples must be hashed artifact payloads") + if sum((manifest_path.parent / n).stat().st_size for n in manifest.files) > min( + reservation.demands.values() + ): + raise ResourceUnavailable("artifact files exceed stage memory reservation") + route = resolve_route(config["route"]) + self._worker = ExternalWorker( + route, start_timeout_s=float(config.get("start_timeout_s", 120)), max_payload_bytes=self._max_io_bytes + ) + self._worker.start() + with np.load(manifest_path.parent / examples_name, allow_pickle=False) as example: + inputs = {name: example[name] for name in example.files} + self._describe(inputs) + self.load_report = self._worker.load( + manifest_path.parent / graph_name, + example_inputs=inputs, + ep_dir=config.get("ep_dir"), + intra_op_num_threads=config.get("intra_op_num_threads"), + profile_prefix=config.get("profile_prefix"), + artifact_files=[manifest_path.parent / name for name in manifest.files], + ) + report = self.load_report + minimum = float(config.get("min_fraction_on_target", 1.0)) + if not 0 < minimum <= 1: + raise ValueError("placement threshold must be in (0, 1]") + if ( + report.requested_provider_missing + or report.fraction_on_target is None + or not math.isfinite(report.fraction_on_target) + or not 0 < report.fraction_on_target <= 1 + or report.target_nodes <= 0 + or report.fraction_on_target < minimum + ): + raise RuntimeError(f"REFUSE_EP_PLACEMENT: {report.summary()}") + allowed = set(config.get("allowed_providers", report.session_providers if route.ep == "cpu" else ())) + if not allowed or set(report.node_counts) - allowed: + raise RuntimeError("REFUSE_EP_PLACEMENT: declare all permitted execution providers") + if report.rss_bytes > min(reservation.demands.values()): + raise ResourceUnavailable("loaded worker RSS exceeds stage memory reservation") + self._input_schema = {name: (value.dtype.str, value.shape) for name, value in inputs.items()} + from omni_stage_contracts.types import file_digest + + self.execution_plan = { + "backend": "external.graph.v1", + "stage_id": self.stage_id, + "worker_generation": self._generation, + "artifact_sha256": file_digest(manifest_path), + "component": manifest.component, + "route": route.to_dict(), + "reserved_bytes": dict(reservation.demands), + "max_io_bytes": self._max_io_bytes, + "placement": report.to_dict(), + "worker": {k: v for k, v in self._worker.hello.items() if k != "token"}, + "evidence": "B", + "stateful": False, + } + except BaseException: + self.shutdown() + raise + + def _describe(self, tensors: dict[str, np.ndarray]) -> tuple[BufferRef, ...]: + refs = tuple( + BufferRef( + name, str(self.stage_id), self._generation, array.dtype.name, tuple(array.shape), int(array.nbytes) + ) + for name, array in tensors.items() + ) + if sum(r.nbytes for r in refs) > self._max_io_bytes: + raise ResourceUnavailable("graph payload exceeds admitted I/O bound") + return refs + + async def add_request_async(self, request_id: str, prompt: Any, params: Any = None) -> None: + self.check_health() + if self._active is not None: + raise ResourceUnavailable("graph stage has an unacknowledged request; capacity is one") + if not isinstance(prompt, dict) or not isinstance(prompt.get("tensors"), dict): + raise ValueError("graph prompt must contain a tensors mapping supplied by its model adapter") + inputs = {name: np.asarray(value) for name, value in prompt["tensors"].items()} + refs = self._describe(inputs) + if {name: (value.dtype.str, value.shape) for name, value in inputs.items()} != self._input_schema: + raise ValueError("input names/dtypes/shapes differ from the placement-verified artifact bucket") + # Copy at submission: application mutation cannot change an admitted call. + inputs = {name: value.copy() for name, value in inputs.items()} + self._epoch += 1 + epoch = self._epoch + request = StageRequest(request_id, self.stage_id, epoch, self._generation, inputs=refs) + self._active = request_id + self._task = asyncio.create_task(self._run(request, inputs), name=f"graph-{self.stage_id}-{request_id}") + + async def _run(self, request: StageRequest, inputs: dict[str, np.ndarray]) -> None: + try: + outputs, timing = await asyncio.to_thread(self._worker.run, inputs) + refs = self._describe(outputs) + if sum(x.nbytes for x in outputs.values()) + sum(x.nbytes for x in inputs.values()) > self._max_io_bytes: + raise ResourceUnavailable("combined graph input/output exceeds admitted I/O bound") + event = StageEvent( + request.request_id, self.stage_id, request.epoch, 1, "tensor", self._generation, refs, terminal=True + ) + output = OmniRequestOutput( + request_id=request.request_id, + stage_id=self.stage_id, + final_output_type=self.final_output_type or "latent", + _custom_output={"tensors": outputs, "stage_event": dataclasses.asdict(event)}, + metrics={"graph_timing": timing.to_dict()}, + ) + except asyncio.CancelledError: + raise + except Exception as exc: + output = OmniRequestOutput.from_error(request.request_id, str(exc)) + output.stage_id = self.stage_id + if not self._closed and self._epoch == request.epoch: + loop = asyncio.get_running_loop() + output._stage_release = lambda: loop.call_soon_threadsafe( + self.acknowledge, request.request_id, request.epoch, request.worker_generation + ) + self._output = output + + def get_graph_output_nowait(self): + # Delivery is separate from release: the slot stays charged until ACK. + output, self._output = self._output, None + return output + + def acknowledge(self, request_id: str, epoch: int, generation: str) -> None: + if ( + self._active == request_id + and epoch == self._epoch + and generation == self._generation + and (self._task is None or self._task.done()) + ): + self._active = None + self._task = None + + async def abort_requests_async(self, request_ids: list[str]) -> None: + if self._active not in request_ids: + return + self._epoch += 1 + self._output = None + # A completed call needs no worker kill. Otherwise cancellation retires + # this single-session worker; the runtime will report it unavailable. + if self._task is not None and not self._task.done(): + drained = await asyncio.to_thread(self._worker.terminate) + self._closed = True + done, _ = await asyncio.wait({self._task}, timeout=5) + self._ledger.release(self._reservation, drained=drained and bool(done) and not self._task.cancelled()) + self._active = None + + def check_health(self) -> None: + if self._closed: + from vllm.v1.engine.exceptions import EngineDeadError + + raise EngineDeadError() + + async def collective_rpc_async(self, method, timeout=None, args=(), kwargs=None): + raise NotImplementedError(f"graph backend does not implement collective RPC {method}") + + def shutdown(self) -> None: + self._closed = True + self._output = None + drained = self._worker is None or self._worker.terminate() + task = self._task + if task is not None and not task.done(): + # Process exit does not release Python input/transport references. + # Keep them charged until the call coroutine actually unwinds. + self._ledger.release(self._reservation, drained=False) + task.add_done_callback( + lambda done: self._ledger.release(self._reservation, drained=drained and not done.cancelled()) + ) + else: + self._ledger.release(self._reservation, drained=drained) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 55df201e5..bbf05880c 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -1191,8 +1191,12 @@ class Orchestrator: # Shared catch so a dead replica on either poll path is # evicted rather than tearing down every stage (#4285). try: - if pool.stage_type == "diffusion": - diffusion_output = pool.poll_diffusion_output(replica_id) + if pool.stage_type in ("diffusion", "graph"): + diffusion_output = ( + pool.poll_graph_output(replica_id) + if pool.stage_type == "graph" + else pool.poll_diffusion_output(replica_id) + ) if diffusion_output is None: continue @@ -1282,7 +1286,11 @@ class Orchestrator: got = False for replica_id in pool.available_replica_ids(): try: - output = pool.poll_diffusion_output(replica_id) + output = ( + pool.poll_graph_output(replica_id) + if pool.stage_type == "graph" + else pool.poll_diffusion_output(replica_id) + ) except EngineDeadError as e: await ready_q.put(("error", stage_id, replica_id, e)) continue @@ -1317,7 +1325,7 @@ class Orchestrator: # No live replica yet; nothing to attach to. A later # reconcile picks the stage up once one registers. continue - if stage_type == "diffusion": + if stage_type in ("diffusion", "graph"): live_pools.add(stage_id) existing_poller = pollers.get(stage_id) if existing_poller is None or existing_poller.done(): @@ -1500,6 +1508,8 @@ class Orchestrator: for output in outputs: req_state = self.request_states.get(output.request_id) if req_state is None: + if hasattr(output, "release_stage_buffers"): + output.release_stage_buffers() logger.warning( "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)", output.request_id, @@ -1524,7 +1534,11 @@ class Orchestrator: ) stage_metrics.pipeline_timings = dict(req_state.pipeline_timings) - await self._route_output(stage_id, replica_id, output, req_state, stage_metrics) + try: + await self._route_output(stage_id, replica_id, output, req_state, stage_metrics) + finally: + if pool.stage_type == "graph" and stage_id not in req_state.final_output_stage_ids: + output.release_stage_buffers() async def _handle_stage_error(self, stage_id: int, output: Any) -> None: """Emit a frontend-visible error and clean up request state.""" @@ -1834,6 +1848,10 @@ class Orchestrator: try: if abort: abort_outputs = await self._abort_request_ids(cleanup_ids) + graph_gate = getattr(self, "graph_request_gate", None) + if graph_gate is not None: + for request_id in cleanup_ids: + graph_gate.release(graph_gate.current(request_id)) self._release_request_bindings(cleanup_ids) for request_id in cleanup_ids: self._pd_kv_params.pop(request_id, None) @@ -2513,6 +2531,17 @@ class Orchestrator: requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) _t_submit_start = _time.perf_counter() + if next_pool.stage_type == "graph": + if is_streaming_session: + raise ValueError("external graph requires a complete upstream payload") + transform = next_client.custom_process_input_func + if transform is None: + raise ValueError("graph dependency requires a model-owned custom_process_input_func") + graph_prompt = transform(source_outputs, req_state.prompt, requires_multimodal_data) + await next_pool.submit_initial(req_id, req_state, graph_prompt) + req_state.stage_submit_ts[next_logical] = _time.time() + return + if next_pool.stage_type == "diffusion": # Gate: never dispatch with an incomplete CFG bundle. Checked # non-destructively BEFORE popping — a pop-then-redefer would lose diff --git a/vllm_omni/engine/resource_ledger.py b/vllm_omni/engine/resource_ledger.py new file mode 100644 index 000000000..00b410165 --- /dev/null +++ b/vllm_omni/engine/resource_ledger.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Atomic reservations shared by native and external stage startup. + +Capacities are admission ceilings for this controller, captured before loads. +They are not fresh `available` samples: resident allocations remain charged, +so sampling available memory cannot subtract our own allocation a second time. +Refresh ceilings only with an explicit, reconciled controller budget. +""" + +from __future__ import annotations + +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType + + +class ResourceUnavailable(RuntimeError): # noqa: N818 - public admission result, like PlanNotAdmitted + pass + + +@dataclass(frozen=True) +class Reservation: + owner: str + demands: Mapping[str, int] + + +class ResourceLedger: + def __init__(self, capacities: Mapping[str, int]) -> None: + if not capacities or any(type(n) is not int or n < 0 for n in capacities.values()): + raise ValueError("explicit nonnegative memory ceilings are required") + self.capacities = MappingProxyType(dict(capacities)) + self._active: dict[str, Reservation] = {} + self._quarantined: set[str] = set() + self._lock = threading.RLock() + + def reserve(self, owner: str, demands: Mapping[str, int]) -> Reservation: + with self._lock: + if not owner or owner in self._active: + raise ValueError(f"reservation owner missing or already active: {owner}") + if not demands or any(type(n) is not int or n < 0 for n in demands.values()): + raise ValueError("nonnegative memory demands are required") + for pool, count in demands.items(): + if pool not in self.capacities: + raise ResourceUnavailable(f"unknown memory pool: {pool}") + used = sum(r.demands.get(pool, 0) for r in self._active.values()) + if used + count > self.capacities[pool]: + raise ResourceUnavailable(f"{pool}: {used} reserved + {count} requested > {self.capacities[pool]}") + reservation = Reservation(owner, MappingProxyType(dict(demands))) + self._active[owner] = reservation + return reservation + + def release(self, reservation: Reservation, *, drained: bool) -> bool: + with self._lock: + if self._active.get(reservation.owner) is not reservation: + return False # stale token must never release a newer allocation + if not drained: + self._quarantined.add(reservation.owner) + return False + del self._active[reservation.owner] + self._quarantined.discard(reservation.owner) + return True + + def snapshot(self) -> dict: + with self._lock: + return { + "capacities": dict(self.capacities), + "reserved": {p: sum(r.demands.get(p, 0) for r in self._active.values()) for p in self.capacities}, + "owners": list(self._active), + "quarantined": sorted(self._quarantined), + } + + +class GraphRequestGate: + """One admitted pipeline request, held through the final consumer ACK. + + Stage reservations already include its payload allowance. Rejecting excess + work before the ingress queue keeps slow consumers from accumulating inputs + or error outputs, and leaves the existing control queue available to abort. + """ + + def __init__(self): + self._lock = threading.Lock() + self._active = None + + def acquire(self, request_id): + with self._lock: + if self._active is not None: + raise ResourceUnavailable("graph pipeline capacity is one; consume or cancel the active request first") + ticket = (request_id, object()) + self._active = ticket + return ticket + + def current(self, request_id): + with self._lock: + return self._active if self._active is not None and self._active[0] == request_id else None + + def release(self, ticket): + with self._lock: + if ticket is not None and self._active is ticket: + self._active = None diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 2a3c26177..a9a7860ac 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -71,6 +71,14 @@ class StagePoolClient(StageClient, Protocol): ) -> Any: ... +class StagePoolGraphClient(StagePoolClient, Protocol): + async def add_request_async(self, request_id: str, prompt: Any, params: Any = None) -> None: ... + + def get_graph_output_nowait(self) -> Any | None: ... + + def acknowledge(self, request_id: str, epoch: int, generation: str) -> None: ... + + class StagePoolLLMClient(StagePoolClient, Protocol): """Pool-facing API for LLM-style stages.""" diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 1a4838387..15b3ab00e 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -12,12 +12,6 @@ out of StageEngineCoreClient into reusable functions. from __future__ import annotations import copy -# [edge-infer W1] fcntl does not exist on Windows; this module is in the -# import closure of AsyncOmni, so the bare import makes the engine -# unimportable before anything can be configured. The locking logic is -# sound and torch-free -- only the primitive changes. The shim keeps the -# BlockingIOError contract this file's stale-lock path depends on. -from vllm_omni._filelock_compat import flock_exclusive_nb, funlock import importlib import json import multiprocessing as mp @@ -41,6 +35,12 @@ from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.executor import Executor +# [edge-infer W1] fcntl does not exist on Windows; this module is in the +# import closure of AsyncOmni, so the bare import makes the engine +# unimportable before anything can be configured. The locking logic is +# sound and torch-free -- only the primitive changes. The shim keeps the +# BlockingIOError contract this file's stale-lock path depends on. +from vllm_omni._filelock_compat import flock_exclusive_nb, funlock from vllm_omni.config.omni_config import ( _CACHE_STAGE_ENGINE_FIELD_MAP, _DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP, @@ -549,7 +549,7 @@ class StageMetadata: """Lightweight stage attributes extracted from stage_config.""" stage_id: int - stage_type: Literal["llm", "diffusion"] + stage_type: Literal["llm", "diffusion", "graph"] engine_output_type: str | None is_comprehension: bool requires_multimodal_data: bool @@ -599,7 +599,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: engine-argument and stage-init consumers together. """ stage_id: int = stage_config.stage_id - stage_type: Literal["llm", "diffusion"] = _get_attr_or_item(stage_config, "stage_type", "llm") + stage_type: Literal["llm", "diffusion", "graph"] = _get_attr_or_item(stage_config, "stage_type", "llm") engine_args = stage_config.engine_args _apply_rocm_attention_backend(engine_args, stage_type) @@ -711,8 +711,11 @@ def extract_stage_metadata_from_omni_stage_config( layout, engine-argument, remote-diffusion, and platform setup paths still require the legacy StageConfig/OmegaConf shape. """ - stage_type: Literal["llm", "diffusion"] = "diffusion" if stage_config.stage_type == StageType.DIFFUSION else "llm" - sampling_params_cls = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams + stage_type: Literal["llm", "diffusion", "graph"] = { + StageType.DIFFUSION: "diffusion", + StageType.GRAPH: "graph", + }.get(stage_config.stage_type, "llm") + sampling_params_cls = OmniDiffusionSamplingParams if stage_type == "diffusion" else SamplingParams sampling_params: OmniSamplingParams = sampling_params_cls( **(stage_config.model_config.default_sampling_params or {}) ) @@ -738,7 +741,7 @@ def extract_stage_metadata_from_omni_stage_config( return StageMetadata( stage_id=stage_config.stage_id, - stage_type="llm", + stage_type=stage_type, engine_output_type=stage_config.engine_output_type, is_comprehension=stage_config.is_comprehension, requires_multimodal_data=stage_config.requires_multimodal_data, @@ -1574,7 +1577,7 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer) -def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str: +def device_init_lock_path(device_id: int, lock_dir: str | None = None) -> str: """Return the per-physical-device initialization lock file path. Shared by the orchestrator-side ``acquire_device_locks`` (legacy full-init @@ -1585,7 +1588,9 @@ def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str: unlink this path: two holders of the same pathname on different inodes do not conflict. The PID written into the file is diagnostic only. """ - return os.path.join(lock_dir, f"vllm_omni_device_{device_id}_init.lock") + from vllm_omni.host.paths import device_lock_directory + + return os.path.join(lock_dir or device_lock_directory(), f"vllm_omni_device_{device_id}_init.lock") def _open_existing_lock_file(lock_file: str) -> tuple[int, bool] | None: @@ -1637,6 +1642,12 @@ def open_device_lock_file(lock_file: str) -> tuple[int, bool]: very failure this avoids. ``link`` also fails cleanly if another process wins the race, which keeps the "first creator wins" inode stable. """ + if os.name == "nt": + # Windows byte-range locks require a writable handle. O_CREAT without + # O_TRUNC atomically opens/creates the stable file; there is no fchmod + # or POSIX cross-user mode publication on this host. + return os.open(lock_file, os.O_RDWR | os.O_CREAT, 0o600), True + existing = _open_existing_lock_file(lock_file) if existing is not None: return existing diff --git a/vllm_omni/engine/stage_phase_lock.py b/vllm_omni/engine/stage_phase_lock.py index 44493f649..7b5365bf4 100644 --- a/vllm_omni/engine/stage_phase_lock.py +++ b/vllm_omni/engine/stage_phase_lock.py @@ -76,6 +76,7 @@ def _unlock(fd: int) -> None: funlock(fd) + logger = init_logger(__name__) # Default upper bound on how long a single phase will wait for its device locks @@ -144,7 +145,7 @@ class DevicePhaseLock: device_ids: list[int], *, timeout_s: float = _DEFAULT_LOCK_TIMEOUT_S, - lock_dir: str = "/tmp", + lock_dir: str | None = None, ) -> None: self._device_ids = sorted(device_ids) self._timeout_s = timeout_s diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 0ab6c1da7..5a24a5fb3 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -990,6 +990,14 @@ class StagePool: if self.stage_type == "diffusion": params = OmniDiffusionSamplingParams.from_params(params) submit_kwargs = dict(submit_kwargs or {}) + if self.stage_type == "graph": + replica_id = await self._pick_or_select(request_id, affinity_request_id=affinity_request_id) + try: + await self.clients[replica_id].add_request_async(request_id, request, params) + except BaseException: + self.release_binding(request_id) + raise + return replica_id if self.stage_type == "diffusion": if isinstance(request, list): raise ValueError( @@ -1050,6 +1058,8 @@ class StagePool: prompt_text: Any = None, ) -> int: """Submit a streaming update to an already admitted request.""" + if self.stage_type == "graph": + raise ValueError("external.graph.v1 accepts complete graphs, not streaming updates") params = req_state.sampling_params_list[self.stage_id] if self.stage_type == "diffusion": params = OmniDiffusionSamplingParams.from_params(params) @@ -1207,6 +1217,12 @@ class StagePool: return None return cast(StagePoolDiffusionClient, raw_client).get_diffusion_output_nowait() + def poll_graph_output(self, replica_id: int) -> Any | None: + if not self.is_replica_available(replica_id): + return None + client = self.clients[replica_id] + return None if client is None else client.get_graph_output_nowait() + # ---- Stage-local control plane ---- async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index ae72f0870..549a29818 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -155,6 +155,9 @@ class StageRuntime: # device groups can initialize concurrently. self._replica_launch_lock = threading.Lock() self._init_visible_devices_baseline: str | None = None + self.resource_ledger = None + self._resource_reservations = {} + self._resource_started = set() @staticmethod def _client_addresses_from_zmq(addresses: Any) -> dict[str, str]: @@ -221,18 +224,21 @@ class StageRuntime: return collected @staticmethod - def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None: + def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> set[int]: """Best-effort shutdown for attached clients after init failure.""" + drained = set() for client in reversed(list(clients)): if client is None: continue try: client.shutdown() + drained.add(id(client)) except Exception as cleanup_error: logger.warning( "[StageRuntime] Failed to shutdown initialized client after init failure: %s", cleanup_error, ) + return drained def initialize(self) -> None: """Run the full stage initialization sequence.""" @@ -241,6 +247,7 @@ class StageRuntime: plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans } try: + self._reserve_stage_resources(stage_plans) self._before_initialize_stage_replicas(stage_plans) initialized_clients = self._initialize_stage_replicas(stage_plans, self._stage_init_timeout) initialized_clients_by_stage = initialized_clients @@ -259,16 +266,80 @@ class StageRuntime: "[StageRuntime] Stage initialization failed; shutting down %s initialized client(s)", len(cleanup_clients), ) - self._shutdown_initialized_clients(cleanup_clients) + drained = self._shutdown_initialized_clients(cleanup_clients) or set() + if self.resource_ledger is not None: + for client in cleanup_clients: + if client.stage_type != "graph": + reservation = self._resource_reservations.get((client.stage_id, client.replica_id)) + if reservation is not None: + self.resource_ledger.release(reservation, drained=id(client) in drained) self._cleanup_after_initialize_failure() + self._release_unstarted_resources(initialized_clients_by_stage) raise exc + def _reserve_stage_resources(self, stage_plans) -> None: + """Explicit controller ceilings are mandatory for mixed graph plans. + + Native stage budgets join the same transaction; absent estimates fail + closed rather than treating native allocations as free memory. + """ + from vllm_omni.engine.resource_ledger import ResourceLedger + + replicas = [r for plan in stage_plans for r in plan.replicas] + if not any(r.metadata.stage_type == "graph" for r in replicas): + return + # One final consumer owns the ingress ticket. Fan-out or multiple final + # outputs would need reference-counted leases before any early ACK can + # release that ticket. Refuse those topologies in the v1 graph backend. + previous_stage = None + for index, plan in enumerate(stage_plans): + metadata = plan.replicas[0].metadata + expected_sources = [] if previous_stage is None else [previous_stage] + if ( + list(metadata.engine_input_source or []) != expected_sources + or metadata.final_output != (index == len(stage_plans) - 1) + ): + raise ValueError("graph v1 requires a linear chain with exactly one final output") + previous_stage = metadata.stage_id + budgets = [r.stage_cfg.engine_args.get("resource_budget") for r in replicas] + if any(not b for b in budgets): + raise ValueError("every stage in a graph pipeline must declare resource_budget capacities and demands") + capacities = dict(budgets[0]["capacities"]) + if any(dict(b["capacities"]) != capacities for b in budgets): + raise ValueError("all stages must agree on physical memory ceilings") + self.resource_ledger = ResourceLedger(capacities) + try: + for replica, budget in zip(replicas, budgets): + key = (replica.metadata.stage_id, replica.replica_id) + self._resource_reservations[key] = self.resource_ledger.reserve(str(key), dict(budget["demands"])) + except BaseException: + for reservation in self._resource_reservations.values(): + self.resource_ledger.release(reservation, drained=True) + self._resource_reservations.clear() + raise + + def _release_unstarted_resources(self, initialized) -> None: + if self.resource_ledger is None: + return + for (stage_id, replica_id), reservation in self._resource_reservations.items(): + clients = initialized.get(stage_id, []) + if replica_id >= len(clients) or clients[replica_id] is None: + # Graph factories release or quarantine their own failed load. + if reservation.owner not in self.resource_ledger.snapshot()["quarantined"]: + self.resource_ledger.release( + reservation, drained=(stage_id, replica_id) not in self._resource_started + ) + def shutdown(self) -> None: for pool in self.stage_pools: for client in pool.clients: if client is not None and hasattr(client, "shutdown"): try: client.shutdown() + if self.resource_ledger is not None and client.stage_type != "graph": + reservation = self._resource_reservations.get((client.stage_id, client.replica_id)) + if reservation is not None: + self.resource_ledger.release(reservation, drained=True) except Exception: logger.warning("[StageRuntime] client shutdown failed", exc_info=True) if self._stage_init_executor is not None: @@ -374,7 +445,7 @@ class StageRuntime: stage_vllm_config = None executor_class = None engine_args_dict = None - if base_metadata.stage_type != "diffusion": + if base_metadata.stage_type not in ("diffusion", "graph"): # The stable adapter entry point still receives the same # legacy stage object as replica planning. Its implementation # switches only at the coordinated RFC #4021 cutover. @@ -384,6 +455,7 @@ class StageRuntime: stage_connector_spec=stage_connector_spec, cli_tokenizer=self._tokenizer, ) + engine_args_dict.pop("resource_budget", None) inject_omni_kv_connector_config( engine_args_dict, omni_kv_connector, @@ -597,8 +669,7 @@ class StageRuntime: if local and all(self._cpu_placed(r) for r in local): import psutil - self._run_host_admission(stage_plans, check_host_admission, - lambda: int(psutil.virtual_memory().available)) + self._run_host_admission(stage_plans, check_host_admission, lambda: int(psutil.virtual_memory().available)) return def _resolve(replica: ReplicaInitPlan) -> list[int] | AdmissionExempt | None: @@ -718,6 +789,8 @@ class StageRuntime: def _init_group_key_override(self, replica: ReplicaInitPlan) -> str | None: """Key that bypasses device-overlap grouping, or ``None`` to use it.""" + if replica.metadata.stage_type == "graph": + return "inline:graph" if replica.launch_mode == "local" and replica.metadata.stage_type == "diffusion": # Local diffusion process spawning must stay on the orchestrator # thread. Keep all local diffusion replicas in one sequential group. @@ -736,6 +809,18 @@ class StageRuntime: plan: ReplicaInitPlan, stage_init_timeout: int, ) -> StagePoolClient: + if plan.metadata.stage_type == "graph": + if type(self) is not StageRuntime or plan.launch_mode != "local": + raise ValueError("graph backends support one local device only") + from vllm_omni.engine.backends import create_graph_client + + return create_graph_client( + plan.metadata, + plan.stage_cfg.engine_args, + self.resource_ledger, + self._resource_reservations[(plan.metadata.stage_id, plan.replica_id)], + ) + self._resource_started.add((plan.metadata.stage_id, plan.replica_id)) if plan.launch_mode == "remote": return self._initialize_remote_replica(plan, stage_init_timeout) if plan.metadata.stage_type == "diffusion": @@ -985,7 +1070,7 @@ class StageRuntime: clients: list[StagePoolClient] = [client for client in replica_clients if client is not None] stage_vllm_config = None output_processor = None - if plan.replicas[0].metadata.stage_type != "diffusion": + if plan.replicas[0].metadata.stage_type not in ("diffusion", "graph"): stage_vllm_config = plan.replicas[0].stage_vllm_config if stage_vllm_config is None: raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config") diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 0012039e0..595c19612 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -918,7 +918,10 @@ class AsyncOmni(EngineClient, OmniBase): stage_id, getattr(output_to_yield, "final_output_type", None), ) - yield output_to_yield + try: + yield output_to_yield + finally: + output_to_yield.release_stage_buffers() # The Orchestrator sets "finished" when the final stage is done if result.finished: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 28bbe9219..4d7526daa 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -40,7 +40,7 @@ class Omni(OmniBase): sp = copy.deepcopy(params) stage_meta = self.engine.get_stage_metadata(stage_id) if ( - stage_meta.stage_type != "diffusion" + stage_meta.stage_type not in ("diffusion", "graph") and hasattr(sp, "output_kind") and sp.output_kind != RequestOutputKind.DELTA ): @@ -191,7 +191,10 @@ class Omni(OmniBase): final_stage_id_for_e2e=req_final_stage_ids[req_id], ) if output_to_yield is not None: - yield output_to_yield + try: + yield output_to_yield + finally: + output_to_yield.release_stage_buffers() if isinstance(msg, OutputMessage) and msg.finished: active_reqs.discard(req_id) diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 281cbc55c..16d416988 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -191,6 +191,10 @@ def shm_write_bytes(payload: bytes, name: str | None = None) -> dict[str, Any]: Caller should close the segment; the receiver should unlink. """ + if os.name == "nt": + from vllm_omni.host.shared_memory import owner + + return owner.write(payload, name) try: shm = _shm.SharedMemory(create=True, size=len(payload), name=name) except FileExistsError: @@ -218,6 +222,10 @@ def shm_write_bytes(payload: bytes, name: str | None = None) -> dict[str, Any]: def shm_read_bytes(meta: dict[str, Any]) -> bytes: """Read bytes from SharedMemory by meta {name,size} and cleanup.""" + if os.name == "nt": + from vllm_omni.host.shared_memory import read_buffer + + return read_buffer(meta) shm = _shm.SharedMemory(name=meta["name"]) # type: ignore[index] mv = memoryview(shm.buf) data = bytes(mv[: meta["size"]]) diff --git a/vllm_omni/host/__init__.py b/vllm_omni/host/__init__.py new file mode 100644 index 000000000..8b5255b87 --- /dev/null +++ b/vllm_omni/host/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host mechanics, independent of accelerator and model selection.""" diff --git a/vllm_omni/host/paths.py b/vllm_omni/host/paths.py new file mode 100644 index 000000000..8d96da4bb --- /dev/null +++ b/vllm_omni/host/paths.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-local paths, independent of the checkout and model location.""" + +import os +import tempfile + + +def device_lock_directory() -> str: + # A rooted POSIX path resolves against the current UNC share on Windows. + # Opening /tmp from a WSL checkout can block inside the filesystem driver. + # All same-user controllers/workers use the same native temporary directory. + return tempfile.gettempdir() if os.name == "nt" else "/tmp" diff --git a/vllm_omni/host/process.py b/vllm_omni/host/process.py new file mode 100644 index 000000000..b09690cd3 --- /dev/null +++ b/vllm_omni/host/process.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Portable worker ownership; termination reports whether resources drained.""" + +from __future__ import annotations + +import os +import signal +import subprocess + + +def spawn_worker(argv: list[str], **kwargs) -> subprocess.Popen: + if os.name == "nt": + kwargs.setdefault("creationflags", subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP) + else: + kwargs.setdefault("start_new_session", True) + return subprocess.Popen(argv, **kwargs) + + +def terminate_worker(process: subprocess.Popen, timeout: float = 5.0) -> bool: + """Kill descendants and parent, retaining the caller's resource charge on failure.""" + import psutil + + try: + parent = psutil.Process(process.pid) + children = parent.children(recursive=True) + except psutil.NoSuchProcess: + process.wait(timeout=timeout) + if os.name == "nt": + return False # absent parent cannot prove that descendants exited + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + return True + except OSError: + return False + return False # an orphaned group was found; retain its charge until verified + except psutil.Error: + return False + for child in reversed(children): + try: + child.kill() + except psutil.NoSuchProcess: + pass + except psutil.Error: + return False + try: + if os.name == "nt": + process.kill() + else: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=timeout) + _, alive = psutil.wait_procs(children, timeout=timeout) + return not alive + except (OSError, subprocess.TimeoutExpired): + return False diff --git a/vllm_omni/host/routes.py b/vllm_omni/host/routes.py new file mode 100644 index 000000000..fdf2a7751 --- /dev/null +++ b/vllm_omni/host/routes.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deployment-owned worker routes, independent of model and physical device.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path + +from vllm_omni.edge.local.external.launch import Route, is_wsl, resolve + + +def resolve_route(spec: str | Mapping) -> Route: + """Names retain compatibility; explicit profiles have no machine-local defaults.""" + if isinstance(spec, str): + return resolve(spec) + required = {"name", "interpreter", "worker", "ep", "os_domain"} + if not required <= set(spec) or set(spec) - required: + raise ValueError(f"explicit route requires exactly {sorted(required)}") + if spec["os_domain"] not in {"windows", "posix"}: + raise ValueError("route os_domain must be windows or posix") + windows = spec["os_domain"] == "windows" + interpreter = Path(os.path.expandvars(str(spec["interpreter"]))).expanduser() + worker = Path(os.path.expandvars(str(spec["worker"]))).expanduser() + reason = "" + if windows and os.name != "nt" and not is_wsl(): + reason = "Windows route requires Windows or WSL interop on this machine" + elif not windows and os.name == "nt": + reason = "POSIX route needs a POSIX controller; no implicit remote/WSL launch" + elif not interpreter.is_file() or not worker.is_file(): + reason = "configured interpreter or worker script is missing" + return Route(str(spec["name"]), str(interpreter), str(worker), str(spec["ep"]), windows, not reason, reason) diff --git a/vllm_omni/host/shared_memory.py b/vllm_omni/host/shared_memory.py new file mode 100644 index 000000000..c61dc12f5 --- /dev/null +++ b/vllm_omni/host/shared_memory.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bounded, single-consumer shared buffers with explicit Windows ownership. + +Windows destroys a named mapping when its last handle closes. Keep a producer +handle until a generation-specific acknowledgement arrives from the reader. +The header also records payload length: opening a Windows mapping can report a +page-rounded allocation size. POSIX's existing unlink-based path is unchanged. +""" + +from __future__ import annotations + +import hashlib +import struct +import threading +import uuid +from multiprocessing import shared_memory +from pathlib import Path + +_HEADER = struct.Struct("<8s16sQ") +_MAGIC = b"OMNISHM1" + + +def _ack_path(name: str, generation: bytes) -> Path: + from vllm_omni.windows.paths import shm_dir + + identity = hashlib.sha256(name.encode() + generation).hexdigest() + return Path(shm_dir()) / f"ack_{identity}" + + +class SharedBufferOwner: + def __init__(self, max_bytes: int = 64 << 20, max_buffers: int = 256): + self.max_bytes, self.max_buffers = max_bytes, max_buffers + self._owned = {} + self._lock = threading.RLock() + self._collector = None + + @property + def retained_bytes(self) -> int: + with self._lock: + return sum(size for _, _, size in self._owned.values()) + + @property + def retained_names(self) -> frozenset[str]: + with self._lock: + return frozenset(self._owned) + + def collect(self) -> None: + with self._lock: + for name, (_, generation, _) in list(self._owned.items()): + if _ack_path(name, generation).exists(): + self.release(name) + + def release(self, name: str) -> None: + with self._lock: + item = self._owned.get(name) + if item is not None: + segment, generation, _ = item + segment.close() + segment.unlink() # no-op on Windows; useful for contract tests + self._owned.pop(name) + _ack_path(name, generation).unlink(missing_ok=True) + + def _collect_loop(self) -> None: + event = threading.Event() + while not event.wait(0.05): + try: + self.collect() + except OSError: + # Keep ownership charged if acknowledgement cleanup fails. + continue + + def write(self, payload: bytes, name: str | None = None) -> dict: + with self._lock: + self.collect() + size = _HEADER.size + len(payload) + if len(self._owned) >= self.max_buffers or self.retained_bytes + size > self.max_bytes: + raise BufferError("Windows shared-buffer ownership limit exceeded") + generation = uuid.uuid4().bytes + segment = shared_memory.SharedMemory(create=True, size=size, name=name) + try: + segment.buf[: _HEADER.size] = _HEADER.pack(_MAGIC, generation, len(payload)) + segment.buf[_HEADER.size : size] = payload + except BaseException: + segment.close() + segment.unlink() + raise + self._owned[segment.name] = (segment, generation, size) + if self._collector is None: + self._collector = threading.Thread(target=self._collect_loop, daemon=True, name="omni-shm-owner") + self._collector.start() + return {"name": segment.name, "size": len(payload), "generation": generation.hex()} + + +def read_buffer(meta: dict) -> bytes: + segment = shared_memory.SharedMemory(name=meta["name"]) + try: + magic, generation, size = _HEADER.unpack(segment.buf[: _HEADER.size]) + if magic != _MAGIC or size > segment.size - _HEADER.size: + raise ValueError("invalid Windows shared-buffer header") + if meta.get("generation", generation.hex()) != generation.hex(): + raise ValueError("stale Windows shared-buffer generation") + payload = bytes(segment.buf[_HEADER.size : _HEADER.size + size]) + finally: + segment.close() + _ack_path(meta["name"], generation).touch() + return payload + + +owner = SharedBufferOwner() diff --git a/vllm_omni/host/windows_process.py b/vllm_omni/host/windows_process.py new file mode 100644 index 000000000..dc6e61054 --- /dev/null +++ b/vllm_omni/host/windows_process.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Standalone native-Windows process identity and bounded tree retirement. + +This file is executed by the worker interpreter from WSL; it imports no engine +or third-party packages. A launcher PID is never treated as a Windows PID. +""" + +from __future__ import annotations + +import ctypes +import json +import os +import subprocess +import sys +from ctypes import wintypes + + +def _api(): + kernel = ctypes.WinDLL("kernel32", use_last_error=True) + kernel.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel.OpenProcess.restype = wintypes.HANDLE + kernel.CloseHandle.argtypes = [wintypes.HANDLE] + kernel.GetProcessTimes.argtypes = [wintypes.HANDLE] + [ctypes.POINTER(wintypes.FILETIME)] * 4 + kernel.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] + return kernel + + +def _created(kernel, handle) -> int: + times = [wintypes.FILETIME() for _ in range(4)] + if not kernel.GetProcessTimes(handle, *(ctypes.byref(t) for t in times)): + raise ctypes.WinError(ctypes.get_last_error()) + return (times[0].dwHighDateTime << 32) | times[0].dwLowDateTime + + +def identity() -> dict: + kernel = _api() + handle = kernel.OpenProcess(0x1000, False, os.getpid()) # QUERY_LIMITED_INFORMATION + if not handle: + raise ctypes.WinError(ctypes.get_last_error()) + try: + return {"pid": os.getpid(), "created_filetime": _created(kernel, handle)} + finally: + kernel.CloseHandle(handle) + + +def retire(pid: int, created_filetime: int) -> bool: + kernel = _api() + handle = kernel.OpenProcess(0x1000 | 0x100000, False, pid) # query + synchronize + if not handle: + return False # parent absence alone says nothing about its descendants + try: + if _created(kernel, handle) != created_filetime: + return False # stale identity must not kill a reused PID + result = subprocess.run( + ["taskkill.exe", "/PID", str(pid), "/T", "/F"], + capture_output=True, + timeout=10, + creationflags=subprocess.CREATE_NO_WINDOW, + ) + return result.returncode == 0 and kernel.WaitForSingleObject(handle, 5000) == 0 + finally: + kernel.CloseHandle(handle) + + +if __name__ == "__main__": + try: + print(json.dumps({"drained": retire(int(sys.argv[1]), int(sys.argv[2]))})) + except Exception as exc: + print(json.dumps({"drained": False, "error": str(exc)})) + sys.exit(1) diff --git a/vllm_omni/host/worker_bootstrap.py b/vllm_omni/host/worker_bootstrap.py new file mode 100644 index 000000000..bc0d1041b --- /dev/null +++ b/vllm_omni/host/worker_bootstrap.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Standalone worker bootstrap: record OS identity without importing Omni.""" + +import json +import os +import runpy +import sys +from pathlib import Path + +if os.name == "nt": + from windows_process import identity + + os.environ["VLLM_OMNI_WORKER_IDENTITY"] = json.dumps(identity()) +else: + os.environ.pop("VLLM_OMNI_WORKER_IDENTITY", None) + +worker = sys.argv.pop(1) +sys.argv[0] = worker +sys.path.insert(0, str(Path(worker).resolve().parent)) +runpy.run_path(worker, run_name="__main__") diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index e74931a20..6c712caba 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -96,6 +96,7 @@ _REQUEST_OUTPUT_CONTENT_ATTRS = ( # Omni-specific content copied when the wrapped stage output is itself an # OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline). _OMNI_CONTENT_ATTRS = ( + "_stage_release", "images", "latents", "trajectory_latents", @@ -170,6 +171,13 @@ class OmniRequestOutput(RequestOutput): trajectory_decoded: list | None = None _multimodal_output: dict[str, Any] = field(default_factory=dict) _custom_output: dict[str, Any] = field(default_factory=dict) + _stage_release: Any = field(default=None, repr=False, compare=False) + + def release_stage_buffers(self) -> None: + """Acknowledge consumption of an externally owned stage payload.""" + release, self._stage_release = self._stage_release, None + if release is not None: + release() # profiling data stage_durations: dict[str, float] = field(default_factory=dict) diff --git a/vllm_omni/windows/paths.py b/vllm_omni/windows/paths.py index 75b16383f..0200edaf8 100644 --- a/vllm_omni/windows/paths.py +++ b/vllm_omni/windows/paths.py @@ -20,7 +20,7 @@ def shm_dir() -> str: use. Callers that used to write ``f"/dev/shm/..."`` go through this so the path is the same on every platform that has one and a real one elsewhere. """ - if os.path.isdir(_POSIX_SHM): + if sys.platform != "win32" and os.path.isdir(_POSIX_SHM): return _POSIX_SHM base = os.environ.get("LOCALAPPDATA") or tempfile.gettempdir() path = Path(base) / "vllm_omni" / "shm" diff --git a/vllm_omni/windows/shims.py b/vllm_omni/windows/shims.py index a60b4055b..e1d01d811 100644 --- a/vllm_omni/windows/shims.py +++ b/vllm_omni/windows/shims.py @@ -18,6 +18,7 @@ built on injectable primitives so the logic is testable on Linux. from __future__ import annotations import asyncio +import errno import os import sys import time @@ -88,8 +89,10 @@ class FcntlShim: self._at_lock_offset(fd, lambda: self._locking(fd, self._nblck, 1)) return except OSError as exc: + if exc.errno not in (errno.EACCES, errno.EAGAIN, errno.EDEADLK): + raise if non_blocking: - raise BlockingIOError(str(exc)) from exc + raise BlockingIOError(exc.errno, str(exc)) from exc self._sleep(self._poll_interval) def lockf(self, fd: int, cmd: int, len: int = 0, start: int = 0, whence: int = 0) -> None: # noqa: A002 diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index c1325f101..27e2ad1e4 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -2035,7 +2035,11 @@ class OmniGPUModelRunner(GPUModelRunner): accepted: frozenset[str] | None try: - parameters = inspect.signature(type(self.model).forward).parameters + # CUDA graph wrappers delegate attributes to their wrapped model; + # their class does not define ``forward``. Inspect the resolved + # callable so eager and captured native models use the same filter. + forward = getattr(self.model, "forward", self.model) + parameters = inspect.signature(forward).parameters except (TypeError, ValueError): accepted = None # not introspectable; assume it copes, as before else: