Add admitted InternVLA policy stages and update device matrix

benchmarks/edge_harness/experiments/probe_omni_internvla_entrypoint.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Check a bounded InternVLA action request through public AsyncOmni."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import hashlib
+import json
+import tempfile
+import time
+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()
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    for name in ("model-dir", "cosmos-dir", "processor-dir", "python-bin", "log-file", "output-report"):
+        parser.add_argument(f"--{name}", type=Path, required=True)
+    parser.add_argument("--graph-file", type=Path)
+    parser.add_argument("--placement", choices=("cpu", "radeon-cosmos"), required=True)
+    args = parser.parse_args()
+    if args.placement == "radeon-cosmos" and args.graph_file is None:
+        parser.error("Radeon placement requires graph-file")
+
+    import numpy as np
+    import torch
+    import yaml
+
+    from vllm_omni.diffusion.models.internvla_a1_whole_pipeline import INTERNVLA_A1_WHOLE_POLICY_PIPELINE
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+    model = args.model_dir.resolve(strict=True)
+    cosmos = args.cosmos_dir.resolve(strict=True)
+    processor = args.processor_dir.resolve(strict=True)
+    python = args.python_bin.absolute()
+    graph = args.graph_file.resolve(strict=True) if args.graph_file else None
+    files = {
+        "python": python, "model": model / "model.safetensors",
+        "model_config": model / "config.json", "train_config": model / "train_config.json",
+        "stats": model / "stats.json", "cosmos_encoder": cosmos / "encoder.safetensors",
+        "cosmos_decoder": cosmos / "decoder.safetensors",
+        "processor_tokenizer": processor / "tokenizer.json",
+        "processor_config": processor / "preprocessor_config.json",
+    }
+    if graph is not None:
+        files["graph"] = graph
+    backend = {
+        "name": "external.internvla.policy.v1", "placement": args.placement,
+        "python_bin": str(python), "expected_torch": str(torch.__version__),
+        "model_dir": str(model), "cosmos_dir": str(cosmos), "processor_dir": str(processor),
+        "graph_file": str(graph) if graph else None,
+        "artifact_sha256": {name: sha256(path) for name, path in files.items()},
+        "log_file": str(args.log_file), "memory_overhead_bytes": 8 << 30,
+        "max_input_bytes": 8 << 20, "max_action_bytes": 1 << 20,
+        "start_timeout_s": 180, "request_timeout_s": 60,
+    }
+    report = {"entrypoint": "AsyncOmni.generate", "placement": args.placement,
+              "pipeline": INTERNVLA_A1_WHOLE_POLICY_PIPELINE.model_type}
+    engine = None
+    try:
+        with tempfile.TemporaryDirectory(prefix="omni-internvla-") as directory:
+            deployment = Path(directory) / "deploy.yaml"
+            deployment.write_text(yaml.safe_dump({
+                "pipeline": INTERNVLA_A1_WHOLE_POLICY_PIPELINE.model_type,
+                "async_chunk": False,
+                "stages": [{"stage_id": 0, "backend": backend,
+                            "resource_budget": {"capacities": {"host_ram": 30 << 30},
+                                                "demands": {"host_ram": 16 << 30}}}],
+            }), encoding="utf-8")
+            started = time.perf_counter()
+            engine = AsyncOmni(
+                model=str(model), deploy_config=str(deployment),
+                stage_init_timeout=180, init_timeout=240,
+            )
+            report["startup_s"] = time.perf_counter() - started
+            images = [np.zeros((1, 2, 3, 224, 224), dtype=np.float32) for _ in range(3)]
+            images[0][:, :, 0, 56:168, 56:168] = 1
+            images[1][:, :, 1, 56:168, 56:168] = .5
+            prompt = {
+                **{f"image{i}": images[i] for i in range(3)},
+                **{f"mask{i}": np.ones((1,), dtype=np.bool_) for i in range(3)},
+                "state": np.zeros((1, 32), dtype=np.float32),
+                "noise": np.zeros((1, 50, 32), dtype=np.float32),
+                "task": "Place the marker pen in its holder.",
+                "observation_timestamp_ns": time.time_ns(),
+            }
+            outputs = []
+            started = time.perf_counter()
+            async for output in engine.generate(prompt, request_id="public-internvla-1"):
+                if output.error:
+                    raise RuntimeError(output.error)
+                actions = np.asarray(output.custom_output["actions"])
+                outputs.append({
+                    "request_id": output.request_id,
+                    "action_shape": list(actions.shape),
+                    "action_sha256": hashlib.sha256(actions.tobytes()).hexdigest(),
+                    "finite": bool(np.isfinite(actions).all()),
+                    "metadata": output.custom_output.get("action_metadata"),
+                    "stage_event": output.custom_output.get("stage_event"),
+                })
+            report["request_wall_s"] = time.perf_counter() - started
+            report["outputs"] = outputs
+            assert len(outputs) == 1 and outputs[0]["request_id"] == "public-internvla-1"
+            assert outputs[0]["action_shape"] == [1, 50, 32] and outputs[0]["finite"]
+            assert outputs[0]["metadata"]["control_ready"] is False
+            report["status"] = "passed"
+    except BaseException as exc:
+        report["status"] = "failed"
+        report["error"] = f"{type(exc).__name__}: {exc}"
+        raise
+    finally:
+        if engine is not None:
+            engine.shutdown()
+        args.output_report.parent.mkdir(parents=True, exist_ok=True)
+        args.output_report.write_text(json.dumps(report, indent=2, default=str) + "\n", encoding="utf-8")
+        print(json.dumps(report, indent=2, default=str))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

benchmarks/edge_harness/experiments/probe_omni_internvla_policy.py

@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+"""Profile an admitted InternVLA action stage under Omni StageRuntime."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import hashlib
+import json
+import math
+import time
+from pathlib import Path
+from types import SimpleNamespace
+
+
+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 rank(values: list[float], fraction: float) -> float:
+    return sorted(values)[math.ceil(len(values) * fraction) - 1]
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    for name in ("model-dir", "cosmos-dir", "processor-dir", "python-bin", "log-file", "output-report"):
+        parser.add_argument(f"--{name}", type=Path, required=True)
+    parser.add_argument("--graph-file", type=Path)
+    parser.add_argument("--placement", choices=("cpu", "radeon-cosmos"), required=True)
+    parser.add_argument("--output-actions", type=Path)
+    parser.add_argument("--reference-actions", type=Path)
+    parser.add_argument("--capacity-gib", type=int, default=30)
+    parser.add_argument("--reserve-gib", type=int, default=16)
+    parser.add_argument("--warmups", type=int, default=1)
+    parser.add_argument("--repeats", type=int, default=20)
+    parser.add_argument("--admission-refusal", action="store_true")
+    parser.add_argument("--abort-check", action="store_true")
+    args = parser.parse_args()
+    if args.warmups < 0 or args.repeats < 1 or args.reserve_gib < 1 or args.capacity_gib < args.reserve_gib:
+        parser.error("invalid profile counts or explicit memory budget")
+    if args.placement == "radeon-cosmos" and args.graph_file is None:
+        parser.error("Radeon placement requires --graph-file")
+
+    import numpy as np
+
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.diffusion.models.internvla_a1_whole_pipeline import INTERNVLA_A1_WHOLE_POLICY_PIPELINE
+    from vllm_omni.engine.stage_runtime import StageRuntime
+
+    model = args.model_dir.resolve(strict=True)
+    cosmos = args.cosmos_dir.resolve(strict=True)
+    processor = args.processor_dir.resolve(strict=True)
+    python = args.python_bin.absolute()
+    graph = args.graph_file.resolve(strict=True) if args.graph_file else None
+    files = {
+        "python": python, "model": model / "model.safetensors",
+        "model_config": model / "config.json", "train_config": model / "train_config.json",
+        "stats": model / "stats.json", "cosmos_encoder": cosmos / "encoder.safetensors",
+        "cosmos_decoder": cosmos / "decoder.safetensors",
+        "processor_tokenizer": processor / "tokenizer.json",
+        "processor_config": processor / "preprocessor_config.json",
+    }
+    if graph is not None:
+        files["graph"] = graph
+    hashes = {name: sha256(path) for name, path in files.items()}
+    import torch
+
+    backend = {
+        "name": "external.internvla.policy.v1", "placement": args.placement,
+        "python_bin": str(python), "expected_torch": str(torch.__version__),
+        "model_dir": str(model), "cosmos_dir": str(cosmos), "processor_dir": str(processor),
+        "graph_file": str(graph) if graph else None, "artifact_sha256": hashes,
+        "log_file": str(args.log_file), "memory_overhead_bytes": 8 << 30,
+        "max_input_bytes": 8 << 20, "max_action_bytes": 1 << 20,
+        "start_timeout_s": 180, "request_timeout_s": 60,
+    }
+    budget = {"capacities": {"host_ram": args.capacity_gib << 30},
+              "demands": {"host_ram": args.reserve_gib << 30}}
+    deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(
+        stage_id=0, backend=backend, resource_budget=budget,
+    )])
+    configs = [item.to_omegaconf() for item in merge_pipeline_deploy(
+        INTERNVLA_A1_WHOLE_POLICY_PIPELINE, deploy,
+    )]
+    runtime = StageRuntime(configs, "local-internvla-policy", "", stage_init_timeout=180, async_chunk=False)
+    report = {
+        "scope": "real InternVLA Place_Markpen checkpoint via bounded Omni whole-policy graph stage; synthetic patterned observations/noise; no robot-task quality claim",
+        "placement": args.placement, "artifact_sha256": hashes, "budget": budget,
+        "warmups": args.warmups, "repeats": args.repeats,
+    }
+    try:
+        started = time.perf_counter()
+        if args.admission_refusal:
+            try:
+                runtime.initialize()
+            except Exception as exc:
+                if "exceed reservation" not in str(exc):
+                    raise
+                report["admission_refusal"] = f"{type(exc).__name__}: {exc}"
+                report["status"] = "passed"
+                return
+            raise AssertionError("InternVLA admitted the insufficient memory reservation")
+        runtime.initialize()
+        report["startup_s"] = time.perf_counter() - started
+        pool = runtime.stage_pools[0]
+        report["execution_plan_start"] = dict(pool.stage_client.execution_plan)
+        state = SimpleNamespace(sampling_params_list=[None])
+        images = [np.zeros((1, 2, 3, 224, 224), dtype=np.float32) for _ in range(3)]
+        images[0][:, :, 0, 56:168, 56:168] = 1.0
+        images[1][:, :, 1, 56:168, 56:168] = .5
+        prompt = {
+            **{f"image{i}": images[i] for i in range(3)},
+            **{f"mask{i}": np.ones((1,), dtype=np.bool_) for i in range(3)},
+            "state": np.zeros((1, 32), dtype=np.float32),
+            "noise": np.zeros((1, 50, 32), dtype=np.float32),
+            "task": "Place the marker pen in its holder.",
+        }
+
+        async def request_one(request_id: str):
+            prompt["observation_timestamp_ns"] = time.time_ns()
+            started = time.perf_counter()
+            await pool.submit_initial(request_id, state, prompt)
+            deadline = time.monotonic() + 75
+            while True:
+                output = pool.poll_graph_output(0)
+                if output is not None:
+                    break
+                if time.monotonic() >= deadline:
+                    raise TimeoutError("InternVLA Omni action request timed out")
+                await asyncio.sleep(.005)
+            try:
+                if output.error:
+                    raise RuntimeError(output.error)
+                values = np.asarray(output.custom_output["actions"])
+                event = output.custom_output["stage_event"]
+                metadata = output.custom_output["action_metadata"]
+                if values.dtype != np.float32 or values.shape != (1, 50, 32) or not np.isfinite(values).all():
+                    raise RuntimeError("InternVLA returned invalid actions")
+                if not event["terminal"] or event["kind"] != "action" or metadata["control_ready"]:
+                    raise RuntimeError("InternVLA action event or metadata differs from contract")
+                return {"wall_s": time.perf_counter() - started,
+                        "action_sha256": hashlib.sha256(values.tobytes()).hexdigest(),
+                        "stage_event": event, "action_metadata": metadata,
+                        "metrics": output.metrics, "actions": values.copy()}
+            finally:
+                output.release_stage_buffers()
+                await asyncio.sleep(0)
+
+        report["warmup_results"] = []
+        for index in range(args.warmups):
+            row = await request_one(f"internvla-warmup-{index}")
+            row.pop("actions")
+            report["warmup_results"].append(row)
+        measured = []
+        report["measured"] = measured
+        first_actions = None
+        for index in range(args.repeats):
+            row = await request_one(f"internvla-measured-{index}")
+            values = row.pop("actions")
+            if first_actions is None:
+                first_actions = values
+            measured.append(row)
+        assert first_actions is not None
+        report["nearest_rank_p50_wall_s"] = rank([r["wall_s"] for r in measured], .5)
+        report["nearest_rank_p95_wall_s"] = rank([r["wall_s"] for r in measured], .95)
+        report["unique_action_hashes"] = sorted({r["action_sha256"] for r in measured})
+        if args.output_actions:
+            args.output_actions.parent.mkdir(parents=True, exist_ok=True)
+            np.save(args.output_actions, first_actions)
+            report["output_actions"] = str(args.output_actions)
+        if args.reference_actions:
+            reference = np.load(args.reference_actions, allow_pickle=False)
+            if reference.shape != first_actions.shape:
+                raise RuntimeError("reference action shape differs")
+            difference = first_actions - reference
+            report["reference_comparison"] = {
+                "max_abs": float(np.max(np.abs(difference))),
+                "relative_l2": float(np.linalg.norm(difference) / np.linalg.norm(reference)),
+                "cosine": float(np.dot(first_actions.ravel(), reference.ravel()) / (
+                    np.linalg.norm(first_actions) * np.linalg.norm(reference)
+                )),
+            }
+        report["execution_plan_end"] = dict(pool.stage_client.execution_plan)
+        if args.abort_check:
+            prompt["observation_timestamp_ns"] = time.time_ns()
+            await pool.submit_initial("internvla-abort", state, prompt)
+            await asyncio.sleep(.05)
+            await pool.abort_requests(["internvla-abort"])
+            await asyncio.sleep(.1)
+            report["abort_check"] = {
+                "stale_output": pool.poll_graph_output(0) is not None,
+                "ledger": runtime.resource_ledger.snapshot(),
+            }
+            if report["abort_check"]["stale_output"]:
+                raise RuntimeError("InternVLA cancellation leaked a stale action")
+        report["status"] = "passed"
+    except BaseException as exc:
+        report["status"] = "failed"
+        report["error"] = f"{type(exc).__name__}: {exc}"
+        raise
+    finally:
+        runtime.shutdown()
+        if runtime.resource_ledger is not None:
+            report["ledger_after_shutdown"] = runtime.resource_ledger.snapshot()
+        args.output_report.parent.mkdir(parents=True, exist_ok=True)
+        args.output_report.write_text(json.dumps(report, indent=2, default=str) + "\n", encoding="utf-8")
+        print(json.dumps({key: value for key, value in report.items() if key != "measured"}, indent=2, default=str))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())