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.
116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""A stage worker that speaks the protocol and never touches a device.
|
|
|
|
Selected with ``VLLM_OMNI_EXTERNAL_WORKER_ORT_CPU``, so the client, the
|
|
framing, the placement gate and every refusal path are exercised on a machine
|
|
with no AMD hardware, no onnxruntime and no Windows. What it replies is read
|
|
from the JSON file named by ``VLLM_OMNI_FAKE_WORKER_JSON``:
|
|
|
|
``load`` body merged into the load reply -- set ``node_counts`` and
|
|
``fraction_on_target`` to stage a declined graph.
|
|
``run`` ``{"scale": float}``; outputs are the inputs times it.
|
|
``die_on`` op name after which the worker exits without replying, for the
|
|
"worker vanished mid-request" path.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
_HERE = Path(__file__).resolve()
|
|
_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)
|
|
sys.modules["_fake_protocol"] = proto
|
|
spec.loader.exec_module(proto)
|
|
|
|
|
|
def main() -> int:
|
|
host, port, token = sys.argv[2], int(sys.argv[4]), sys.argv[6]
|
|
config = {}
|
|
config_path = os.environ.get("VLLM_OMNI_FAKE_WORKER_JSON")
|
|
if config_path and Path(config_path).is_file():
|
|
config = json.loads(Path(config_path).read_text())
|
|
|
|
import socket
|
|
|
|
sock = socket.create_connection((host, port), timeout=30)
|
|
proto.send_message(
|
|
sock,
|
|
proto.OP_HELLO,
|
|
{
|
|
"token": config.get("token", token),
|
|
"executable": sys.executable,
|
|
"python": sys.version,
|
|
"platform": sys.platform,
|
|
"pid": os.getpid(),
|
|
"onnxruntime": "fake",
|
|
"numpy": np.__version__,
|
|
"available_providers": ["CPUExecutionProvider"],
|
|
"rss_bytes": 1024,
|
|
},
|
|
)
|
|
|
|
while True:
|
|
op, body, tensors = proto.recv_message(sock)
|
|
if config.get("die_on") == op:
|
|
sock.close()
|
|
return 0
|
|
if op == proto.OP_LOAD:
|
|
reply = {
|
|
"ep": body.get("ep", "cpu"),
|
|
"session_providers": ["CPUExecutionProvider"],
|
|
"available_providers": ["CPUExecutionProvider"],
|
|
"node_counts": {"CPUExecutionProvider": 4},
|
|
"total_nodes": 4,
|
|
"target_nodes": 4,
|
|
"fraction_on_target": 1.0,
|
|
"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()],
|
|
"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,
|
|
proto.OP_OK,
|
|
{"run_s": 0.001},
|
|
{f"{k}_out": (v * scale) for k, v in tensors.items()},
|
|
)
|
|
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": {},
|
|
},
|
|
)
|
|
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"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|