tzhouam
723c550209
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.
527 lines
20 KiB
Python
527 lines
20 KiB
Python
"""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()
|