383 lines
20 KiB
Python
383 lines
20 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Bounded whole-policy InternVLA stage on CPU, CUDA or CPU plus Radeon Cosmos.
|
|
|
|
The existing Omni diffusion policy owns action semantics. StageRuntime owns
|
|
shared-RAM admission, one in-flight request, terminal output and cancellation.
|
|
The isolated worker is retired on an in-flight abort; no state is replayed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import dataclasses
|
|
import io
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
from omni_stage_contracts import BufferRef, StageEvent, StageRequest
|
|
|
|
from vllm_omni.engine.resource_ledger import ResourceUnavailable
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
from .crisp_tts import CrispTTSStageClient, _local_port, _sha256
|
|
|
|
|
|
def _post_actions(url: str, body: bytes, *, timeout: float, max_bytes: int) -> bytes:
|
|
request = urllib.request.Request(
|
|
url, data=body, headers={"Content-Type": "application/octet-stream"}, method="POST"
|
|
)
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
if response.status != 200:
|
|
raise RuntimeError(f"InternVLA worker returned HTTP {response.status}")
|
|
payload = response.read(max_bytes + 1)
|
|
if len(payload) > max_bytes:
|
|
raise ResourceUnavailable("InternVLA action response exceeds admitted bound")
|
|
return payload
|
|
|
|
|
|
class InternVLAStageClient(CrispTTSStageClient):
|
|
"""One complete action request under Omni's graph-stage control path."""
|
|
|
|
def __init__(self, metadata, config: dict, 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._proc: subprocess.Popen | None = None
|
|
self._log_stream = None
|
|
self._closed = False
|
|
self._active: str | None = None
|
|
self._task: asyncio.Task | None = None
|
|
self._output: OmniRequestOutput | None = None
|
|
self._epoch = 0
|
|
self._memory_pool = str(config.get("memory_pool", "host_ram"))
|
|
self._max_input_bytes = int(config.get("max_input_bytes", 8 << 20))
|
|
self._max_action_bytes = int(config.get("max_action_bytes", 1 << 20))
|
|
self._request_timeout_s = float(config.get("request_timeout_s", 60))
|
|
self._port = int(config.get("port") or _local_port())
|
|
self._placement = str(config.get("placement", ""))
|
|
self._cuda_demand_bytes = reservation.demands.get("cuda:0", 0)
|
|
try:
|
|
if (
|
|
self._placement not in {"cpu", "cuda", "radeon-cosmos"} or self._max_input_bytes <= 0
|
|
or self._max_action_bytes <= 0 or self._request_timeout_s <= 0
|
|
or not 0 < self._port < 65536
|
|
):
|
|
raise ValueError("invalid InternVLA placement, request bounds, timeout or port")
|
|
if self._memory_pool not in reservation.demands:
|
|
raise ResourceUnavailable("InternVLA shared-RAM pool absent from stage reservation")
|
|
demand = reservation.demands[self._memory_pool]
|
|
if self._placement == "cuda" and reservation.demands.get("cuda:0", 0) <= 0:
|
|
raise ResourceUnavailable("InternVLA CUDA placement requires an explicit cuda:0 reservation")
|
|
if self._max_input_bytes + self._max_action_bytes > demand:
|
|
raise ResourceUnavailable("InternVLA I/O bound exceeds stage reservation")
|
|
|
|
self._python = Path(config["python_bin"]).absolute()
|
|
if not self._python.is_file():
|
|
raise FileNotFoundError(f"InternVLA interpreter absent: {self._python}")
|
|
self._model_dir = Path(config["model_dir"]).resolve(strict=True)
|
|
self._processor_dir = Path(config["processor_dir"]).resolve(strict=True)
|
|
self._cosmos_dir = Path(config["cosmos_dir"]).resolve(strict=True)
|
|
self._graph = Path(config["graph_file"]).resolve(strict=True) if self._placement == "radeon-cosmos" else None
|
|
self._worker = Path(__file__).with_name("internvla_worker.py").resolve(strict=True)
|
|
artifacts = {
|
|
"python": self._python,
|
|
"model": self._model_dir / "model.safetensors",
|
|
"model_config": self._model_dir / "config.json",
|
|
"train_config": self._model_dir / "train_config.json",
|
|
"stats": self._model_dir / "stats.json",
|
|
"cosmos_encoder": self._cosmos_dir / "encoder.safetensors",
|
|
"cosmos_decoder": self._cosmos_dir / "decoder.safetensors",
|
|
"processor_tokenizer": self._processor_dir / "tokenizer.json",
|
|
"processor_config": self._processor_dir / "preprocessor_config.json",
|
|
}
|
|
if self._graph is not None:
|
|
artifacts["graph"] = self._graph
|
|
hashes = dict(config["artifact_sha256"])
|
|
if set(hashes) != set(artifacts):
|
|
raise ValueError("InternVLA artifact manifest has missing or unexpected entries")
|
|
for name, path in artifacts.items():
|
|
if _sha256(path) != str(hashes[name]).lower():
|
|
raise ValueError(f"InternVLA {name} differs from declared artifact hash")
|
|
overhead = int(config["memory_overhead_bytes"])
|
|
artifact_bytes = sum(path.stat().st_size for name, path in artifacts.items() if name != "python")
|
|
if overhead <= 0 or artifact_bytes + overhead > demand:
|
|
raise ResourceUnavailable("InternVLA weights, graph and state/workspace/headroom exceed reservation")
|
|
cuda_min_bytes = 0
|
|
if self._placement == "cuda":
|
|
cuda_min_bytes = (artifacts["model"].stat().st_size
|
|
+ artifacts["cosmos_encoder"].stat().st_size + overhead)
|
|
if cuda_min_bytes > self._cuda_demand_bytes:
|
|
raise ResourceUnavailable("InternVLA CUDA weights and workspace/headroom exceed reservation")
|
|
|
|
self._log_path = Path(config["log_file"]).resolve()
|
|
self._log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._log_stream = self._log_path.open("wb")
|
|
env = os.environ.copy()
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
env.update({
|
|
"PYTHONPATH": str(repo_root) + os.pathsep + env.get("PYTHONPATH", ""),
|
|
"CUDA_VISIBLE_DEVICES": "0" if self._placement == "cuda" else "",
|
|
"HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1",
|
|
"OMP_NUM_THREADS": "8", "MKL_NUM_THREADS": "8", "PYTHONUTF8": "1",
|
|
"TOKENIZERS_PARALLELISM": "false",
|
|
})
|
|
command = [
|
|
str(self._python), "-u", str(self._worker),
|
|
"--model-dir", str(self._model_dir),
|
|
"--processor-dir", str(self._processor_dir),
|
|
"--cosmos-dir", str(self._cosmos_dir),
|
|
"--placement", self._placement,
|
|
"--port", str(self._port), "--threads", "8",
|
|
"--max-input-bytes", str(self._max_input_bytes),
|
|
]
|
|
if self._graph is not None:
|
|
command += ["--graph", str(self._graph), "--graph-sha256", hashes["graph"]]
|
|
flags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
|
|
self._proc = subprocess.Popen(
|
|
command, stdout=self._log_stream, stderr=subprocess.STDOUT,
|
|
env=env, creationflags=flags,
|
|
)
|
|
base = f"http://127.0.0.1:{self._port}"
|
|
deadline = time.monotonic() + float(config.get("start_timeout_s", 180))
|
|
while True:
|
|
if self._proc.poll() is not None:
|
|
raise RuntimeError(f"InternVLA worker exited during load; see {self._log_path}")
|
|
try:
|
|
with urllib.request.urlopen(base + "/health", timeout=1) as response:
|
|
if response.status == 200:
|
|
break
|
|
except (urllib.error.URLError, TimeoutError, OSError):
|
|
pass
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError(f"InternVLA worker did not become ready; see {self._log_path}")
|
|
time.sleep(.2)
|
|
with urllib.request.urlopen(base + "/props", timeout=5) as response:
|
|
props = json.load(response)
|
|
expected_external = self._placement == "radeon-cosmos"
|
|
observed_external = props.get("external_load")
|
|
if (
|
|
props.get("placement") != self._placement
|
|
or Path(props.get("model_dir", "")).resolve() != self._model_dir
|
|
or Path(props.get("processor_dir", "")).resolve() != self._processor_dir
|
|
or Path(props.get("cosmos_dir", "")).resolve() != self._cosmos_dir
|
|
or props.get("policy_device") != ("cuda" if self._placement == "cuda" else "cpu")
|
|
or props.get("runtime_mode") != "real_checkpoint_loaded"
|
|
or props.get("policy_dtype") != "bfloat16"
|
|
or props.get("cosmos_dtype") != ("float32" if expected_external else "bfloat16")
|
|
or props.get("action_shape") != [1, 50, 32]
|
|
or props.get("action_mode") != "delta"
|
|
or props.get("control_ready") is not False
|
|
or props.get("torch") != str(config["expected_torch"])
|
|
or bool(observed_external) != expected_external
|
|
or (expected_external and observed_external.get("device_name") != "AMD Radeon(TM) 890M Graphics")
|
|
or (self._placement == "cuda" and (
|
|
props.get("cuda_device_name") != config.get("expected_cuda_device_name")
|
|
or props.get("cuda_device_index") != 0
|
|
or not 0 < props.get("cuda_reserved_after_load_bytes", 0) <= reservation.demands["cuda:0"]
|
|
))
|
|
):
|
|
raise RuntimeError("InternVLA worker model, precision or device differs from plan")
|
|
import psutil
|
|
|
|
root = psutil.Process(self._proc.pid)
|
|
descendants = [root, *root.children(recursive=True)]
|
|
self._loaded_process_tree = []
|
|
for process in descendants:
|
|
try:
|
|
memory = process.memory_full_info()
|
|
self._loaded_process_tree.append({
|
|
"pid": process.pid, "rss_bytes": int(memory.rss),
|
|
"private_bytes": int(getattr(memory, "private", memory.rss)),
|
|
})
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
self._loaded_rss_bytes = sum(row["rss_bytes"] for row in self._loaded_process_tree)
|
|
if self._loaded_rss_bytes > demand:
|
|
raise ResourceUnavailable("loaded InternVLA process tree exceeds stage reservation")
|
|
self._base_url = base
|
|
self.execution_plan = {
|
|
"backend": "external.internvla.policy.v1", "stage_id": self.stage_id,
|
|
"worker_generation": self._generation, "worker_pid": self._proc.pid,
|
|
"worker_script_sha256": _sha256(self._worker),
|
|
"artifact_sha256": hashes, "placement": self._placement,
|
|
"worker_props": props, "loaded_rss_bytes": self._loaded_rss_bytes,
|
|
"loaded_process_tree": self._loaded_process_tree,
|
|
"reserved_bytes": dict(reservation.demands),
|
|
"memory_overhead_bytes": overhead,
|
|
"cuda_weights_and_headroom_min_bytes": cuda_min_bytes,
|
|
"max_input_bytes": self._max_input_bytes,
|
|
"max_action_bytes": self._max_action_bytes,
|
|
"request_capacity": 1,
|
|
"stateful_session": "InternVLA policy worker-owned; complete action chunk per request",
|
|
"evidence": "B",
|
|
}
|
|
except BaseException:
|
|
self.shutdown()
|
|
raise
|
|
|
|
@staticmethod
|
|
def _encode_prompt(request_id: str, prompt: Any) -> bytes:
|
|
if not isinstance(prompt, dict):
|
|
raise ValueError("InternVLA prompt must contain camera histories and state")
|
|
expected = {"image0", "image1", "image2", "mask0", "mask1", "mask2",
|
|
"state", "task", "noise", "observation_timestamp_ns"}
|
|
if set(prompt) != expected:
|
|
raise ValueError("InternVLA observation fields differ from the declared contract")
|
|
arrays: dict[str, Any] = {}
|
|
for i in range(3):
|
|
image = np.asarray(prompt[f"image{i}"])
|
|
mask = np.asarray(prompt[f"mask{i}"])
|
|
if (
|
|
image.dtype != np.float32 or image.shape != (1, 2, 3, 224, 224)
|
|
or not np.isfinite(image).all() or image.min() < 0 or image.max() > 1
|
|
):
|
|
raise ValueError("camera history must be finite normalized float32 [1,2,3,224,224]")
|
|
if mask.dtype != np.bool_ or mask.shape != (1,):
|
|
raise ValueError("camera mask must be bool [1]")
|
|
arrays[f"image{i}"] = image
|
|
arrays[f"mask{i}"] = mask
|
|
for name, shape in (("state", (1, 32)), ("noise", (1, 50, 32))):
|
|
value = np.asarray(prompt[name])
|
|
if value.dtype != np.float32 or value.shape != shape or not np.isfinite(value).all():
|
|
raise ValueError(f"{name} must be finite float32 {shape}")
|
|
arrays[name] = value
|
|
task = prompt["task"]
|
|
timestamp = prompt["observation_timestamp_ns"]
|
|
if not isinstance(task, str) or len(task.encode("utf-8")) > 4096:
|
|
raise ValueError("task must be a bounded string")
|
|
if type(timestamp) is not int or timestamp <= 0:
|
|
raise ValueError("observation timestamp must be positive Unix nanoseconds")
|
|
arrays.update(task=np.array(task), observation_timestamp_ns=np.int64(timestamp),
|
|
request_id=np.array(request_id))
|
|
with io.BytesIO() as output:
|
|
np.savez(output, **arrays)
|
|
return output.getvalue()
|
|
|
|
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("InternVLA stage has an unacknowledged request; capacity is one")
|
|
body = self._encode_prompt(request_id, prompt)
|
|
if len(body) > self._max_input_bytes:
|
|
raise ResourceUnavailable("InternVLA observation exceeds admitted input bound")
|
|
self._epoch += 1
|
|
request = StageRequest(request_id, self.stage_id, self._epoch, self._generation)
|
|
self._active = request_id
|
|
self._task = asyncio.create_task(self._run(request, body), name=f"internvla-{request_id}")
|
|
|
|
async def _run(self, request: StageRequest, body: bytes) -> None:
|
|
try:
|
|
started = time.perf_counter()
|
|
data = await asyncio.to_thread(
|
|
_post_actions, self._base_url + "/v1/actions", body,
|
|
timeout=self._request_timeout_s, max_bytes=self._max_action_bytes,
|
|
)
|
|
with np.load(io.BytesIO(data), allow_pickle=False) as response:
|
|
if set(response.files) != {"actions", "observation_timestamp_ns",
|
|
"generation_timestamp_ns", "request_id", "worker_wall_s",
|
|
"cuda_peak_reserved_bytes"}:
|
|
raise ValueError("InternVLA response fields differ from declared action contract")
|
|
actions = np.asarray(response["actions"])
|
|
observation_ns = int(response["observation_timestamp_ns"].item())
|
|
generation_ns = int(response["generation_timestamp_ns"].item())
|
|
returned_id = str(response["request_id"].item())
|
|
worker_wall_s = float(response["worker_wall_s"].item())
|
|
cuda_peak_reserved = int(response["cuda_peak_reserved_bytes"].item())
|
|
if (returned_id != request.request_id or actions.dtype != np.float32
|
|
or actions.shape != (1, 50, 32) or not np.isfinite(actions).all()
|
|
or generation_ns < observation_ns):
|
|
raise ValueError("InternVLA worker returned invalid, stale or misrouted actions")
|
|
if self._placement == "cuda" and not 0 < cuda_peak_reserved <= self._cuda_demand_bytes:
|
|
raise ResourceUnavailable("InternVLA CUDA peak reservation exceeded the admitted VRAM budget")
|
|
if self._placement != "cuda" and cuda_peak_reserved != 0:
|
|
raise ValueError("InternVLA non-CUDA worker reported CUDA allocation")
|
|
ref = BufferRef(
|
|
"actions", str(self.stage_id), self._generation, "float32", tuple(actions.shape), int(actions.nbytes)
|
|
)
|
|
event = StageEvent(
|
|
request.request_id, self.stage_id, request.epoch, 1, "action", self._generation,
|
|
(ref,), terminal=True,
|
|
)
|
|
output = OmniRequestOutput(
|
|
request_id=request.request_id, stage_id=self.stage_id,
|
|
final_output_type="action",
|
|
_custom_output={
|
|
"actions": actions.copy(), "stage_event": dataclasses.asdict(event),
|
|
"action_metadata": {
|
|
"observation_timestamp_ns": observation_ns,
|
|
"generation_timestamp_ns": generation_ns,
|
|
"observation_age_s": (generation_ns - observation_ns) / 1e9,
|
|
"action_mode": "delta", "action_shape": [1, 50, 32],
|
|
"action_units": "unverified", "joint_order": "unverified",
|
|
"action_step_s": None, "control_ready": False,
|
|
},
|
|
},
|
|
metrics={"policy_wall_s": time.perf_counter() - started,
|
|
"worker_wall_s": worker_wall_s,
|
|
"cuda_peak_reserved_bytes": cuda_peak_reserved},
|
|
)
|
|
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 _terminate(self) -> bool:
|
|
"""Retire both the policy server and its possible DirectML child."""
|
|
import psutil
|
|
|
|
children = []
|
|
try:
|
|
if self._proc is not None:
|
|
children = psutil.Process(self._proc.pid).children(recursive=True)
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
except psutil.Error:
|
|
return False
|
|
for child in reversed(children):
|
|
try:
|
|
child.terminate()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
except psutil.Error:
|
|
return False
|
|
# The policy process is the worker's parent and may be the one that
|
|
# reaps it. Retire the parent before deciding a transient zombie means
|
|
# the shared-RAM reservation must remain quarantined.
|
|
parent_drained = super()._terminate()
|
|
_, alive = psutil.wait_procs(children, timeout=5)
|
|
for child in alive:
|
|
try:
|
|
if child.status() != psutil.STATUS_ZOMBIE:
|
|
child.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
except psutil.Error:
|
|
return False
|
|
_, alive = psutil.wait_procs(alive, timeout=5)
|
|
for child in alive:
|
|
try:
|
|
if child.status() != psutil.STATUS_ZOMBIE:
|
|
return False
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
except psutil.Error:
|
|
return False
|
|
return parent_drained
|