260 lines
13 KiB
Python
260 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""InternVLA action worker on CPU/CUDA with an optional Radeon Cosmos encoder.
|
|
|
|
The policy remains the existing Omni diffusion implementation. Requests carry
|
|
explicit normalized camera histories and state; the worker owns policy state
|
|
and, for the hybrid route, Omni's existing external DirectML graph worker.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import time
|
|
import traceback
|
|
import zipfile
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
from pathlib import Path
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for block in iter(lambda: source.read(8 << 20), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--processor-dir", type=Path, required=True)
|
|
parser.add_argument("--cosmos-dir", type=Path, required=True)
|
|
parser.add_argument("--placement", choices=("cpu", "cuda", "radeon-cosmos"), required=True)
|
|
parser.add_argument("--graph", type=Path)
|
|
parser.add_argument("--graph-sha256")
|
|
parser.add_argument("--port", type=int, required=True)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--max-input-bytes", type=int, default=8 << 20)
|
|
args = parser.parse_args()
|
|
if not 0 < args.port < 65536 or not 0 < args.threads <= 24 or args.max_input_bytes <= 0:
|
|
parser.error("invalid port, threads or input bound")
|
|
if args.placement == "radeon-cosmos" and (args.graph is None or not args.graph_sha256):
|
|
parser.error("Radeon route requires a graph and hash")
|
|
expected_visible = "0" if args.placement == "cuda" else ""
|
|
if os.environ.get("CUDA_VISIBLE_DEVICES") != expected_visible:
|
|
raise RuntimeError("InternVLA CUDA visibility differs from the placement plan")
|
|
if os.environ.get("HF_HUB_OFFLINE") != "1":
|
|
raise RuntimeError("InternVLA worker requires offline checkpoint loading")
|
|
|
|
model_dir = args.model_dir.resolve(strict=True)
|
|
processor_dir = args.processor_dir.resolve(strict=True)
|
|
cosmos_dir = args.cosmos_dir.resolve(strict=True)
|
|
os.environ["INTERNVLA_A1_COSMOS_DIR"] = str(cosmos_dir)
|
|
os.environ["INTERNVLA_A1_PROCESSOR_DIR"] = str(processor_dir)
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
if args.placement == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError("InternVLA CUDA placement requested but CUDA is unavailable")
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.registry import initialize_model
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
torch.set_num_threads(args.threads)
|
|
policy_device = "cuda" if args.placement == "cuda" else "cpu"
|
|
config = OmniDiffusionConfig(
|
|
model=str(model_dir), model_class_name="InternVLAA1Pipeline", dtype=torch.bfloat16,
|
|
custom_pipeline_args={
|
|
"device": policy_device, "dtype": "bfloat16", "compile_model": False,
|
|
"enable_regional_compile": False, "enable_warmup": False,
|
|
"strict_load": True, "processor_model_name": str(processor_dir),
|
|
},
|
|
)
|
|
started = time.perf_counter()
|
|
pipeline = initialize_model(config)
|
|
policy_load_s = time.perf_counter() - started
|
|
if pipeline.runtime_mode() != "real_checkpoint_loaded":
|
|
raise RuntimeError("InternVLA real checkpoint did not load")
|
|
if {parameter.device.type for parameter in pipeline.policy.parameters()} != {policy_device}:
|
|
raise RuntimeError("InternVLA policy parameters differ from the placement plan")
|
|
|
|
external = None
|
|
external_report = None
|
|
if args.placement == "radeon-cosmos":
|
|
from vllm_omni.edge.local.external.client import ExternalWorker
|
|
from vllm_omni.edge.local.external.launch import ROUTE_TORCH_DML, resolve
|
|
|
|
graph = args.graph.resolve(strict=True)
|
|
if _sha256(graph) != args.graph_sha256.lower():
|
|
raise RuntimeError("Radeon Cosmos graph hash differs from plan")
|
|
route = resolve(ROUTE_TORCH_DML)
|
|
if not route.available:
|
|
raise RuntimeError(f"DirectML worker route unavailable: {route.reason}")
|
|
external = ExternalWorker(route, max_payload_bytes=args.max_input_bytes)
|
|
try:
|
|
external.start()
|
|
external_report = external.load(
|
|
graph, example_inputs={"pixels": np.zeros((6, 3, 256, 256), dtype=np.float32)}
|
|
)
|
|
if (
|
|
external_report.device_name != "AMD Radeon(TM) 890M Graphics"
|
|
or external_report.placement_granularity != "output_device"
|
|
or external_report.fraction_on_target != 1.0
|
|
or tuple(external_report.outputs[0]["shape"]) != (6, 16, 32, 32)
|
|
):
|
|
raise RuntimeError("Radeon Cosmos output placement/shape differs from plan")
|
|
|
|
class RemoteEncoder(torch.nn.Module):
|
|
def forward(self, pixels: torch.Tensor) -> torch.Tensor:
|
|
if pixels.device.type != "cpu" or tuple(pixels.shape) != (6, 3, 256, 256):
|
|
raise ValueError("Cosmos input differs from fixed six-frame CPU contract")
|
|
output, _ = external.run({"pixels": pixels.detach().float().contiguous().numpy()})
|
|
latent = output["out_0"]
|
|
if latent.shape != (6, 16, 32, 32) or not np.isfinite(latent).all():
|
|
raise RuntimeError("Radeon Cosmos returned invalid latent")
|
|
return torch.from_numpy(latent).to(pixels.dtype)
|
|
|
|
pipeline.policy.model.cosmos._enc_model = RemoteEncoder()
|
|
except BaseException:
|
|
external.close()
|
|
raise
|
|
|
|
props = {
|
|
"model_dir": str(model_dir), "processor_dir": str(processor_dir),
|
|
"cosmos_dir": str(cosmos_dir), "placement": args.placement,
|
|
"torch": torch.__version__, "policy_dtype": "bfloat16",
|
|
"cosmos_dtype": "float32" if external is not None else "bfloat16",
|
|
"policy_device": policy_device, "runtime_mode": pipeline.runtime_mode(),
|
|
"cuda_device_name": torch.cuda.get_device_name(0) if policy_device == "cuda" else None,
|
|
"cuda_device_index": torch.cuda.current_device() if policy_device == "cuda" else None,
|
|
"cuda_allocated_after_load_bytes": torch.cuda.memory_allocated(0) if policy_device == "cuda" else None,
|
|
"cuda_reserved_after_load_bytes": torch.cuda.memory_reserved(0) if policy_device == "cuda" else None,
|
|
"policy_load_s": policy_load_s, "action_shape": [1, 50, 32],
|
|
"image_shape": [1, 2, 3, 224, 224], "state_shape": [1, 32],
|
|
"action_mode": "delta", "action_units": "unverified",
|
|
"action_joint_order": "unverified", "action_step_s": None,
|
|
"control_ready": False,
|
|
"external_load": external_report.to_dict() if external_report else None,
|
|
}
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def _json(self, code: int, value: dict) -> None:
|
|
body = json.dumps(value).encode("utf-8")
|
|
self.send_response(code)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def do_GET(self) -> None:
|
|
if self.path == "/health":
|
|
self._json(200, {"status": "ok"})
|
|
elif self.path == "/props":
|
|
self._json(200, props)
|
|
else:
|
|
self._json(404, {"error": "unknown endpoint"})
|
|
|
|
def do_POST(self) -> None:
|
|
if self.path != "/v1/actions":
|
|
self._json(404, {"error": "unknown endpoint"})
|
|
return
|
|
size = int(self.headers.get("Content-Length", "0"))
|
|
if not 0 < size <= args.max_input_bytes:
|
|
self._json(413, {"error": "observation exceeds admitted input bound"})
|
|
return
|
|
try:
|
|
raw = self.rfile.read(size)
|
|
with zipfile.ZipFile(io.BytesIO(raw)) as archive:
|
|
if (len(archive.infolist()) != 11 or
|
|
any(info.compress_type != zipfile.ZIP_STORED for info in archive.infolist()) or
|
|
sum(info.file_size for info in archive.infolist()) > args.max_input_bytes):
|
|
raise ValueError("observation archive exceeds uncompressed bound")
|
|
with np.load(io.BytesIO(raw), allow_pickle=False) as data:
|
|
expected = {
|
|
"image0", "image1", "image2", "mask0", "mask1", "mask2",
|
|
"state", "task", "noise", "observation_timestamp_ns", "request_id",
|
|
}
|
|
if set(data.files) != expected:
|
|
raise ValueError("observation fields differ from the declared contract")
|
|
images = [np.asarray(data[f"image{i}"]) for i in range(3)]
|
|
masks = [np.asarray(data[f"mask{i}"]) for i in range(3)]
|
|
state = np.asarray(data["state"])
|
|
noise = np.asarray(data["noise"])
|
|
task = str(data["task"].item())
|
|
request_id = str(data["request_id"].item())
|
|
observation_ns = int(data["observation_timestamp_ns"].item())
|
|
if len(task.encode("utf-8")) > 4096 or not request_id or observation_ns <= 0:
|
|
raise ValueError("invalid task, request identity or observation time")
|
|
for image in images:
|
|
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]")
|
|
for mask in masks:
|
|
if mask.dtype != np.bool_ or mask.shape != (1,):
|
|
raise ValueError("camera mask must be bool [1]")
|
|
if state.dtype != np.float32 or state.shape != (1, 32) or not np.isfinite(state).all():
|
|
raise ValueError("state must be finite float32 [1,32]")
|
|
if noise.dtype != np.float32 or noise.shape != (1, 50, 32) or not np.isfinite(noise).all():
|
|
raise ValueError("noise must be finite float32 [1,50,32]")
|
|
batch = {"observation.state": torch.from_numpy(state.copy()).to(device=policy_device, dtype=torch.bfloat16),
|
|
"observation.task": [task]}
|
|
for i in range(3):
|
|
batch[f"observation.images.image{i}"] = torch.from_numpy(images[i].copy()).to(device=policy_device, dtype=torch.bfloat16)
|
|
batch[f"observation.images.image{i}_mask"] = torch.from_numpy(masks[i].copy()).to(policy_device)
|
|
started = time.perf_counter()
|
|
result = pipeline.forward(DiffusionRequestBatch(requests=[
|
|
OmniDiffusionRequest(
|
|
prompt="",
|
|
sampling_params=OmniDiffusionSamplingParams(extra_args={
|
|
"batch_inputs": batch, "noise": torch.from_numpy(noise.copy()).to(policy_device),
|
|
"decode_image": False,
|
|
}),
|
|
request_id=request_id,
|
|
)
|
|
]))
|
|
if result.error:
|
|
raise RuntimeError(result.error)
|
|
actions = result.output["payload"]["actions"]
|
|
if actions.device.type != policy_device or tuple(actions.shape) != (1, 50, 32):
|
|
raise RuntimeError("policy action device or shape differs from plan")
|
|
values = actions.detach().float().cpu().contiguous().numpy()
|
|
if policy_device == "cuda":
|
|
torch.cuda.synchronize(0)
|
|
cuda_peak_reserved = torch.cuda.max_memory_reserved(0) if policy_device == "cuda" else 0
|
|
if not np.isfinite(values).all():
|
|
raise RuntimeError("policy returned nonfinite actions")
|
|
with io.BytesIO() as out:
|
|
np.savez(out, actions=values, observation_timestamp_ns=np.int64(observation_ns),
|
|
generation_timestamp_ns=np.int64(time.time_ns()), request_id=np.array(request_id),
|
|
worker_wall_s=np.float64(time.perf_counter() - started),
|
|
cuda_peak_reserved_bytes=np.int64(cuda_peak_reserved))
|
|
body = out.getvalue()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/octet-stream")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
except Exception as exc:
|
|
traceback.print_exc()
|
|
self._json(500, {"error": f"{type(exc).__name__}: {exc}"})
|
|
|
|
try:
|
|
with HTTPServer(("127.0.0.1", args.port), Handler) as server:
|
|
print(f"internvla-worker ready port={args.port} props={json.dumps(props)}", flush=True)
|
|
server.serve_forever(poll_interval=0.1)
|
|
finally:
|
|
if external is not None:
|
|
external.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|