diff --git a/.gitignore b/.gitignore index dad7e719..256048bc 100755 --- a/.gitignore +++ b/.gitignore @@ -158,6 +158,7 @@ INSTALL_HYS.md AGENTS.md _version.py.mcp.json telefuser/_version.py +!telefuser/pipelines/lingbot_vla_v2/assets/*.json # LingBot regression example assets !examples/data/lingbot_world_fast/image.jpg !examples/data/lingbot_world_fast/poses.npy diff --git a/benchmarks/telefuser_aiperf/README.md b/benchmarks/telefuser_aiperf/README.md index ae261644..93f93258 100644 --- a/benchmarks/telefuser_aiperf/README.md +++ b/benchmarks/telefuser_aiperf/README.md @@ -92,6 +92,31 @@ Available batch configs: | `configs/video_generation_rate.yaml` | Poisson-arrival load | | `configs/video_generation_wan21_i2v_480p_compare.yaml` | Fixed Wan2.1 I2V comparison | +## LingBot-VLA v2 Structured Actions + +Start the native VLA service from its isolated model environment, then run the AIPerf workload from the repository +root: + +```bash +bash benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh +``` + +The repository-owned `telefuser_vla_structured` endpoint and `telefuser_structured_http` transport submit +`POST /v1/tasks/structured`, poll `GET /v1/tasks/{task_id}/status`, and pass request latency, throughput, success, +trace, and server metric facts into AIPerf's normal warmup and aggregation pipeline. Defaults are two excluded warmup +requests followed by 20 measured requests at concurrency one. Override them without changing the checked-in config: + +```bash +TELEFUSER_AIPERF_REQUESTS=100 \ +TELEFUSER_AIPERF_CONCURRENCY=2 \ + bash benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh +``` + +Each terminal result is required to contain a finite `50x55` action chunk and the frozen structured result fields. +The adapter retains an action hash, bounds, dimensions, verification status, target inference time, and peak memory; +it does not copy full action arrays or Base64 cameras into AIPerf response records. This validates service execution +and normalized action structure, not physical robot control semantics. + ## LingBot-World v2 Streaming The v2 pipeline expects the following files below `TF_MODEL_ZOO_PATH`: @@ -264,10 +289,12 @@ AIPerf environment first, then run the checks from the repository root: PYTHONPATH=benchmarks/telefuser_aiperf \ .venv-aiperf/bin/python -m pytest \ benchmarks/telefuser_aiperf/tests/test_livekit_adapter.py \ - benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py + benchmarks/telefuser_aiperf/tests/test_sglang_adapter.py \ + benchmarks/telefuser_aiperf/tests/test_vla_structured.py bash -n \ scripts/setup_aiperf.sh \ benchmarks/telefuser_aiperf/scripts/run_stream_bench.sh \ - benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh + benchmarks/telefuser_aiperf/scripts/run_sglang_lingbot_world_v2_4gpu.sh \ + benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh ``` diff --git a/benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml b/benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml new file mode 100644 index 00000000..3053385b --- /dev/null +++ b/benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml @@ -0,0 +1,49 @@ +# yaml-language-server: $schema=../../aiperf/src/aiperf/config/schema/aiperf-config.schema.json + +schemaVersion: "2.0" + +randomSeed: 42 + +benchmark: + model: lingbot-vla-v2-6b + + endpoint: + url: ${TELEFUSER_AIPERF_URL:http://127.0.0.1:18080} + type: telefuser_vla_structured + transport: telefuser_structured_http + timeout: ${TELEFUSER_AIPERF_TIMEOUT:120} + + tokenizer: + name: builtin + + dataset: + type: file + path: ./benchmarks/telefuser_aiperf/data/vla_structured.jsonl + format: single_turn + sampling: sequential + + warmup: + type: concurrency + concurrency: 1 + requests: ${TELEFUSER_AIPERF_WARMUP_REQUESTS:2} + excludeFromResults: true + + profiling: + type: concurrency + concurrency: ${TELEFUSER_AIPERF_CONCURRENCY:1} + requests: ${TELEFUSER_AIPERF_REQUESTS:20} + duration: ${TELEFUSER_AIPERF_DURATION:3600} + gracePeriod: ${TELEFUSER_AIPERF_GRACE_PERIOD:120} + + artifacts: + dir: ./artifacts/telefuser_aiperf/vla_structured + summary: [json] + records: [jsonl] + showTraceTiming: true + trace: true + + serverMetrics: + enabled: ${TELEFUSER_AIPERF_SERVER_METRICS:true} + urls: + - ${TELEFUSER_AIPERF_METRICS_URL:http://127.0.0.1:18080/v1/service/metrics} + formats: [json, csv] diff --git a/benchmarks/telefuser_aiperf/data/vla_structured.jsonl b/benchmarks/telefuser_aiperf/data/vla_structured.jsonl new file mode 100644 index 00000000..eb2fbaff --- /dev/null +++ b/benchmarks/telefuser_aiperf/data/vla_structured.jsonl @@ -0,0 +1 @@ +{"text":"pick up the object","image":"examples/data/101235-video-720_0.png","extra":{"state":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],"seed":7}} diff --git a/benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh b/benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh new file mode 100755 index 00000000..adf86b15 --- /dev/null +++ b/benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${ROOT_DIR}" + +CONFIG_PATH="${1:-benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml}" +SERVER_URL="${TELEFUSER_AIPERF_URL:-http://127.0.0.1:18080}" +HEALTH_URL="${TELEFUSER_AIPERF_HEALTH_URL:-${SERVER_URL}/v1/service/ready}" +DEFAULT_PYTHON="${ROOT_DIR}/.venv-aiperf/bin/python" +ADAPTER_ROOT="${ROOT_DIR}/benchmarks/telefuser_aiperf" +AIPERF_PYTHON="${TELEFUSER_AIPERF_PYTHON:-${DEFAULT_PYTHON}}" + +if [[ ! -x "${AIPERF_PYTHON}" ]]; then + echo "The isolated AIPerf environment is unavailable. Run: bash scripts/setup_aiperf.sh" >&2 + exit 1 +fi + +if command -v curl >/dev/null 2>&1; then + echo "Checking TeleFuser VLA readiness: ${HEALTH_URL}" + curl --noproxy '*' --fail --silent --show-error "${HEALTH_URL}" >/dev/null +fi + +export PYTHONPATH="${ADAPTER_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +exec "${AIPERF_PYTHON}" -m telefuser_aiperf.cli profile --config "${CONFIG_PATH}" diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py index 7834d47c..2a3b3473 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/__init__.py @@ -6,6 +6,12 @@ from telefuser_aiperf.adapter import TeleFuserLiveKitAdapter from telefuser_aiperf.sglang_adapter import SGLangRealtimeAdapter +from telefuser_aiperf.vla_structured import ( + ENDPOINT_METADATA, + TRANSPORT_METADATA, + TeleFuserStructuredHttpTransport, + TeleFuserVlaStructuredEndpoint, +) def register_adapters(*, replace: bool = False) -> None: @@ -23,4 +29,43 @@ def register_adapters(*, replace: bool = False) -> None: ) -__all__ = ["SGLangRealtimeAdapter", "TeleFuserLiveKitAdapter", "register_adapters"] +def register_plugins(*, replace: bool = False) -> None: + """Register repository-owned AIPerf batch endpoint and transport plugins.""" + from aiperf.plugin import plugins + from aiperf.plugin.enums import EndpointType, TransportType + + if "telefuser_vla_structured" not in EndpointType: + EndpointType.register("TELEFUSER_VLA_STRUCTURED", "telefuser_vla_structured") + if "telefuser_structured_http" not in TransportType: + TransportType.register("TELEFUSER_STRUCTURED_HTTP", "telefuser_structured_http") + + definitions = ( + ( + "endpoint", + "telefuser_vla_structured", + TeleFuserVlaStructuredEndpoint, + ENDPOINT_METADATA, + ), + ( + "transport", + "telefuser_structured_http", + TeleFuserStructuredHttpTransport, + TRANSPORT_METADATA, + ), + ) + for category, name, plugin_class, metadata in definitions: + if plugins.has_entry(category, name): + if not replace: + continue + plugins.unregister(category, name) + plugins.register(category, name, plugin_class, metadata=metadata) + + +__all__ = [ + "SGLangRealtimeAdapter", + "TeleFuserLiveKitAdapter", + "TeleFuserStructuredHttpTransport", + "TeleFuserVlaStructuredEndpoint", + "register_adapters", + "register_plugins", +] diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py index 3ad1f4be..de77e0a8 100644 --- a/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/cli.py @@ -6,9 +6,10 @@ def main() -> None: """Register TeleFuser adapters and delegate to the AIPerf CLI.""" - from telefuser_aiperf import register_adapters + from telefuser_aiperf import register_adapters, register_plugins register_adapters() + register_plugins() from aiperf.cli import app diff --git a/benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py b/benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py new file mode 100644 index 00000000..a7a64697 --- /dev/null +++ b/benchmarks/telefuser_aiperf/telefuser_aiperf/vla_structured.py @@ -0,0 +1,385 @@ +"""AIPerf endpoint and HTTP polling transport for TeleFuser VLA actions.""" + +from __future__ import annotations + +import asyncio +import base64 +import binascii +import hashlib +import math +import struct +import time +from dataclasses import dataclass +from typing import Any +from urllib.parse import quote, urlsplit, urlunsplit + +import orjson +from aiperf.common.exceptions import NotInitializedError +from aiperf.common.models import ( + BaseResponseData, + ErrorDetails, + InferenceServerResponse, + ParsedResponse, + RequestInfo, + RequestRecord, + TextResponse, +) +from aiperf.endpoints.base_endpoint import BaseEndpoint +from aiperf.plugin.schema.schemas import TransportMetadata +from aiperf.transports.aiohttp_transport import AioHttpTransport + +_RESULT_FIELDS = frozenset( + { + "canonical_normalized_actions", + "horizon", + "action_dim", + "checkpoint_variant", + "policy_verified", + "verification_status", + } +) +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) +_EXPECTED_HORIZON = 50 +_EXPECTED_ACTION_DIM = 55 +_POLL_INTERVAL_SECONDS = 0.05 + + +@dataclass(slots=True) +class VlaActionResponseData(BaseResponseData): + """Validated summary of one VLA action chunk.""" + + task_id: str + horizon: int + action_dim: int + value_count: int + sha256_float64_le: str + checkpoint_variant: str + policy_verified: bool + verification_status: str + inference_time_s: float | None = None + peak_memory_mb: float | None = None + + +def _finite_number(value: Any, *, name: str, allow_none: bool = False) -> float | None: + if value is None and allow_none: + return None + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"{name} must be a finite number") + normalized = float(value) + if not math.isfinite(normalized): + raise ValueError(f"{name} must be a finite number") + return normalized + + +def validate_state(state: Any) -> list[float]: + """Validate the canonical 14-dimensional VLA state vector.""" + if not isinstance(state, list) or len(state) != 14: + raise ValueError("VLA state must contain exactly 14 values") + return [float(_finite_number(value, name="VLA state value")) for value in state] + + +def image_content_to_base64(content: str) -> str: + """Normalize an AIPerf image data URL to raw validated base64.""" + if content.lower().startswith(("http://", "https://")): + raise ValueError("TeleFuser VLA camera inputs must be inline image data, not URLs") + encoded = content + if content.startswith("data:"): + try: + header, encoded = content.split(",", 1) + except ValueError as error: + raise ValueError("VLA image data URL is missing a comma") from error + if ";base64" not in header.lower() or not header.lower().startswith("data:image/"): + raise ValueError("VLA camera input must be a base64 image data URL") + try: + decoded = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as error: + raise ValueError("VLA camera input is not valid base64") from error + if not decoded: + raise ValueError("VLA camera input is empty") + return encoded + + +def build_vla_payload( + instruction: str, + image_content: str, + *, + extra: dict[str, Any] | None, +) -> dict[str, Any]: + """Build the stable TeleFuser structured action request body.""" + if not instruction.strip(): + raise ValueError("VLA instruction must not be empty") + parameters = dict(extra or {}) + unsupported = sorted(set(parameters).difference({"state", "seed"})) + if unsupported: + raise ValueError(f"Unsupported VLA request fields: {', '.join(unsupported)}") + if "state" not in parameters: + raise ValueError("VLA dataset entry must provide state in extra") + state = validate_state(parameters["state"]) + seed = parameters.get("seed", 7) + if isinstance(seed, bool) or not isinstance(seed, int): + raise ValueError("VLA seed must be an integer") + image_base64 = image_content_to_base64(image_content) + return { + "task": "vla_action", + "instruction": instruction, + "state": state, + "camera_high": image_base64, + "camera_left_wrist": image_base64, + "camera_right_wrist": image_base64, + "seed": seed, + } + + +def summarize_action_result(result: Any) -> dict[str, Any]: + """Validate a 50 x 55 normalized action chunk and return bounded facts.""" + if not isinstance(result, dict) or set(result) != set(_RESULT_FIELDS): + observed = sorted(result) if isinstance(result, dict) else type(result).__name__ + raise ValueError(f"VLA result fields changed: {observed}") + actions = result.get("canonical_normalized_actions") + if not isinstance(actions, list) or len(actions) != _EXPECTED_HORIZON: + raise ValueError(f"VLA action horizon must be {_EXPECTED_HORIZON}") + if result.get("horizon") != _EXPECTED_HORIZON or result.get("action_dim") != _EXPECTED_ACTION_DIM: + raise ValueError("VLA action dimension metadata changed") + + digest = hashlib.sha256() + value_count = 0 + minimum = math.inf + maximum = -math.inf + for row_index, row in enumerate(actions): + if not isinstance(row, list) or len(row) != _EXPECTED_ACTION_DIM: + raise ValueError(f"VLA action row {row_index} must contain {_EXPECTED_ACTION_DIM} values") + for raw_value in row: + value = float(_finite_number(raw_value, name="VLA action value")) + digest.update(struct.pack(" str: + """Build the native task status URL on the same origin as submission.""" + parsed = urlsplit(submit_url) + path = f"/v1/tasks/{quote(task_id, safe='')}/status" + return urlunsplit((parsed.scheme, parsed.netloc, path, "", "")) + + +class TeleFuserVlaStructuredEndpoint(BaseEndpoint): + """Format and parse TeleFuser's native VLA structured API.""" + + def format_payload(self, request_info: RequestInfo) -> dict[str, Any]: + if not request_info.turns: + raise ValueError("TeleFuser VLA endpoint requires one dataset turn") + turn = request_info.turns[-1] + if not turn.texts or not turn.texts[0].contents: + raise ValueError("TeleFuser VLA endpoint requires one instruction") + if not turn.images or not turn.images[0].contents: + raise ValueError("TeleFuser VLA endpoint requires one camera image") + merged_extra = dict(request_info.model_endpoint.endpoint.extra or {}) + merged_extra.update(turn.extra_body or {}) + return build_vla_payload( + turn.texts[0].contents[0], + turn.images[0].contents[0], + extra=merged_extra, + ) + + def parse_response(self, response: InferenceServerResponse) -> ParsedResponse | None: + body = response.get_json() + if not isinstance(body, dict) or body.get("status") != "completed": + return None + summary = body.get("action_summary") + if not isinstance(summary, dict): + raise ValueError("completed VLA benchmark response has no action_summary") + data = VlaActionResponseData( + task_id=str(body["task_id"]), + horizon=int(summary["horizon"]), + action_dim=int(summary["action_dim"]), + value_count=int(summary["value_count"]), + sha256_float64_le=str(summary["sha256_float64_le"]), + checkpoint_variant=str(summary["checkpoint_variant"]), + policy_verified=bool(summary["policy_verified"]), + verification_status=str(summary["verification_status"]), + inference_time_s=body.get("inference_time_s"), + peak_memory_mb=body.get("peak_memory_mb"), + ) + return ParsedResponse(perf_ns=response.perf_ns, data=data, metadata={"media_type": "structured"}) + + +class TeleFuserStructuredHttpTransport(AioHttpTransport): + """HTTP JSON transport for TeleFuser submit/poll structured tasks.""" + + @classmethod + def metadata(cls) -> TransportMetadata: + return TransportMetadata(transport_type="telefuser_structured_http", url_schemes=[]) + + @staticmethod + def _parse_json_record(record: RequestRecord, context: str) -> tuple[dict[str, Any], TextResponse] | ErrorDetails: + if record.error: + return record.error + if not record.responses or not isinstance(record.responses[0], TextResponse): + return ErrorDetails(type="VlaStructuredError", message=f"No JSON response from {context}", code=500) + response = record.responses[0] + try: + body = orjson.loads(response.text) + except orjson.JSONDecodeError: + return ErrorDetails(type="VlaStructuredError", message=f"Invalid JSON from {context}", code=500) + if not isinstance(body, dict): + return ErrorDetails(type="VlaStructuredError", message=f"Non-object JSON from {context}", code=500) + return body, response + + async def send_request( + self, + request_info: RequestInfo, + payload: dict[str, Any], + *, + first_token_callback: Any = None, + ) -> RequestRecord: + """Submit one action task, poll terminal state, and retain bounded facts.""" + del first_token_callback + if self.aiohttp_client is None: + raise NotInitializedError("AioHttpClient not initialized") + start_ns = time.perf_counter_ns() + headers = self.build_headers(request_info) + responses: list[TextResponse] = [] + + def make_record(error: ErrorDetails | None = None, status: int | None = None) -> RequestRecord: + return RequestRecord( + request_info=request_info, + request_headers=headers, + start_perf_ns=start_ns, + end_perf_ns=time.perf_counter_ns(), + responses=responses, + error=error, + status=status, + ) + + try: + submit_url = self.build_url(request_info) + submitted = await self.aiohttp_client.post_request(submit_url, orjson.dumps(payload), headers) + parsed_submit = self._parse_json_record(submitted, "VLA task submission") + if isinstance(parsed_submit, ErrorDetails): + return make_record(error=parsed_submit, status=submitted.status) + submit_body, submit_response = parsed_submit + task_id = submit_body.get("task_id") + if not isinstance(task_id, str) or not task_id: + return make_record( + error=ErrorDetails( + type="VlaStructuredError", + message="VLA submission returned no task_id", + code=500, + ) + ) + responses.append( + TextResponse( + perf_ns=submit_response.perf_ns, + text=orjson.dumps({"task_id": task_id, "status": "pending"}).decode(), + content_type="application/json", + ) + ) + + status_url = build_task_status_url(submit_url, task_id) + timeout = request_info.model_endpoint.endpoint.timeout + deadline = time.monotonic() + timeout if timeout > 0 else math.inf + while time.monotonic() < deadline: + polled = await self.aiohttp_client.get_request(status_url, headers) + parsed_poll = self._parse_json_record(polled, "VLA task status") + if isinstance(parsed_poll, ErrorDetails): + return make_record(error=parsed_poll, status=polled.status) + status_body, status_response = parsed_poll + status = status_body.get("status") or status_body.get("task_status") + if status not in _TERMINAL_STATUSES: + await asyncio.sleep(_POLL_INTERVAL_SECONDS) + continue + if status != "completed": + return make_record( + error=ErrorDetails( + type="VlaStructuredError", + message=f"VLA task {task_id} ended with {status}: {status_body.get('error')}", + code=500, + ), + status=polled.status, + ) + action_summary = summarize_action_result(status_body.get("result")) + bounded_status = { + "task_id": task_id, + "status": "completed", + "inference_time_s": _finite_number( + status_body.get("inference_time_s"), name="inference_time_s", allow_none=True + ), + "peak_memory_mb": _finite_number( + status_body.get("peak_memory_mb"), name="peak_memory_mb", allow_none=True + ), + "action_summary": action_summary, + } + responses.append( + TextResponse( + perf_ns=status_response.perf_ns, + text=orjson.dumps(bounded_status).decode(), + content_type="application/json", + ) + ) + return make_record(status=200) + return make_record( + error=ErrorDetails( + type="TimeoutError", + message=f"VLA task {task_id} timed out after {timeout:g}s", + code=504, + ), + status=504, + ) + except asyncio.CancelledError: + raise + except Exception as error: + return make_record(error=ErrorDetails.from_exception(error)) + + +ENDPOINT_METADATA = { + "endpoint_path": "/v1/tasks/structured", + "supports_streaming": False, + "tokenizes_input": False, + "produces_tokens": False, + "supports_images": True, + "requires_polling": True, + "requires_form_data": False, + "metrics_title": "TeleFuser VLA Structured Metrics", + "service_kind": "telefuser_vla", +} + +TRANSPORT_METADATA = { + "transport_type": "telefuser_structured_http", + "url_schemes": [], +} + + +__all__ = [ + "ENDPOINT_METADATA", + "TRANSPORT_METADATA", + "TeleFuserStructuredHttpTransport", + "TeleFuserVlaStructuredEndpoint", + "VlaActionResponseData", + "build_task_status_url", + "build_vla_payload", + "image_content_to_base64", + "summarize_action_result", + "validate_state", +] diff --git a/benchmarks/telefuser_aiperf/tests/test_vla_structured.py b/benchmarks/telefuser_aiperf/tests/test_vla_structured.py new file mode 100644 index 00000000..1ca58d81 --- /dev/null +++ b/benchmarks/telefuser_aiperf/tests/test_vla_structured.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import base64 +import math + +import orjson +import pytest +from aiperf.common.models import TextResponse +from aiperf.plugin import plugins +from telefuser_aiperf import register_plugins +from telefuser_aiperf.vla_structured import ( + TeleFuserStructuredHttpTransport, + TeleFuserVlaStructuredEndpoint, + VlaActionResponseData, + build_task_status_url, + build_vla_payload, + summarize_action_result, +) + + +def _action_result(value: float = 0.25) -> dict: + return { + "canonical_normalized_actions": [[value] * 55 for _ in range(50)], + "horizon": 50, + "action_dim": 55, + "checkpoint_variant": "base", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + } + + +def test_registration_uses_aiperf_endpoint_and_transport_plugins() -> None: + register_plugins(replace=True) + + endpoint_class = plugins.get_class("endpoint", "telefuser_vla_structured") + transport_class = plugins.get_class("transport", "telefuser_structured_http") + + assert endpoint_class is TeleFuserVlaStructuredEndpoint + assert transport_class is TeleFuserStructuredHttpTransport + assert plugins.get_endpoint_metadata("telefuser_vla_structured").requires_polling is True + + +def test_build_vla_payload_reuses_inline_image_for_three_cameras() -> None: + encoded = base64.b64encode(b"image bytes").decode() + + payload = build_vla_payload( + "pick up the object", + f"data:image/png;base64,{encoded}", + extra={"state": [0.0] * 14, "seed": 7}, + ) + + assert payload["task"] == "vla_action" + assert payload["camera_high"] == encoded + assert payload["camera_left_wrist"] == encoded + assert payload["camera_right_wrist"] == encoded + assert payload["state"] == [0.0] * 14 + + +@pytest.mark.parametrize( + "extra,match", + [ + ({"state": [0.0] * 13}, "exactly 14"), + ({"state": [0.0] * 13 + [math.nan]}, "finite"), + ({"state": [0.0] * 14, "unknown": 1}, "Unsupported"), + ], +) +def test_build_vla_payload_rejects_contract_drift(extra: dict, match: str) -> None: + encoded = base64.b64encode(b"image bytes").decode() + + with pytest.raises(ValueError, match=match): + build_vla_payload("instruction", encoded, extra=extra) + + +def test_summarize_action_result_validates_shape_and_omits_full_actions() -> None: + summary = summarize_action_result(_action_result()) + + assert summary["horizon"] == 50 + assert summary["action_dim"] == 55 + assert summary["value_count"] == 2750 + assert summary["minimum"] == 0.25 + assert len(summary["sha256_float64_le"]) == 64 + assert "canonical_normalized_actions" not in summary + + +def test_summarize_action_result_rejects_non_finite_action() -> None: + result = _action_result() + result["canonical_normalized_actions"][0][0] = math.inf + + with pytest.raises(ValueError, match="finite"): + summarize_action_result(result) + + +def test_endpoint_parses_bounded_completed_response() -> None: + endpoint = object.__new__(TeleFuserVlaStructuredEndpoint) + summary = summarize_action_result(_action_result()) + response = TextResponse( + perf_ns=123, + content_type="application/json", + text=orjson.dumps( + { + "task_id": "task-1", + "status": "completed", + "inference_time_s": 0.65, + "peak_memory_mb": None, + "action_summary": summary, + } + ).decode(), + ) + + parsed = endpoint.parse_response(response) + + assert parsed is not None + assert isinstance(parsed.data, VlaActionResponseData) + assert parsed.data.task_id == "task-1" + assert parsed.data.value_count == 2750 + assert parsed.metadata == {"media_type": "structured"} + + +def test_task_status_url_uses_native_structured_route() -> None: + assert ( + build_task_status_url("http://127.0.0.1:18080/v1/tasks/structured", "task id") + == "http://127.0.0.1:18080/v1/tasks/task%20id/status" + ) diff --git a/benchmarks/telefuser_aiperf/vla_structured_contract.yaml b/benchmarks/telefuser_aiperf/vla_structured_contract.yaml new file mode 100644 index 00000000..81ca0cc8 --- /dev/null +++ b/benchmarks/telefuser_aiperf/vla_structured_contract.yaml @@ -0,0 +1,39 @@ +contract_version: v1 +name: telefuser_lingbot_vla_v2_structured +mode: structured_action +implementation: telefuser +model_family: lingbot_vla_v2 +model: lingbot-vla-v2-6b +supported_tasks: + - vla_action +transport: http_polling +endpoint: + submit_path: /v1/tasks/structured + status_path: /v1/tasks/{task_id}/status + protocol: telefuser_structured_task +request_encoding: + content_type: application/json + parameters: + instruction: text + state: extra.state + camera_high: image + camera_left_wrist: image + camera_right_wrist: image + seed: extra.seed +result_delivery: + terminal_status: completed + result_field: result + action_field: canonical_normalized_actions + expected_shape: [50, 55] +workload: + warmup_requests: 2 + profile_requests: 20 + concurrency: 1 +metrics: + - request_latency + - request_throughput + - success_rate + - server_metrics +artifacts: + config: benchmarks/telefuser_aiperf/configs/vla_structured_e2e.yaml + dataset: benchmarks/telefuser_aiperf/data/vla_structured.jsonl diff --git a/docs/en/benchmark_aiperf.md b/docs/en/benchmark_aiperf.md index 94c1d143..d7dc2067 100644 --- a/docs/en/benchmark_aiperf.md +++ b/docs/en/benchmark_aiperf.md @@ -2,8 +2,8 @@ TeleFuser exposes raw target-side facts; AIPerf owns workload execution, aggregation, resource collection, artifacts, GreptimeDB history, and visualization. The checked-in integration covers batch video generation through the -OpenAI-compatible `/v1/videos` API, TeleFuser LingBot streaming through LiveKit, and SGLang LingBot streaming through -its native realtime WebSocket endpoint. +OpenAI-compatible `/v1/videos` API, LingBot-VLA structured actions through native HTTP task polling, TeleFuser LingBot +streaming through LiveKit, and SGLang LingBot streaming through its native realtime WebSocket endpoint. AIPerf's stream runner and result schema are transport-neutral. The LiveKit adapter is maintained by TeleFuser, loads from source at process startup, and produces AIPerf's standard session results. The contract records WebRTC as @@ -80,6 +80,7 @@ parity comparisons. See the benchmark README for model, GPU, port, and executabl |---|---|---| | TeleFuser runtime | TeleFuser | Emit synchronized phase, chunk, runtime, cache, and environment facts | | Batch target adapter | AIPerf | Convert `/v1/videos` HTTP events into the standard request timeline | +| VLA structured adapter | TeleFuser | Validate action results and convert native submit/poll events into bounded AIPerf records | | LiveKit source adapter | TeleFuser | Convert room, track, status, metrics, and control events into session results | | SGLang source adapter | TeleFuser | Convert MessagePack frames, chunk timings, and camera events into session results | | Aggregation and history | AIPerf | Apply warmup, percentiles, throughput, artifacts, GreptimeDB, and visualization | diff --git a/docs/en/service.md b/docs/en/service.md index 0f956890..26be39e3 100644 --- a/docs/en/service.md +++ b/docs/en/service.md @@ -152,7 +152,7 @@ telefuser serve /path/to/pipeline --task i2v [OPTIONS] | Parameter | Shortcut | Type | Default | Description | |-----------|----------|------|---------|-------------| | `pipe_path` | | string | **Required** | Positional path to the pipeline Python file | -| `--task` | `-t` | choice | `i2v` | Task type: t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr | +| `--task` | `-t` | choice | `i2v` | Task type: t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr, vla_action | | `--port` | `-p` | int | `8000` | Server port | | `--host` | | string | `127.0.0.1` | Server host address | | `--cache-dir` | `-c` | string | `work_dirs/server_cache` | Cache directory | @@ -258,6 +258,7 @@ telefuser serve --help | `i2i` | Image-to-Image: Generate image from input image and prompt | | `s2v` | Speech-to-Video: Generate video from speech | | `vsr` | Video Super-Resolution: Upscale an input video | +| `vla_action` | Structured VLA canonical action inference | ### Environment Variables @@ -1102,6 +1103,17 @@ telefuser serve ./pipeline.py --task t2v --- +### Structured JSON Tasks + +Pipelines that return JSON instead of image or video artifacts can declare a task contract with +`media_type="structured"`. Submit those tasks to `POST /v1/tasks/structured`; the existing status, cancellation, +queue, pool, and metrics endpoints remain unchanged. The pipeline entrypoint must return a finite JSON object. On +completion, `GET /v1/tasks/{task_id}/status` exposes that object under `result` together with +`inference_time_s` and the optional `peak_memory_mb`. + + +Media tasks continue to use `POST /v1/tasks/create` and artifact paths. The structured endpoint rejects media +contracts, and the media endpoint retains its existing request and response format. ## Client SDK ### Installation diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md new file mode 100644 index 00000000..918c0090 --- /dev/null +++ b/examples/lingbot_vla_v2/README.md @@ -0,0 +1,348 @@ +# LingBot-VLA v2 Base Model SDK + +This example loads the official LingBot-VLA v2 6B base checkpoint through TeleFuser and returns its normalized +55-dimensional canonical action chunk. The RobotWin profile is used only to prepare the example observation; the +result is not converted to physical RobotWin actions. + +## Inputs + +- Three RGB cameras in the upstream RobotWin order: high, left wrist, right wrist. +- A raw 14-dimensional RobotWin state. +- A non-empty task string. + +The SDK applies the bundled upstream RobotWin `bounds_99_woclip` statistics and maps the observation into +LingBot's 55-dimensional canonical state. + +## Output + +The pipeline returns `LingBotVlaV2CanonicalActionChunk` with: + +- `canonical_normalized_actions`: `[H, 55]` base-model output. +- `horizon`: action chunk length, normally 50 for the official base config. +- `action_dim`: canonical action dimension, normally 55. +- `checkpoint_variant`: `base`. +- `policy_verified=False` and `verification_status="unverified_official_6b_base"`. + +## Checkpoints + +The VLA directory must contain `model.safetensors.index.json` and every referenced shard. The Qwen3-VL directory +supplies the visual-language backbone configuration and processor. + +## Example + +```bash +python examples/lingbot_vla_v2/lingbot_vla_v2_inference.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --camera-high /data/cam_high.png \ + --camera-left-wrist /data/cam_left_wrist.png \ + --camera-right-wrist /data/cam_right_wrist.png \ + --task "pick up the red block" \ + --state-json '[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]' \ + --output canonical_action_chunk.npz +``` + +The example saves canonical actions and checkpoint metadata in an `.npz` file. The base output must not be sent to +a robot without an embodiment-specific post-training checkpoint, action mapping, and policy validation. + +## Minimal Single-GPU HTTP Service + +The VLA-specific server loads one policy replica and serializes all inference calls on the selected GPU. It does not +use the shared media service, Ray, multi-GPU execution, dynamic batching, or robot control. Start it from the repository +with the isolated VLA environment: + +```bash +.venv-vla/bin/python examples/lingbot_vla_v2/lingbot_vla_v2_server.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --device cuda:0 \ + --host 127.0.0.1 \ + --port 8000 +``` + +The process reports ready only after both the processor and policy have loaded: + +```bash +curl http://127.0.0.1:8000/health +``` + +`POST /v1/vla/actions` accepts raw Base64 or a Base64 data URL for each camera. The state must contain exactly 14 +finite values. For example: + +```bash +.venv-vla/bin/python - <<'PY' +import base64 +from pathlib import Path + +import httpx + + +def encode(path: str) -> str: + return base64.b64encode(Path(path).read_bytes()).decode("ascii") + + +response = httpx.post( + "http://127.0.0.1:8000/v1/vla/actions", + json={ + "task": "pick up the red block", + "state": [0.0] * 14, + "camera_high": encode("/data/cam_high.png"), + "camera_left_wrist": encode("/data/cam_left_wrist.png"), + "camera_right_wrist": encode("/data/cam_right_wrist.png"), + "seed": 7, + }, + timeout=300.0, +) +response.raise_for_status() +print(response.json()) +PY +``` + +The response contains `canonical_normalized_actions`, `horizon`, `action_dim`, `checkpoint_variant`, +`policy_verified`, and `verification_status`. A successful HTTP response confirms service and model execution only; +the normalized base-model output is not a physical robot command. + +## Native TeleFuser Service + +The native service uses the shared `PIPELINE_CONTRACT`, asynchronous task scheduler, pipeline pool, status API, runtime +metrics, and `TFClient`. It keeps the standalone endpoint above as a small debugging path. + +The example resolves checkpoints under the existing `TF_MODEL_ZOO_PATH` layout: + +- `lingbot/lingbot-vla-v2-6b` +- `Qwen3-VL-4B-Instruct` + +Start one replica on one visible GPU: + +```bash +TF_MODEL_ZOO_PATH=/hhb-data/aigc/model_zoo \ + .venv-vla/bin/telefuser serve \ + examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py \ + --task vla_action \ + --parallelism 1 \ + --host 127.0.0.1 \ + --port 18080 +``` + +Submit `POST /v1/tasks/structured` with `task="vla_action"`, an `instruction`, the 14-dimensional `state`, and +the three Base64 camera fields. The creation response contains a task ID. Poll +`GET /v1/tasks/{task_id}/status`; a completed response contains the action payload under `result` and includes +`inference_time_s` and the optional `peak_memory_mb`. + +The unified client handles image encoding, submission, polling, and result extraction: + +```python +from telefuser.client import TFClient + +client = TFClient("http://127.0.0.1:18080") +actions = client.predict_vla_actions( + instruction="pick up the red block", + state=[0.0] * 14, + camera_high_path="/data/cam_high.png", + camera_left_wrist_path="/data/cam_left_wrist.png", + camera_right_wrist_path="/data/cam_right_wrist.png", + seed=7, +) +print(actions["horizon"], actions["action_dim"]) +``` + +For independent replicas, expose one GPU per replica through the existing pipeline pool: + +```bash +CUDA_VISIBLE_DEVICES=0,1 TF_MODEL_ZOO_PATH=/hhb-data/aigc/model_zoo \ + .venv-vla/bin/telefuser serve \ + examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py \ + --task vla_action \ + --parallelism 2 \ + --num-replicas 2 \ + --port 18080 +``` + +This is request-level replication, not tensor parallelism inside one policy replica. The response remains a normalized +base-model canonical action chunk and must not be treated as a physical robot command. + +## Single-GPU Service Benchmark + +Use the VLA-specific benchmark to measure checkpoint construction, first-request latency, steady-state latency, +sequential throughput, process RSS, CUDA allocator peaks, and source-image-size overhead. The pipeline always converts +the three source images to the official `256x256` model input, so source size affects boundary and preprocessing cost, +not the model token shape. + +```bash +CUDA_VISIBLE_DEVICES=0 .venv-vla/bin/python \ + tools/validation/benchmark_lingbot_vla_v2_service.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --image examples/data/lingbot_world_fast/image.jpg \ + --image-sizes 256x256,640x480,1280x720 \ + --warmup 1 \ + --runs 20 \ + --output work_dirs/vla_service_benchmark/report.json +``` + +The native service moves the policy to its target GPU and runs one synthetic fixed-shape warmup before readiness. It +also keeps the allocator cache between requests. The report records construction and startup warmup separately, while +the first accepted request represents a ready replica. The default `service-thread` execution mode matches the native +service runner's fixed worker thread; use `--execution-mode direct` only to measure the in-process pipeline ceiling. +Shutdown still offloads the policy explicitly. + +## Native Structured API Validation + +Use the VLA-specific HTTP validator after the native service reports ready. This is the structured-output counterpart +to the model-specific direct and AIPerf workloads used by the video and LingBot-World integrations: it exercises the +real TeleFuser HTTP boundary, asynchronous scheduler, task status polling, pipeline pool, and result serialization. +It emits raw request facts and aggregate latency distributions to a JSON artifact; it does not add a VLA-specific +service interface or change shared metric semantics. + +Run a single-replica smoke and latency check: + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_structured_service.py \ + --base-url http://127.0.0.1:18080 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --warmup 1 \ + --requests 20 \ + --concurrency 1 \ + --output work_dirs/vla_service_validation/smoke_20.json +``` + +When the target was started with two independent replicas, validate request-level concurrency with: + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_structured_service.py \ + --base-url http://127.0.0.1:18080 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --warmup 2 \ + --requests 100 \ + --concurrency 2 \ + --output work_dirs/vla_service_validation/two_replica_100.json +``` + +Use duration mode for a bounded soak. Workers use closed-loop scheduling: each worker submits its next request only +after its previous task reaches a terminal state. + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_structured_service.py \ + --base-url http://127.0.0.1:18080 \ + --camera-high /data/cam_high.png \ + --camera-left-wrist /data/cam_left_wrist.png \ + --camera-right-wrist /data/cam_right_wrist.png \ + --duration-seconds 7200 \ + --concurrency 1 \ + --service-pid \ + --gpu-indexes 0 \ + --resource-interval-seconds 1 \ + --output work_dirs/vla_service_validation/soak_2h.json +``` + +Resource sampling is opt-in and local-only. `--service-pid` must identify the parent `telefuser serve` process; its +replica descendants are discovered on every sample. RSS is summed across that process tree, while `nvidia-smi` +process memory is grouped by physical GPU index. For a two-replica service on physical GPUs 0 and 1, pass +`--gpu-indexes 0,1`. Omitting `--service-pid` keeps remote-service validation lightweight and does not invoke +`nvidia-smi`. Reports retain bounded raw samples plus distributions and first/last 10% trends for latency, RSS, and +per-GPU process memory. + +The validator freezes the current structured contract. Requests contain exactly `task`, `instruction`, `state`, the +three camera fields, and optional `seed`. Action results contain exactly `canonical_normalized_actions`, `horizon`, +`action_dim`, `checkpoint_variant`, `policy_verified`, and `verification_status`. Safe additive task-status metadata +remains allowed, but status responses must not echo the three Base64 camera fields. + +The command exits nonzero if readiness or contract checks fail, any measured request fails, task IDs are duplicated, +or the queue is not drained at the end. Each successful record validates the expected `50x55` finite action tensor +and retains only statistics and a float64 action fingerprint. Full actions and Base64 camera contents are deliberately +excluded from the artifact. `--max-records` bounds retained per-request samples during long runs while aggregate +latency and success counters still cover the complete run. `--max-resource-samples` independently bounds retained +resource samples. + +For fault handling, run the independent validator against a ready service: + +```bash +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_service_faults.py \ + --base-url http://127.0.0.1:18080 \ + --image examples/data/lingbot_world_fast/image.jpg +``` + +It checks missing cameras, invalid state size, invalid Base64, and cancellation. Replica termination is opt-in and +requires a disposable two-replica service: add `--service-pid ` and +`--kill-replica-gpu-index `. The tool only selects a GPU compute process inside that parent process +tree, sends `SIGTERM`, and verifies one-replica capacity degradation plus a subsequent valid `50x55` response. It does +not promise automatic replica restart. + +The same structured API is available through the repository-owned AIPerf workload. Install the pinned isolated +AIPerf environment once, then run the workload while the native service is ready: + +```bash +bash scripts/setup_aiperf.sh +bash benchmarks/telefuser_aiperf/scripts/run_vla_structured_bench.sh +``` + +AIPerf excludes the configured warmup, aggregates request latency, throughput, success, traces, and server metrics, +and writes normal AIPerf artifacts. The adapter strictly validates the action contract but retains only bounded action +facts, not full arrays or Base64 inputs. Passing either validator proves serving and normalized action structure, not +embodiment-specific control semantics. + +## TeleFuser Regression Baseline + +The validation capture runs through the public loader and pipeline, then records preprocessing tensors, fixed initial +noise, every flow-matching `x_t` and velocity step, and the final canonical action. Run it twice before changing VLA +model code to establish and verify a strict local baseline: + +```bash +.venv-vla/bin/python tools/validation/capture_lingbot_vla_v2_telefuser.py \ + --model-root /hhb-data/aigc/model_zoo/lingbot/lingbot-vla-v2-6b \ + --qwen3vl-root /hhb-data/aigc/model_zoo/Qwen3-VL-4B-Instruct \ + --camera-high /data/cam_high.png \ + --camera-left-wrist /data/cam_left_wrist.png \ + --camera-right-wrist /data/cam_right_wrist.png \ + --task "pick up the red block" \ + --state-json '[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]' \ + --seed 7 \ + --deterministic-moe \ + --output work_dirs/vla_regression/baseline_seed7.npz + +# Repeat the same command with: +# --output work_dirs/vla_regression/replay_seed7.npz + +.venv-vla/bin/python tools/validation/run_lingbot_vla_v2_parity.py \ + --reference work_dirs/vla_regression/baseline_seed7.npz \ + --candidate work_dirs/vla_regression/replay_seed7.npz \ + --profile strict \ + --output work_dirs/vla_regression/strict_report.json +``` + +Each `.npz` has a same-name `.json` sidecar containing the checkpoint, processor, input, runtime, and tensor contract +metadata. The default checkpoint identity is a fast filename-and-size manifest. Add `--full-checkpoint-hash` when a +content hash of every checkpoint shard is required. Keep generated artifacts under `work_dirs`; do not commit them. + +This is a TeleFuser regression check, not upstream parity. It detects changes to the current implementation but does +not establish equivalence with the official repository. + +## Official Upstream Parity + +The strict upstream baseline pins `Robbyant/lingbot-vla-v2` at commit +`be27333c9b5f2663b0ec33f069dd7dfd67fa32b5`. Keep the checkout, uv environment, cache, and artifacts under +`work_dirs`; Git ignores them. Create the isolated runtime with: + +```bash +mkdir -p work_dirs/.uv-cache-upstream work_dirs/.uv-tmp-upstream +UV_CACHE_DIR="$PWD/work_dirs/.uv-cache-upstream" TMPDIR="$PWD/work_dirs/.uv-tmp-upstream" uv venv work_dirs/.venv-lingbot-upstream --python .venv-vla/bin/python +UV_CACHE_DIR="$PWD/work_dirs/.uv-cache-upstream" TMPDIR="$PWD/work_dirs/.uv-tmp-upstream" uv pip install --python work_dirs/.venv-lingbot-upstream/bin/python -r tools/validation/requirements-lingbot-vla-v2-upstream.txt +UV_CACHE_DIR="$PWD/work_dirs/.uv-cache-upstream" TMPDIR="$PWD/work_dirs/.uv-tmp-upstream" uv pip install --python work_dirs/.venv-lingbot-upstream/bin/python --no-deps "lerobot @ https://github.com/huggingface/lerobot/archive/refs/tags/v0.4.2.tar.gz" +git clone https://github.com/Robbyant/lingbot-vla-v2 work_dirs/lingbot-vla-v2-upstream +git -C work_dirs/lingbot-vla-v2-upstream checkout be27333c9b5f2663b0ec33f069dd7dfd67fa32b5 +``` + +Generate the reference with `capture_lingbot_vla_v2_upstream.py` in the upstream uv environment and the candidate +with `capture_lingbot_vla_v2_telefuser.py` in `.venv-vla`. Pass identical model, processor, camera, task, state, seed, +and device arguments to both commands, add `--deterministic-moe`, and pass `--upstream-root` to the upstream command. +Then compare them with the strict comparator shown above. Generated artifacts belong in `work_dirs/vla_upstream_parity`. + +This is a minimal inference-parity runtime, not a LeRobot training environment. The upstream setup itself combines +LeRobot 0.4.2 metadata constraints with versions outside those constraints, so LeRobot is installed with `--no-deps`; +the capture import and end-to-end run are the runtime checks. + +The official code hard-codes FlashAttention during construction. The upstream capture replaces that selection only +inside its validation process so both sides use eager attention on the Python 3.10.12 / PyTorch 2.11 stack. Production +inference keeps the upstream Triton MoE path through `telefuser.ops`; strict capture uses `--deterministic-moe` because +the upstream kernel uses atomic accumulation and is not bitwise repeatable across separate processes. Artifact metadata +records both `attention_backend` and `moe_backend`, and the comparator rejects mixed-backend artifacts. diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py new file mode 100644 index 00000000..608e9bff --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py @@ -0,0 +1,93 @@ +"""Run the LingBot-VLA v2 base checkpoint with a RobotWin observation adapter.""" + +from __future__ import annotations + +import json + +import click +import numpy as np + +from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_CAMERA_KEYS, + LingBotVlaV2Observation, + LingBotVlaV2Pipeline, +) +from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline + + +def get_pipeline( + model_root: str, + qwen3vl_root: str, + device: str = "cuda", +) -> LingBotVlaV2Pipeline: + """Load the official 6B checkpoint and Qwen3-VL processor.""" + return get_lingbot_vla_v2_pipeline(model_root, qwen3vl_root, device=device) + + +@click.command() +@click.option("--model-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--qwen3vl-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--camera-high", required=True, type=click.Path(exists=True, dir_okay=False)) +@click.option("--camera-left-wrist", required=True, type=click.Path(exists=True, dir_okay=False)) +@click.option("--camera-right-wrist", required=True, type=click.Path(exists=True, dir_okay=False)) +@click.option("--task", required=True) +@click.option("--state-json", required=True, help="Raw 14-D RobotWin state as a JSON list") +@click.option("--output", default="canonical_action_chunk.npz", type=click.Path(dir_okay=False)) +@click.option("--seed", default=None, type=int) +@click.option("--device", default="cuda") +def main( + model_root: str, + qwen3vl_root: str, + camera_high: str, + camera_left_wrist: str, + camera_right_wrist: str, + task: str, + state_json: str, + output: str, + seed: int | None, + device: str, +) -> None: + """Predict and save a normalized canonical action chunk.""" + try: + state = json.loads(state_json) + except json.JSONDecodeError as error: + raise click.BadParameter("state-json must be valid JSON") from error + if not isinstance(state, list) or len(state) != 14: + raise click.BadParameter("state-json must decode to a 14-element JSON list") + observation = LingBotVlaV2Observation( + task=task, + state=state, + images=dict( + zip( + ROBOTWIN_CAMERA_KEYS, + (camera_high, camera_left_wrist, camera_right_wrist), + strict=True, + ) + ), + ) + pipeline = get_pipeline( + model_root, + qwen3vl_root, + device=device, + ) + try: + chunk = pipeline(observation, seed=seed) + arrays = { + "canonical_normalized_actions": chunk.canonical_normalized_actions.numpy(), + "horizon": np.asarray(chunk.horizon), + "action_dim": np.asarray(chunk.action_dim), + "checkpoint_variant": np.asarray(chunk.checkpoint_variant), + "policy_verified": np.asarray(chunk.policy_verified), + "verification_status": np.asarray(chunk.verification_status), + } + np.savez(output, **arrays) + click.echo( + f"Saved {chunk.horizon}-step normalized canonical action chunk to {output}; " + f"policy status: {chunk.verification_status}" + ) + finally: + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py new file mode 100644 index 00000000..deb6b3e2 --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py @@ -0,0 +1,115 @@ +"""Native TeleFuser service contract for LingBot-VLA v2 action inference.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2Pipeline +from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline +from telefuser.pipelines.lingbot_vla_v2.service import ( + LingBotVlaV2ActionRequest, + predict_lingbot_vla_v2_action, +) + +TF_MODEL_ZOO_PATH = Path(os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo")).expanduser() + +PPL_CONFIG = { + "model_root": str(TF_MODEL_ZOO_PATH / "lingbot" / "lingbot-vla-v2-6b"), + "qwen3vl_root": str(TF_MODEL_ZOO_PATH / "Qwen3-VL-4B-Instruct"), + "device": "cuda:0", + "max_image_bytes": 10 * 1024 * 1024, +} + +PIPELINE_CONTRACT = { + "contract_version": "v1", + "pipeline_name": "lingbot_vla_v2_6b_base", + "supported_tasks": ["vla_action"], + "supported_media_types": ["structured"], + "execution_mode": "serial_single_pipeline", + "effective_max_concurrent_tasks": 1, + "entrypoints": { + "get_pipeline": "get_pipeline", + "run_with_file": "run_structured", + }, + "task_contracts": { + "vla_action": { + "media_type": "structured", + "required_inputs": ["camera_high", "camera_left_wrist", "camera_right_wrist"], + "optional_inputs": [], + "parameters": { + "instruction": { + "type": "string", + "required": True, + "description": "Robot instruction.", + }, + "state": { + "type": "array", + "required": True, + "description": "Raw 14-dimensional RobotWin state.", + }, + "camera_high": { + "type": "string", + "required": True, + "description": "Base64-encoded high camera image.", + }, + "camera_left_wrist": { + "type": "string", + "required": True, + "description": "Base64-encoded left wrist camera image.", + }, + "camera_right_wrist": { + "type": "string", + "required": True, + "description": "Base64-encoded right wrist camera image.", + }, + "seed": { + "type": "integer", + "required": False, + "default": None, + "description": "Optional deterministic inference seed.", + }, + }, + } + }, +} + + +def get_pipeline(parallelism: int = 1) -> LingBotVlaV2Pipeline: + """Load one policy replica for the native TeleFuser service.""" + if parallelism != 1: + raise ValueError("LingBot-VLA v2 supports parallelism=1 per replica; use --num-replicas for a pipeline pool") + return get_lingbot_vla_v2_pipeline( + PPL_CONFIG["model_root"], + PPL_CONFIG["qwen3vl_root"], + device=PPL_CONFIG["device"], + warmup=True, + ) + + +def run_structured( + pipeline: LingBotVlaV2Pipeline, + instruction: str, + state: list[float], + camera_high: str, + camera_left_wrist: str, + camera_right_wrist: str, + seed: int | None = None, + **_: Any, +) -> dict[str, Any]: + """Return one JSON-serializable canonical normalized action chunk.""" + request = LingBotVlaV2ActionRequest( + task=instruction, + state=state, + camera_high=camera_high, + camera_left_wrist=camera_left_wrist, + camera_right_wrist=camera_right_wrist, + seed=seed, + ) + response = predict_lingbot_vla_v2_action( + pipeline, + request, + max_image_bytes=int(PPL_CONFIG["max_image_bytes"]), + ) + return response.model_dump(mode="json") diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_server.py new file mode 100644 index 00000000..22dfecbc --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_server.py @@ -0,0 +1,38 @@ +"""Start a minimal single-GPU LingBot-VLA v2 HTTP service.""" + +from __future__ import annotations + +import click +import uvicorn + +from telefuser.pipelines.lingbot_vla_v2.service import LingBotVlaV2ServiceConfig, create_lingbot_vla_v2_app + + +@click.command() +@click.option("--model-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--qwen3vl-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--device", default="cuda:0", show_default=True) +@click.option("--host", default="127.0.0.1", show_default=True) +@click.option("--port", default=8000, show_default=True, type=click.IntRange(1, 65535)) +@click.option("--max-image-mb", default=10, show_default=True, type=click.IntRange(1, 100)) +def main( + model_root: str, + qwen3vl_root: str, + device: str, + host: str, + port: int, + max_image_mb: int, +) -> None: + """Load one policy replica and serve normalized canonical actions.""" + config = LingBotVlaV2ServiceConfig( + model_root=model_root, + qwen3vl_root=qwen3vl_root, + device=device, + max_image_bytes=max_image_mb * 1024 * 1024, + ) + app = create_lingbot_vla_v2_app(config) + uvicorn.run(app, host=host, port=port, workers=1) + + +if __name__ == "__main__": + main() diff --git a/telefuser/client/tf_client.py b/telefuser/client/tf_client.py index f71add4c..a44a06c2 100644 --- a/telefuser/client/tf_client.py +++ b/telefuser/client/tf_client.py @@ -95,10 +95,22 @@ TASK_S2V = "s2v" TASK_VSR = "vsr" TASK_EDIT = "edit" +TASK_VLA_ACTION = "vla_action" VIDEO_TASKS = (TASK_T2V, TASK_I2V, TASK_FL2V, TASK_VC, TASK_S2V, TASK_VSR) IMAGE_TASKS = (TASK_T2I, TASK_I2I, TASK_EDIT) -VALID_TASK_TYPES = (TASK_T2V, TASK_I2V, TASK_FL2V, TASK_VC, TASK_T2I, TASK_I2I, TASK_S2V, TASK_VSR, TASK_EDIT) +VALID_TASK_TYPES = ( + TASK_T2V, + TASK_I2V, + TASK_FL2V, + TASK_VC, + TASK_T2I, + TASK_I2I, + TASK_S2V, + TASK_VSR, + TASK_EDIT, + TASK_VLA_ACTION, +) # ── Constants: Aspect Ratios ──────────────────────────────────────────────── @@ -248,6 +260,74 @@ def create_task(self, task_type: str, **params: Any) -> Dict[str, Any]: except requests.RequestException as e: raise TaskCreationError(f"Task creation request failed: {e}") from e + def create_structured_task(self, task_type: str, **params: Any) -> Dict[str, Any]: + """Create a task whose pipeline contract returns a JSON result.""" + payload = {"task": task_type, **params} + try: + response = self._session.post( + f"{self.base_url}/v1/tasks/structured", + json=payload, + timeout=self.timeout, + ) + response.raise_for_status() + return response.json() + except requests.HTTPError as error: + raise TaskCreationError( + f"Structured task creation failed (HTTP {error.response.status_code}): {error.response.text}" + ) from error + except requests.RequestException as error: + raise TaskCreationError(f"Structured task creation request failed: {error}") from error + + def create_vla_action_task( + self, + instruction: str, + state: List[float], + camera_high_path: str, + camera_left_wrist_path: str, + camera_right_wrist_path: str, + seed: Optional[int] = None, + ) -> Dict[str, Any]: + """Create a LingBot-VLA v2 canonical action inference task.""" + return self.create_structured_task( + TASK_VLA_ACTION, + instruction=instruction, + state=state, + camera_high=self._encode_file_input(camera_high_path), + camera_left_wrist=self._encode_file_input(camera_left_wrist_path), + camera_right_wrist=self._encode_file_input(camera_right_wrist_path), + seed=seed, + ) + + def predict_vla_actions( + self, + instruction: str, + state: List[float], + camera_high_path: str, + camera_left_wrist_path: str, + camera_right_wrist_path: str, + seed: Optional[int] = None, + timeout: int = 300, + poll_interval: float = 0.5, + ) -> Dict[str, Any]: + """Run LingBot-VLA v2 inference and return its structured action result.""" + created = self.create_vla_action_task( + instruction=instruction, + state=state, + camera_high_path=camera_high_path, + camera_left_wrist_path=camera_left_wrist_path, + camera_right_wrist_path=camera_right_wrist_path, + seed=seed, + ) + status = self.wait_for_completion( + created["task_id"], + timeout=timeout, + poll_interval=poll_interval, + ) + result = status.get("result") + if not isinstance(result, dict): + raise TaskFailedError(f"Task {created['task_id']} completed without a structured result") + return result + # ── Video task creation methods ────────────────────────────────────────── def create_t2v_task( diff --git a/telefuser/entrypoints/cli/main.py b/telefuser/entrypoints/cli/main.py index 67465027..bf69be6d 100644 --- a/telefuser/entrypoints/cli/main.py +++ b/telefuser/entrypoints/cli/main.py @@ -31,7 +31,7 @@ def main(): "-t", default="i2v", type=click.Choice(TaskType.values(), case_sensitive=False), - help="Task type (t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr)", + help="Task type (t2v, i2v, fl2v, vc, t2i, i2i, s2v, vsr, vla_action)", ) @click.option("--port", "-p", default=8000, type=int, help="Server port") @click.option("--host", default="127.0.0.1", type=str, help="Server host") diff --git a/telefuser/kernel/triton/lingbot_vla_v2_moe.py b/telefuser/kernel/triton/lingbot_vla_v2_moe.py new file mode 100644 index 00000000..eb851d58 --- /dev/null +++ b/telefuser/kernel/triton/lingbot_vla_v2_moe.py @@ -0,0 +1,265 @@ +# ruff: noqa: E741 +"""Triton grouped-MoE kernels for LingBot-VLA v2. + +Internal implementation; callers should use ``telefuser.ops.lingbot_vla_v2_moe``. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _zero_i32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.int32), mask=offs < N) + + +@triton.jit +def _zero_fp32_kernel(out_ptr, N: tl.constexpr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + tl.store(out_ptr + offs, tl.zeros((BLOCK,), dtype=tl.float32), mask=offs < N) + + +@triton.jit +def _moe_pack_selected_kernel( + selected_ptr, + route_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + T: tl.constexpr, + TOPK: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row = tl.program_id(0) + slots = tl.arange(0, BLOCK_K) + mask = slots < TOPK + experts = tl.load(selected_ptr + row * TOPK + slots, mask=mask, other=0).to(tl.int32) + pos = tl.atomic_add(counts_ptr + experts, 1, sem="relaxed", mask=mask) + store_mask = mask & (pos < MAX_ROUTES) + tl.store(rows_ptr + experts * MAX_ROUTES + pos, row, mask=store_mask) + tl.store(slots_ptr + experts * MAX_ROUTES + pos, slots, mask=store_mask) + + +@triton.jit +def _moe_gate_up_grouped_kernel( + x_ptr, + gate_ptr, + up_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + route_ptr, + inter_ptr, + T: tl.constexpr, + D: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + I: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_I: tl.constexpr, + BLOCK_D: tl.constexpr, +): + expert = tl.program_id(0) + bid_m = tl.program_id(1) + bid_i = tl.program_id(2) + count = tl.load(counts_ptr + expert).to(tl.int32) + start_m = bid_m * BLOCK_M + if start_m >= count: + return + route_idx = start_m + tl.arange(0, BLOCK_M) + offs_i = bid_i * BLOCK_I + tl.arange(0, BLOCK_I) + offs_d = tl.arange(0, BLOCK_D) + valid_m = route_idx < count + rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + acc_g = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) + acc_u = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32) + for d0 in range(0, D, BLOCK_D): + ds = d0 + offs_d + x = tl.load( + x_ptr + rows[:, None] * D + ds[None, :], + mask=valid_m[:, None] & (ds[None, :] < D), + other=0.0, + ) + gw = tl.load( + gate_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], + mask=(offs_i[None, :] < I) & (ds[:, None] < D), + other=0.0, + ) + uw = tl.load( + up_ptr + (expert * I + offs_i[None, :]) * D + ds[:, None], + mask=(offs_i[None, :] < I) & (ds[:, None] < D), + other=0.0, + ) + acc_g += tl.dot(x, gw) + acc_u += tl.dot(x, uw) + route = tl.load(route_ptr + rows * TOPK + slots, mask=valid_m, other=0.0).to(tl.float32) + silu = acc_g * (1.0 / (1.0 + tl.exp(-acc_g))) + val = silu * acc_u * route[:, None] + tl.store( + inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + offs_i[None, :]), + val.to(inter_ptr.dtype.element_ty), + mask=valid_m[:, None] & (offs_i[None, :] < I), + ) + + +@triton.jit +def _moe_down_grouped_kernel( + inter_ptr, + down_ptr, + counts_ptr, + rows_ptr, + slots_ptr, + out_ptr, + T: tl.constexpr, + D: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + I: tl.constexpr, + MAX_ROUTES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_I: tl.constexpr, +): + expert = tl.program_id(0) + bid_m = tl.program_id(1) + bid_d = tl.program_id(2) + count = tl.load(counts_ptr + expert).to(tl.int32) + start_m = bid_m * BLOCK_M + if start_m >= count: + return + route_idx = start_m + tl.arange(0, BLOCK_M) + offs_d = bid_d * BLOCK_D + tl.arange(0, BLOCK_D) + offs_i = tl.arange(0, BLOCK_I) + valid_m = route_idx < count + rows = tl.load(rows_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + slots = tl.load(slots_ptr + expert * MAX_ROUTES + route_idx, mask=valid_m, other=0).to(tl.int32) + acc = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32) + for i0 in range(0, I, BLOCK_I): + is_ = i0 + offs_i + x = tl.load( + inter_ptr + ((rows[:, None] * TOPK + slots[:, None]) * I + is_[None, :]), + mask=valid_m[:, None] & (is_[None, :] < I), + other=0.0, + ) + w = tl.load( + down_ptr + (expert * D + offs_d[None, :]) * I + is_[:, None], + mask=(offs_d[None, :] < D) & (is_[:, None] < I), + other=0.0, + ) + acc += tl.dot(x, w) + tl.atomic_add( + out_ptr + rows[:, None] * D + offs_d[None, :], + acc, + sem="relaxed", + mask=valid_m[:, None] & (offs_d[None, :] < D), + ) + + +def robby_moe_forward( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + workspace: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Inference-only grouped MoE path migrated from robbyvla_infer _moe.""" + if hidden_states.ndim != 2: + raise ValueError(f"hidden_states must be 2D, got {tuple(hidden_states.shape)}") + if selected_experts.ndim != 2 or routing_weights.ndim != 2: + raise ValueError("selected_experts and routing_weights must be 2D") + if not hidden_states.is_cuda: + raise ValueError("robby_moe_forward requires CUDA tensors") + + T, D = hidden_states.shape + E, I, weight_d = gate_weight.shape + top_k = selected_experts.shape[1] + if weight_d != D or up_weight.shape != gate_weight.shape or down_weight.shape != (E, D, I): + raise ValueError( + "Unexpected MoE weight shapes: " + f"hidden={tuple(hidden_states.shape)} gate={tuple(gate_weight.shape)} " + f"up={tuple(up_weight.shape)} down={tuple(down_weight.shape)}" + ) + + max_routes = T * top_k + if workspace is None: + counts = torch.empty((E,), device=hidden_states.device, dtype=torch.int32) + rows = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) + slots = torch.empty((E, max_routes), device=hidden_states.device, dtype=torch.int32) + inter = torch.empty((T, top_k, I), device=hidden_states.device, dtype=hidden_states.dtype) + out = torch.empty((T, D), device=hidden_states.device, dtype=torch.float32) + else: + counts = workspace["counts"] + rows = workspace["rows"] + slots = workspace["slots"] + inter = workspace["inter"] + out = workspace["out"] + + selected_i32 = selected_experts.to(torch.int32).contiguous() + route = routing_weights.contiguous() + + _zero_i32_kernel[(1,)](counts, E, BLOCK=triton.next_power_of_2(E), num_warps=1) + _moe_pack_selected_kernel[(T,)]( + selected_i32, + route, + counts, + rows, + slots, + T, + top_k, + max_routes, + BLOCK_K=triton.next_power_of_2(top_k), + num_warps=1, + ) + _moe_gate_up_grouped_kernel[(E, triton.cdiv(max_routes, 16), triton.cdiv(I, 32))]( + hidden_states, + gate_weight, + up_weight, + counts, + rows, + slots, + route, + inter, + T, + D, + E, + top_k, + I, + max_routes, + BLOCK_M=16, + BLOCK_I=32, + BLOCK_D=64, + num_warps=4, + ) + _zero_fp32_kernel[(triton.cdiv(out.numel(), 1024),)]( + out, + out.numel(), + BLOCK=1024, + num_warps=4, + ) + _moe_down_grouped_kernel[(E, triton.cdiv(max_routes, 16), triton.cdiv(D, 64))]( + inter, + down_weight, + counts, + rows, + slots, + out, + T, + D, + E, + top_k, + I, + max_routes, + BLOCK_M=16, + BLOCK_D=64, + BLOCK_I=64, + num_warps=4, + ) + return out.reshape_as(hidden_states) diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py new file mode 100644 index 00000000..653b0d95 --- /dev/null +++ b/telefuser/models/lingbot_vla_v2.py @@ -0,0 +1,1493 @@ +"""Native LingBot-VLA v2 policy and flow-matching implementation. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +# Copyright 2026 Robbyant Team and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, List, Literal, Optional, Union + +import einops +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from transformers import AutoConfig, AutoTokenizer, PreTrainedModel, PretrainedConfig +from transformers.cache_utils import Cache +from transformers.modeling_flash_attention_utils import is_flash_attn_available +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm +from transformers.models.qwen3_vl.modeling_qwen3_vl import apply_rotary_pos_emb +from transformers.utils import logging + +from telefuser.models.lingbot_vla_v2_loader import ( + LingBotVLAWeightLoader, + LingBotVlaV2StateDictConverter, + TaskTokenDepthHead, + block_suffix_to_fv_, + build_block_mask, + create_sinusoidal_pos_embedding, + flex_attention_forward, + flex_attention_with_block_mask, + make_att_2d_masks, + our_eager_attention_forward, + prefix_query_segments, + prefix_query_token_spans, +) +from telefuser.models.lingbot_vla_v2_moe import ( + FixQwen2RMSNorm, + Qwen2ForCausalLM, + Qwen2FusedExperts, + Qwen2TokenMoeBlock, +) +from telefuser.models.lingbot_vla_v2_qwen import ( + Qwen3VLForConditionalGeneration, + Qwen3VLPreTrainedModel, + Qwen3VLTextModel, +) + +try: + from dinov3.hub.backbones import dinov3_vitb16 +except Exception: + dinov3_vitb16 = None + +logger = logging.get_logger(__name__) + + +class LingbotVLAConfig(PretrainedConfig): + """Configuration class for Lingbot-VLA. + This is the configuration class to store the configuration of a [`Lingbot-VLA`]. + """ + + model_type = "lingbotvla" + is_composition = True + + def __init__( + self, + vlm_repo_id: Optional[str] = None, + expert_vision_path: Optional[str] = None, + tokenizer_path: Optional[str] = None, + post_training: bool = False, + adanorm_time: bool = False, + split_gate_liner: bool = False, + nosplit_gate_liner: bool = False, + separate_time_proj: bool = False, + final_norm_adanorm: bool = False, + enable_expert_vision: bool = False, + expert_vision_type: Optional[str] = None, + freeze_vision_encoder: bool = False, + incremental_training: bool = False, + depth_incremental_training: bool = False, + reinit_mismatched_weights: bool = False, + action_dim: int = 14, + max_action_dim: int = 14, + max_state_dim: int = 14, + chunk_size: int = 50, + vlm_causal: bool = False, + tokenizer_max_length: int = 48, + loss_type: str = "fm", + norm_qkv: bool = False, + align_params: Optional[Dict[str, Any]] = None, + use_compile: bool = False, + use_moe: bool = False, + token_moe_layers: Optional[list] = None, + token_num_experts: int = 32, + token_top_k: int = 1, + token_moe_intermediate_size: int = 256, + token_shared_intermediate_size: int = 256, + bias_update_speed: float = 0.001, + sequence_wise_loss_coeff: float = 0.001, + sequence_wise_mode: str = "per_sequence", + router_z_loss_coeff: float = 0.0, + router_activation: str = "softmax", + routed_scaling_factor: float = 1.0, + use_shared_expert_gate: bool = True, + moe_implementation: Optional[Literal[None, "eager", "fused"]] = None, + use_robby_moe_kernel: bool = False, + split_fused_experts_from_decoder_fsdp: bool = False, + expert_hidden_size: int = 768, + expert_intermediate_size: int = 2752, + action_num_attention_heads: int = 16, + action_num_key_value_heads: int = 2, + action_head_dim: int = 128, + action_fp32: bool = False, + use_qwen3_chat_template: bool = False, + return_image_grid_thw: bool = False, + qwen3vl_use_vision_boundaries: bool = False, + precompute_grid_thw: bool = False, + use_qwen3_fixed_grid_cache: bool = False, + use_lm_head: bool = False, + vocab_size: int = 0, + vit_attn_implementation: str = "flash_attention_2", + attention_implementation: str = "flex", + train_expert_only: bool = False, + train_state_proj: bool = True, + **kwargs, + ): + super().__init__() + if moe_implementation is None: + moe_implementation = kwargs.pop("_moe_implementation", None) + self.architectures = ["LingbotVlaPolicy"] + self.train_state_proj = train_state_proj + self.train_expert_only = train_expert_only + self.use_cache = False + self.attention_implementation = attention_implementation + self.num_steps = 10 + self.n_obs_steps = 1 + + assert not (split_gate_liner and nosplit_gate_liner), ( + "split_gate_liner and nosplit_gate_liner can not be both True." + ) + + self.vlm_repo_id = vlm_repo_id + self.expert_vision_path = expert_vision_path + self.tokenizer_path = tokenizer_path + self.post_training = post_training + self.adanorm_time = adanorm_time + self.split_gate_liner = split_gate_liner + self.nosplit_gate_liner = nosplit_gate_liner + self.enable_expert_vision = enable_expert_vision + self.expert_vision_type = expert_vision_type + self.incremental_training = incremental_training + self.depth_incremental_training = depth_incremental_training + self.reinit_mismatched_weights = reinit_mismatched_weights + self.norm_qkv = norm_qkv + self.use_compile = use_compile + self.loss_type = loss_type + self.separate_time_proj = separate_time_proj + self.final_norm_adanorm = final_norm_adanorm + self.freeze_vision_encoder = freeze_vision_encoder + self.tokenizer_max_length = tokenizer_max_length + self.action_dim = action_dim + self.max_action_dim = max_action_dim + self.max_state_dim = max_state_dim + self.chunk_size = chunk_size + self.n_action_steps = chunk_size + self.vlm_causal = vlm_causal + self.align_params = align_params + self.use_moe = use_moe + if self.use_moe: + self.token_moe_layers = token_moe_layers + self.token_num_experts = token_num_experts + self.token_top_k = token_top_k + self.token_moe_intermediate_size = token_moe_intermediate_size + self.token_shared_intermediate_size = token_shared_intermediate_size + self.bias_update_speed = bias_update_speed + self.sequence_wise_loss_coeff = sequence_wise_loss_coeff + self.sequence_wise_mode = sequence_wise_mode + self.router_z_loss_coeff = router_z_loss_coeff + self.router_activation = router_activation + self.routed_scaling_factor = routed_scaling_factor + self.use_shared_expert_gate = use_shared_expert_gate + self.moe_implementation = moe_implementation + self.use_robby_moe_kernel = use_robby_moe_kernel + if moe_implementation is not None: + if moe_implementation not in ("eager", "fused"): + raise ValueError(f"Invalid moe_implementation: {moe_implementation}") + self._moe_implementation = moe_implementation + self.split_fused_experts_from_decoder_fsdp = split_fused_experts_from_decoder_fsdp + self.expert_hidden_size = expert_hidden_size + self.expert_intermediate_size = expert_intermediate_size + self.action_num_attention_heads = action_num_attention_heads + self.action_num_key_value_heads = action_num_key_value_heads + self.action_head_dim = action_head_dim + self.action_fp32 = action_fp32 + self.use_qwen3_chat_template = use_qwen3_chat_template + self.return_image_grid_thw = return_image_grid_thw + self.qwen3vl_use_vision_boundaries = qwen3vl_use_vision_boundaries + self.precompute_grid_thw = precompute_grid_thw + self.use_qwen3_fixed_grid_cache = use_qwen3_fixed_grid_cache + self.use_lm_head = use_lm_head + if vocab_size == 0: + if vlm_repo_id and "paligemma" in vlm_repo_id.lower(): + self.vocab_size = 257216 + elif vlm_repo_id and "qwen" in vlm_repo_id.lower(): + self.vocab_size = 151936 + else: + self.vocab_size = 257152 + else: + self.vocab_size = vocab_size + self.vit_attn_implementation = vit_attn_implementation + + +class LingbotVLAV2Config(LingbotVLAConfig): + def __init__(self, **kwargs): + kwargs.setdefault("attention_implementation", "flex_cached") + kwargs.setdefault("vit_attn_implementation", "flash_attention_2") + kwargs.setdefault("action_num_attention_heads", 32) + kwargs.setdefault("action_num_key_value_heads", 8) + kwargs.setdefault("action_head_dim", 128) + kwargs.setdefault("expert_hidden_size", 768) + kwargs.setdefault("use_qwen3_chat_template", True) + kwargs.setdefault("return_image_grid_thw", True) + kwargs.setdefault("qwen3vl_use_vision_boundaries", True) + kwargs.setdefault("use_qwen3_fixed_grid_cache", True) + super().__init__(**kwargs) + self.architectures = ["LingbotVlaV2Policy"] + self.vlm_family = "qwen3_vl" + + +ConfigClass = [LingbotVLAConfig, LingbotVLAV2Config] +__all__ = ["LingbotVLAConfig", "LingbotVLAV2Config"] + + +class AdaRMSNorm(nn.Module): + def __init__(self, hidden_size, cond_dim, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.gamma = nn.Linear(cond_dim, hidden_size) + self.beta = nn.Linear(cond_dim, hidden_size) + + nn.init.zeros_(self.gamma.weight) + nn.init.zeros_(self.gamma.bias) + nn.init.zeros_(self.beta.weight) + nn.init.zeros_(self.beta.bias) + + def forward(self, hidden_states, cond): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + hidden_states = self.weight * hidden_states + gamma = self.gamma(cond).unsqueeze(1) + beta = self.beta(cond).unsqueeze(1) + hidden_states = (1 + gamma.to(torch.float32)) * hidden_states + beta.to(torch.float32) + return hidden_states.to(input_dtype) + + +class FixAdaRMSNorm(AdaRMSNorm): + def forward(self, hidden_states, cond): + return super().forward(hidden_states, cond.float()) + + +def replace_lnorm_with_adanorm(module, hidden_size, cond_dim, final_norm_adanorm): + for name, child in module.named_children(): + if isinstance(child, Qwen2RMSNorm) and "q_layernorm" not in name and "k_layernorm" not in name: + setattr(module, name, AdaRMSNorm(hidden_size, cond_dim)) + elif ( + final_norm_adanorm + and isinstance(child, FixQwen2RMSNorm) + and "q_layernorm" not in name + and "k_layernorm" not in name + ): + setattr(module, name, FixAdaRMSNorm(hidden_size, cond_dim)) + else: + replace_lnorm_with_adanorm(child, hidden_size, cond_dim, final_norm_adanorm) + + +class FlowMatchingBase(nn.Module): + def init_depth_heads(self, config): + self.llm_image_token_size = config["llm"]["image_token_size"] + self.llm_image_input_size = config["llm"]["image_input_size"] + self.depth_token_size = config["depth"]["token_size"] + self.depth_input_size = config["depth"]["input_size"] + self.align_type = config.get("mode", None) + self.model_type = config["depth"]["model_type"] + if self.align_type != "query": + raise ValueError(f"Only query depth alignment is supported, got {self.align_type!r}.") + if self.model_type != "MoRGBD": + raise ValueError(f"Only MoRGBD depth distillation is supported, got {self.model_type!r}.") + self.use_future_depth = (config.get("depth") or {}).get("use_future_depth", False) + self.block_future_depth_to_action = (config.get("depth") or {}).get("block_future_depth_to_action", False) + self.detach_future_depth_image_feats = bool((config.get("depth") or {}).get("detach_future_image_feats", False)) + self.use_future_video = bool(config.get("use_future_video", False)) + self.use_future_video_patch = False + self.use_current_video_patch = False + self.use_current_shared_task_proj = False + self.use_future_video_cls = False + self.use_shared_future_task_proj = False + self.future_video_share_future_depth_query = False + self.num_task_tokens = config["num_task_tokens"] + assert config["depth"]["num_backbone_tokens"] % self.num_task_tokens == 0 + self.depth_align_embs = nn.Parameter( + torch.randn(config["depth"]["num_backbone_tokens"], config["llm"]["dim_out"]) + ) + + self.depth_align_head = TaskTokenDepthHead(config["depth"], llm_hidden_size=config["llm"]["dim_out"]).to( + dtype=torch.bfloat16 + ) + + if self.use_future_depth: + self.future_depth_align_embs = nn.Parameter( + torch.randn(config["depth"]["num_backbone_tokens"], config["llm"]["dim_out"]) + ) + + self.future_depth_align_head = TaskTokenDepthHead( + config["depth"], llm_hidden_size=config["llm"]["dim_out"] + ).to(dtype=torch.bfloat16) + + def init_video_heads(self, config): + if self.align_type != "query": + raise ValueError("future-video alignment is only supported for query align mode.") + + video_config = dict(config.get("depth", {})) + video_config.update(config.get("video", {})) + required_keys = ("num_backbone_tokens", "dim_out", "num_layers", "num_heads", "dim_head", "ff_mult") + missing = [key for key in required_keys if key not in video_config] + if missing: + raise ValueError(f"video align config missing required keys: {missing}") + self.use_future_video_patch = bool(video_config.get("use_patch_loss", True)) + self.use_current_video_patch = bool(video_config.get("use_current_patch_loss", False)) + if self.use_current_video_patch and not self.use_future_video_patch: + raise ValueError( + "align_params.video.use_current_patch_loss=True requires align_params.video.use_patch_loss=True." + ) + self.use_current_shared_task_proj = bool( + video_config.get("use_current_shared_task_proj", self.use_current_video_patch) + ) + if self.use_current_shared_task_proj and not self.use_current_video_patch: + raise ValueError( + "align_params.video.use_current_shared_task_proj=True requires " + "align_params.video.use_current_patch_loss=True." + ) + self.use_future_video_cls = bool(video_config.get("use_cls_loss", False)) + self.future_video_share_future_depth_query = bool(video_config.get("share_future_depth_query", False)) + self.use_shared_future_task_proj = bool(video_config.get("use_shared_future_task_proj", False)) + if self.use_shared_future_task_proj and not self.use_future_video_patch: + raise ValueError( + "align_params.video.use_shared_future_task_proj=True requires align_params.video.use_patch_loss=True." + ) + if self.use_shared_future_task_proj and not self.future_video_share_future_depth_query: + raise ValueError( + "align_params.video.use_shared_future_task_proj=True requires " + "align_params.video.share_future_depth_query=True." + ) + if self.future_video_share_future_depth_query: + if not self.use_future_depth: + raise ValueError( + "align_params.video.share_future_depth_query=True requires " + "align_params.depth.use_future_depth=True." + ) + if int(video_config["num_backbone_tokens"]) != int(config["depth"]["num_backbone_tokens"]): + raise ValueError( + "future-video shared query requires video.num_backbone_tokens to match depth.num_backbone_tokens." + ) + + self.block_suffix_to_future_video = bool(video_config.get("block_suffix_to_future_video", False)) + self.future_video_context_mode = str(video_config.get("context_mode", "img_query")).lower() + if self.future_video_context_mode not in ("img_query", "query_only"): + raise ValueError( + "future-video context_mode must be 'img_query' or 'query_only', " + f"got {self.future_video_context_mode!r}." + ) + if self.use_future_video_patch: + if self.use_current_video_patch: + self.current_video_align_embs = nn.Parameter( + torch.randn(video_config["num_backbone_tokens"], config["llm"]["dim_out"]) + ) + if self.use_current_shared_task_proj: + self.current_shared_task_proj = nn.Linear( + config["llm"]["dim_out"] * 2, + config["llm"]["dim_out"], + ) + self.current_video_align_head = TaskTokenDepthHead( + video_config, llm_hidden_size=config["llm"]["dim_out"] + ).to(dtype=torch.bfloat16) + + if not self.future_video_share_future_depth_query or self.use_shared_future_task_proj: + self.future_video_align_embs = nn.Parameter( + torch.randn(video_config["num_backbone_tokens"], config["llm"]["dim_out"]) + ) + if self.use_shared_future_task_proj: + self.future_shared_task_proj = nn.Linear( + config["llm"]["dim_out"] * 2, + config["llm"]["dim_out"], + ) + self.future_video_align_head = TaskTokenDepthHead( + video_config, llm_hidden_size=config["llm"]["dim_out"] + ).to(dtype=torch.bfloat16) + + if self.use_future_video_cls: + self.future_video_cls_align_emb = nn.Embedding(1, config["llm"]["dim_out"]) + self.future_video_cls_head = nn.Sequential( + nn.LayerNorm(config["llm"]["dim_out"]), + nn.Linear(config["llm"]["dim_out"], video_config["dim_out"]), + ).to(dtype=torch.bfloat16) + + def _future_depth_token_count(self): + return self.num_task_tokens if getattr(self, "use_future_depth", False) else 0 + + def _future_video_own_token_count(self): + if not getattr(self, "use_future_video", False): + return 0 + count = 1 if getattr(self, "use_future_video_cls", False) else 0 + if getattr(self, "use_future_video_patch", True) and not getattr( + self, "future_video_share_future_depth_query", False + ): + count += self.num_task_tokens + return count + + def _future_video_own_span(self, hidden_states): + own_count = self._future_video_own_token_count() + future_depth_count = self._future_depth_token_count() + end = hidden_states.shape[1] - future_depth_count + start = end - own_count + return start, end + + def _future_depth_task_tokens(self, hidden_states): + if not getattr(self, "use_future_depth", False): + raise ValueError("future-depth query tokens are not enabled.") + return hidden_states[:, -self.num_task_tokens :, :] + + def _future_video_cls_task_tokens(self, hidden_states): + if not getattr(self, "use_future_video_cls", False): + return None + start, _ = self._future_video_own_span(hidden_states) + return hidden_states[:, start : start + 1, :] + + def _future_video_patch_task_tokens(self, hidden_states): + if getattr(self, "future_video_share_future_depth_query", False): + return self._future_depth_task_tokens(hidden_states) + start, end = self._future_video_own_span(hidden_states) + if getattr(self, "use_future_video_cls", False): + start += 1 + return hidden_states[:, start:end, :] + + def _current_depth_task_tokens(self, hidden_states, num_images=3): + chunk_size = self.llm_image_token_size * self.llm_image_token_size + image_token_len = chunk_size + (2 if getattr(self.config, "qwen3vl_use_vision_boundaries", False) else 0) + if getattr(self, "use_future_depth", False): + start = num_images * image_token_len + return hidden_states[:, start : start + self.num_task_tokens, :] + end = hidden_states.shape[1] - self._future_video_own_token_count() + start = end - self.num_task_tokens + return hidden_states[:, start:end, :] + + def _future_video_query_span(self, prefix_len): + if not getattr(self, "use_future_video", False): + return prefix_len, prefix_len + future_depth_count = self._future_depth_token_count() + own_count = self._future_video_own_token_count() + end = prefix_len - future_depth_count + return end - own_count, end + + def _block_suffix_to_future_video_(self, att_2d_masks, suffix_row_start, prefix_len): + start, end = self._future_video_query_span(prefix_len) + if end <= start: + return att_2d_masks + att_2d_masks[:, suffix_row_start:, start:end] = False + return att_2d_masks + + def _block_suffix_to_future_video_if_enabled_( + self, + att_2d_masks, + suffix_row_start, + prefix_len, + ): + if not getattr(self, "block_suffix_to_future_video", False): + return att_2d_masks + return self._block_suffix_to_future_video_( + att_2d_masks, + suffix_row_start=suffix_row_start, + prefix_len=prefix_len, + ) + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, (nn.Linear, nn.Conv3d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + if module.weight is not None: + module.weight.data.fill_(1.0) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Qwen2FusedExperts): + module.initializer_range = std + module.reset_parameters() + reset_post_init = getattr(module, "_reset_post_init_parameters", None) + if reset_post_init is not None: + reset_post_init() + + @staticmethod + def _fp32_linear(module, x): + """Compute linear layer in fp32 regardless of module's current parameter dtype.""" + return F.linear(x.float(), module.weight.float(), module.bias.float() if module.bias is not None else None) + + def embed_suffix( + self, state, noisy_actions, timestep + ): # (torch.Size([state_bs, 32]), torch.Size([1, state_bs*50, 32]), torch.Size([1])) + bsize = state.shape[0] # state_bs = img_bs + device = state.device + dtype = state.dtype + _fp32 = getattr(self.config, "action_fp32", False) + # embed state + state_emb = self._fp32_linear(self.state_proj, state) if _fp32 else self.state_proj(state) + + # embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1] + time_emb = create_sinusoidal_pos_embedding( # 1, 1024 + timestep, # torch.Size([1])) + self.config.proj_width, # 1024 + min_period=4e-3, + max_period=4.0, + device=device, + ) + time_emb = time_emb.type(dtype=dtype) + + time_emb_ori = time_emb + + # Fuse timestep + action information using an MLP + action_emb = ( + self._fp32_linear(self.action_in_proj, noisy_actions) if _fp32 else self.action_in_proj(noisy_actions) + ) # torch.Size([1, state_bs*50, 1024]) + time_emb = einops.repeat(time_emb, "b d -> b n d", n=action_emb.shape[1]) # [1, 1024] -> [1, state_bs*50, 1024] + action_time_emb = torch.cat([action_emb, time_emb], dim=-1) # [1, state_bs*50, 2048] + + action_time_emb = ( + self._fp32_linear(self.action_time_mlp_in, action_time_emb) + if _fp32 + else self.action_time_mlp_in(action_time_emb) + ) + action_time_emb = F.silu(action_time_emb) # swish == silu + action_time_emb = ( + self._fp32_linear(self.action_time_mlp_out, action_time_emb) + if _fp32 + else self.action_time_mlp_out(action_time_emb) + ) # [1, state_bs*50, 1024] + action_time_dim = action_time_emb.shape[1] + + embs = torch.cat([state_emb[:, None], action_time_emb], dim=1) + pad_masks = torch.ones((bsize, action_time_dim + 1), device=device, dtype=torch.bool) + + # Set attention masks for suffix tokens so that prefix tokens cannot attend to suffix tokens. + # And state token cannot attend action tokens. + # Action tokens use a bidirectional attention. + att_masks = torch.zeros((bsize, action_time_dim + 1), device=device, dtype=torch.bool) + att_masks[:, :2] = True + + return time_emb_ori, embs, pad_masks, att_masks + + +class QwenvlWithExpertV2Config(PretrainedConfig): + model_type = "QwenvlWithExpertV2Model" + + def __init__( + self, + freeze_vision_encoder: bool = False, + train_expert_only: bool = False, + vocab_size: int = 0, + use_lm_head: bool = False, + attention_implementation: str = "flex_cached", + tokenizer_path: str | None = None, + enable_expert_vision: bool = False, + expert_vision_type: str | None = None, + use_cache: bool = False, + expert_hidden_size: int = 768, + expert_intermediate_size: int = 2752, + action_num_attention_heads: int = 32, + action_num_key_value_heads: int = 8, + action_head_dim: int = 128, + **kwargs, + ): + self.freeze_vision_encoder = freeze_vision_encoder + self.train_expert_only = train_expert_only + self.attention_implementation = attention_implementation + self.tokenizer_path = tokenizer_path + self.enable_expert_vision = enable_expert_vision + self.expert_vision_type = expert_vision_type + self.vocab_size = vocab_size + self.use_lm_head = use_lm_head + self.action_num_attention_heads = action_num_attention_heads + self.action_num_key_value_heads = action_num_key_value_heads + self.action_head_dim = action_head_dim + num_layers = 36 + + self.qwen_expert_config = CONFIG_MAPPING["qwen2"]( + attention_dropout=0.0, + bos_token_id=151643, + eos_token_id=151645, + hidden_act="silu", + hidden_size=expert_hidden_size, + head_dim=action_head_dim, + initializer_range=0.02, + intermediate_size=expert_intermediate_size, + max_position_embeddings=32768, + max_window_layers=21, + model_type="qwen2", + num_attention_heads=action_num_attention_heads, + num_hidden_layers=num_layers, + num_key_value_heads=action_num_key_value_heads, + rms_norm_eps=1e-06, + rope_theta=1000000.0, + sliding_window=32768, + tie_word_embeddings=True, + torch_dtype="bfloat16", + transformers_version="4.57.3", + use_cache=use_cache, + use_sliding_window=False, + vocab_size=151936, + ) + print( + "=====Action Expert V2 init " + f"{num_layers} Layers, hidden={expert_hidden_size}, " + f"q_heads={action_num_attention_heads}, kv_heads={action_num_key_value_heads}, " + f"head_dim={action_head_dim}.=====" + ) + super().__init__(**kwargs) + + +class QwenvlWithExpertV2Model(PreTrainedModel): + config_class = QwenvlWithExpertV2Config + + def __init__(self, config: QwenvlWithExpertV2Config, eval=False): + super().__init__(config=config) + self.config = config + vlm_config = AutoConfig.from_pretrained(self.config.tokenizer_path, local_files_only=True) + if self.config.vocab_size not in (0, 257152): + vlm_config.text_config.vocab_size = self.config.vocab_size + base_attn_implementation = "flash_attention_2" if is_flash_attn_available() else "eager" + vision_attn_implementation = self.config.vit_attn_implementation + if vision_attn_implementation == "flash_attention_2" and not is_flash_attn_available(): + logger.warning_once("flash-attn is unavailable; using eager attention for Qwen3-VL vision") + vision_attn_implementation = "eager" + vlm_config._attn_implementation = base_attn_implementation + vlm_config.text_config._attn_implementation = base_attn_implementation + vlm_config.vision_config._attn_implementation = vision_attn_implementation + self.qwenvl = Qwen3VLForConditionalGeneration._from_config(vlm_config) + if self.config.use_lm_head: + self.qwenvl.tie_weights() + + self.config.qwen_expert_config._attn_implementation = base_attn_implementation + self.qwen_expert = Qwen2ForCausalLM._from_config(self.config.qwen_expert_config, eval=eval) + + if getattr(self.config, "adanorm_time", False): + replace_lnorm_with_adanorm( + self.qwen_expert, + self.config.qwen_expert_config.hidden_size, + self.config.qwen_expert_config.hidden_size, + config.final_norm_adanorm, + ) + + self._install_moe_blocks() + self.pos_embeds = None + self.position_embeddings = None + self.cu_seqlens = None + self.visual_split_sizes = None + self.visual_max_seqlen = None + self._cached_image_grid_signature = None + + del self.qwen_expert.model.embed_tokens + if self.config.enable_expert_vision: + if dinov3_vitb16 is None: + raise ImportError("dinov3 is required when enable_expert_vision=True") + if "dinov3_vitb16" in self.config.expert_vision_type: + self.expert_visual = dinov3_vitb16(pretrained=False) + self.expert_visual_mlp = nn.Sequential( + nn.Linear(self.expert_visual.embed_dim, self.expert_visual.embed_dim * 2), + nn.GELU(), + nn.Linear(self.expert_visual.embed_dim * 2, self.config.qwen_expert_config.hidden_size), + ) + + self.attention_interface = self.get_attention_interface() + + def _apply(self, fn): + super()._apply(fn) + for name in ("pos_embeds", "position_embeddings", "cu_seqlens"): + value = getattr(self, name, None) + if isinstance(value, torch.Tensor): + setattr(self, name, fn(value)) + elif isinstance(value, tuple): + setattr( + self, + name, + tuple(fn(item) if isinstance(item, torch.Tensor) else item for item in value), + ) + return self + + def _install_moe_blocks(self): + if not getattr(self.config, "use_moe", False): + return + bias_update_speed = getattr(self.config, "bias_update_speed", 0.001) + hidden_size = self.config.qwen_expert_config.hidden_size + token_moe_layers = getattr(self.config, "token_moe_layers", None) or [] + + _moe_impl = getattr(self.config, "_moe_implementation", None) + + if token_moe_layers: + token_config = CONFIG_MAPPING["qwen2_moe"]( + num_experts=getattr(self.config, "token_num_experts", 32), + num_experts_per_tok=getattr(self.config, "token_top_k", 1), + norm_topk_prob=True, + hidden_size=hidden_size, + moe_intermediate_size=getattr(self.config, "token_moe_intermediate_size", 256), + shared_expert_intermediate_size=getattr(self.config, "token_shared_intermediate_size", 256), + output_router_logits=False, + ) + token_config.bias_update_speed = bias_update_speed + token_config._moe_implementation = _moe_impl + token_config.router_activation = getattr(self.config, "router_activation", "softmax") + token_config.routed_scaling_factor = getattr(self.config, "routed_scaling_factor", 1.0) + token_config.use_shared_expert_gate = getattr(self.config, "use_shared_expert_gate", True) + token_config.use_robby_moe_kernel = getattr(self.config, "use_robby_moe_kernel", False) + for idx in token_moe_layers: + self.qwen_expert.model.layers[idx].mlp = Qwen2TokenMoeBlock(token_config) + + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor, + ): + precompute_grid_thw = getattr(self.config, "precompute_grid_thw", False) + grid_signature = tuple(image_grid_thw.detach().to(device="cpu").reshape(-1).tolist()) + cache_miss = self.position_embeddings is None or self._cached_image_grid_signature != grid_signature + if precompute_grid_thw and cache_miss: + ( + self.pos_embeds, + self.position_embeddings, + self.cu_seqlens, + self.visual_split_sizes, + self.visual_max_seqlen, + ) = self.qwenvl.visual.preprcess_grid_thw(grid_thw=image_grid_thw) + self._cached_image_grid_signature = grid_signature + image_embeds, deepstack_image_embeds = self.qwenvl.visual( + pixel_values, + grid_thw=image_grid_thw, + pos_embeds=self.pos_embeds, + position_embeddings=self.position_embeddings, + cu_seqlens=self.cu_seqlens, + max_seqlen=self.visual_max_seqlen, + ) + split_sizes = self.visual_split_sizes + if split_sizes is None: + split_sizes = (image_grid_thw.prod(-1) // self.qwenvl.visual.spatial_merge_size**2).tolist() + image_chunks = list(torch.split(image_embeds, split_sizes)) + deepstack_chunks = [ + list(torch.split(deepstack_embeds, split_sizes)) for deepstack_embeds in deepstack_image_embeds + ] + image_embeds = torch.stack(image_chunks, dim=0) + deepstack_image_embeds = [torch.stack(chunks, dim=0) for chunks in deepstack_chunks] + return image_embeds, deepstack_image_embeds + + def embed_image(self, image: torch.Tensor, image_grid_thw: torch.LongTensor): + return self.get_image_features( + image, + image_grid_thw=image_grid_thw, + ) + + def embed_language_tokens(self, tokens: torch.Tensor): + return self.qwenvl.model.language_model.embed_tokens(tokens) + + def embed_special_token(self, token_id: int, batch: int, count: int, device, dtype): + token = torch.tensor([token_id], device=device, dtype=torch.long) + emb = self.embed_language_tokens(token).to(dtype=dtype) + return emb.view(1, 1, 1, -1).expand(batch, count, 1, -1) + + def build_prefix_position_ids(self, input_ids, attention_mask, image_grid_thw=None, video_grid_thw=None): + position_ids, _ = self.qwenvl.model.get_rope_index( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + ) + return position_ids + + def apply_mrope(self, query_states, key_states, position_ids): + position_embeddings = self.qwenvl.model.language_model.rotary_emb(query_states, position_ids) + return apply_rotary_pos_emb(query_states, key_states, *position_embeddings, unsqueeze_dim=2) + + def handle_kv_cache( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + use_cache: Optional[bool] = None, + fill_kv_cache: Optional[bool] = None, + ): + if use_cache: + if past_key_values is None: + past_key_values = {} + if fill_kv_cache: + past_key_values[layer_idx] = {"key_states": key_states, "value_states": value_states} + else: + key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1) + value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1) + return key_states, value_states, past_key_values + + def _apply_deepstack(self, hidden_states, layer_idx, visual_pos_masks, deepstack_visual_embeds): + if ( + deepstack_visual_embeds is not None + and visual_pos_masks is not None + and layer_idx < len(deepstack_visual_embeds) + ): + hidden_states = self.qwenvl.model.language_model._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + return hidden_states + + def forward( + self, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + vlm_position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[List[torch.FloatTensor], Cache]] = None, + inputs_embeds: List[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + fill_kv_cache: Optional[bool] = None, + ada_cond: List[torch.FloatTensor] = None, + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, + ): + models = [self.qwenvl.model.language_model, self.qwen_expert.model] + num_layers = self.qwenvl.config.text_config.num_hidden_layers + action_num_layers = self.config.qwen_expert_config.num_hidden_layers + router_logits_list = [] + + assert action_num_layers == num_layers, ( + "Action expert and VLM must have the same number of layers " + f"(got action={action_num_layers}, vlm={num_layers})." + ) + + for layer_idx in range(num_layers): + query_states = [] + key_states = [] + value_states = [] + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + continue + if i == 1: + q, k, v = models[i].layers[layer_idx](hidden_states, compute_kqv=True, ada_cond=ada_cond) + else: + q, k, v = models[i].layers[layer_idx](hidden_states, compute_kqv=True) + query_states.append(q.float()) + key_states.append(k.float()) + value_states.append(v.float()) + + query_states = torch.cat(query_states, dim=1) + key_states = torch.cat(key_states, dim=1) + value_states = torch.cat(value_states, dim=1) + query_states, key_states = self.apply_mrope(query_states, key_states, position_ids) + key_states, value_states, past_key_values = self.handle_kv_cache( + key_states, + value_states, + layer_idx, + past_key_values=past_key_values, + use_cache=use_cache, + fill_kv_cache=fill_kv_cache, + ) + if self.config.attention_implementation == "flex_cached": + if layer_idx == 0: + _full_len = query_states.shape[1] + _full_block_mask = build_block_mask( + attention_mask, + self.qwenvl.config.text_config.num_attention_heads, + _full_len, + _full_len, + ) + att_output = flex_attention_with_block_mask( + query_states, key_states, value_states, _full_block_mask, query_states.shape[1] + ) + else: + att_output = self.attention_interface(query_states, key_states, value_states, attention_mask) + + outputs_embeds = [] + start = 0 + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + outputs_embeds.append(None) + continue + end = start + hidden_states.shape[1] + if i == 1: + out_emb, router_logits = models[i].layers[layer_idx]( + hidden_states, + att_output, + start, + end, + output_atten=True, + ada_cond=ada_cond, + ) + if router_logits is not None: + router_logits_list.append(router_logits) + else: + out_emb = models[i].layers[layer_idx](hidden_states, att_output, start, end, output_atten=True) + out_emb = self._apply_deepstack(out_emb, layer_idx, visual_pos_masks, deepstack_visual_embeds) + outputs_embeds.append(out_emb) + start = end + inputs_embeds = outputs_embeds + + outputs_embeds = [] + for i, hidden_states in enumerate(inputs_embeds): + if hidden_states is None: + outputs_embeds.append(None) + elif self.config.final_norm_adanorm and i == 1: + out_emb, _ = models[i].norm(hidden_states, ada_cond) + outputs_embeds.append(out_emb) + else: + outputs_embeds.append(models[i].norm(hidden_states)) + return outputs_embeds, past_key_values, router_logits_list + + def get_attention_interface(self): + if self.config.attention_implementation == "flex": + print("=====Using Flex Attn=====") + return flex_attention_forward + if self.config.attention_implementation == "flex_cached": + print("=====Using Flex Cached (prebuilt BlockMask) Attn=====") + return flex_attention_forward + if self.config.attention_implementation == "eager": + print("=====Using Eager Attn=====") + return our_eager_attention_forward + raise ValueError(f"Invalid attention implementation: {self.config.attention_implementation}") + + +class FlowMatchingV2(FlowMatchingBase): + def __init__(self, config, eval): + nn.Module.__init__(self) + self.config = config + qwenvl_with_export_config = QwenvlWithExpertV2Config( + freeze_vision_encoder=self.config.freeze_vision_encoder, + train_expert_only=self.config.train_expert_only, + vocab_size=getattr(self.config, "vocab_size", 0), + use_lm_head=getattr(self.config, "use_lm_head", False), + attention_implementation=self.config.attention_implementation, + tokenizer_path=self.config.tokenizer_path, + enable_expert_vision=self.config.enable_expert_vision, + expert_vision_type=self.config.expert_vision_type, + use_cache=getattr(self.config, "use_cache", True), + expert_hidden_size=getattr(self.config, "expert_hidden_size", 768), + expert_intermediate_size=getattr(self.config, "expert_intermediate_size", 2752), + action_num_attention_heads=getattr(self.config, "action_num_attention_heads", 32), + action_num_key_value_heads=getattr(self.config, "action_num_key_value_heads", 8), + action_head_dim=getattr(self.config, "action_head_dim", 128), + ) + for name in [ + "adanorm_time", + "final_norm_adanorm", + "precompute_grid_thw", + "vit_attn_implementation", + "use_moe", + "bias_update_speed", + "token_moe_layers", + "token_num_experts", + "token_top_k", + "token_moe_intermediate_size", + "token_shared_intermediate_size", + "router_activation", + "routed_scaling_factor", + "use_shared_expert_gate", + "use_robby_moe_kernel", + "_moe_implementation", + ]: + if hasattr(config, name): + setattr(qwenvl_with_export_config, name, getattr(config, name)) + self.qwenvl_with_expert = QwenvlWithExpertV2Model(qwenvl_with_export_config, eval) + self.config.proj_width = qwenvl_with_export_config.qwen_expert_config.hidden_size + self.config.initializer_range = getattr(qwenvl_with_export_config.qwen_expert_config, "initializer_range", None) + + self.state_proj = nn.Linear(self.config.max_state_dim, self.config.proj_width) + self.action_in_proj = nn.Linear(self.config.max_action_dim, self.config.proj_width) + self.action_out_proj = nn.Linear(self.config.proj_width, self.config.max_action_dim) + self.action_time_mlp_in = nn.Linear(self.config.proj_width * 2, self.config.proj_width) + self.action_time_mlp_out = nn.Linear(self.config.proj_width, self.config.proj_width) + + self.config.align_params = getattr(self.config, "align_params", None) or {} + if self.config.align_params != {}: + self.steps = 0 + self.use_depth_align = True + self.init_depth_heads(self.config.align_params) + self.use_future_video = self.config.align_params.get("use_future_video", False) + if self.use_future_video: + self.init_video_heads(self.config.align_params) + else: + self.use_depth_align = False + self.use_future_video = False + self.use_future_video_patch = False + self.use_current_video_patch = False + self.use_current_shared_task_proj = False + self.use_future_video_cls = False + self.use_shared_future_task_proj = False + self.future_video_share_future_depth_query = False + self.block_future_depth_to_action = False + + def embed_prefix( + self, + images, + img_masks, + lang_tokens, + lang_masks, + image_grid_thw=None, + ): + if image_grid_thw is None: + raise ValueError("LingbotVlaV2Policy requires image_grid_thw from the Qwen3-VL image processor.") + bsize = images.shape[0] + device = images.device + if images.ndim == 3: + bsize = 1 + num_images = images.shape[0] + else: + num_images = images.shape[1] if images.ndim >= 4 else 1 + if images.ndim == 4: + images = einops.rearrange(images, "b n l d -> (b n) l d") + elif images.ndim == 5: + images = einops.rearrange(images, "b n c h w -> (b n) c h w") + if image_grid_thw.ndim == 3: + flat_grid_thw = einops.rearrange(image_grid_thw, "b n d -> (b n) d") + else: + flat_grid_thw = image_grid_thw + + img_emb, deepstack_embs = self.qwenvl_with_expert.embed_image( + images, + flat_grid_thw, + ) + embed_dtype = img_emb.dtype + num_patch = img_emb.shape[1] + img_emb = einops.rearrange(img_emb, "(b n) l d -> b n l d", b=bsize, n=num_images) + deepstack_embs = [einops.rearrange(x, "(b n) l d -> b n l d", b=bsize, n=num_images) for x in deepstack_embs] + if img_masks.ndim == 1: + img_masks = img_masks.unsqueeze(0) + + cfg = self.qwenvl_with_expert.qwenvl.config + visual_token_id = cfg.image_token_id + + if getattr(self.config, "qwen3vl_use_vision_boundaries", True): + start_emb = self.qwenvl_with_expert.embed_special_token( + cfg.vision_start_token_id, bsize, num_images, device, embed_dtype + ) + end_emb = self.qwenvl_with_expert.embed_special_token( + cfg.vision_end_token_id, bsize, num_images, device, embed_dtype + ) + img_chunks = torch.cat([start_emb, img_emb, end_emb], dim=2) + image_token_len = num_patch + 2 + image_pad_masks = einops.repeat(img_masks, "b n -> b n l", l=image_token_len) + image_visual_masks = torch.zeros_like(image_pad_masks) + image_visual_masks[:, :, 1 : 1 + num_patch] = einops.repeat(img_masks, "b n -> b n l", l=num_patch) + fake_image_ids = torch.full( + (bsize, num_images, image_token_len), + visual_token_id, + dtype=torch.long, + device=device, + ) + fake_image_ids[:, :, 0] = cfg.vision_start_token_id + fake_image_ids[:, :, -1] = cfg.vision_end_token_id + else: + img_chunks = img_emb + image_token_len = num_patch + image_pad_masks = einops.repeat(img_masks, "b n -> b n l", l=image_token_len) + image_visual_masks = image_pad_masks + fake_image_ids = torch.full( + (bsize, num_images, image_token_len), + visual_token_id, + dtype=torch.long, + device=device, + ) + + img_emb = einops.rearrange(img_chunks, "b n l d -> b (n l) d") + image_pad_masks = einops.rearrange(image_pad_masks, "b n l -> b (n l)") + visual_pos_masks = einops.rearrange(image_visual_masks, "b n l -> b (n l)") + fake_image_ids = einops.rearrange(fake_image_ids, "b n l -> b (n l)") + + lang_emb = self.qwenvl_with_expert.embed_language_tokens(lang_tokens).to(dtype=embed_dtype) + + if self.use_depth_align and self.align_type == "query": + + def _get_align_tokens(tokens): + tk_weights = tokens.view(self.num_task_tokens, tokens.shape[0] // self.num_task_tokens, tokens.shape[1]) + tk_weights = tk_weights.mean(dim=1) + return tk_weights + + align_pad_masks = torch.ones(bsize, self.num_task_tokens, device=device, dtype=lang_masks.dtype) + fake_align_ids = torch.full( + (bsize, self.num_task_tokens), cfg.text_config.eos_token_id, dtype=torch.long, device=device + ) + + current_task = _get_align_tokens(self.depth_align_embs) + if ( + getattr(self, "use_future_video", False) + and getattr(self, "use_current_video_patch", False) + and getattr(self, "use_current_shared_task_proj", False) + ): + current_video_task = _get_align_tokens(self.current_video_align_embs) + current_task = self.current_shared_task_proj(torch.cat([current_task, current_video_task], dim=-1)) + align_embs = current_task.repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) + parts = [img_emb] + masks = [image_pad_masks] + input_ids = [fake_image_ids] + visual_masks = [visual_pos_masks] + + def _append( + tokens, + token_masks, + token_ids, + token_visual_masks=None, + ): + parts.append(tokens) + masks.append(token_masks) + input_ids.append(token_ids) + if token_visual_masks is None: + token_visual_masks = torch.zeros_like(token_masks) + visual_masks.append(token_visual_masks) + + future_align_embs = None + if self.use_future_depth: + future_task = _get_align_tokens(self.future_depth_align_embs) + if ( + getattr(self, "use_future_video", False) + and getattr(self, "use_future_video_patch", True) + and getattr(self, "future_video_share_future_depth_query", False) + and getattr(self, "use_shared_future_task_proj", False) + ): + future_video_task = _get_align_tokens(self.future_video_align_embs) + future_task = self.future_shared_task_proj(torch.cat([future_task, future_video_task], dim=-1)) + future_align_embs = future_task.repeat(img_emb.size(0), 1, 1).to(img_emb.device, img_emb.dtype) + + if ( + not self.use_future_depth + and getattr(self, "use_future_video", False) + and getattr(self, "future_video_share_future_depth_query", False) + ): + raise ValueError("share_future_depth_query=True requires depth.use_future_depth=True.") + + for segment_name in prefix_query_segments( + use_depth_align=True, + use_future_depth=self.use_future_depth, + use_future_video=getattr(self, "use_future_video", False), + use_future_video_cls=getattr(self, "use_future_video_cls", False), + use_future_video_patch=getattr(self, "use_future_video_patch", True), + future_video_share_future_depth_query=getattr( + self, + "future_video_share_future_depth_query", + False, + ), + ): + if segment_name == "language": + _append( + lang_emb, + lang_masks, + lang_tokens.to(device), + ) + elif segment_name == "current_depth": + _append(align_embs, align_pad_masks, fake_align_ids) + elif segment_name == "future_video_cls": + future_video_cls_align_emb = self.future_video_cls_align_emb.weight.repeat( + img_emb.size(0), 1, 1 + ).to(img_emb.device, img_emb.dtype) + cls_align_pad_masks = torch.ones( + bsize, + 1, + device=device, + dtype=lang_masks.dtype, + ) + fake_cls_align_ids = torch.full( + (bsize, 1), + cfg.text_config.eos_token_id, + dtype=torch.long, + device=device, + ) + _append(future_video_cls_align_emb, cls_align_pad_masks, fake_cls_align_ids) + elif segment_name == "future_video": + future_video_align_embs = ( + _get_align_tokens(self.future_video_align_embs) + .repeat(img_emb.size(0), 1, 1) + .to(img_emb.device, img_emb.dtype) + ) + _append(future_video_align_embs, align_pad_masks, fake_align_ids) + elif segment_name == "future_depth": + _append(future_align_embs, align_pad_masks, fake_align_ids) + else: + raise ValueError(f"Unsupported prefix query segment: {segment_name}") + + embs = torch.cat(parts, dim=1) + pad_masks = torch.cat(masks, dim=1) + prefix_input_ids = torch.cat(input_ids, dim=1) + full_visual_pos_masks = torch.cat(visual_masks, dim=1) + else: + embs = torch.cat([img_emb, lang_emb], dim=1) + pad_masks = torch.cat([image_pad_masks, lang_masks], dim=1) + prefix_input_ids = torch.cat([fake_image_ids, lang_tokens.to(device)], dim=1) + full_visual_pos_masks = torch.cat([visual_pos_masks, torch.zeros_like(lang_masks)], dim=1) + + if getattr(self.config, "vlm_causal", False): + att_masks = torch.ones((bsize, embs.shape[1]), device=device, dtype=torch.bool) + else: + att_masks = torch.zeros((bsize, embs.shape[1]), device=device, dtype=torch.bool) + + flat_img_masks = einops.rearrange(img_masks, "b n -> (b n)") + rope_grid_thw = flat_grid_thw[flat_img_masks] + if rope_grid_thw.numel() == 0: + rope_grid_thw = flat_grid_thw[:1] + prefix_position_ids = self.qwenvl_with_expert.build_prefix_position_ids( + prefix_input_ids, + pad_masks.long(), + image_grid_thw=rope_grid_thw, + video_grid_thw=None, + ) + filtered_deepstack = [] + img_visual_only = einops.repeat(img_masks, "b n -> b n l", l=num_patch) + for deepstack in deepstack_embs: + filtered_deepstack.append(deepstack[img_visual_only]) + + result = ( + embs, + pad_masks, + att_masks, + prefix_position_ids, + full_visual_pos_masks, + filtered_deepstack, + ) + return result + + def _build_full_position_ids(self, prefix_position_ids, prefix_pad_masks, suffix_pad_masks): + valid_prefix_pos = prefix_position_ids.masked_fill(~prefix_pad_masks.unsqueeze(0), 0) + prefix_offsets = valid_prefix_pos.amax(dim=(0, 2)) + 1 + suffix_1d = prefix_offsets[:, None] + torch.cumsum(suffix_pad_masks.long(), dim=1) - 1 + suffix_1d = suffix_1d.masked_fill(~suffix_pad_masks, 1) + suffix_position_ids = suffix_1d.unsqueeze(0).expand(3, -1, -1) + return torch.cat([prefix_position_ids, suffix_position_ids], dim=-1) + + def _current_depth_task_tokens(self, hidden_states, num_images=3): + query_spans = prefix_query_token_spans( + prefix_len=hidden_states.shape[1], + num_task_tokens=self.num_task_tokens, + use_depth_align=True, + use_future_depth=getattr(self, "use_future_depth", False), + use_future_video=getattr(self, "use_future_video", False), + use_future_video_cls=getattr(self, "use_future_video_cls", False), + use_future_video_patch=getattr(self, "use_future_video_patch", True), + future_video_share_future_depth_query=getattr( + self, + "future_video_share_future_depth_query", + False, + ), + ) + start, end = query_spans["current_depth"] + return hidden_states[:, start:end, :] + + def forward(self, *args, **kwargs): + """Reject the upstream training API in the inference-only model.""" + del args, kwargs + raise RuntimeError("LingBot-VLA v2 is inference-only; use sample_actions()") + + def sample_actions( + self, + images, + img_masks, + lang_tokens, + lang_masks, + state, + noise=None, + image_grid_thw=None, + ) -> Tensor: + """Do a full Qwen3-VL inference forward and compute the action.""" + bsize = state.shape[0] + device = state.device + dtype = state.dtype + + if noise is None: + actions_shape = ( + bsize, + self.config.n_action_steps, + self.config.max_action_dim, + ) + noise = torch.randn(actions_shape, device=device, dtype=dtype) + + ( + prefix_embs, + prefix_pad_masks, + prefix_att_masks, + prefix_position_ids, + visual_pos_masks, + deepstack_visual_embeds, + ) = self.embed_prefix( + images, + img_masks, + lang_tokens, + lang_masks, + image_grid_thw=image_grid_thw, + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + + _, past_key_values, _ = self.qwenvl_with_expert.forward( + attention_mask=prefix_att_2d_masks, + position_ids=prefix_position_ids, + vlm_position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=self.config.use_cache, + fill_kv_cache=True, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + + dt = torch.tensor(-1.0 / self.config.num_steps, dtype=dtype, device=device) + x_t = noise + time = torch.tensor(1.0, dtype=dtype, device=device) + count = 0 + predict_velocity_fn = self.predict_velocity + if getattr(self, "_use_compile_predict_velocity", False): + predict_velocity_fn = getattr(self, "_compiled_predict_velocity", None) + if predict_velocity_fn is None: + predict_velocity_fn = torch.compile( + self.predict_velocity, + fullgraph=False, + dynamic=False, + options={"triton.cudagraphs": False}, + ) + self._compiled_predict_velocity = predict_velocity_fn + + while time >= -dt / 2: + count += 1 + expanded_time = time.expand(bsize) + v_t = predict_velocity_fn( + state, + prefix_pad_masks, + past_key_values, + x_t, + expanded_time, + prefix_position_ids=prefix_position_ids, + ) + + x_t += dt * v_t + time += dt + logger.debug("Denoised actions in %d steps", count) + return x_t + + def predict_velocity( + self, + state, + prefix_pad_masks, + past_key_values, + x_t, + timestep, + prefix_position_ids=None, + ): + """Predict velocity at time t using cached Qwen3-VL prefix states.""" + if prefix_position_ids is None: + raise ValueError("FlowMatchingV2.predict_velocity requires Qwen3-VL prefix_position_ids.") + + time_embs, suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix( + state, + x_t, + timestep, + ) + + suffix_len = suffix_pad_masks.shape[1] + batch_size = prefix_pad_masks.shape[0] + prefix_len = prefix_pad_masks.shape[1] + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand( + batch_size, + suffix_len, + prefix_len, + ) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + if self.block_future_depth_to_action: + # Query rows here are all suffix (state/action), so row start is 0. + full_att_2d_masks = block_suffix_to_fv_( + full_att_2d_masks, + suffix_row_start=0, + prefix_len=prefix_len, + num_task_tokens=self.num_task_tokens, + ) + full_att_2d_masks = self._block_suffix_to_future_video_if_enabled_( + full_att_2d_masks, + suffix_row_start=0, + prefix_len=prefix_len, + ) + + full_position_ids = self._build_full_position_ids( + prefix_position_ids, + prefix_pad_masks, + suffix_pad_masks, + ) + position_ids = full_position_ids[:, :, -suffix_len:] + + outputs_embeds, _, _ = self.qwenvl_with_expert.forward( + attention_mask=full_att_2d_masks, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=self.config.use_cache, + fill_kv_cache=False, + ada_cond=time_embs if getattr(self.config, "adanorm_time", False) else None, + ) + suffix_out = outputs_embeds[1] + suffix_out = suffix_out[:, -self.config.n_action_steps :] + if getattr(self.config, "action_fp32", False): + v_t = self._fp32_linear(self.action_out_proj, suffix_out) + else: + if suffix_out.dtype != self.action_out_proj.weight.dtype: + suffix_out = suffix_out.to(self.action_out_proj.weight.dtype) + v_t = self.action_out_proj(suffix_out) + return v_t + + +class LingbotVlaV2Policy(PreTrainedModel): + config_class = LingbotVLAV2Config + name = "torch_lingbot_vla_v2" + _no_split_modules = ["Qwen2DecoderLayer", "FixQwen2RMSNorm", "FixAdaRMSNorm"] + + @classmethod + def get_weight_loader(cls): + return LingBotVLAWeightLoader() + + def __init__(self, config: LingbotVLAV2Config, eval: bool = True): + if not eval: + raise ValueError("LingBot-VLA v2 only supports inference mode") + super().__init__(config) + self.config = config + self.language_tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, local_files_only=True) + self.model = FlowMatchingV2(config, eval) + if not getattr(self.config, "use_lm_head", False): + del self.model.qwenvl_with_expert.qwenvl.lm_head + del self.model.qwenvl_with_expert.qwen_expert.lm_head + self.requires_grad_(False) + self.eval() + self.reset() + torch.set_float32_matmul_precision("high") + + def reset(self): + return None + + def forward(self, *args, **kwargs): + """Reject the upstream training API in the inference-only model.""" + del args, kwargs + raise RuntimeError("LingBot-VLA v2 is inference-only; use sample_actions()") + + def sample_actions(self, *args, **kwargs) -> Tensor: + return self.model.sample_actions(*args, **kwargs) + + +ModelClass = LingbotVlaV2Policy + +__all__ = [ + "LingbotVlaV2Policy", + "Qwen3VLForConditionalGeneration", + "Qwen3VLTextModel", + "Qwen3VLPreTrainedModel", + "Qwen2ForCausalLM", +] +# __V2_END__ + + +class LingBotVlaV2Model(LingbotVlaV2Policy): + """TeleFuser-native entry point preserving official checkpoint key names.""" + + name = "lingbot_vla_v2" + + def __init__(self, config, eval=True): + super().__init__(config=config, eval=eval) + + @staticmethod + def state_dict_converter(**kwargs): + return LingBotVlaV2StateDictConverter(**kwargs) + + +# __WRAPPER_END__ diff --git a/telefuser/models/lingbot_vla_v2_loader.py b/telefuser/models/lingbot_vla_v2_loader.py new file mode 100644 index 00000000..8f74fbc0 --- /dev/null +++ b/telefuser/models/lingbot_vla_v2_loader.py @@ -0,0 +1,940 @@ +"""Native utility, alignment, and checkpoint support for LingBot-VLA v2. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +import math + +import einops +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from packaging.version import Version +from torch import Tensor + +# from xformers.ops import memory_efficient_attention + + +def find_next_divisible_by_8_numpy(n: np.ndarray) -> np.ndarray: + """ + Finds the smallest integers greater than each element in a NumPy array 'n' + that are divisible by 8. Assumes non-negative integers. + + Args: + n: A NumPy array of integers. + + Returns: + A NumPy array containing the smallest integers greater than each input element + that are divisible by 8. + """ + remainder = n % 8 + # Calculate the amount to add: 0 if already divisible, otherwise 8 - remainder + # np.where is efficient for conditional operations on arrays + amount_to_add = np.where(remainder == 0, 8, 8 - remainder) + return n + amount_to_add + + +def create_sinusoidal_pos_embedding( + time: torch.tensor, + dimension: int, + min_period: float, + max_period: float, + device="cpu", +) -> Tensor: + """Computes sine-cosine positional embedding vectors for scalar positions.""" + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + + if time.ndim != 1: + raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") + + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float32, device=device) + period = min_period * (max_period / min_period) ** fraction + + # Compute the outer product + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None] + pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + return pos_emb + + +def make_att_2d_masks(pad_masks, att_masks): + """Copied from big_vision. + + Tokens can attend to valid inputs tokens which have a cumulative mask_ar + smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to + setup several types of attention, for example: + + [[1 1 1 1 1 1]]: pure causal attention. + + [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between + themselves and the last 3 tokens have a causal attention. The first + entry could also be a 1 without changing behaviour. + + [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a + block can attend all previous blocks and all tokens on the same block. + + Args: + input_mask: bool[B, N] true if its part of the input, false if padding. + mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on + it and 0 where it shares the same attention mask as the previous token. + """ + if att_masks.ndim != 2: + raise ValueError(att_masks.ndim) + if pad_masks.ndim != 2: + raise ValueError(pad_masks.ndim) + + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + att_2d_masks = att_2d_masks & pad_2d_masks + return att_2d_masks + + +def prefix_query_segments( + use_depth_align, + use_future_depth, + use_future_video=False, + use_future_video_cls=False, + use_future_video_patch=True, + future_video_share_future_depth_query=False, +): + """Return prefix segment order after the image block. + + Task-specific query tokens are always placed after language tokens. Current + task queries precede future task queries; future-depth remains the last + query segment so the existing suffix-to-future-depth blocking can keep using + the tail span. + """ + segments = ["language"] + if not use_depth_align: + return tuple(segments) + + segments.append("current_depth") + if use_future_video: + if use_future_video_cls: + segments.append("future_video_cls") + if use_future_video_patch and not future_video_share_future_depth_query: + segments.append("future_video") + if use_future_depth: + segments.append("future_depth") + return tuple(segments) + + +def prefix_query_token_spans( + prefix_len, + num_task_tokens, + use_depth_align, + use_future_depth, + use_future_video=False, + use_future_video_cls=False, + use_future_video_patch=True, + future_video_share_future_depth_query=False, +): + """Return [start, end) spans for non-language task query segments.""" + counts = { + "current_depth": num_task_tokens, + "future_video_cls": 1, + "future_video": num_task_tokens, + "future_depth": num_task_tokens, + } + ordered = prefix_query_segments( + use_depth_align=use_depth_align, + use_future_depth=use_future_depth, + use_future_video=use_future_video, + use_future_video_cls=use_future_video_cls, + use_future_video_patch=use_future_video_patch, + future_video_share_future_depth_query=future_video_share_future_depth_query, + ) + query_segments = [name for name in ordered if name != "language"] + cursor = prefix_len - sum(counts[name] for name in query_segments) + spans = {} + for name in query_segments: + count = counts[name] + spans[name] = (cursor, cursor + count) + cursor += count + return spans + + +def fv_col_span(prefix_len, num_task_tokens, use_cls, use_patch): + """Return [start, end) of a tail query block inside the prefix. + + This legacy helper is still used for future-depth tail blocking in V2. + New prefix layout code should prefer prefix_query_token_spans(), which also + handles current-depth and separate future-video spans. + """ + fv_len = (1 if use_cls else 0) + (num_task_tokens if use_patch else 0) + return prefix_len - fv_len, prefix_len + + +def block_suffix_to_fv_( + att_2d_masks, suffix_row_start, prefix_len, num_task_tokens, use_cls=False, use_patch=True, drop_mask=None +): + """In-place mask out the suffix-to-future-video attention edge. + + `make_att_2d_masks`' cumsum scheme cannot express "a query cannot see a + segment that precedes it", so we zero the rectangular [suffix rows, FV cols] + block on the already-built 2D mask instead of touching mask_ar. + + att_2d_masks: bool[B, Q, K], True == visible. `suffix_row_start` is the first + query row belonging to the suffix: prefix_len in the square training mask, + 0 in the suffix-only inference mask. Leaves FV -> img/lang rows untouched so + the distillation query still reads the current observation. + + `drop_mask`: optional bool[B], True where this sample's suffix must NOT see + FV. None == block every sample (hard mask). Used for per-sample stochastic + masking (FV-attention dropout): keep = visible iff not dropped, applied via + broadcast multiply so it stays a static graph under torch.compile. + """ + fv_start, fv_end = fv_col_span(prefix_len, num_task_tokens, use_cls, use_patch) + if fv_end <= fv_start: + return att_2d_masks + if drop_mask is None: + att_2d_masks[:, suffix_row_start:, fv_start:fv_end] = False + else: + # keep[b] = True where the sample is NOT dropped -> AND keeps those rows + # visible and zeros the dropped ones, with no data-dependent indexing. + keep = (~drop_mask).view(-1, 1, 1) + block = att_2d_masks[:, suffix_row_start:, fv_start:fv_end] + att_2d_masks[:, suffix_row_start:, fv_start:fv_end] = block & keep + return att_2d_masks + + +def resize_with_pad(img, width, height, pad_value=-1): + # assume no-op when width height fits already + if img.ndim != 4: + raise ValueError(f"(b,c,h,w) expected, but {img.shape}") + + cur_height, cur_width = img.shape[2:] + + ratio = max(cur_width / width, cur_height / height) + resized_height = int(cur_height / ratio) + resized_width = int(cur_width / ratio) + resized_img = F.interpolate(img, size=(resized_height, resized_width), mode="bilinear", align_corners=False) + + pad_height = max(0, int(height - resized_height)) + pad_width = max(0, int(width - resized_width)) + + # pad on left and top of image + padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value) + return padded_img + + +def our_eager_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor, +): + """ + Performs eager attention, optimized with torch.einsum. + + Args: + query_states: Query tensor of shape [batch_size, seq_len, num_attention_heads, head_dim]. + key_states: Key tensor of shape [batch_size, seq_len, num_key_value_heads, head_dim]. + value_states: Value tensor of shape [batch_size, seq_len, num_key_value_heads, head_dim]. + attention_mask: Attention mask tensor, typically + [batch_size, 1, seq_len, seq_len] or [batch_size, seq_len, seq_len]. + + Returns: + Output tensor of shape [batch_size, seq_len, num_attention_heads * head_dim]. + """ + bsize, seq_len, num_att_heads, head_dim = query_states.shape + num_key_value_heads = key_states.shape[2] + num_key_value_groups = num_att_heads // num_key_value_heads + + key_states = einops.repeat(key_states, "b l h d -> b l (h g) d", g=num_key_value_groups) + value_states = einops.repeat(value_states, "b l h d -> b l (h g) d", g=num_key_value_groups) + + query_states_permuted = torch.einsum("blhd->bhld", query_states) + key_states_permuted = torch.einsum("blhd->bhld", key_states) + + att_weights = torch.einsum("bhqd,bhkd->bhqk", query_states_permuted, key_states_permuted) + att_weights *= head_dim**-0.5 + + big_neg = -2.3819763e38 + masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg) + + probs = nn.functional.softmax(masked_att_weights, dim=-1) + probs = probs.to(dtype=value_states.dtype) + + value_states_permuted = torch.einsum("blhd->bhld", value_states) # [B, H, L_v, D] + att_output = torch.einsum("bhqk,bhkv->bhqv", probs, value_states_permuted) # [B, H, L_q, D] + att_output = torch.einsum("bhld->blhd", att_output) # [B, L, H, D] + att_output = att_output.reshape(bsize, seq_len, num_att_heads * head_dim) + + return att_output + + +# @torch.jit.script +def apply_rope( + x: torch.Tensor, + positions: torch.Tensor, + max_wavelength: float = 10_000.0, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Applies RoPE positions [B, L] to x [B, L, H, D].""" + original_dtype = x.dtype # bf16 + d = x.shape[-1] + d_half = d // 2 + device = x.device + + # Cast input to compute_dtype for all internal operations + x_casted = x.to(dtype) + positions_casted = positions.to(dtype) + + freq_exponents = (2.0 / d) * torch.arange(d_half, dtype=dtype, device=device) + timescale = max_wavelength**freq_exponents + radians = torch.einsum("bl,h->blh", positions_casted, 1.0 / timescale) # fp32 -> bf16 + + radians = radians[..., None, :] # [B, L, 1, D_half] + + sin = torch.sin(radians) # bf16 + cos = torch.cos(radians) # bf16 + + x1, x2 = x_casted.split(d_half, dim=-1) # fp32 + + res = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) # fp32 + + return res.to(original_dtype) # bf16 + + +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +FLEX_SPARSE_BLOCK_SIZE = 128 +FLEX_KERNEL_OPTIONS = {"BLOCK_M": 32, "BLOCK_N": 64, "num_warps": 4, "num_stages": 2} + +if Version(torch.__version__) > Version("2.5.0"): + # Ffex attention is only available from torch 2.5 onwards + from torch.nn.attention.flex_attention import ( + _mask_mod_signature, + _round_up_to_multiple, + create_block_mask, + create_mask, + flex_attention, + ) + + +# @torch.compile(dynamic=False) +def flex_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: torch.Tensor, + scaling=None, +): + """ + This is defined out of classes to make compile happy. + """ + batch_size, seq_len, num_att_heads, head_dim = query_states.shape + original_dtype = query_states.dtype + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + query_states = query_states.to(torch.float32) + key_states = key_states.to(torch.float32) + value_states = value_states.to(torch.float32) + + causal_mask = attention_mask + if causal_mask is not None: + causal_mask = causal_mask[:, None, :, : key_states.shape[2]] + + if causal_mask.shape[1] == 1 and query_states.shape[1] > 1: + causal_mask = causal_mask.expand(-1, query_states.shape[1], -1, -1) + + def precomputed_mask_factory(precomputed_mask: torch.Tensor) -> _mask_mod_signature: + def mask_mod(b, h, q_idx, kv_idx): + # Danger zone: if b,h,q_idx,kv_idx exceed the shape, device-side assert occurs. + return precomputed_mask[b][h][q_idx][kv_idx] + + return mask_mod + + b_mask, h_mask, q_len, kv_len = causal_mask.shape # The shape of your mask + # ipdb.set_trace() + block_size = FLEX_SPARSE_BLOCK_SIZE + q_len_rounded = _round_up_to_multiple(q_len, block_size) + kv_len_rounded = _round_up_to_multiple(kv_len, block_size) + + # *CRITICAL* we do need to expand here, else we get a CUDA index error + + pad_q = q_len_rounded - q_len + pad_k = kv_len_rounded - kv_len + + if pad_q > 0: + query_states = F.pad(query_states, (0, 0, 0, pad_q), value=0.0) # [B, H, q_len_rounded, D] + if pad_k > 0: + key_states = F.pad(key_states, (0, 0, 0, pad_k), value=0.0) + value_states = F.pad(value_states, (0, 0, 0, pad_k), value=0.0) + padded_causal_mask = F.pad(causal_mask, (0, pad_k, 0, pad_q), value=0.0) + mask_mod_fn_orig = precomputed_mask_factory(padded_causal_mask) + + mask_4d = create_mask( + mod_fn=mask_mod_fn_orig, + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, + device=causal_mask.device, + ) + + mask_mod_fn_padded = precomputed_mask_factory(mask_4d) + block_mask = create_block_mask( + mask_mod=mask_mod_fn_padded, + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, + BLOCK_SIZE=block_size, + device=causal_mask.device, + _compile=False, + ) + + # mask is applied inside the kernel, ideally more efficiently than score_mod. + attn_output, attention_weights = flex_attention( + query_states, + key_states, + value_states, + block_mask=block_mask, + enable_gqa=True, # because we shaped query/key states for GQA + scale=head_dim**-0.5 if scaling is None else scaling, + return_lse=True, + kernel_options=FLEX_KERNEL_OPTIONS, + ) + attn_output = attn_output[:, :, :seq_len, :].to(dtype=original_dtype) + attn_output = attn_output.transpose(1, 2).contiguous() # [B, Q_LEN, H, head_dim] + attn_output = attn_output.reshape( + batch_size, + -1, + attn_output.shape[2] * attn_output.shape[3], # merges [H, head_dim] + ) + return attn_output + + +@torch.compiler.disable +def build_block_mask( + attention_mask_3d: torch.Tensor, + num_heads: int, + q_len: int, + kv_len: int, + block_size: int = FLEX_SPARSE_BLOCK_SIZE, +): + """ + Build a reusable BlockMask from a 3D attention mask [B, Q, KV]. + This allocates the dense 4D mask once; the returned BlockMask can be reused across layers. + """ + from torch.nn.attention.flex_attention import ( + _round_up_to_multiple, + create_block_mask, + create_mask, + ) + + causal_mask = attention_mask_3d[:, None, :, :].expand(-1, num_heads, -1, -1).contiguous() + b_mask, h_mask = causal_mask.shape[0], causal_mask.shape[1] + + q_len_rounded = _round_up_to_multiple(q_len, block_size) + kv_len_rounded = _round_up_to_multiple(kv_len, block_size) + + pad_q = q_len_rounded - q_len + pad_k = kv_len_rounded - kv_len + padded_mask = F.pad(causal_mask, (0, pad_k, 0, pad_q), value=0.0) + + def precomputed_mask_factory(precomputed_mask: torch.Tensor): + def mask_mod(b, h, q_idx, kv_idx): + return precomputed_mask[b][h][q_idx][kv_idx] + + return mask_mod + + mask_4d = create_mask( + mod_fn=precomputed_mask_factory(padded_mask), + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, + device=causal_mask.device, + ) + + block_mask = create_block_mask( + mask_mod=precomputed_mask_factory(mask_4d), + B=b_mask, + H=h_mask, + Q_LEN=q_len_rounded, + KV_LEN=kv_len_rounded, + BLOCK_SIZE=block_size, + device=causal_mask.device, + _compile=False, + ) + return block_mask + + +def flex_attention_with_block_mask( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + block_mask, + seq_len: int, + scaling=None, +): + """ + Run flex_attention with a pre-built BlockMask (no create_mask allocation per call). + """ + batch_size = query_states.shape[0] + head_dim = query_states.shape[3] + original_dtype = query_states.dtype + + query_states = query_states.transpose(1, 2).to(torch.float32) + key_states = key_states.transpose(1, 2).to(torch.float32) + value_states = value_states.transpose(1, 2).to(torch.float32) + + q_len_rounded = block_mask.shape[-2] if hasattr(block_mask, "shape") else query_states.shape[2] + kv_len_rounded = block_mask.shape[-1] if hasattr(block_mask, "shape") else key_states.shape[2] + + pad_q = q_len_rounded - query_states.shape[2] + pad_k = kv_len_rounded - key_states.shape[2] + + if pad_q > 0: + query_states = F.pad(query_states, (0, 0, 0, pad_q), value=0.0) + if pad_k > 0: + key_states = F.pad(key_states, (0, 0, 0, pad_k), value=0.0) + value_states = F.pad(value_states, (0, 0, 0, pad_k), value=0.0) + + attn_output, _ = flex_attention( + query_states, + key_states, + value_states, + block_mask=block_mask, + enable_gqa=True, + scale=head_dim**-0.5 if scaling is None else scaling, + return_lse=True, + kernel_options=FLEX_KERNEL_OPTIONS, + ) + attn_output = attn_output[:, :, :seq_len, :].to(dtype=original_dtype) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(batch_size, -1, attn_output.shape[2] * attn_output.shape[3]) + return attn_output + + +# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +import torch + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + # (bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + batch_size, latent_length, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(batch_size, latent_length, -1) + + return self.to_out(out) + + +class TaskTokenResampler(nn.Module): + def __init__( + self, + dim_in=768, + dim_mid=1024, + dim_head=64, + dim_out=1024, + num_layers=8, + num_queries=8, + num_heads=16, + ff_mult=4, + ): + super().__init__() + + self.num_queries = num_queries + self.proj_in1 = nn.Linear(dim_in, dim_mid) + self.proj_in2 = nn.Linear(dim_in, dim_mid) + self.proj_out = nn.Linear(dim_mid, dim_out) + self.norm_out = nn.LayerNorm(dim_out) + + self.layers = nn.ModuleList([]) + for _ in range(num_layers): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim_mid, dim_head=dim_head, heads=num_heads), + FeedForward(dim=dim_mid, mult=ff_mult), + ] + ) + ) + + def forward(self, x, queries): + queries = self.proj_in1(queries) + x = self.proj_in2(x) + + for attn, ff in self.layers: + queries = attn(x, queries) + queries + queries = ff(queries) + queries + + queries = self.proj_out(queries) + queries = self.norm_out(queries) + return queries + + +class TaskTokenDepthHead(nn.Module): + def __init__( + self, + proj_config=None, + llm_hidden_size=4096, + use_intermediate_depth=False, + ): + super(TaskTokenDepthHead, self).__init__() + + self.projector = TaskTokenResampler( + dim_in=llm_hidden_size, + dim_mid=llm_hidden_size, + dim_head=proj_config["dim_head"], + dim_out=proj_config["dim_out"], + num_layers=proj_config["num_layers"], + num_heads=proj_config["num_heads"], + num_queries=proj_config["num_backbone_tokens"], + ff_mult=proj_config["ff_mult"], + ) + + def forward(self, llm_feats, queries): + queries = self.projector(llm_feats, queries) + return queries + + +import json +from copy import deepcopy +from pathlib import Path +from typing import Any + +from transformers import AutoConfig + + +class LingBotVLAWeightLoader: + """Minimal native weight-name mapper retained for model compatibility.""" + + def get_vlm_submodule(self, model): + return model.model.qwenvl_with_expert.qwenvl + + def get_expert_vision_submodule(self, model): + return getattr(model.model.qwenvl_with_expert, "expert_visual", None) + + def map_ckpt_key(self, key, load_vlm_only=False, post_training=False): + if key.startswith("expert_visual.") and not post_training: + return "model.qwenvl_with_expert." + key + if load_vlm_only: + return "model.qwenvl_with_expert.qwenvl." + key + return key + + +OFFICIAL_6B_MODEL_CONFIG: dict[str, Any] = { + "post_training": False, + "adanorm_time": True, + "moe_implementation": "fused", + "use_robby_moe_kernel": True, + "attention_implementation": "eager", + "precompute_grid_thw": True, + "vlm_causal": True, + "use_moe": True, + "token_moe_layers": list(range(36)), + "token_num_experts": 32, + "token_top_k": 4, + "token_moe_intermediate_size": 512, + "token_shared_intermediate_size": 704, + "bias_update_speed": 0.0, + "sequence_wise_mode": "per_sequence", + "sequence_wise_loss_coeff": 1e-3, + "router_z_loss_coeff": 1e-4, + "router_activation": "sigmoid", + "routed_scaling_factor": 4.0, + "use_shared_expert_gate": False, + "freeze_vision_encoder": False, + "tokenizer_max_length": 72, + "loss_type": "L1_fm", + "action_dim": 55, + "max_action_dim": 55, + "max_state_dim": 55, + "align_params": { + "mode": "query", + "num_task_tokens": 8, + "depth_loss_weight": 0.004, + "future_depth_loss_weight": 0.004, + "use_future_video": True, + "llm": { + "dim_out": 2560, + "image_token_size": 8, + "image_input_size": 224, + }, + "depth": { + "model_type": "MoRGBD", + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "token_size": 16, + "dim_out": 1024, + "input_size": 224, + "use_future_depth": True, + "block_future_depth_to_action": True, + "future_depth_head_type": "resampler", + "detach_future_image_feats": True, + }, + "video": { + "attention_mode": "flex_block_causal", + "input_size": 256, + "block_suffix_to_future_video": True, + "share_future_depth_query": True, + "use_shared_future_task_proj": True, + "use_current_shared_task_proj": True, + "num_future_frames": 1, + "use_warmup_frame": True, + "effective_fps": 1.0, + "n_blocks": 1, + "cls_pool": "last", + "detach_image_feats": True, + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "dim_out": 1024, + "future_video_loss_weight": 0.004, + "use_smooth_l1_loss": False, + "use_mse_loss": True, + "mse_loss_weight": 1.0, + "use_patch_loss": True, + "use_current_patch_loss": True, + "use_cosine_loss": False, + "cosine_loss_weight": 0.2, + "use_cls_loss": False, + "cls_loss_type": "mse", + "cls_loss_weight": 0.2, + }, + }, +} + + +def resolve_lingbot_vla_v2_checkpoint(model_path: str | Path) -> Path: + path = Path(model_path).expanduser().resolve() + if path.is_file(): + if path.name != "model.safetensors.index.json": + raise ValueError(f"Expected model.safetensors.index.json, got: {path}") + return path + if not path.is_dir(): + raise FileNotFoundError(f"LingBot-VLA v2 model path does not exist: {path}") + index_path = path / "model.safetensors.index.json" + if not index_path.is_file(): + raise FileNotFoundError(f"Missing sharded checkpoint index: {index_path}") + return index_path + + +def resolve_lingbot_vla_v2_shards(model_path: str | Path) -> list[str]: + index_path = resolve_lingbot_vla_v2_checkpoint(model_path) + with index_path.open("r", encoding="utf-8") as handle: + index = json.load(handle) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Invalid safetensors index without weight_map: {index_path}") + + shard_paths = [index_path.parent / name for name in sorted(set(weight_map.values()))] + missing = [str(path) for path in shard_paths if not path.is_file()] + if missing: + raise FileNotFoundError(f"Missing LingBot-VLA v2 checkpoint shards: {missing}") + return [str(path) for path in shard_paths] + + +def build_official_6b_config( + qwen3vl_path: str | Path, + *, + checkpoint_variant: str = "base", + checkpoint_path: str | Path | None = None, +): + from telefuser.models.lingbot_vla_v2 import LingbotVLAV2Config + + if checkpoint_variant != "base": + raise ValueError(f"Unsupported LingBot-VLA v2 checkpoint variant: {checkpoint_variant!r}") + + qwen_path = Path(qwen3vl_path).expanduser().resolve() + qwen_config = AutoConfig.from_pretrained(str(qwen_path), local_files_only=True) + if not hasattr(qwen_config, "text_config") or not hasattr(qwen_config, "vision_config"): + raise ValueError( + "LingBot-VLA v2 requires the local Qwen3-VL-4B-Instruct architecture/tokenizer " + f"directory; this is not a complete Qwen3-VL directory: {qwen_path}" + ) + + text_config = qwen_config.text_config + expected_architecture = {"hidden_size": 2560, "num_hidden_layers": 36} + mismatches = { + key: (expected, getattr(text_config, key, None)) + for key, expected in expected_architecture.items() + if getattr(text_config, key, None) != expected + } + if mismatches: + raise ValueError( + "LingBot-VLA v2 6B was trained with Qwen3-VL-4B-Instruct; " + f"the supplied architecture is incompatible: {mismatches}" + ) + + values = deepcopy(OFFICIAL_6B_MODEL_CONFIG) + values["tokenizer_path"] = str(qwen_path) + config = LingbotVLAV2Config(**values) + for key in ( + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "rms_norm_eps", + "rope_theta", + "vocab_size", + "max_position_embeddings", + "hidden_act", + "tie_word_embeddings", + ): + if hasattr(text_config, key): + setattr(config, key, getattr(text_config, key)) + config.vision_config = qwen_config.vision_config + config.tokenizer_path = str(qwen_path) + config.use_cache = True + config.attention_implementation = "eager" + config.checkpoint_variant = checkpoint_variant + config.checkpoint_path = None if checkpoint_path is None else str(Path(checkpoint_path).expanduser().resolve()) + config.policy_verified = False + config.verification_status = "unverified_official_6b_base" + return config + + +def validate_official_6b_checkpoint(state_dict): + gate = "model.qwenvl_with_expert.qwen_expert.model.layers.0.mlp.experts.gate_proj" + last_gate = "model.qwenvl_with_expert.qwen_expert.model.layers.35.mlp.experts.gate_proj" + expected = (32, 512, 768) + for key in (gate, last_gate): + if key not in state_dict: + raise ValueError(f"Missing official LingBot-VLA v2 weight: {key}") + if tuple(state_dict[key].shape) != expected: + raise ValueError(f"Unexpected shape for {key}: expected {expected}, got {tuple(state_dict[key].shape)}") + + +class LingBotVlaV2StateDictConverter: + def __init__( + self, + qwen3vl_path: str | Path, + checkpoint_variant: str = "base", + checkpoint_path: str | Path | None = None, + ): + self.qwen3vl_path = Path(qwen3vl_path) + self.checkpoint_variant = checkpoint_variant + self.checkpoint_path = checkpoint_path + + def from_official(self, state_dict): + validate_official_6b_checkpoint(state_dict) + config = build_official_6b_config( + self.qwen3vl_path, + checkpoint_variant=self.checkpoint_variant, + checkpoint_path=self.checkpoint_path, + ) + return state_dict, {"config": config, "eval": True} + + def from_diffusers(self, state_dict): + del state_dict + raise ValueError("LingBot-VLA v2 does not provide a Diffusers checkpoint") + + +def load_lingbot_vla_v2( + module_manager, + model_path: str | Path, + qwen3vl_path: str | Path, + *, + torch_dtype=torch.bfloat16, + device=None, + checkpoint_variant: str = "base", +): + from telefuser.models.lingbot_vla_v2 import LingBotVlaV2Model + + checkpoint_path = resolve_lingbot_vla_v2_checkpoint(model_path).parent + shard_paths = resolve_lingbot_vla_v2_shards(checkpoint_path) + module_manager.load_model( + shard_paths, + device=device, + torch_dtype=torch_dtype, + low_cpu_mem_usage=True, + name="lingbot_vla_v2", + model_class=LingBotVlaV2Model, + model_resource="official", + converter_kwargs={ + "qwen3vl_path": str(qwen3vl_path), + "checkpoint_variant": checkpoint_variant, + "checkpoint_path": str(checkpoint_path), + }, + ) + return module_manager.fetch_module("lingbot_vla_v2") diff --git a/telefuser/models/lingbot_vla_v2_moe.py b/telefuser/models/lingbot_vla_v2_moe.py new file mode 100644 index 00000000..b8c8e91e --- /dev/null +++ b/telefuser/models/lingbot_vla_v2_moe.py @@ -0,0 +1,548 @@ +"""Native fused-MoE action expert used by LingBot-VLA v2. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + +import torch + +from telefuser.ops.lingbot_vla_v2_moe import robby_moe_forward + + +def fused_moe_forward( + module, + num_experts, + routing_weights, + selected_experts, + hidden_states, + fc1_1_weight, + fc1_2_weight, + fc2_weight, +): + """Single-device PyTorch fallback for the fused 3D expert layout.""" + del module + output = torch.zeros_like(hidden_states) + for expert_id in range(num_experts): + routes = (selected_experts == expert_id).nonzero(as_tuple=False) + if routes.numel() == 0: + continue + token_ids = routes[:, 0] + route_ids = routes[:, 1] + expert_input = hidden_states.index_select(0, token_ids) + gate = torch.nn.functional.linear(expert_input, fc1_1_weight[expert_id]) + up = torch.nn.functional.linear(expert_input, fc1_2_weight[expert_id]) + intermediate = torch.nn.functional.silu(gate) * up + expert_output = torch.nn.functional.linear(intermediate, fc2_weight[expert_id]) + weights = routing_weights[token_ids, route_ids].unsqueeze(-1) + output.index_add_(0, token_ids, expert_output * weights) + return output + + +from typing import Optional, Tuple + +import torch.nn.functional as F +from torch import nn +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config +from transformers.processing_utils import Unpack +from transformers.utils import auto_docstring, logging + + +def _update_moe_runtime_stats(block, routing_weights, selected_experts): + """Update MoE runtime buffers outside torch.compile graphs.""" + with torch.no_grad(): + if routing_weights is not None and hasattr(block, "avg_topk_sigmoid_score"): + avg_score = routing_weights.detach().float().mean() + block.avg_topk_sigmoid_score.copy_( + avg_score.reshape_as(block.avg_topk_sigmoid_score).to( + device=block.avg_topk_sigmoid_score.device, + dtype=block.avg_topk_sigmoid_score.dtype, + ) + ) + + if hasattr(block, "tokens_per_expert"): + counts = F.one_hot( + selected_experts.detach().reshape(-1), + num_classes=block.num_experts, + ).sum(dim=0) + block.tokens_per_expert.add_( + counts.to( + device=block.tokens_per_expert.device, + dtype=block.tokens_per_expert.dtype, + ) + ) + + +from transformers.models.qwen2.modeling_qwen2 import ( + PreTrainedModel, + Qwen2Attention, + Qwen2MLP, + Qwen2RMSNorm, + Qwen2RotaryEmbedding, +) +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2ForCausalLM as _Qwen2ForCausalLM, +) +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Model as _Qwen2Model, +) + +logger = logging.get_logger(__name__) +# from transformers.models.mistral.modeling_mistral import MistralMLP + + +# Modified from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Qwen2Moe +class Qwen2MoeRoutedExpertMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class Qwen2MoeSharedExpertMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class Qwen2FusedExperts(nn.Module): + """Fused expert module: stores E experts' weights as 3D tensors for group_gemm. + + Shape convention matches nn.Linear(in, out).weight = [out, in]: + gate_proj: [E, intermediate_size, hidden_size] + up_proj: [E, intermediate_size, hidden_size] + down_proj: [E, hidden_size, intermediate_size] + + The forward() method runs the full fused_moe computation. This is critical + for FSDP2: calling self.experts(...) triggers FSDP2's forward pre-hook to + unshard the expert params on ep_fsdp_mesh BEFORE they are used by kernels. + """ + + def __init__(self, num_experts, hidden_size, intermediate_size, initializer_range=0.02): + super().__init__() + self.num_experts = num_experts + self.intermediate_size = intermediate_size + self.initializer_range = initializer_range + self.gate_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.up_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.down_proj = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.register_buffer("_gate_up_proj_cache", None, persistent=False) + self._gate_up_proj_cache_key = None + self._robby_moe_workspace = None + self._robby_moe_workspace_key = None + self.reset_parameters() + + def reset_parameters(self): + nn.init.normal_(self.gate_proj, mean=0.0, std=self.initializer_range) + nn.init.normal_(self.up_proj, mean=0.0, std=self.initializer_range) + nn.init.normal_(self.down_proj, mean=0.0, std=self.initializer_range) + self.clear_inference_cache() + + def clear_inference_cache(self): + self._gate_up_proj_cache = None + self._gate_up_proj_cache_key = None + self._robby_moe_workspace = None + self._robby_moe_workspace_key = None + + def _get_robby_moe_workspace(self, hidden_states, top_k): + if self.training or torch.is_grad_enabled() or not hidden_states.is_cuda: + return None + num_tokens, hidden_size = hidden_states.shape + key = ( + num_tokens, + int(top_k), + self.num_experts, + hidden_size, + self.intermediate_size, + hidden_states.dtype, + hidden_states.device, + ) + if self._robby_moe_workspace is None or self._robby_moe_workspace_key != key: + max_routes = num_tokens * int(top_k) + self._robby_moe_workspace = { + "counts": torch.empty((self.num_experts,), device=hidden_states.device, dtype=torch.int32), + "rows": torch.empty((self.num_experts, max_routes), device=hidden_states.device, dtype=torch.int32), + "slots": torch.empty((self.num_experts, max_routes), device=hidden_states.device, dtype=torch.int32), + "inter": torch.empty( + (num_tokens, int(top_k), self.intermediate_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ), + "out": torch.empty((num_tokens, hidden_size), device=hidden_states.device, dtype=torch.float32), + } + self._robby_moe_workspace_key = key + return self._robby_moe_workspace + + def forward(self, module, num_experts, routing_weights, selected_experts, hidden_states): + """Run fused_moe_forward with FSDP2-managed weights. + + Must be called via self.experts(...) so FSDP2 unshards params first. + """ + return fused_moe_forward( + module=module, + num_experts=num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_states, + fc1_1_weight=self.gate_proj, + fc1_2_weight=self.up_proj, + fc2_weight=self.down_proj, + ) + + +class FixQwen2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + FixQwen2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + # print(f'self.weight dtype is {self.weight.dtype}') + input_dtype = hidden_states.dtype + # print(f'input_dtype is {input_dtype}') + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + # print(f'hidden_states dtype is {hidden_states.dtype}') + # print(f'output dtype is {(self.weight * hidden_states.to(input_dtype)).dtype}') + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen2TokenMoeBlock(nn.Module): + """Token-level routing MoE block with all-to-all computation for torch.compile compatibility.""" + + def __init__(self, config): + super().__init__() + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.norm_topk_prob = config.norm_topk_prob + + # Loss-free balancing support. With zero correction bias this is + # equivalent to unbiased top-k selection; the optimizer pre-hook updates + # the bias when bias_update_speed > 0. + self.register_buffer( + "e_score_correction_bias", + torch.zeros(config.num_experts), + persistent=True, + ) + self.register_buffer( + "tokens_per_expert", + torch.zeros(config.num_experts, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "last_tokens_per_expert", + torch.zeros(config.num_experts, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "avg_topk_sigmoid_score", + torch.zeros(1, dtype=torch.float32), + persistent=False, + ) + + # gating (per-token) + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + + # EP/fused support: choose expert storage based on moe_implementation + self._moe_implementation = getattr(config, "_moe_implementation", None) or "eager" + self._use_robby_moe_kernel = bool(getattr(config, "use_robby_moe_kernel", False)) + if self._moe_implementation == "fused": + self.experts = Qwen2FusedExperts( + self.num_experts, + config.hidden_size, + config.moe_intermediate_size, + initializer_range=getattr(config, "initializer_range", 0.02), + ) + else: + self.experts = nn.ModuleList( + [ + Qwen2MoeRoutedExpertMLP(config, intermediate_size=config.moe_intermediate_size) + for _ in range(self.num_experts) + ] + ) + + self.shared_expert = Qwen2MoeSharedExpertMLP(config, intermediate_size=config.shared_expert_intermediate_size) + self._router_activation = getattr(config, "router_activation", "softmax") + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self._use_shared_expert_gate = getattr(config, "use_shared_expert_gate", True) + if self._use_shared_expert_gate: + self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Token-level routing with all-to-all computation for torch.compile compatibility.""" + batch_size, sequence_length, hidden_dim = hidden_states.shape + + # Token-level routing: each token individually + hidden_flat = hidden_states.reshape(-1, hidden_dim) # (B*T, D) + # Gate in true fp32 (autocast disabled): bf16 gate logits can flip top-k + # selection on near-equal scores -> routing jitter / rotating dead experts. + # cf. VideoPretrain lumos/moe/router.py TokenChoiceTopKRouter. + with torch.amp.autocast(hidden_flat.device.type, enabled=False): + router_logits = F.linear(hidden_flat.float(), self.gate.weight.float()) # (B*T, num_experts) + + if self._router_activation == "sigmoid": + routing_scores = router_logits.sigmoid() + else: + routing_scores = F.softmax(router_logits, dim=1, dtype=torch.float) + + scores_for_choice = routing_scores + self.e_score_correction_bias.unsqueeze(0) + _, selected_experts = torch.topk(scores_for_choice, self.top_k, dim=-1) + routing_weights = routing_scores.gather(1, selected_experts) + if self.training: + _update_moe_runtime_stats(self, routing_weights, selected_experts) + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20) + if self.routed_scaling_factor != 1.0: + routing_weights = routing_weights * self.routed_scaling_factor + routing_weights = routing_weights.to(hidden_states.dtype) + + # Expert computation: fused (group_gemm) or eager (per-expert loop) + if self._moe_implementation == "fused": + use_robby_moe = ( + self._use_robby_moe_kernel + and robby_moe_forward is not None + and hidden_flat.is_cuda + and not self.training + and not torch.is_grad_enabled() + ) + if use_robby_moe: + try: + final_hidden_states = robby_moe_forward( + hidden_flat, + routing_weights, + selected_experts, + self.experts.gate_proj, + self.experts.up_proj, + self.experts.down_proj, + workspace=self.experts._get_robby_moe_workspace( + hidden_flat, + selected_experts.shape[1], + ), + ) + except Exception as exc: + logger.warning_once(f"robby_moe_forward failed, falling back to fused_moe_forward: {exc}") + final_hidden_states = self.experts( + module=self, + num_experts=self.num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_flat, + ) + else: + final_hidden_states = self.experts( + module=self, + num_experts=self.num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_flat, + ) + else: + # Original eager path: every expert processes all tokens + expert_outputs = torch.stack( + [expert(hidden_flat) for expert in self.experts], dim=0 + ) # (num_experts, B*T, D) + expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).float() # (B*T, top_k, num_experts) + weights = ( + (expert_mask * routing_weights.unsqueeze(-1).float()).sum(dim=1).to(hidden_states.dtype) + ) # (B*T, num_experts) + final_hidden_states = torch.einsum("ebd,be->bd", expert_outputs, weights) # (B*T, D) + + # Shared expert: applied to all tokens (fixed shape) + if final_hidden_states.dtype != hidden_flat.dtype: + final_hidden_states = final_hidden_states.to(hidden_flat.dtype) + shared_expert_output = self.shared_expert(hidden_flat) + if self._use_shared_expert_gate: + shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_flat)) * shared_expert_output + final_hidden_states = final_hidden_states + shared_expert_output + + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +class Qwen2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen2Attention(config=config, layer_idx=layer_idx) + self.mlp = Qwen2MLP(config) + self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + if config.use_sliding_window and config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + + def forward( + self, + hidden_states: torch.Tensor, + att_output: Optional[torch.Tensor] = None, + start: Optional[int] = 0, + end: Optional[int] = 0, + compute_kqv: bool = False, + output_atten: bool = False, + ada_cond: Optional[torch.Tensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + # Ensure input dtypes match weight dtype (needed for gradient checkpointing + # recomputation where autocast context is lost) + param_dtype = self.self_attn.q_proj.weight.dtype + hidden_states = hidden_states.to(param_dtype) + if att_output is not None: + att_output = att_output.to(param_dtype) + if ada_cond is not None: + ada_cond = ada_cond.to(param_dtype) + + if compute_kqv: + if ada_cond is not None: + hidden_states = self.input_layernorm(hidden_states, ada_cond) + else: + hidden_states = self.input_layernorm(hidden_states) + hidden_shape = (*hidden_states.shape[:-1], -1, self.self_attn.head_dim) + + query_state = self.self_attn.q_proj(hidden_states).view(hidden_shape) + key_state = self.self_attn.k_proj(hidden_states).view(hidden_shape) + value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) + + return query_state, key_state, value_state + + elif output_atten: + if att_output.dtype != self.self_attn.o_proj.weight.dtype: + att_output = att_output.to(self.self_attn.o_proj.weight.dtype) + out_emb = self.self_attn.o_proj(att_output[:, start:end]) + + # first residual + out_emb += hidden_states + after_first_residual = out_emb.clone() + if ada_cond is not None: + out_emb = self.post_attention_layernorm(out_emb, ada_cond) + else: + out_emb = self.post_attention_layernorm(out_emb) + out_emb = self.mlp(out_emb) + # Handle MoE block returning (hidden_states, router_logits) + router_logits = None + if isinstance(out_emb, tuple): + out_emb, router_logits = out_emb + # second residual + out_emb += after_first_residual + + return out_emb, router_logits + + else: + raise ValueError( + f"Invalid operation compute_kqv={compute_kqv} and output_atten={output_atten} " + "with Qwen2DecoderLayer in LingBot-VLA" + ) + + +@auto_docstring +class Qwen2PreTrainedModel(PreTrainedModel): + config: Qwen2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen2DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Qwen2DecoderLayer, + "attentions": Qwen2Attention, + } + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Qwen2FusedExperts): + module.initializer_range = std + module.reset_parameters() + + +class Qwen2Model(Qwen2PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`] + + Args: + config: Qwen2Config + """ + + get_input_embeddings = _Qwen2Model.get_input_embeddings + set_input_embeddings = _Qwen2Model.set_input_embeddings + forward = _Qwen2Model.forward + + def __init__(self, config: Qwen2Config, eval=False): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = FixQwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen2RotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + if eval: + self._init_weights = lambda module: None + self.post_init() + + +class Qwen2ForCausalLM(Qwen2PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + get_input_embeddings = _Qwen2ForCausalLM.get_input_embeddings + set_input_embeddings = _Qwen2ForCausalLM.set_input_embeddings + get_output_embeddings = _Qwen2ForCausalLM.get_output_embeddings + set_output_embeddings = _Qwen2ForCausalLM.set_output_embeddings + forward = _Qwen2ForCausalLM.forward + set_decoder = _Qwen2ForCausalLM.set_decoder + get_decoder = _Qwen2ForCausalLM.get_decoder + + def __init__(self, config, eval): + super().__init__(config) + self.model = Qwen2Model(config, eval) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() diff --git a/telefuser/models/lingbot_vla_v2_qwen.py b/telefuser/models/lingbot_vla_v2_qwen.py new file mode 100644 index 00000000..47cc54cb --- /dev/null +++ b/telefuser/models/lingbot_vla_v2_qwen.py @@ -0,0 +1,331 @@ +"""Native Qwen vision-language layers used by LingBot-VLA v2. + +Adapted from the Apache-2.0 licensed LingBot-VLA v2 implementation. +""" + + +# Qwen3-VL implementation used by LingBot-VLA v2. + +from types import MethodType +from typing import Callable, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.generation import GenerationMixin +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLForConditionalGeneration as _Qwen3VLForConditionalGeneration, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLModel as _Qwen3VLModel, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLPreTrainedModel as _Qwen3VLPreTrainedModel, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextAttention, + Qwen3VLTextMLP, + Qwen3VLTextRMSNorm, + Qwen3VLTextRotaryEmbedding, + Qwen3VLVisionMLP, + Qwen3VLVisionModel, + apply_rotary_pos_emb_vision, + eager_attention_forward, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextModel as _Qwen3VLTextModel, +) +from transformers.processing_utils import Unpack +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +class Qwen3VLPreTrainedModel(_Qwen3VLPreTrainedModel): + def _init_weights(self, module): + return + + +class Qwen3VLVisionAttention(nn.Module): + def __init__(self, config: Qwen3VLVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + max_seqlen: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + if self.config._attn_implementation == "flash_attention_2": + if max_seqlen is None: + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) + out_fp32_atten = False + if key_states.dtype == torch.float32: + out_fp32_atten = True + query_states = query_states.to(torch.bfloat16) + key_states = key_states.to(torch.bfloat16) + value_states = value_states.to(torch.bfloat16) + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + if out_fp32_atten: + attn_output = attn_output.to(torch.float32) + else: + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) + ] + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Qwen3VLVisionBlock(GradientCheckpointingLayer): + def __init__(self, config, attn_implementation: str = "sdpa") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen3VLVisionAttention(config=config) + self.mlp = Qwen3VLVisionMLP(config=config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class Qwen3VLTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx) + self.mlp = Qwen3VLTextMLP(config) + self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + att_output: Optional[torch.Tensor] = None, + start: Optional[int] = 0, + end: Optional[int] = 0, + compute_kqv: bool = False, + output_atten: bool = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + param_dtype = self.self_attn.q_proj.weight.dtype + hidden_states = hidden_states.to(param_dtype) + if att_output is not None: + att_output = att_output.to(param_dtype) + + if compute_kqv: + hidden_states = self.input_layernorm(hidden_states) + hidden_shape = (*hidden_states.shape[:-1], -1, self.self_attn.head_dim) + query_state = self.self_attn.q_norm(self.self_attn.q_proj(hidden_states).view(hidden_shape)) + key_state = self.self_attn.k_norm(self.self_attn.k_proj(hidden_states).view(hidden_shape)) + value_state = self.self_attn.v_proj(hidden_states).view(hidden_shape) + return query_state, key_state, value_state + + if output_atten: + if att_output.dtype != self.self_attn.o_proj.weight.dtype: + att_output = att_output.to(self.self_attn.o_proj.weight.dtype) + out_emb = self.self_attn.o_proj(att_output[:, start:end]) + out_emb += hidden_states + after_first_residual = out_emb.clone() + out_emb = self.post_attention_layernorm(out_emb) + out_emb = self.mlp(out_emb) + out_emb += after_first_residual + return out_emb + + position_embeddings = kwargs.pop("position_embeddings", None) + attention_mask = kwargs.pop("attention_mask", None) + if position_embeddings is not None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + raise ValueError( + f"Invalid operation compute_kqv={compute_kqv} and output_atten={output_atten} " + "with Qwen3VLTextDecoderLayer in LingBot-VLA" + ) + + +class Qwen3VLTextModel(_Qwen3VLTextModel): + def __init__(self, config: Qwen3VLTextConfig): + Qwen3VLPreTrainedModel.__init__(self, config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config) + self.gradient_checkpointing = False + self.post_init() + + +class Qwen3VLModel(_Qwen3VLModel): + def __init__(self, config: Qwen3VLConfig): + Qwen3VLPreTrainedModel.__init__(self, config) + self.visual = Qwen3VLVisionModel._from_config(config.vision_config) + self.visual.blocks = nn.ModuleList([Qwen3VLVisionBlock(config.vision_config) for _ in self.visual.blocks]) + self.visual.forward = MethodType(forward_without_grid_thw, self.visual) + self.visual.preprcess_grid_thw = MethodType(preprcess_grid_thw, self.visual) + self.language_model = Qwen3VLTextModel._from_config(config.text_config) + self.rope_deltas = None + self.post_init() + + +class Qwen3VLForConditionalGeneration(_Qwen3VLForConditionalGeneration, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + config_class = Qwen3VLConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"] + + def __init__(self, config): + Qwen3VLPreTrainedModel.__init__(self, config) + self.model = Qwen3VLModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + +@torch.compiler.disable +def preprcess_grid_thw(self, grid_thw: torch.Tensor): + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + seq_len = int(torch.prod(grid_thw, dim=1).sum().item()) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + split_sizes = (grid_thw.prod(-1) // self.spatial_merge_size**2).tolist() + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) + return None, position_embeddings, cu_seqlens, split_sizes, max_seqlen + + +def forward_without_grid_thw( + self, + hidden_states: torch.Tensor, + grid_thw: torch.Tensor = None, + pos_embeds: Optional[torch.Tensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + **kwargs, +) -> torch.Tensor: + hidden_states = self.patch_embed(hidden_states) + + if pos_embeds is None or position_embeddings is None or cu_seqlens is None or max_seqlen is None: + pos_embeds, position_embeddings, cu_seqlens, _, max_seqlen = self.preprcess_grid_thw(grid_thw) + if pos_embeds is None: + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + + hidden_states = hidden_states + pos_embeds + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len, -1) + + deepstack_feature_lists = [] + for layer_num, blk in enumerate(self.blocks): + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + max_seqlen=max_seqlen, + **kwargs, + ) + if layer_num in self.deepstack_visual_indexes: + deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)]( + hidden_states + ) + deepstack_feature_lists.append(deepstack_feature) + + hidden_states = self.merger(hidden_states) + return hidden_states, deepstack_feature_lists diff --git a/telefuser/ops/lingbot_vla_v2_moe.py b/telefuser/ops/lingbot_vla_v2_moe.py new file mode 100644 index 00000000..f734ebdc --- /dev/null +++ b/telefuser/ops/lingbot_vla_v2_moe.py @@ -0,0 +1,33 @@ +"""Compile-aware LingBot-VLA v2 MoE operation dispatch.""" + +from __future__ import annotations + +import torch + + +def robby_moe_forward( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + workspace: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Run the optional Triton grouped-MoE path for VLA eager inference.""" + if torch.compiler.is_compiling(): + raise RuntimeError("LingBot-VLA v2 Triton MoE is disabled during torch.compile") + if hidden_states.device.type != "cuda": + raise RuntimeError("LingBot-VLA v2 Triton MoE requires CUDA tensors") + + from telefuser.kernel.triton.lingbot_vla_v2_moe import robby_moe_forward as _triton_robby_moe_forward + + return _triton_robby_moe_forward( + hidden_states, + routing_weights, + selected_experts, + gate_weight, + up_weight, + down_weight, + workspace=workspace, + ) diff --git a/telefuser/pipelines/lingbot_vla_v2/__init__.py b/telefuser/pipelines/lingbot_vla_v2/__init__.py new file mode 100644 index 00000000..de38aed1 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py @@ -0,0 +1,19 @@ +"""TeleFuser pipeline components for LingBot-VLA v2 action inference.""" + +from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation +from .pipeline import LingBotVlaV2CanonicalActionChunk, LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig +from .policy import LingBotVlaV2PolicyStage +from .robot_profile import ROBOTWIN_CAMERA_KEYS, LingBotVlaV2ActionChunk, RobotWinProfile + +__all__ = [ + "LingBotVlaV2ActionChunk", + "LingBotVlaV2CanonicalActionChunk", + "LingBotVlaV2InputProcessor", + "LingBotVlaV2Inputs", + "LingBotVlaV2Observation", + "LingBotVlaV2Pipeline", + "LingBotVlaV2PipelineConfig", + "LingBotVlaV2PolicyStage", + "ROBOTWIN_CAMERA_KEYS", + "RobotWinProfile", +] diff --git a/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json b/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json new file mode 100644 index 00000000..71b222fb --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json @@ -0,0 +1,229 @@ +{ + "norm_stats": { + "action.arm.position": { + "mean": [ + -0.2395261526107788, + 1.1349077224731445, + 0.8105292320251465, + -0.31164029240608215, + 0.055165957659482956, + -0.03837483748793602, + 0.21789361536502838, + 1.0841481685638428, + 0.7913455963134766, + -0.32266682386398315, + -0.017583254724740982, + 0.03598335385322571 + ], + "std": [ + 0.41590750217437744, + 1.003921389579773, + 0.7768228650093079, + 0.6747837662696838, + 0.2715570628643036, + 0.6217551827430725, + 0.3499184548854828, + 1.024870753288269, + 0.7947020530700684, + 0.6865031123161316, + 0.24511374533176422, + 0.6167412400245667 + ], + "q01": [ + -1.0185421916961674, + -0.0010411963462829688, + -0.004309860050678266, + -1.5662350454330445, + -0.6512136519670486, + -2.2326875198364258, + -0.1715596118927003, + -0.003369329285621614, + -0.0018556645691394785, + -1.6451744033813476, + -1.0230259281158447, + -1.6478794967651362 + ], + "q99": [ + 0.17221172838211007, + 2.601926616668701, + 2.450952765509486, + 1.3516903750896456, + 1.2373228998184205, + 1.6001025575637815, + 0.9952441711425788, + 2.6186830965638164, + 2.453483357307315, + 1.2904379455566408, + 0.875431350231171, + 2.2640067550659193 + ], + "q02": [ + -0.948497843456269, + -0.0010411963462829688, + -0.0020969149172306023, + -1.474721703195572, + -0.39062818100452423, + -1.6791497331619256, + -0.0868722405433644, + -0.003369329285621614, + -0.0007056228727102265, + -1.5527206142425536, + -0.8707946321487425, + -1.5263084461212157 + ], + "q98": [ + 0.14699576301574702, + 2.5201579942703245, + 2.2860883530676364, + 1.221171345996857, + 1.0305539935112003, + 1.4689324659347545, + 0.9262396463394165, + 2.5411415680885314, + 2.298227728289366, + 1.154620874786377, + 0.577619640159607, + 1.8101414993286138 + ] + }, + "action.effector.position": { + "mean": [ + 0.664304256439209, + 0.6785873174667358 + ], + "std": [ + 0.45511099696159363, + 0.45013949275016785 + ], + "q01": [ + -1e-10, + -1e-10 + ], + "q99": [ + 0.99980000009996, + 0.99980000009996 + ], + "q02": [ + -1e-10, + -1e-10 + ], + "q98": [ + 0.99980000009996, + 0.99980000009996 + ] + }, + "observation.state.arm.position": { + "mean": [ + -0.2384551614522934, + 1.1301639080047607, + 0.8070681095123291, + -0.31032508611679077, + 0.05487748235464096, + -0.037838324904441833, + 0.21649977564811707, + 1.0786004066467285, + 0.78729248046875, + -0.3211615979671478, + -0.017434170469641685, + 0.03533728048205376 + ], + "std": [ + 0.41534423828125, + 1.0045241117477417, + 0.7767844200134277, + 0.6732552647590637, + 0.2708289623260498, + 0.619656503200531, + 0.34926003217697144, + 1.024977207183838, + 0.7943728566169739, + 0.6845870018005371, + 0.24417006969451904, + 0.6140097379684448 + ], + "q01": [ + -1.0185421916961674, + -0.0010411963462829688, + -0.004309860050678266, + -1.56473482670784, + -0.6483812011957168, + -2.224817314338684, + -0.1715596118927003, + -0.003369329285621614, + -0.0018556645691394785, + -1.6435380531311035, + -1.022286941242218, + -1.645177917861938 + ], + "q99": [ + 0.17221172838211007, + 2.601926616668701, + 2.4498462929427625, + 1.350190156364441, + 1.234490449047089, + 1.6001025575637815, + 0.9952441711425788, + 2.616768490922451, + 2.4511832739144563, + 1.2879834201812748, + 0.8739533764839176, + 2.26130517616272 + ], + "q02": [ + -0.948497843456269, + -0.0010411963462829688, + -0.0020969149172306023, + -1.4724713751077652, + -0.38921195561885824, + -1.6712795276641845, + -0.0837356712341304, + -0.003369329285621614, + -0.0007056228727102265, + -1.5510842639923097, + -0.8685776715278626, + -1.520905288314819 + ], + "q98": [ + 0.14419398908615033, + 2.5201579942703245, + 2.2838754079341888, + 1.2204212366342548, + 1.0291377681255343, + 1.4663090641021732, + 0.9231030770301825, + 2.5392269624471666, + 2.2959276448965076, + 1.152166349411011, + 0.575402679538727, + 1.7993351837158205 + ] + }, + "observation.state.effector.position": { + "mean": [ + 0.6655541062355042, + 0.6796996593475342 + ], + "std": [ + 0.45465800166130066, + 0.4496912956237793 + ], + "q01": [ + -1e-10, + -1e-10 + ], + "q99": [ + 0.99980000009996, + 0.99980000009996 + ], + "q02": [ + -1e-10, + -1e-10 + ], + "q98": [ + 0.99980000009996, + 0.99980000009996 + ] + } + }, + "count": 6062592 +} \ No newline at end of file diff --git a/telefuser/pipelines/lingbot_vla_v2/data.py b/telefuser/pipelines/lingbot_vla_v2/data.py new file mode 100644 index 00000000..5aedda3f --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/data.py @@ -0,0 +1,163 @@ +"""RobotWin input preparation for LingBot-VLA v2.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np +import torch +from PIL import Image +from torchvision.transforms.v2 import Resize + +from .robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile + +ImageInput = Image.Image | np.ndarray | torch.Tensor | str | Path + + +@dataclass(frozen=True) +class LingBotVlaV2Observation: + """One RobotWin observation accepted by the public SDK.""" + + task: str + state: torch.Tensor | Sequence[float] + images: Mapping[str, ImageInput] + + +@dataclass(frozen=True) +class LingBotVlaV2Inputs: + """Tensor contract consumed by ``LingBotVlaV2PolicyStage``.""" + + images: torch.Tensor + img_masks: torch.Tensor + lang_tokens: torch.Tensor + lang_masks: torch.Tensor + state: torch.Tensor + image_grid_thw: torch.Tensor + + +def _image_to_chw_uint8(image: ImageInput) -> torch.Tensor: + """Convert one RGB image to the format used by the upstream processor.""" + if isinstance(image, (str, Path)): + with Image.open(image) as opened: + image = np.asarray(opened.convert("RGB")) + elif isinstance(image, Image.Image): + image = np.asarray(image.convert("RGB")) + if isinstance(image, np.ndarray): + image = torch.from_numpy(np.asarray(image).copy()) + if not isinstance(image, torch.Tensor): + raise TypeError(f"unsupported image type: {type(image)!r}") + image = image.detach().to(device="cpu") + if image.ndim != 3: + raise ValueError(f"each image must have three dimensions, got {tuple(image.shape)}") + if image.shape[0] == 3: + chw = image + elif image.shape[-1] == 3: + chw = image.permute(2, 0, 1) + else: + raise ValueError(f"each image must have three RGB channels, got {tuple(image.shape)}") + + if chw.dtype == torch.uint8: + return chw.contiguous() + chw = chw.to(dtype=torch.float32) + if not torch.isfinite(chw).all(): + raise ValueError("images must contain only finite values") + if chw.numel() and float(chw.max()) <= 2.0 and float(chw.min()) >= 0.0: + chw = chw * 255.0 + return chw.round().clamp_(0, 255).to(dtype=torch.uint8).contiguous() + + +class LingBotVlaV2InputProcessor: + """Prepare RobotWin images, task text, and canonical state tensors.""" + + def __init__( + self, + processor: Any, + model_config: Any, + robot_profile: RobotWinProfile, + *, + image_size: int = 256, + ) -> None: + if processor is None or not hasattr(processor, "image_processor") or not hasattr(processor, "tokenizer"): + raise TypeError("LingBot-VLA v2 requires a Qwen3-VL AutoProcessor") + self.processor = processor + self.robot_profile = robot_profile + if image_size <= 0: + raise ValueError(f"image_size must be positive, got {image_size}") + self.image_size = int(image_size) + self.image_resize = Resize((self.image_size, self.image_size), antialias=True) + self.max_state_dim = int(getattr(model_config, "max_state_dim", 55)) + self.tokenizer_max_length = int(getattr(model_config, "tokenizer_max_length", 72)) + if self.max_state_dim != robot_profile.canonical_dim: + raise ValueError( + f"model max_state_dim is {self.max_state_dim}, RobotWin requires {robot_profile.canonical_dim}" + ) + + def _process_images(self, images: Mapping[str, ImageInput]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + missing = [key for key in ROBOTWIN_CAMERA_KEYS if key not in images] + if missing: + raise ValueError(f"RobotWin observation is missing camera keys: {missing}") + + processed_images: list[torch.Tensor] = [] + grids: list[torch.Tensor] = [] + for key in self.robot_profile.camera_keys: + image = self.image_resize(_image_to_chw_uint8(images[key]).to(dtype=torch.float32)) + output = self.processor.image_processor(image) + pixels = output["pixel_values"] if isinstance(output, dict) else output.pixel_values + grid = output.get("image_grid_thw") if isinstance(output, dict) else getattr(output, "image_grid_thw", None) + pixels = torch.as_tensor(pixels) + grid = None if grid is None else torch.as_tensor(grid) + if pixels.ndim == 3 and pixels.shape[0] == 1: + pixels = pixels.squeeze(0) + if pixels.ndim != 2: + raise ValueError(f"Qwen3-VL image processor must return [patches, features], got {tuple(pixels.shape)}") + if grid is None or grid.numel() < 3: + raise ValueError("Qwen3-VL image processor must return image_grid_thw") + processed_images.append(pixels) + grids.append(grid.reshape(-1, 3)[0].to(dtype=torch.long)) + + first_shape = processed_images[0].shape + if any(image.shape != first_shape for image in processed_images[1:]): + shapes = [tuple(image.shape) for image in processed_images] + raise ValueError(f"all RobotWin cameras must produce equal patch shapes, got {shapes}") + return ( + torch.stack(processed_images, dim=0).unsqueeze(0), + torch.ones(1, len(processed_images), dtype=torch.bool), + torch.stack(grids, dim=0).unsqueeze(0), + ) + + def _process_language(self, task: str) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(task, str) or not task.strip(): + raise ValueError("task must be a non-empty string") + tokenizer = self.processor.tokenizer + rendered = tokenizer.apply_chat_template( + [{"role": "user", "content": task}], + tokenize=False, + add_generation_prompt=False, + ) + tokens = tokenizer( + [rendered], + padding="max_length", + padding_side="right", + truncation=True, + max_length=self.tokenizer_max_length, + return_tensors="pt", + ) + return tokens["input_ids"], tokens["attention_mask"].to(dtype=torch.bool) + + def prepare(self, observation: LingBotVlaV2Observation) -> LingBotVlaV2Inputs: + """Prepare one RobotWin observation for model inference.""" + if not isinstance(observation, LingBotVlaV2Observation): + raise TypeError("observation must be a LingBotVlaV2Observation") + image_tensors, image_masks, image_grid_thw = self._process_images(observation.images) + state = self.robot_profile.normalize_state(observation.state).unsqueeze(0) + lang_tokens, lang_masks = self._process_language(observation.task) + return LingBotVlaV2Inputs( + images=image_tensors, + img_masks=image_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + state=state, + image_grid_thw=image_grid_thw, + ) diff --git a/telefuser/pipelines/lingbot_vla_v2/pipeline.py b/telefuser/pipelines/lingbot_vla_v2/pipeline.py new file mode 100644 index 00000000..c7653075 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/pipeline.py @@ -0,0 +1,122 @@ +"""BasePipeline integration for LingBot-VLA v2 base-model inference.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager + +from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation +from .policy import LingBotVlaV2PolicyStage +from .robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile + + +@dataclass +class LingBotVlaV2PipelineConfig: + """Runtime configuration for one LingBot-VLA v2 pipeline replica.""" + + policy_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + robot_profile: RobotWinProfile = field(default_factory=RobotWinProfile.default) + image_size: int = 256 + enable_metrics: bool = False + + +@dataclass(frozen=True) +class LingBotVlaV2CanonicalActionChunk: + """Normalized canonical actions produced by the base checkpoint.""" + + canonical_normalized_actions: torch.Tensor + horizon: int + action_dim: int + checkpoint_variant: str = "base" + policy_verified: bool = False + verification_status: str = "unverified_official_6b_base" + + +class LingBotVlaV2Pipeline(BasePipeline): + """Single-replica LingBot-VLA v2 canonical action SDK.""" + + # The service owns one fixed-shape resident policy. Per-request GC and + # allocator cache eviction add latency without releasing model weights. + clear_memory_after_call = False + + def _get_stages(self) -> list: + return [self.policy_stage] + + def init(self, module_manager: ModuleManager, config: LingBotVlaV2PipelineConfig) -> None: + self._model_info = module_manager.get_model_info() + self.config = config + policy = module_manager.fetch_module("lingbot_vla_v2") + processor = module_manager.fetch_module("lingbot_vla_v2_processor") + if policy is None or processor is None: + raise RuntimeError("LingBot-VLA v2 requires policy and lingbot_vla_v2_processor modules") + self.input_processor = LingBotVlaV2InputProcessor( + processor, + policy.config, + config.robot_profile, + image_size=config.image_size, + ) + self.policy_stage = LingBotVlaV2PolicyStage("policy", module_manager, config.policy_config) + if config.enable_metrics: + self.enable_metrics() + + @torch.inference_mode() + def predict( + self, + inputs: LingBotVlaV2Inputs, + seed: int | None = None, + ) -> LingBotVlaV2CanonicalActionChunk: + """Run prepared tensors and return normalized canonical actions.""" + actions = self.policy_stage.process(inputs, seed=seed) + if actions.shape[0] != 1: + raise RuntimeError(f"LingBot-VLA v2 pipeline expects batch size 1, got {actions.shape[0]}") + canonical_actions = actions[0] + policy_config = self.policy_stage.policy.config + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=canonical_actions, + horizon=int(canonical_actions.shape[0]), + action_dim=int(canonical_actions.shape[1]), + checkpoint_variant=str(getattr(policy_config, "checkpoint_variant", "base")), + policy_verified=bool(getattr(policy_config, "policy_verified", False)), + verification_status=str(getattr(policy_config, "verification_status", "unverified_official_6b_base")), + ) + + @torch.inference_mode() + def __call__( + self, + observation: LingBotVlaV2Observation, + seed: int | None = None, + ) -> LingBotVlaV2CanonicalActionChunk: + """Predict one normalized canonical action chunk.""" + return self.predict(self.input_processor.prepare(observation), seed=seed) + + def prepare_for_inference(self) -> None: + """Move the policy to its target device before the service becomes ready.""" + if not self.policy_stage.onload_models_flag: + self.policy_stage.onload_models() + self.policy_stage.onload_models_flag = True + + @torch.inference_mode() + def warmup(self) -> None: + """Initialize fixed-shape CUDA kernels before accepting service requests.""" + self.prepare_for_inference() + image_size = self.input_processor.image_size + image = torch.zeros(3, image_size, image_size, dtype=torch.uint8) + self( + LingBotVlaV2Observation( + task="warm up the policy", + state=[0.0] * 14, + images={key: image for key in ROBOTWIN_CAMERA_KEYS}, + ), + seed=0, + ) + + def close(self) -> None: + """Release policy device memory.""" + if hasattr(self, "policy_stage"): + self.policy_stage.offload_models() + self.policy_stage.onload_models_flag = False diff --git a/telefuser/pipelines/lingbot_vla_v2/policy.py b/telefuser/pipelines/lingbot_vla_v2/policy.py new file mode 100644 index 00000000..51a5b09e --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/policy.py @@ -0,0 +1,73 @@ +"""TeleFuser stage for LingBot-VLA v2 flow-matching action inference.""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.metrics import with_metrics + +from .data import LingBotVlaV2Inputs + + +class LingBotVlaV2PolicyStage(BaseStage): + """Run Qwen3-VL prefix encoding and the complete 10-step action sampler.""" + + def __init__(self, name: str, module_manager: ModuleManager, runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, runtime_config) + self.policy = module_manager.fetch_module("lingbot_vla_v2") + if self.policy is None: + raise RuntimeError("ModuleManager does not contain 'lingbot_vla_v2'") + self.model_names = ["policy"] + self._validate_parallelism() + + def _validate_parallelism(self) -> None: + parallel_config = self.model_runtime_config.parallel_config + if getattr(parallel_config, "world_size", 1) != 1: + raise ValueError("LingBot-VLA v2 currently supports one GPU per pipeline replica") + + @with_model_offload(["policy"]) + @torch.inference_mode() + @with_metrics + def process(self, inputs: LingBotVlaV2Inputs, seed: int | None = None) -> torch.Tensor: + """Return a CPU float32 normalized action chunk with shape ``[1, H, 55]``.""" + device = self.device + dtype = self.torch_dtype + tensors = { + "images": inputs.images.to(device=device, dtype=dtype), + "img_masks": inputs.img_masks.to(device=device), + "lang_tokens": inputs.lang_tokens.to(device=device), + "lang_masks": inputs.lang_masks.to(device=device), + "state": inputs.state.to(device=device, dtype=dtype), + "image_grid_thw": inputs.image_grid_thw.to(device=device, dtype=torch.long), + } + noise = None + if seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + config = self.policy.config + noise = torch.randn( + tensors["state"].shape[0], + int(config.n_action_steps), + int(config.max_action_dim), + device=device, + dtype=dtype, + generator=generator, + ) + actions = self.policy.sample_actions(**tensors, noise=noise) + if not isinstance(actions, torch.Tensor) or actions.ndim != 3: + raise RuntimeError(f"LingBot-VLA v2 policy returned an invalid action tensor: {type(actions)!r}") + config = self.policy.config + expected_shape = ( + tensors["state"].shape[0], + int(config.n_action_steps), + int(config.max_action_dim), + ) + if tuple(actions.shape) != expected_shape: + raise RuntimeError( + f"LingBot-VLA v2 policy returned shape {tuple(actions.shape)}, expected {expected_shape}" + ) + if not torch.isfinite(actions).all(): + raise RuntimeError("LingBot-VLA v2 policy returned non-finite actions") + return actions.detach().to(device="cpu", dtype=torch.float32) diff --git a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py new file mode 100644 index 00000000..e460b88d --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py @@ -0,0 +1,170 @@ +"""RobotWin feature mapping for LingBot-VLA v2 inference.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Mapping, Sequence + +import torch + +ROBOTWIN_CAMERA_KEYS = ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", +) +ROBOTWIN_STATE_DIM = 14 +CANONICAL_DIM = 55 +ARM_SLICE = slice(0, 12) +EFFECTOR_SLICE = slice(28, 30) + + +@dataclass(frozen=True) +class LingBotVlaV2ActionChunk: + """Structured RobotWin action chunk returned by the SDK.""" + + fields: Mapping[str, torch.Tensor] + raw_actions: torch.Tensor + action_mask: torch.Tensor + horizon: int + robot_profile: str = "robotwin" + policy_verified: bool = False + verification_status: str = "unverified_official_6b_base" + canonical_normalized_actions: torch.Tensor | None = None + + +class RobotWinProfile: + """Map RobotWin observations and actions to LingBot's canonical space.""" + + name = "robotwin" + camera_keys = ROBOTWIN_CAMERA_KEYS + canonical_dim = CANONICAL_DIM + raw_state_dim = ROBOTWIN_STATE_DIM + _REQUIRED_STATS = ( + "observation.state.arm.position", + "observation.state.effector.position", + "action.arm.position", + "action.effector.position", + ) + + def __init__(self, norm_stats: Mapping[str, Mapping[str, object]]) -> None: + self._stats = { + key: { + stat_name: torch.as_tensor(stat_value, dtype=torch.float64) for stat_name, stat_value in values.items() + } + for key, values in norm_stats.items() + } + self._validate_stats() + + @classmethod + def from_json(cls, path: str | Path) -> "RobotWinProfile": + """Load RobotWin normalization statistics from an upstream-format JSON file.""" + payload = json.loads(Path(path).read_text(encoding="utf-8")) + norm_stats = payload.get("norm_stats") + if not isinstance(norm_stats, dict): + raise ValueError("RobotWin normalization file must contain a norm_stats object") + return cls(norm_stats) + + @classmethod + def default(cls) -> "RobotWinProfile": + """Load the RobotWin statistics bundled with TeleFuser.""" + path = Path(__file__).with_name("assets") / "robotwin_norm_stats.json" + return cls.from_json(path) + + @property + def action_mask(self) -> torch.Tensor: + """Return the canonical dimensions used by RobotWin actions.""" + mask = torch.zeros(self.canonical_dim, dtype=torch.bool) + mask[ARM_SLICE] = True + mask[EFFECTOR_SLICE] = True + return mask + + def normalize_state(self, raw_state: torch.Tensor | Sequence[float]) -> torch.Tensor: + """Convert one raw 14-D RobotWin state to normalized canonical 55-D space.""" + state = torch.as_tensor(raw_state, dtype=torch.float32, device="cpu") + if state.shape != (self.raw_state_dim,): + raise ValueError(f"RobotWin state must have shape ({self.raw_state_dim},), got {tuple(state.shape)}") + if not torch.isfinite(state).all(): + raise ValueError("RobotWin state must contain only finite values") + + arm = torch.cat((state[0:6], state[7:13])) + effector = state[[6, 13]] + canonical = torch.zeros(self.canonical_dim, dtype=torch.float32) + canonical[ARM_SLICE] = self._normalize("observation.state.arm.position", arm) + canonical[EFFECTOR_SLICE] = self._normalize("observation.state.effector.position", effector) + return canonical + + def structure_actions( + self, + canonical_normalized_actions: torch.Tensor, + *, + include_canonical: bool = False, + ) -> LingBotVlaV2ActionChunk: + """Convert a normalized canonical action chunk to RobotWin action fields.""" + actions = torch.as_tensor(canonical_normalized_actions, dtype=torch.float32, device="cpu") + if actions.ndim == 3: + if actions.shape[0] != 1: + raise ValueError("RobotWin structured output currently supports a single observation") + actions = actions[0] + if actions.ndim != 2 or actions.shape[-1] != self.canonical_dim: + raise ValueError( + f"canonical actions must have shape [H,{self.canonical_dim}] or [1,H,{self.canonical_dim}], " + f"got {tuple(actions.shape)}" + ) + if not torch.isfinite(actions).all(): + raise ValueError("canonical actions must contain only finite values") + + arm = self._unnormalize("action.arm.position", actions[:, ARM_SLICE]) + effector = self._unnormalize("action.effector.position", actions[:, EFFECTOR_SLICE]) + raw = torch.empty(actions.shape[0], self.raw_state_dim, dtype=torch.float32) + raw[:, 0:6] = arm[:, 0:6] + raw[:, 6] = effector[:, 0] + raw[:, 7:13] = arm[:, 6:12] + raw[:, 13] = effector[:, 1] + fields = MappingProxyType( + { + "action.arm.position": arm, + "action.effector.position": effector, + "action": raw, + } + ) + return LingBotVlaV2ActionChunk( + fields=fields, + raw_actions=raw, + action_mask=self.action_mask, + horizon=int(actions.shape[0]), + canonical_normalized_actions=actions.clone() if include_canonical else None, + ) + + def _validate_stats(self) -> None: + expected_dims = { + "observation.state.arm.position": 12, + "observation.state.effector.position": 2, + "action.arm.position": 12, + "action.effector.position": 2, + } + missing = [key for key in self._REQUIRED_STATS if key not in self._stats] + if missing: + raise ValueError(f"RobotWin normalization statistics are missing keys: {missing}") + for key, expected_dim in expected_dims.items(): + values = self._stats[key] + for stat_name in ("q01", "q99"): + if stat_name not in values or values[stat_name].shape != (expected_dim,): + shape = None if stat_name not in values else tuple(values[stat_name].shape) + raise ValueError( + f"RobotWin statistic {key}.{stat_name} must have shape ({expected_dim},), got {shape}" + ) + + def _normalize(self, key: str, value: torch.Tensor) -> torch.Tensor: + low = self._stats[key]["q01"] + high = self._stats[key]["q99"] + normalized = (value.to(dtype=torch.float64) - low) / (high - low + 1e-6) * 2.0 - 1.0 + return normalized.to(dtype=value.dtype) + + def _unnormalize(self, key: str, value: torch.Tensor) -> torch.Tensor: + low = self._stats[key]["q01"] + high = self._stats[key]["q99"] + unnormalized = (value.to(dtype=torch.float64) + 1.0) / 2.0 * (high - low + 1e-6) + low + return unnormalized.to(dtype=value.dtype) diff --git a/telefuser/pipelines/lingbot_vla_v2/runtime.py b/telefuser/pipelines/lingbot_vla_v2/runtime.py new file mode 100644 index 00000000..6e3f93a2 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/runtime.py @@ -0,0 +1,52 @@ +"""Runtime construction for single-replica LingBot-VLA v2 inference.""" + +from __future__ import annotations + +import torch +from transformers import AutoProcessor + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2 + +from .pipeline import LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig + + +def get_lingbot_vla_v2_pipeline( + model_root: str, + qwen3vl_root: str, + device: str = "cuda:0", + *, + warmup: bool = False, +) -> LingBotVlaV2Pipeline: + """Load one official 6B base checkpoint replica for inference.""" + target_device = torch.device(device) + if target_device.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError(f"CUDA device {device!r} was requested, but CUDA is unavailable") + device_index = target_device.index or 0 + if device_index >= torch.cuda.device_count(): + raise ValueError( + f"CUDA device index {device_index} is unavailable; visible device count is {torch.cuda.device_count()}" + ) + target_device = torch.device("cuda", device_index) + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + processor = AutoProcessor.from_pretrained(qwen3vl_root, local_files_only=True, padding_side="right") + manager = ModuleManager(torch_dtype=dtype, device="cpu") + manager.add_module(processor, "lingbot_vla_v2_processor", path=qwen3vl_root) + load_lingbot_vla_v2(manager, model_root, qwen3vl_root, torch_dtype=dtype) + pipeline = LingBotVlaV2Pipeline(device=str(target_device), torch_dtype=dtype) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig( + device_type=target_device.type, + device_id=target_device.index or 0, + torch_dtype=dtype, + ), + ), + ) + pipeline.prepare_for_inference() + if warmup: + pipeline.warmup() + return pipeline diff --git a/telefuser/pipelines/lingbot_vla_v2/service.py b/telefuser/pipelines/lingbot_vla_v2/service.py new file mode 100644 index 00000000..f0f78b38 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/service.py @@ -0,0 +1,206 @@ +"""Minimal single-GPU HTTP service for LingBot-VLA v2 action inference.""" + +from __future__ import annotations + +import base64 +import binascii +import io +import math +import threading +from collections.abc import Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Protocol + +from PIL import Image, UnidentifiedImageError +from fastapi import FastAPI, HTTPException +from fastapi.concurrency import run_in_threadpool +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from .data import LingBotVlaV2Observation +from .pipeline import LingBotVlaV2CanonicalActionChunk +from .robot_profile import ROBOTWIN_CAMERA_KEYS +from .runtime import get_lingbot_vla_v2_pipeline + + +@dataclass(frozen=True) +class LingBotVlaV2ServiceConfig: + """Configuration for one process-local LingBot-VLA v2 replica.""" + + model_root: str + qwen3vl_root: str + device: str = "cuda:0" + max_image_bytes: int = 10 * 1024 * 1024 + + def __post_init__(self) -> None: + if self.max_image_bytes <= 0: + raise ValueError("max_image_bytes must be positive") + + +class LingBotVlaV2ActionRequest(BaseModel): + """One RobotWin observation encoded for the HTTP boundary.""" + + model_config = ConfigDict(extra="forbid") + + task: str = Field(min_length=1) + state: list[float] = Field(min_length=14, max_length=14) + camera_high: str = Field(min_length=1) + camera_left_wrist: str = Field(min_length=1) + camera_right_wrist: str = Field(min_length=1) + seed: int | None = None + + @field_validator("task") + @classmethod + def validate_task(cls, value: str) -> str: + """Reject whitespace-only instructions.""" + value = value.strip() + if not value: + raise ValueError("task must be a non-empty string") + return value + + @field_validator("state") + @classmethod + def validate_state(cls, value: list[float]) -> list[float]: + """Reject non-finite robot state values.""" + if not all(math.isfinite(item) for item in value): + raise ValueError("state must contain only finite values") + return value + + +class LingBotVlaV2ActionResponse(BaseModel): + """Normalized canonical action chunk returned by the base checkpoint.""" + + canonical_normalized_actions: list[list[float]] + horizon: int + action_dim: int + checkpoint_variant: str + policy_verified: bool + verification_status: str + + +class LingBotVlaV2HealthResponse(BaseModel): + """Readiness state for the process-local model replica.""" + + status: str + model: str + device: str + policy_verified: bool + + +class _Pipeline(Protocol): + def __call__( + self, + observation: LingBotVlaV2Observation, + seed: int | None = None, + ) -> LingBotVlaV2CanonicalActionChunk: ... + + def close(self) -> None: ... + + +PipelineFactory = Callable[[LingBotVlaV2ServiceConfig], _Pipeline] + + +def _default_pipeline_factory(config: LingBotVlaV2ServiceConfig) -> _Pipeline: + return get_lingbot_vla_v2_pipeline(config.model_root, config.qwen3vl_root, device=config.device, warmup=True) + + +def _decode_image(value: str, *, max_image_bytes: int) -> Image.Image: + payload = value.strip() + if payload.startswith("data:"): + header, separator, payload = payload.partition(",") + if not separator or ";base64" not in header.lower(): + raise ValueError("image data URLs must use base64 encoding") + max_encoded_length = 4 * ((max_image_bytes + 2) // 3) + if len(payload) > max_encoded_length: + raise ValueError(f"decoded image must not exceed {max_image_bytes} bytes") + try: + decoded = base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError) as error: + raise ValueError("image must be valid base64") from error + if not decoded or len(decoded) > max_image_bytes: + raise ValueError(f"decoded image must contain 1 to {max_image_bytes} bytes") + try: + with Image.open(io.BytesIO(decoded)) as image: + return image.convert("RGB").copy() + except (UnidentifiedImageError, OSError) as error: + raise ValueError("decoded payload must be a supported image") from error + + +def predict_lingbot_vla_v2_action( + pipeline: _Pipeline, + request: LingBotVlaV2ActionRequest, + *, + max_image_bytes: int, +) -> LingBotVlaV2ActionResponse: + """Decode one request and return the canonical normalized action chunk.""" + encoded_images = (request.camera_high, request.camera_left_wrist, request.camera_right_wrist) + images = { + key: _decode_image(value, max_image_bytes=max_image_bytes) + for key, value in zip(ROBOTWIN_CAMERA_KEYS, encoded_images, strict=True) + } + observation = LingBotVlaV2Observation(task=request.task, state=request.state, images=images) + chunk = pipeline(observation, seed=request.seed) + return LingBotVlaV2ActionResponse( + canonical_normalized_actions=chunk.canonical_normalized_actions.tolist(), + horizon=chunk.horizon, + action_dim=chunk.action_dim, + checkpoint_variant=chunk.checkpoint_variant, + policy_verified=chunk.policy_verified, + verification_status=chunk.verification_status, + ) + + +class LingBotVlaV2Service: + """Serialize requests through one loaded policy replica.""" + + def __init__(self, pipeline: _Pipeline, config: LingBotVlaV2ServiceConfig) -> None: + self.pipeline = pipeline + self.config = config + self._inference_lock = threading.Lock() + + def predict(self, request: LingBotVlaV2ActionRequest) -> LingBotVlaV2ActionResponse: + """Decode one request and run it on the process-local replica.""" + with self._inference_lock: + return predict_lingbot_vla_v2_action(self.pipeline, request, max_image_bytes=self.config.max_image_bytes) + + def close(self) -> None: + """Release model resources during application shutdown.""" + self.pipeline.close() + + +def create_lingbot_vla_v2_app( + config: LingBotVlaV2ServiceConfig, + *, + pipeline_factory: PipelineFactory = _default_pipeline_factory, +) -> FastAPI: + """Create a FastAPI application backed by exactly one policy replica.""" + + @asynccontextmanager + async def lifespan(app: FastAPI): + service = LingBotVlaV2Service(pipeline_factory(config), config) + app.state.lingbot_vla_v2_service = service + try: + yield + finally: + service.close() + + app = FastAPI(title="LingBot VLA v2", version="1", lifespan=lifespan) + + @app.get("/health", response_model=LingBotVlaV2HealthResponse) + async def health() -> LingBotVlaV2HealthResponse: + return LingBotVlaV2HealthResponse( + status="ready", + model="lingbot-vla-v2-6b-base", + device=config.device, + policy_verified=False, + ) + + @app.post("/v1/vla/actions", response_model=LingBotVlaV2ActionResponse) + async def predict(request: LingBotVlaV2ActionRequest) -> LingBotVlaV2ActionResponse: + service: LingBotVlaV2Service = app.state.lingbot_vla_v2_service + try: + return await run_in_threadpool(service.predict, request) + except (TypeError, ValueError) as error: + raise HTTPException(status_code=422, detail=str(error)) from error + + return app diff --git a/telefuser/service/api/__init__.py b/telefuser/service/api/__init__.py index 3660471d..4da4d21a 100644 --- a/telefuser/service/api/__init__.py +++ b/telefuser/service/api/__init__.py @@ -17,6 +17,8 @@ OutputFormat, StopTaskResponse, StopTaskStatus, + StructuredTaskRequest, + StructuredTaskResponse, TaskRequest, TaskResponse, TaskStatus, @@ -29,6 +31,8 @@ "RateLimitMiddleware", "LoggingMiddleware", "setup_middleware", + "StructuredTaskRequest", + "StructuredTaskResponse", "TaskRequest", "TaskResponse", "StopTaskResponse", diff --git a/telefuser/service/api/api_server.py b/telefuser/service/api/api_server.py index f1398506..76096e1f 100644 --- a/telefuser/service/api/api_server.py +++ b/telefuser/service/api/api_server.py @@ -18,7 +18,7 @@ from ..core.file_service import FileService from ..core.task_manager import TaskManager from ..core.task_processor import AsyncTaskProcessor -from ..core.task_service import MediaGenerationService +from ..core.task_service import MediaGenerationService, StructuredInferenceService from . import routers from .task_application_service import TaskApplicationService @@ -58,6 +58,7 @@ def __init__( self.file_service: FileService | None = None self.inference_service: PipelineService | None = None self.media_service: MediaGenerationService | None = None + self.structured_service: StructuredInferenceService | None = None self.task_app_service = TaskApplicationService(self) self.cache_service: Any | None = None self.max_queue_size = max_queue_size @@ -188,6 +189,7 @@ async def ensure_task_processor_running(self) -> None: return if self.task_processor.is_running: + self.task_processor.notify_task_available() await self.ensure_artifact_cleanup_running() return @@ -196,8 +198,10 @@ async def ensure_task_processor_running(self) -> None: logger.warning("Task processor is not initialized; task will remain pending until services are ready") return if self.task_processor.is_running: + self.task_processor.notify_task_available() return await self.task_processor.start() + self.task_processor.notify_task_available() await self.ensure_artifact_cleanup_running() async def ensure_artifact_cleanup_running(self) -> None: @@ -338,9 +342,11 @@ def initialize_services( cache_service=cache_service, cache_adapter=cache_adapter, ) + self.structured_service = StructuredInferenceService(inference_service) self.task_processor = AsyncTaskProcessor( task_manager=self.task_manager, media_service=self.media_service, + structured_service=self.structured_service, max_concurrent=self.max_concurrent_tasks, ) diff --git a/telefuser/service/api/routers/tasks.py b/telefuser/service/api/routers/tasks.py index 437264d7..e48beb62 100644 --- a/telefuser/service/api/routers/tasks.py +++ b/telefuser/service/api/routers/tasks.py @@ -19,7 +19,7 @@ from telefuser.service_types import MediaType, StopTaskStatus from telefuser.utils.logging import logger -from ..schema import StopTaskResponse, TaskRequest, TaskResponse +from ..schema import StopTaskResponse, StructuredTaskRequest, StructuredTaskResponse, TaskRequest, TaskResponse from ..task_contract_runtime import match_task_candidates if TYPE_CHECKING: @@ -40,6 +40,11 @@ async def create_task(message: TaskRequest) -> TaskResponse: """Create a new generation task.""" return await routes.create_task(message) + @new_router.post("/structured", response_model=StructuredTaskResponse) + async def create_structured_task(message: StructuredTaskRequest) -> StructuredTaskResponse: + """Create a task whose pipeline contract declares a structured result.""" + return await routes.create_structured_task(message) + @new_router.post("/form", response_model=TaskResponse) async def create_task_form( request: Request, @@ -99,6 +104,19 @@ async def check_image_path(image_name: str) -> None: logger.error(f"Failed to create task: {e}") raise HTTPException(status_code=500, detail=str(e)) + async def create_structured_task(self, message: StructuredTaskRequest) -> StructuredTaskResponse: + """Create a structured inference task without allocating an artifact path.""" + try: + return await self.api.task_app_service.submit_structured( + message, + explicit_fields=set(getattr(message, "model_fields_set", set())), + ) + except HTTPException: + raise + except Exception as error: + logger.error(f"Failed to create structured task: {error}") + raise HTTPException(status_code=500, detail=str(error)) from error + async def list_tasks(self) -> dict: """List all tasks.""" return self.api.task_manager.get_all_tasks() diff --git a/telefuser/service/api/schema.py b/telefuser/service/api/schema.py index fa23bbd7..e61fa25f 100644 --- a/telefuser/service/api/schema.py +++ b/telefuser/service/api/schema.py @@ -65,6 +65,23 @@ class TaskStatusMessage(BaseModel): task_id: str = Field(..., description="Task ID") +class StructuredTaskRequest(BaseModel): + """Request model for JSON-serializable inference results.""" + + model_config = ConfigDict(extra="allow") + + task_id: str = Field(default_factory=generate_task_id, description="Task ID (auto-generated)") + task: str = Field(..., description="Structured task type declared by the pipeline contract") + + @field_validator("task") + @classmethod + def validate_task(cls: type["StructuredTaskRequest"], value: str) -> str: + return validate_task_name_format(value) + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + class TaskResponse(BaseModel): """Response model for task creation.""" @@ -80,3 +97,10 @@ class StopTaskResponse(BaseModel): stop_status: StopTaskStatus reason: str + + +class StructuredTaskResponse(BaseModel): + """Response returned when a structured task is accepted.""" + + task_id: str + task_status: TaskStatus diff --git a/telefuser/service/api/task_application_service.py b/telefuser/service/api/task_application_service.py index 6d53c820..70b16d08 100644 --- a/telefuser/service/api/task_application_service.py +++ b/telefuser/service/api/task_application_service.py @@ -15,7 +15,7 @@ from telefuser.service.core.pipeline_contract import infer_media_type_for_task from telefuser.service_types import MediaType, TaskStatus -from .schema import TaskRequest, TaskResponse +from .schema import StructuredTaskRequest, StructuredTaskResponse, TaskRequest, TaskResponse from .task_contract_runtime import apply_task_contract_defaults, validate_required_task_parameters if TYPE_CHECKING: @@ -56,6 +56,34 @@ async def submit( except RuntimeError as exc: raise HTTPException(status_code=503, detail=str(exc)) + async def submit_structured( + self, + message: StructuredTaskRequest, + *, + explicit_fields: set[str], + ensure_processing: bool = True, + ) -> StructuredTaskResponse: + """Validate and enqueue a task whose result is returned as JSON.""" + try: + self.api.validate_task_supported(message.task) + contract = self.api.get_task_contract(message.task) + media_type = str((contract or {}).get("media_type") or infer_media_type_for_task(message.task)) + if media_type != MediaType.STRUCTURED.value: + raise HTTPException( + status_code=400, + detail=f"Task '{message.task}' does not declare a structured result contract", + ) + apply_task_contract_defaults(message, task_contract=contract, explicit_fields=explicit_fields) + validate_required_task_parameters(message, task_contract=contract) + + task_id = self.api.task_manager.create_task(message) + message.task_id = task_id + if ensure_processing: + await self.api.ensure_task_processor_running() + return StructuredTaskResponse(task_id=task_id, task_status=TaskStatus.PENDING) + except RuntimeError as exc: + raise HTTPException(status_code=503, detail=str(exc)) + def validate_output_path(self, message: TaskRequest) -> None: file_service = self.api.file_service if file_service is None: diff --git a/telefuser/service/api/task_contract_runtime.py b/telefuser/service/api/task_contract_runtime.py index 74e2a9c8..35fa8a42 100644 --- a/telefuser/service/api/task_contract_runtime.py +++ b/telefuser/service/api/task_contract_runtime.py @@ -132,6 +132,8 @@ def _build_default_output_path(message: Any) -> str: return "" media_type = infer_media_type_for_task(task) + if media_type == "structured": + return "" if media_type == "image": output_format = getattr(message, "output_format", "png") or "png" return f"{task_id}.{output_format}" diff --git a/telefuser/service/core/pipeline_contract.py b/telefuser/service/core/pipeline_contract.py index 747f2ffe..5a72de9f 100644 --- a/telefuser/service/core/pipeline_contract.py +++ b/telefuser/service/core/pipeline_contract.py @@ -10,6 +10,7 @@ VIDEO_TASKS = frozenset({"t2v", "i2v", "fl2v", "vc", "s2v", "vsr"}) IMAGE_TASKS = frozenset({"t2i", "i2i", "edit"}) +STRUCTURED_TASKS = frozenset({"vla_action"}) TASK_NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]{1,31}$") @@ -257,6 +258,8 @@ def _derive_media_types_from_tasks(tasks: tuple[str, ...]) -> list[str]: media_types.append("video") if any(task in IMAGE_TASKS for task in tasks): media_types.append("image") + if any(task in STRUCTURED_TASKS for task in tasks): + media_types.append("structured") if not media_types: media_types.append("unknown") return media_types @@ -336,4 +339,6 @@ def infer_media_type_for_task(task: str) -> str: return "image" if is_video_task(task): return "video" + if task in STRUCTURED_TASKS: + return "structured" return "video" diff --git a/telefuser/service/core/pipeline_pool.py b/telefuser/service/core/pipeline_pool.py index e5081852..74b982fb 100644 --- a/telefuser/service/core/pipeline_pool.py +++ b/telefuser/service/core/pipeline_pool.py @@ -207,8 +207,11 @@ async def acquire(self) -> AsyncIterator[ReplicaHandle]: continue handle = self._handles[idx] - if handle._dead: - self._evict_replica(idx, "pre-existing dead state") + process = getattr(handle, "process", None) + process_dead = process is not None and not process.is_alive() + if handle._dead or process_dead: + reason = "process exited" if process_dead else "pre-existing dead state" + self._evict_replica(idx, reason) continue break diff --git a/telefuser/service/core/pipeline_runner.py b/telefuser/service/core/pipeline_runner.py index dbc5dcda..09d5fddf 100644 --- a/telefuser/service/core/pipeline_runner.py +++ b/telefuser/service/core/pipeline_runner.py @@ -144,6 +144,8 @@ async def shutdown(self) -> None: await self._pipeline.astop() elif hasattr(self._pipeline, "stop"): await asyncio.to_thread(self._pipeline.stop) + elif hasattr(self._pipeline, "close"): + await asyncio.to_thread(self._pipeline.close) self._started = False diff --git a/telefuser/service/core/replica_worker.py b/telefuser/service/core/replica_worker.py index abb7de60..61945089 100644 --- a/telefuser/service/core/replica_worker.py +++ b/telefuser/service/core/replica_worker.py @@ -197,13 +197,16 @@ def _forward_cancel_fn( cancel_event: mp_stdlib.Event, stop_event: threading.Event, forwarder_exit: threading.Event, + forwarder_wake: threading.Event, forwarder_done: threading.Event, ) -> None: """Forward main-process stop_event to subprocess cancel_event (polled).""" while not forwarder_exit.is_set(): - if stop_event.wait(timeout=0.5): + if stop_event.is_set(): cancel_event.set() break + forwarder_wake.wait(timeout=0.5) + forwarder_wake.clear() forwarder_done.set() @@ -236,24 +239,32 @@ async def run_task( self.cancel_event.clear() forwarder_exit = threading.Event() + forwarder_wake = threading.Event() forwarder_done = threading.Event() forwarder = threading.Thread( target=_forward_cancel_fn, - args=(self.cancel_event, stop_event, forwarder_exit, forwarder_done), + args=(self.cancel_event, stop_event, forwarder_exit, forwarder_wake, forwarder_done), daemon=True, ) forwarder.start() loop = asyncio.get_running_loop() - self.conn.send(("task", task_data, timeout_s, output_root)) - ipc_timeout = (timeout_s or 600) + _TASK_IPC_MARGIN_S try: - result = await loop.run_in_executor(None, self._recv_with_health_check, ipc_timeout) + if not self.process.is_alive(): + self._dead = True + raise ReplicaDeadError(f"Replica {self.replica_id} process is not alive") + try: + self.conn.send(("task", task_data, timeout_s, output_root)) + result = await loop.run_in_executor(None, self._recv_with_health_check, ipc_timeout) + except (EOFError, OSError) as error: + self._dead = True + raise ReplicaDeadError(f"Replica {self.replica_id} IPC failed: {error}") from error finally: forwarder_exit.set() + forwarder_wake.set() forwarder_done.wait(2.0) if result is None: diff --git a/telefuser/service/core/task_manager.py b/telefuser/service/core/task_manager.py index c04c1334..3258efe7 100644 --- a/telefuser/service/core/task_manager.py +++ b/telefuser/service/core/task_manager.py @@ -30,6 +30,7 @@ class TaskInfo: output_path: str | None = None peak_memory_mb: float | None = None inference_time_s: float | None = None + result: dict[str, Any] | None = None stop_event: threading.Event = field(default_factory=threading.Event) thread: threading.Thread | None = None @@ -127,6 +128,7 @@ def complete_task( *, peak_memory_mb: float | None = None, inference_time_s: float | None = None, + result: dict[str, Any] | None = None, ) -> None: """Mark task as completed with metrics.""" with self._lock: @@ -151,6 +153,7 @@ def complete_task( task.inference_time_s = inference_time_s if inference_time_s is not None else duration get_service_metrics().record_task_completed(duration) task.peak_memory_mb = peak_memory_mb + task.result = result def fail_task(self, task_id: str, error: str) -> None: """Mark task as failed with metrics.""" @@ -227,6 +230,8 @@ def get_task_status(self, task_id: str) -> dict[str, Any] | None: "peak_memory_mb": task.peak_memory_mb, "inference_time_s": task.inference_time_s, } + if task.result is not None: + status["result"] = task.result status.update(self._serialize_task_message(task.message)) return status diff --git a/telefuser/service/core/task_processor.py b/telefuser/service/core/task_processor.py index 16fe8ecf..18577478 100644 --- a/telefuser/service/core/task_processor.py +++ b/telefuser/service/core/task_processor.py @@ -11,8 +11,9 @@ from telefuser.utils.logging import logger +from ..api.schema import StructuredTaskRequest from .task_manager import TaskManager, TaskStatus -from .task_service import MediaGenerationService +from .task_service import MediaGenerationService, StructuredInferenceService class AsyncTaskProcessor: @@ -27,13 +28,15 @@ def __init__( task_manager: TaskManager, media_service: MediaGenerationService, max_concurrent: int = 1, + structured_service: StructuredInferenceService | None = None, ) -> None: """Initialize the async task processor.""" self.task_manager = task_manager self.media_service = media_service + self.structured_service = structured_service self.max_concurrent = max_concurrent - self._queue: asyncio.Queue = asyncio.Queue() + self._queue: asyncio.Queue[None] = asyncio.Queue() self._workers: list[asyncio.Task] = [] self._running = False self._stop_event = asyncio.Event() @@ -44,6 +47,10 @@ def is_running(self) -> bool: """Whether the processor workers are running.""" return self._running + def notify_task_available(self) -> None: + """Wake one idle worker after a task is added to the task manager.""" + self._queue.put_nowait(None) + async def start(self) -> None: """Start the task processor workers.""" if self._running: @@ -54,6 +61,7 @@ async def start(self) -> None: self._stop_event.clear() self._loop = asyncio.get_running_loop() + self._queue = asyncio.Queue() for i in range(self.max_concurrent): worker = asyncio.create_task(self._worker_loop(f"worker-{i}"), name=f"task-processor-{i}") self._workers.append(worker) @@ -98,7 +106,7 @@ async def _worker_loop(self, worker_name: str) -> None: task_id = self.task_manager.claim_next_pending_task() if task_id is None: - await asyncio.wait_for(self._stop_event.wait(), timeout=1.0) + await asyncio.wait_for(self._queue.get(), timeout=1.0) continue await self._process_task(task_id) @@ -137,9 +145,16 @@ async def _process_task(self, task_id: str) -> None: return try: - result = await self.media_service.generate_media_with_stop_event( - task_info.message, task_info.stop_event - ) + if isinstance(task_info.message, StructuredTaskRequest): + if self.structured_service is None: + raise RuntimeError("Structured inference service is not initialized") + result = await self.structured_service.execute_with_stop_event( + task_info.message, task_info.stop_event + ) + else: + result = await self.media_service.generate_media_with_stop_event( + task_info.message, task_info.stop_event + ) if result: self.task_manager.complete_task( @@ -147,6 +162,7 @@ async def _process_task(self, task_id: str) -> None: result.output_path, peak_memory_mb=result.peak_memory_mb, inference_time_s=result.inference_time_s, + result=getattr(result, "result", None), ) logger.info(f"Task {task_id} completed successfully") else: diff --git a/telefuser/service/core/task_service.py b/telefuser/service/core/task_service.py index 927f4cad..224a09d5 100644 --- a/telefuser/service/core/task_service.py +++ b/telefuser/service/core/task_service.py @@ -2,14 +2,16 @@ from __future__ import annotations +import json import threading +from dataclasses import dataclass from types import SimpleNamespace from typing import TYPE_CHECKING, Any from telefuser.service_types import MediaType, PipelineRunStatus, TaskStatus, TaskType from telefuser.utils.logging import logger -from ..api.schema import TaskRequest, TaskResponse +from ..api.schema import StructuredTaskRequest, TaskRequest, TaskResponse from ..media.media_base import AudioHandler, ImageHandler, VideoHandler from .file_service import FileService from .pipeline_contract import infer_media_type_for_task @@ -17,6 +19,59 @@ if TYPE_CHECKING: from .pipeline_service import PipelineService + +@dataclass(frozen=True) +class StructuredTaskExecutionResponse: + """Internal result returned to the shared task processor.""" + + task_id: str + result: dict[str, Any] + output_path: None = None + peak_memory_mb: float | None = None + inference_time_s: float | None = None + + +class StructuredInferenceService: + """Execute pipeline tasks that return JSON objects instead of artifacts.""" + + def __init__(self, inference_service: "PipelineService") -> None: + self.inference_service = inference_service + + async def execute_with_stop_event( + self, + message: StructuredTaskRequest, + stop_event: threading.Event, + ) -> StructuredTaskExecutionResponse | None: + """Run one structured task and retain its JSON result in task state.""" + if stop_event.is_set(): + logger.info(f"Task {message.task_id} cancelled before processing") + return None + + task_data = message.model_dump(mode="json") + result = await self.inference_service.run_task_with_stop_event(task_data, stop_event) + if result is None: + if stop_event.is_set(): + return None + raise RuntimeError("Task processing timeout") + if result.get("status") != PipelineRunStatus.SUCCESS: + raise RuntimeError(result.get("message") or "Inference failed") + + payload = result.get("raw") + if not isinstance(payload, dict): + raise RuntimeError("Structured pipeline entrypoint must return a JSON object") + try: + json.dumps(payload, allow_nan=False) + except (TypeError, ValueError) as error: + raise RuntimeError("Structured pipeline result must contain finite JSON-serializable values") from error + + return StructuredTaskExecutionResponse( + task_id=message.task_id, + result=payload, + peak_memory_mb=result.get("peak_memory_mb"), + inference_time_s=result.get("inference_time_s"), + ) + + # Media handlers _image_handler = ImageHandler() _video_handler = VideoHandler() diff --git a/telefuser/service_types.py b/telefuser/service_types.py index f2c98cc6..7dc13dbb 100644 --- a/telefuser/service_types.py +++ b/telefuser/service_types.py @@ -17,7 +17,7 @@ def values(cls) -> list[str]: class TaskType(_StringEnum): - """Supported media generation task types.""" + """Supported inference task types.""" T2V = "t2v" I2V = "i2v" @@ -27,6 +27,7 @@ class TaskType(_StringEnum): I2I = "i2i" S2V = "s2v" VSR = "vsr" + VLA_ACTION = "vla_action" class AspectRatio(_StringEnum): @@ -70,10 +71,11 @@ class StopTaskStatus(_StringEnum): class MediaType(_StringEnum): - """Generated media type.""" + """Inference result type.""" IMAGE = "image" VIDEO = "video" + STRUCTURED = "structured" class PipelineRunStatus(_StringEnum): diff --git a/tests/unit/models/test_lingbot_vla_v2.py b/tests/unit/models/test_lingbot_vla_v2.py new file mode 100644 index 00000000..daa829da --- /dev/null +++ b/tests/unit/models/test_lingbot_vla_v2.py @@ -0,0 +1,70 @@ +from types import SimpleNamespace + +import pytest +import torch + +from telefuser.models import lingbot_vla_v2_loader +from telefuser.models.lingbot_vla_v2 import LingbotVlaV2Policy, QwenvlWithExpertV2Model + + +class _Visual: + spatial_merge_size = 1 + + def __init__(self) -> None: + self.preprocess_calls = 0 + + def preprcess_grid_thw(self, grid_thw: torch.Tensor): + self.preprocess_calls += 1 + token_count = int(grid_thw.prod(dim=-1).sum()) + position_embeddings = (torch.zeros(token_count, 2), torch.ones(token_count, 2)) + cu_seqlens = torch.tensor([0, token_count], dtype=torch.int32) + split_sizes = grid_thw.prod(dim=-1).tolist() + return None, position_embeddings, cu_seqlens, split_sizes, token_count + + def __call__(self, pixel_values: torch.Tensor, **kwargs): + del kwargs + embeddings = torch.zeros(pixel_values.shape[0], 3) + return embeddings, [embeddings.clone()] + + +def _model(visual: _Visual) -> SimpleNamespace: + return SimpleNamespace( + config=SimpleNamespace(precompute_grid_thw=True), + qwenvl=SimpleNamespace(visual=visual), + pos_embeds=None, + position_embeddings=None, + cu_seqlens=None, + visual_split_sizes=None, + visual_max_seqlen=None, + _cached_image_grid_signature=None, + ) + + +def test_image_grid_cache_is_reused_and_invalidated_by_grid_shape() -> None: + visual = _Visual() + model = _model(visual) + first_grid = torch.tensor([[1, 2, 2], [1, 2, 2]]) + second_grid = torch.tensor([[1, 1, 2], [1, 1, 2]]) + + first = QwenvlWithExpertV2Model.get_image_features(model, torch.zeros(8, 6), first_grid) + repeated = QwenvlWithExpertV2Model.get_image_features(model, torch.zeros(8, 6), first_grid.clone()) + changed = QwenvlWithExpertV2Model.get_image_features(model, torch.zeros(4, 6), second_grid) + + assert visual.preprocess_calls == 2 + assert first[0].shape == repeated[0].shape == (2, 4, 3) + assert changed[0].shape == (2, 2, 3) + + +def test_policy_rejects_training_entrypoints() -> None: + with pytest.raises(RuntimeError, match="inference-only"): + LingbotVlaV2Policy.forward(None) + with pytest.raises(ValueError, match="only supports inference mode"): + LingbotVlaV2Policy.__init__(None, SimpleNamespace(), eval=False) + + assert "get_optim_params" not in LingbotVlaV2Policy.__dict__ + assert "get_parallel_plan" not in LingbotVlaV2Policy.__dict__ + + +def test_loader_does_not_expose_training_loss_helpers() -> None: + assert not hasattr(lingbot_vla_v2_loader, "triton_sequence_wise_balance_loss") + assert not hasattr(lingbot_vla_v2_loader, "triton_load_balancing_loss_func") diff --git a/tests/unit/models/test_lingbot_vla_v2_loader.py b/tests/unit/models/test_lingbot_vla_v2_loader.py new file mode 100644 index 00000000..b0e01136 --- /dev/null +++ b/tests/unit/models/test_lingbot_vla_v2_loader.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import json +import sys +from types import SimpleNamespace + +import pytest +import torch + +from telefuser.models.lingbot_vla_v2_loader import ( + build_official_6b_config, + load_lingbot_vla_v2, + resolve_lingbot_vla_v2_shards, + validate_official_6b_checkpoint, +) + + +def test_resolve_lingbot_vla_v2_shards_uses_index_manifest(tmp_path) -> None: + shard_names = ["model-00002-of-00002.safetensors", "model-00001-of-00002.safetensors"] + for name in shard_names: + (tmp_path / name).write_bytes(b"") + index = { + "weight_map": { + "layer.0": shard_names[0], + "layer.1": shard_names[1], + "layer.2": shard_names[0], + } + } + (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index), encoding="utf-8") + + resolved = resolve_lingbot_vla_v2_shards(tmp_path) + + assert resolved == [str(tmp_path / name) for name in sorted(shard_names)] + + +def test_resolve_lingbot_vla_v2_shards_rejects_missing_files(tmp_path) -> None: + index = {"weight_map": {"layer.0": "missing.safetensors"}} + (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index), encoding="utf-8") + + with pytest.raises(FileNotFoundError, match="checkpoint shards"): + resolve_lingbot_vla_v2_shards(tmp_path) + + +def test_validate_official_6b_checkpoint_accepts_expected_gate_shapes() -> None: + prefix = "model.qwenvl_with_expert.qwen_expert.model.layers" + state_dict = { + f"{prefix}.0.mlp.experts.gate_proj": SimpleNamespace(shape=(32, 512, 768)), + f"{prefix}.35.mlp.experts.gate_proj": SimpleNamespace(shape=(32, 512, 768)), + } + + validate_official_6b_checkpoint(state_dict) + + +def test_validate_official_6b_checkpoint_rejects_wrong_shape() -> None: + prefix = "model.qwenvl_with_expert.qwen_expert.model.layers" + state_dict = { + f"{prefix}.0.mlp.experts.gate_proj": SimpleNamespace(shape=(1, 2, 3)), + f"{prefix}.35.mlp.experts.gate_proj": SimpleNamespace(shape=(32, 512, 768)), + } + + with pytest.raises(ValueError, match="Unexpected shape"): + validate_official_6b_checkpoint(state_dict) + + +def test_build_official_6b_config_rejects_non_base_variant(tmp_path) -> None: + with pytest.raises(ValueError, match="Unsupported LingBot-VLA v2 checkpoint variant"): + build_official_6b_config(tmp_path, checkpoint_variant="robotwin") + + +def test_public_loader_routes_official_shards_through_module_manager(tmp_path, monkeypatch) -> None: + shard_names = ["model-00002-of-00002.safetensors", "model-00001-of-00002.safetensors"] + for name in shard_names: + (tmp_path / name).write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"layer.0": shard_names[0], "layer.1": shard_names[1]}}), + encoding="utf-8", + ) + + fake_model_class = type("FakeLingBotVlaV2Model", (), {}) + monkeypatch.setitem( + sys.modules, + "telefuser.models.lingbot_vla_v2", + SimpleNamespace(LingBotVlaV2Model=fake_model_class), + ) + + class _RecordingManager: + def __init__(self) -> None: + self.load_kwargs = None + + def load_model(self, file_path, **kwargs) -> None: + self.load_kwargs = {"file_path": file_path, **kwargs} + + def fetch_module(self, name: str): + return SimpleNamespace(name=name) + + manager = _RecordingManager() + + loaded = load_lingbot_vla_v2( + manager, + tmp_path, + tmp_path / "qwen3vl", + torch_dtype=torch.bfloat16, + device="cpu", + ) + + assert loaded.name == "lingbot_vla_v2" + assert manager.load_kwargs == { + "file_path": [str(tmp_path / name) for name in sorted(shard_names)], + "device": "cpu", + "torch_dtype": torch.bfloat16, + "low_cpu_mem_usage": True, + "name": "lingbot_vla_v2", + "model_class": fake_model_class, + "model_resource": "official", + "converter_kwargs": { + "qwen3vl_path": str(tmp_path / "qwen3vl"), + "checkpoint_variant": "base", + "checkpoint_path": str(tmp_path), + }, + } diff --git a/tests/unit/pipelines/lingbot_vla_v2/__init__.py b/tests/unit/pipelines/lingbot_vla_v2/__init__.py new file mode 100644 index 00000000..c7d9de08 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/__init__.py @@ -0,0 +1 @@ +"""LingBot-VLA v2 pipeline unit tests.""" diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_data.py b/tests/unit/pipelines/lingbot_vla_v2/test_data.py new file mode 100644 index 00000000..3459331a --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_data.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from telefuser.pipelines.lingbot_vla_v2.data import LingBotVlaV2InputProcessor, LingBotVlaV2Observation +from telefuser.pipelines.lingbot_vla_v2.robot_profile import ROBOTWIN_CAMERA_KEYS, RobotWinProfile + + +class _ImageProcessor: + def __init__(self) -> None: + self.values: list[float] = [] + + def __call__(self, image: torch.Tensor) -> dict[str, torch.Tensor]: + assert image.dtype == torch.float32 + assert image.shape == (3, 8, 8) + value = float(image[0, 0, 0]) + self.values.append(value) + return { + "pixel_values": torch.full((4, 6), float(value)), + "image_grid_thw": torch.tensor([[1, 4, 4]]), + } + + +class _Tokenizer: + def __init__(self) -> None: + self.rendered_task: str | None = None + self.padding_side: str | None = None + + def apply_chat_template(self, messages, *, tokenize: bool, add_generation_prompt: bool) -> str: + assert tokenize is False + assert add_generation_prompt is False + self.rendered_task = messages[0]["content"] + return f"chat:{self.rendered_task}" + + def __call__(self, prompts, **kwargs) -> dict[str, torch.Tensor]: + assert prompts == [f"chat:{self.rendered_task}"] + self.padding_side = kwargs["padding_side"] + length = kwargs["max_length"] + return { + "input_ids": torch.arange(length).unsqueeze(0), + "attention_mask": torch.ones(1, length), + } + + +def _processor() -> tuple[LingBotVlaV2InputProcessor, _ImageProcessor, _Tokenizer]: + image_processor = _ImageProcessor() + tokenizer = _Tokenizer() + processor = SimpleNamespace(image_processor=image_processor, tokenizer=tokenizer) + config = SimpleNamespace(max_state_dim=55, tokenizer_max_length=6) + return ( + LingBotVlaV2InputProcessor(processor, config, RobotWinProfile.default(), image_size=8), + image_processor, + tokenizer, + ) + + +def _observation() -> LingBotVlaV2Observation: + images = { + ROBOTWIN_CAMERA_KEYS[0]: np.full((8, 8, 3), 10, dtype=np.uint8), + ROBOTWIN_CAMERA_KEYS[1]: np.full((8, 8, 3), 20, dtype=np.uint8), + ROBOTWIN_CAMERA_KEYS[2]: np.full((8, 8, 3), 30, dtype=np.uint8), + } + return LingBotVlaV2Observation(task="pick up the block", state=[0.0] * 14, images=images) + + +def test_prepare_preserves_robotwin_camera_order_and_tensor_contract() -> None: + processor, image_processor, tokenizer = _processor() + observation = _observation() + + inputs = processor.prepare(observation) + + assert image_processor.values == pytest.approx([10.0, 20.0, 30.0], abs=1e-5) + assert tokenizer.rendered_task == observation.task + assert tokenizer.padding_side == "right" + assert inputs.images.shape == (1, 3, 4, 6) + assert inputs.img_masks.tolist() == [[True, True, True]] + assert inputs.image_grid_thw.shape == (1, 3, 3) + assert inputs.lang_tokens.shape == (1, 6) + assert inputs.lang_masks.dtype == torch.bool + assert torch.equal(inputs.state, processor.robot_profile.normalize_state(observation.state).unsqueeze(0)) + + +def test_prepare_rejects_a_missing_robotwin_camera() -> None: + processor, _, _ = _processor() + observation = _observation() + images = dict(observation.images) + del images[ROBOTWIN_CAMERA_KEYS[1]] + + with pytest.raises(ValueError, match="missing camera keys"): + processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) + + +def test_prepare_scales_unit_float_images_to_uint8() -> None: + processor, image_processor, _ = _processor() + observation = _observation() + images = dict(observation.images) + images[ROBOTWIN_CAMERA_KEYS[0]] = np.full((3, 8, 8), 0.5, dtype=np.float32) + + processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) + + assert image_processor.values[0] == 128 + + +def test_prepare_resizes_each_camera_before_qwen_processing() -> None: + processor, image_processor, _ = _processor() + observation = _observation() + images = { + key: np.full((12, 16, 3), value, dtype=np.uint8) + for key, value in zip(ROBOTWIN_CAMERA_KEYS, (10, 20, 30), strict=True) + } + + processor.prepare(LingBotVlaV2Observation(observation.task, observation.state, images)) + + assert image_processor.values == pytest.approx([10.0, 20.0, 30.0], abs=1e-5) diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py new file mode 100644 index 00000000..8889b11a --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_pipeline.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import torch +from torch import nn + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_vla_v2 import ( + LingBotVlaV2Observation, + LingBotVlaV2Pipeline, + LingBotVlaV2PipelineConfig, +) +from telefuser.pipelines.lingbot_vla_v2.robot_profile import ROBOTWIN_CAMERA_KEYS + + +class _ImageProcessor: + def __call__(self, image: torch.Tensor) -> dict[str, torch.Tensor]: + assert image.shape == (3, 8, 8) + return { + "pixel_values": torch.zeros(4, 6), + "image_grid_thw": torch.tensor([[1, 4, 4]]), + } + + +class _Tokenizer: + def apply_chat_template(self, messages, **kwargs) -> str: + return messages[0]["content"] + + def __call__(self, prompts, **kwargs) -> dict[str, torch.Tensor]: + length = kwargs["max_length"] + return { + "input_ids": torch.zeros(1, length, dtype=torch.long), + "attention_mask": torch.ones(1, length, dtype=torch.long), + } + + +class _Policy(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + max_state_dim=55, + max_action_dim=55, + n_action_steps=4, + tokenizer_max_length=6, + checkpoint_variant="base", + policy_verified=False, + verification_status="unverified_official_6b_base", + ) + self.sample_count = 0 + + def sample_actions(self, **inputs) -> torch.Tensor: + self.sample_count += 1 + assert inputs["state"].shape == (1, 55) + assert inputs["images"].shape == (1, 3, 4, 6) + return torch.zeros(1, self.config.n_action_steps, self.config.max_action_dim, device=self.anchor.device) + + +def test_pipeline_returns_normalized_canonical_action_chunk() -> None: + policy = _Policy() + processor = SimpleNamespace(image_processor=_ImageProcessor(), tokenizer=_Tokenizer()) + manager = ModuleManager(torch_dtype=torch.float32, device="cpu") + manager.add_module(policy, "lingbot_vla_v2") + manager.add_module(processor, "lingbot_vla_v2_processor") + pipeline = LingBotVlaV2Pipeline(device="cpu", torch_dtype=torch.float32) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + image_size=8, + ), + ) + observation = LingBotVlaV2Observation( + task="pick up the block", + state=[0.0] * 14, + images={key: np.zeros((8, 8, 3), dtype=np.uint8) for key in ROBOTWIN_CAMERA_KEYS}, + ) + + try: + chunk = pipeline(observation, seed=7) + finally: + pipeline.close() + + assert chunk.horizon == 4 + assert chunk.action_dim == 55 + assert chunk.canonical_normalized_actions.shape == (4, 55) + assert chunk.checkpoint_variant == "base" + assert chunk.policy_verified is False + assert chunk.verification_status == "unverified_official_6b_base" + + +def test_pipeline_prepares_resident_policy_and_disables_per_call_cache_eviction() -> None: + policy = _Policy() + processor = SimpleNamespace(image_processor=_ImageProcessor(), tokenizer=_Tokenizer()) + manager = ModuleManager(torch_dtype=torch.float32, device="cpu") + manager.add_module(policy, "lingbot_vla_v2") + manager.add_module(processor, "lingbot_vla_v2_processor") + pipeline = LingBotVlaV2Pipeline(device="cpu", torch_dtype=torch.float32) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + image_size=8, + ), + ) + + assert pipeline.clear_memory_after_call is False + assert pipeline.policy_stage.onload_models_flag is False + + pipeline.prepare_for_inference() + assert pipeline.policy_stage.onload_models_flag is True + + pipeline.warmup() + assert policy.sample_count == 1 + + pipeline.close() + assert pipeline.policy_stage.onload_models_flag is False diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py new file mode 100644 index 00000000..f2879e53 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import pytest +import torch + +from telefuser.pipelines.lingbot_vla_v2.robot_profile import RobotWinProfile + + +def _stats() -> dict[str, dict[str, list[float]]]: + return { + "observation.state.arm.position": {"q01": [0.0] * 12, "q99": [2.0] * 12}, + "observation.state.effector.position": {"q01": [-1.0] * 2, "q99": [1.0] * 2}, + "action.arm.position": {"q01": [0.0] * 12, "q99": [2.0] * 12}, + "action.effector.position": {"q01": [-1.0] * 2, "q99": [1.0] * 2}, + } + + +def test_normalize_state_uses_robotwin_joint_order() -> None: + profile = RobotWinProfile(_stats()) + state = torch.arange(14, dtype=torch.float32) / 10.0 + + canonical = profile.normalize_state(state) + + arm = torch.cat((state[0:6], state[7:13])) + effector = state[[6, 13]] + assert canonical.shape == (55,) + expected_arm = (arm.to(torch.float64) / (2.0 + 1e-6) * 2.0 - 1.0).to(torch.float32) + expected_effector = ((effector.to(torch.float64) + 1.0) / (2.0 + 1e-6) * 2.0 - 1.0).to(torch.float32) + assert torch.equal(canonical[0:12], expected_arm) + assert torch.equal(canonical[28:30], expected_effector) + assert torch.count_nonzero(canonical[12:28]) == 0 + assert torch.count_nonzero(canonical[30:]) == 0 + + +def test_structure_actions_reconstructs_raw_robotwin_layout() -> None: + profile = RobotWinProfile(_stats()) + canonical = torch.zeros(1, 3, 55) + + chunk = profile.structure_actions(canonical, include_canonical=True) + + arm = chunk.fields["action.arm.position"] + effector = chunk.fields["action.effector.position"] + assert arm.shape == (3, 12) + assert effector.shape == (3, 2) + assert torch.allclose(arm, torch.full_like(arm, 1.0000005)) + assert torch.allclose(effector, torch.zeros_like(effector), atol=1e-6) + assert torch.equal(chunk.raw_actions[:, 0:6], arm[:, 0:6]) + assert torch.equal(chunk.raw_actions[:, 6], effector[:, 0]) + assert torch.equal(chunk.raw_actions[:, 7:13], arm[:, 6:12]) + assert torch.equal(chunk.raw_actions[:, 13], effector[:, 1]) + assert chunk.horizon == 3 + assert chunk.canonical_normalized_actions is not None + + +def test_action_chunk_is_marked_unverified() -> None: + chunk = RobotWinProfile(_stats()).structure_actions(torch.zeros(2, 55)) + + assert chunk.policy_verified is False + assert chunk.verification_status == "unverified_official_6b_base" + assert chunk.robot_profile == "robotwin" + assert chunk.action_mask.shape == (55,) + assert chunk.action_mask.nonzero().flatten().tolist() == list(range(12)) + [28, 29] + assert chunk.canonical_normalized_actions is None + + +def test_default_profile_loads_bundled_upstream_stats() -> None: + profile = RobotWinProfile.default() + + canonical = profile.normalize_state(torch.zeros(14)) + chunk = profile.structure_actions(torch.zeros(1, 55)) + + assert canonical.shape == (55,) + assert torch.isfinite(canonical).all() + assert chunk.raw_actions.shape == (1, 14) + assert torch.isfinite(chunk.raw_actions).all() + + +def test_profile_rejects_invalid_state_and_action_shapes() -> None: + profile = RobotWinProfile(_stats()) + + with pytest.raises(ValueError, match="state must have shape"): + profile.normalize_state(torch.zeros(13)) + with pytest.raises(ValueError, match="canonical actions must have shape"): + profile.structure_actions(torch.zeros(2, 54)) diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_service.py b/tests/unit/pipelines/lingbot_vla_v2/test_service.py new file mode 100644 index 00000000..a79c33d4 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_service.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import base64 +import io +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import torch +from PIL import Image +from fastapi.testclient import TestClient + +from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2CanonicalActionChunk +from telefuser.pipelines.lingbot_vla_v2.robot_profile import ROBOTWIN_CAMERA_KEYS +from telefuser.pipelines.lingbot_vla_v2.service import ( + LingBotVlaV2ActionRequest, + LingBotVlaV2Service, + LingBotVlaV2ServiceConfig, + create_lingbot_vla_v2_app, +) + + +def _encoded_image(*, data_url: bool = False) -> str: + buffer = io.BytesIO() + Image.new("RGB", (8, 8), color=(10, 20, 30)).save(buffer, format="PNG") + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + return f"data:image/png;base64,{encoded}" if data_url else encoded + + +def _payload() -> dict: + image = _encoded_image() + return { + "task": "pick up the red block", + "state": [0.0] * 14, + "camera_high": image, + "camera_left_wrist": image, + "camera_right_wrist": image, + "seed": 7, + } + + +class _Pipeline: + def __init__(self, *, delay: float = 0.0) -> None: + self.delay = delay + self.closed = False + self.observations = [] + self.seeds = [] + self.active = 0 + self.max_active = 0 + self._counter_lock = threading.Lock() + + def __call__(self, observation, seed=None) -> LingBotVlaV2CanonicalActionChunk: + with self._counter_lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + try: + time.sleep(self.delay) + self.observations.append(observation) + self.seeds.append(seed) + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=torch.zeros(2, 55), + horizon=2, + action_dim=55, + ) + finally: + with self._counter_lock: + self.active -= 1 + + def close(self) -> None: + self.closed = True + + +def _config(**kwargs) -> LingBotVlaV2ServiceConfig: + return LingBotVlaV2ServiceConfig( + model_root="/models/lingbot-vla-v2-6b", + qwen3vl_root="/models/Qwen3-VL-4B-Instruct", + **kwargs, + ) + + +def test_app_serves_health_and_normalized_action_contract() -> None: + pipeline = _Pipeline() + config = _config(device="cuda:3") + app = create_lingbot_vla_v2_app(config, pipeline_factory=lambda received: pipeline) + + with TestClient(app) as client: + health = client.get("/health") + response = client.post("/v1/vla/actions", json=_payload()) + + assert health.status_code == 200 + assert health.json() == { + "status": "ready", + "model": "lingbot-vla-v2-6b-base", + "device": "cuda:3", + "policy_verified": False, + } + assert response.status_code == 200 + body = response.json() + assert body["horizon"] == 2 + assert body["action_dim"] == 55 + assert body["checkpoint_variant"] == "base" + assert body["policy_verified"] is False + assert body["verification_status"] == "unverified_official_6b_base" + assert len(body["canonical_normalized_actions"]) == 2 + assert len(body["canonical_normalized_actions"][0]) == 55 + assert pipeline.seeds == [7] + assert tuple(pipeline.observations[0].images) == ROBOTWIN_CAMERA_KEYS + assert all(image.mode == "RGB" for image in pipeline.observations[0].images.values()) + assert pipeline.closed is True + + +def test_app_accepts_image_data_urls() -> None: + pipeline = _Pipeline() + payload = _payload() + payload["camera_high"] = _encoded_image(data_url=True) + app = create_lingbot_vla_v2_app(_config(), pipeline_factory=lambda received: pipeline) + + with TestClient(app) as client: + response = client.post("/v1/vla/actions", json=payload) + + assert response.status_code == 200 + + +def test_app_rejects_invalid_observations_without_running_policy() -> None: + pipeline = _Pipeline() + app = create_lingbot_vla_v2_app(_config(), pipeline_factory=lambda received: pipeline) + payload = _payload() + payload["camera_high"] = "not-base64" + + with TestClient(app) as client: + invalid_image = client.post("/v1/vla/actions", json=payload) + invalid_state = client.post("/v1/vla/actions", json={**_payload(), "state": [0.0] * 13}) + extra_field = client.post("/v1/vla/actions", json={**_payload(), "output_path": "/tmp/action"}) + + assert invalid_image.status_code == 422 + assert invalid_image.json()["detail"] == "image must be valid base64" + assert invalid_state.status_code == 422 + assert extra_field.status_code == 422 + assert pipeline.observations == [] + + +def test_service_serializes_policy_calls() -> None: + pipeline = _Pipeline(delay=0.02) + service = LingBotVlaV2Service(pipeline, _config()) + request = LingBotVlaV2ActionRequest.model_validate(_payload()) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(service.predict, request) for _ in range(2)] + responses = [future.result() for future in futures] + + assert [response.horizon for response in responses] == [2, 2] + assert pipeline.max_active == 1 + + +def test_service_config_rejects_non_positive_image_limit() -> None: + try: + _config(max_image_bytes=0) + except ValueError as error: + assert str(error) == "max_image_bytes must be positive" + else: + raise AssertionError("expected an invalid image size limit to be rejected") diff --git a/tests/unit/service/test_pipeline_pool.py b/tests/unit/service/test_pipeline_pool.py index ff34e06e..e9c7f1cb 100644 --- a/tests/unit/service/test_pipeline_pool.py +++ b/tests/unit/service/test_pipeline_pool.py @@ -16,11 +16,102 @@ from telefuser.platforms import current_platform from telefuser.service.api.schema import TaskRequest from telefuser.service.core.config import ServerConfig +from telefuser.service.core.pipeline_pool import PipelinePool +from telefuser.service.core.replica_worker import ReplicaDeadError, ReplicaHandle, _forward_cancel_fn from telefuser.service.core.task_manager import TaskManager, TaskStatus _DEVICE_ENV_VAR = current_platform.device_control_env_var +def test_cancel_forwarder_wakes_immediately_on_normal_completion() -> None: + cancel_event = threading.Event() + stop_event = threading.Event() + forwarder_exit = threading.Event() + forwarder_wake = threading.Event() + forwarder_done = threading.Event() + forwarder = threading.Thread( + target=_forward_cancel_fn, + args=(cancel_event, stop_event, forwarder_exit, forwarder_wake, forwarder_done), + daemon=True, + ) + forwarder.start() + + forwarder_exit.set() + forwarder_wake.set() + + assert forwarder_done.wait(0.2) + assert not cancel_event.is_set() + assert not stop_event.is_set() + + +def test_cancel_forwarder_preserves_request_cancellation() -> None: + cancel_event = threading.Event() + stop_event = threading.Event() + stop_event.set() + forwarder_done = threading.Event() + + _forward_cancel_fn( + cancel_event, + stop_event, + threading.Event(), + threading.Event(), + forwarder_done, + ) + + assert cancel_event.is_set() + assert forwarder_done.is_set() + + +def test_replica_handle_converts_broken_pipe_to_dead_replica() -> None: + process = MagicMock() + process.is_alive.return_value = True + connection = MagicMock() + connection.send.side_effect = BrokenPipeError("closed") + handle = ReplicaHandle( + replica_id=0, + process=process, + conn=connection, + cancel_event=threading.Event(), + metadata={}, + ) + + with pytest.raises(ReplicaDeadError, match="IPC failed"): + asyncio.run(handle.run_task({}, threading.Event(), None, None)) + + assert handle._dead is True + + +def test_pipeline_pool_evicts_exited_replica_and_uses_remaining_capacity() -> None: + task_manager = MagicMock() + pool = PipelinePool( + num_replicas=2, + replica_device_ids=[["0"], ["1"]], + security_level_name="NONE", + task_manager=task_manager, + ) + dead_handle = MagicMock() + dead_handle._dead = False + dead_handle.process.is_alive.return_value = False + live_handle = MagicMock() + live_handle._dead = False + live_handle.process.is_alive.return_value = True + pool._handles = [dead_handle, live_handle] + pool._instance_status = ["idle", "idle"] + pool._available.put_nowait(0) + pool._available.put_nowait(1) + + async def scenario() -> None: + async with pool.acquire() as handle: + assert handle is live_handle + + asyncio.run(scenario()) + + assert pool._live_count == 1 + assert pool._instance_status == ["dead", "idle"] + dead_handle.shutdown.assert_called_once() + task_manager.set_max_concurrent_processing.assert_called_once_with(1) + + # ============================================================================ # Test 1: TaskManager atomic claim — concurrent claim, single winner # ============================================================================ diff --git a/tests/unit/service/test_service_smoke.py b/tests/unit/service/test_service_smoke.py index 18f635d6..661f545f 100644 --- a/tests/unit/service/test_service_smoke.py +++ b/tests/unit/service/test_service_smoke.py @@ -88,3 +88,13 @@ def test_openai_video_retrieve_includes_artifact_metadata_smoke(tmp_path: Path) assert data["artifact_id"] == f"local:tasks/{video_id}/outputs/videos/clip.mp4" assert data["artifact_metadata"]["backend"] == "local" assert data["artifact_metadata"]["size_bytes"] == 5 + + +def test_running_task_processor_is_notified_for_new_work(tmp_path: Path) -> None: + server = _make_smoke_server(tmp_path) + server.task_processor = Mock() + server.task_processor.is_running = True + + asyncio.run(server.ensure_task_processor_running()) + + server.task_processor.notify_task_available.assert_called_once_with() diff --git a/tests/unit/service/test_structured_tasks.py b/tests/unit/service/test_structured_tasks.py new file mode 100644 index 00000000..41c7adb4 --- /dev/null +++ b/tests/unit/service/test_structured_tasks.py @@ -0,0 +1,240 @@ +from __future__ import annotations + +import asyncio +import base64 +import io +import threading +import time +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +import torch +from PIL import Image +from fastapi.testclient import TestClient + +from examples.lingbot_vla_v2 import lingbot_vla_v2_native_service +from telefuser.client import TFClient, TaskFailedError +from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2CanonicalActionChunk +from telefuser.service.api.api_server import ApiServer +from telefuser.service.api.schema import StructuredTaskRequest +from telefuser.service.core.task_manager import TaskManager +from telefuser.service.core.task_service import StructuredInferenceService + + +def _encoded_image() -> str: + buffer = io.BytesIO() + Image.new("RGB", (8, 8), color=(10, 20, 30)).save(buffer, format="PNG") + return base64.b64encode(buffer.getvalue()).decode("ascii") + + +def _payload() -> dict: + image = _encoded_image() + return { + "task": "vla_action", + "instruction": "pick up the red block", + "state": [0.0] * 14, + "camera_high": image, + "camera_left_wrist": image, + "camera_right_wrist": image, + "seed": 7, + } + + +class _StructuredPipelineService: + is_running = True + + def supported_tasks(self) -> tuple[str, ...]: + return ("vla_action",) + + def get_task_contract(self, task: str) -> dict: + assert task == "vla_action" + return lingbot_vla_v2_native_service.PIPELINE_CONTRACT["task_contracts"][task] + + async def run_task_with_stop_event(self, task_data, stop_event, **kwargs) -> dict: + assert task_data["instruction"] == "pick up the red block" + assert "output_path" not in task_data + return { + "status": "success", + "raw": { + "canonical_normalized_actions": [[0.0] * 55 for _ in range(2)], + "horizon": 2, + "action_dim": 55, + "checkpoint_variant": "base", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + }, + "peak_memory_mb": 128.0, + "inference_time_s": 0.25, + } + + +def test_structured_route_uses_scheduler_and_exposes_result_metrics(tmp_path: Path) -> None: + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.initialize_services(tmp_path, _StructuredPipelineService()) + + with TestClient(server.get_app()) as client: + created = client.post("/v1/tasks/structured", json=_payload()) + + assert created.status_code == 200 + created_body = created.json() + assert created_body["task_status"] == "pending" + assert "output_path" not in created_body + + deadline = time.monotonic() + 2.0 + while True: + status = client.get(f"/v1/tasks/{created_body['task_id']}/status") + assert status.status_code == 200 + body = status.json() + if body["status"] == "completed": + break + assert time.monotonic() < deadline + time.sleep(0.01) + + assert body["output_path"] is None + assert body["media_type"] == "structured" + assert body["peak_memory_mb"] == 128.0 + assert body["inference_time_s"] == 0.25 + assert body["result"]["horizon"] == 2 + assert len(body["result"]["canonical_normalized_actions"][0]) == 55 + assert "camera_high" not in body + + +def test_structured_route_rejects_media_contract(tmp_path: Path) -> None: + inference_service = _StructuredPipelineService() + inference_service.supported_tasks = lambda: ("t2i",) + inference_service.get_task_contract = lambda task: {"media_type": "image", "parameters": {}} + server = ApiServer(task_manager=TaskManager(), enable_openai_api=False) + server.initialize_services(tmp_path, inference_service) + + with TestClient(server.get_app()) as client: + response = client.post("/v1/tasks/structured", json={"task": "t2i"}) + + assert response.status_code == 400 + assert "does not declare a structured result contract" in response.json()["detail"] + + +def test_structured_service_rejects_non_json_pipeline_result() -> None: + class InvalidService: + async def run_task_with_stop_event(self, task_data, stop_event): + return {"status": "success", "raw": {"value": torch.zeros(1)}} + + service = StructuredInferenceService(InvalidService()) + request = StructuredTaskRequest(task="vla_action") + + with pytest.raises(RuntimeError, match="finite JSON-serializable"): + asyncio.run(service.execute_with_stop_event(request, threading.Event())) + + +def test_native_vla_entrypoint_returns_action_contract() -> None: + class Pipeline: + def __call__(self, observation, seed=None): + assert observation.task == "pick up the red block" + assert seed == 7 + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=torch.zeros(2, 55), + horizon=2, + action_dim=55, + ) + + result = lingbot_vla_v2_native_service.run_structured(Pipeline(), **_payload()) + + assert result["horizon"] == 2 + assert result["action_dim"] == 55 + assert len(result["canonical_normalized_actions"][0]) == 55 + + +def test_unified_client_encodes_vla_inputs_and_returns_result(tmp_path: Path) -> None: + image_path = tmp_path / "camera.png" + Image.new("RGB", (8, 8)).save(image_path) + client = TFClient("http://127.0.0.1:8000") + response = Mock() + response.raise_for_status.return_value = None + response.json.return_value = {"task_id": "task-1", "task_status": "pending"} + client._session.post = Mock(return_value=response) + client.wait_for_completion = Mock(return_value={"status": "completed", "result": {"horizon": 2}}) + + result = client.predict_vla_actions( + instruction="pick up the red block", + state=[0.0] * 14, + camera_high_path=str(image_path), + camera_left_wrist_path=str(image_path), + camera_right_wrist_path=str(image_path), + seed=7, + ) + + assert result == {"horizon": 2} + request = client._session.post.call_args + assert request.args[0].endswith("/v1/tasks/structured") + assert request.kwargs["json"]["task"] == "vla_action" + assert request.kwargs["json"]["camera_high"] == base64.b64encode(image_path.read_bytes()).decode("ascii") + + +def test_unified_client_rejects_missing_structured_result() -> None: + client = TFClient() + client.create_vla_action_task = Mock(return_value={"task_id": "task-1"}) + client.wait_for_completion = Mock(return_value={"status": "completed", "result": None}) + + with pytest.raises(TaskFailedError, match="without a structured result"): + client.predict_vla_actions( + instruction="pick", + state=[0.0] * 14, + camera_high_path="unused", + camera_left_wrist_path="unused", + camera_right_wrist_path="unused", + ) + + +def test_pipeline_pool_preserves_structured_result() -> None: + from telefuser.service.core.pipeline_pool import PipelinePool + + result = {"status": "success", "raw": {"horizon": 2}} + handle = SimpleNamespace( + _dead=False, + run_task=AsyncMock(return_value=result), + shutdown=Mock(), + ) + pool = PipelinePool( + num_replicas=1, + replica_device_ids=[["0"]], + security_level_name="NONE", + ) + pool._handles = [handle] + pool._instance_status = ["idle"] + pool._available.put_nowait(0) + + received = asyncio.run( + pool.run_task_with_stop_event( + {"task": "vla_action"}, + threading.Event(), + ) + ) + + assert received == result + assert pool._instance_status == ["idle"] + + +def test_pipeline_runner_closes_close_only_pipeline() -> None: + from telefuser.service.core.pipeline_runner import PipelineRunner + + class Pipeline: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + pipeline = Pipeline() + + def run_structured(pipeline, **kwargs): + return {"value": 1} + + async def scenario() -> None: + runner = PipelineRunner(pipeline=pipeline, run_with_file=run_structured) + result = await runner.run(task_data={"task": "vla_action"}) + assert result.raw == {"value": 1} + await runner.shutdown() + + asyncio.run(scenario()) + assert pipeline.closed is True diff --git a/tests/unit/service/test_task_runtime.py b/tests/unit/service/test_task_runtime.py index dd0bda2c..5c181d08 100644 --- a/tests/unit/service/test_task_runtime.py +++ b/tests/unit/service/test_task_runtime.py @@ -114,6 +114,28 @@ async def wait_for_cancelled_status() -> None: asyncio.run(scenario()) +def test_async_task_processor_wakes_when_task_becomes_available() -> None: + """A newly submitted task should not wait for the idle polling timeout.""" + + async def scenario() -> None: + task_manager = TaskManager(max_queue_size=10) + media_service = _ControlledMediaService() + media_service.finish.set() + processor = AsyncTaskProcessor(task_manager=task_manager, media_service=media_service, max_concurrent=1) + + await processor.start() + try: + await asyncio.sleep(0) + task_manager.create_task(TaskRequest(task="t2i")) + processor.notify_task_available() + + await asyncio.wait_for(media_service.started.wait(), timeout=0.5) + finally: + await processor.stop() + + asyncio.run(scenario()) + + def test_claim_next_pending_task_atomic_single_winner() -> None: """Two PENDING tasks, single slot: only one is claimed, the second claim returns None.""" task_manager = TaskManager(max_queue_size=10) diff --git a/tests/unit/test_example_registry.py b/tests/unit/test_example_registry.py index a56fb120..ff3ff87f 100644 --- a/tests/unit/test_example_registry.py +++ b/tests/unit/test_example_registry.py @@ -10,6 +10,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[2] EXAMPLES_ROOT = PROJECT_ROOT / "examples" SERVICE_PARITY_EXAMPLES = { + "lingbot_vla_v2/lingbot_vla_v2_native_service.py", "wan_video/wan21_14b_image_to_video_480p_service.py", "wan_video/wan22_14b_image_to_video_distill_h100.py", "lingbot_video/lingbot_video_dense_1_3b.py", @@ -38,6 +39,22 @@ def _declares_service_contract(path: Path) -> bool: return False +def _declared_run_entrypoint(path: Path) -> str: + tree = ast.parse(path.read_text(encoding="utf-8-sig"), filename=str(path)) + for node in tree.body: + targets = ( + node.targets if isinstance(node, ast.Assign) else [node.target] if isinstance(node, ast.AnnAssign) else [] + ) + if not any(isinstance(target, ast.Name) and target.id == "PIPELINE_CONTRACT" for target in targets): + continue + try: + contract = ast.literal_eval(node.value) + except (TypeError, ValueError): + break + return contract.get("entrypoints", {}).get("run_with_file", "run_with_file") + return "run_with_file" + + def test_example_regression_registry_has_runnable_entrypoints() -> None: config = load_config() @@ -78,6 +95,8 @@ def test_all_declared_service_examples_have_cpu_parity_coverage() -> None: assert declared_contract_examples == SERVICE_PARITY_EXAMPLES for script in declared_contract_examples: - symbols = _module_symbols(EXAMPLES_ROOT / script) + script_path = EXAMPLES_ROOT / script + symbols = _module_symbols(script_path) assert "get_pipeline" in symbols, f"{script} is missing get_pipeline()" - assert "run_with_file" in symbols, f"{script} is missing run_with_file()" + run_entrypoint = _declared_run_entrypoint(script_path) + assert run_entrypoint in symbols, f"{script} is missing {run_entrypoint}()" diff --git a/tests/unit/validation/test_lingbot_vla_v2_artifacts.py b/tests/unit/validation/test_lingbot_vla_v2_artifacts.py new file mode 100644 index 00000000..09865983 --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_artifacts.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tools.validation.capture_lingbot_vla_v2_telefuser import TensorCapture, trace_predict_velocity +from tools.validation.run_lingbot_vla_v2_parity import compare_artifacts + + +def _arrays() -> dict[str, np.ndarray]: + arrays = { + "images": np.zeros((1, 3, 2, 2), dtype=np.float32), + "img_masks": np.ones((1, 3), dtype=np.bool_), + "image_grid_thw": np.ones((1, 3, 3), dtype=np.int64), + "lang_tokens": np.arange(4, dtype=np.int64).reshape(1, 4), + "lang_masks": np.ones((1, 4), dtype=np.bool_), + "state": np.zeros((1, 55), dtype=np.float32), + "initial_noise": np.ones((1, 2, 55), dtype=np.float32), + "canonical_normalized_actions": np.zeros((2, 55), dtype=np.float32), + } + for step in range(2): + suffix = f"{step:02d}" + arrays[f"timestep_step_{suffix}"] = np.asarray([1.0 - 0.5 * step], dtype=np.float32) + arrays[f"x_t_step_{suffix}"] = np.full((1, 2, 55), step, dtype=np.float32) + arrays[f"velocity_step_{suffix}"] = np.full((1, 2, 55), step + 0.25, dtype=np.float32) + return arrays + + +def _metadata() -> dict[str, object]: + return { + "schema_version": 1, + "artifact_kind": "telefuser_regression", + "checkpoint_manifest_sha256": "checkpoint", + "processor_manifest_sha256": "processor", + "norm_stats_sha256": "norm-stats", + "input_sha256": "input", + "seed": 7, + "num_steps": 2, + "torch_dtype": "bfloat16", + "attention_backend": "eager", + "moe_backend": "deterministic_torch_reference", + } + + +def _write_artifact( + root: Path, + name: str, + arrays: dict[str, np.ndarray], + metadata: dict[str, object] | None = None, +) -> Path: + path = root / f"{name}.npz" + np.savez(path, **arrays) + payload = dict(_metadata() if metadata is None else metadata) + payload["arrays"] = { + key: { + "shape": list(array.shape), + "original_dtype": str(array.dtype), + "stored_dtype": str(array.dtype), + } + for key, array in arrays.items() + } + path.with_suffix(".json").write_text( + json.dumps(payload), + encoding="utf-8", + ) + return path + + +def test_compare_artifacts_accepts_a_complete_strict_replay(tmp_path: Path) -> None: + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", _arrays()) + + report = compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + assert report["passed"] is True + assert report["first_failed_step"] is None + assert len(report["results"]) == 14 + + +def test_compare_artifacts_requires_every_preprocessing_array(tmp_path: Path) -> None: + arrays = _arrays() + del arrays["image_grid_thw"] + reference = _write_artifact(tmp_path, "reference", arrays) + candidate = _write_artifact(tmp_path, "candidate", _arrays()) + + with pytest.raises(ValueError, match="image_grid_thw"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_compare_artifacts_requires_contiguous_sampling_steps(tmp_path: Path) -> None: + arrays = _arrays() + del arrays["velocity_step_01"] + reference = _write_artifact(tmp_path, "reference", arrays) + candidate = _write_artifact(tmp_path, "candidate", _arrays()) + + with pytest.raises(ValueError, match="velocity steps"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_compare_artifacts_reports_the_first_failed_step(tmp_path: Path) -> None: + candidate_arrays = _arrays() + candidate_arrays["velocity_step_01"][0, 0, 0] += 0.5 + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", candidate_arrays) + + report = compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + assert report["passed"] is False + assert report["first_failed_step"] == 1 + failed = [item for item in report["results"] if not item["passed"]] + assert [item["key"] for item in failed] == ["velocity_step_01"] + assert failed[0]["mismatch_count"] == 1 + + +def test_compare_artifacts_rejects_non_finite_values(tmp_path: Path) -> None: + candidate_arrays = _arrays() + candidate_arrays["x_t_step_00"][0, 0, 0] = np.nan + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", candidate_arrays) + + report = compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + assert report["passed"] is False + assert report["first_failed_step"] == 0 + + +def test_compare_artifacts_rejects_different_artifact_identity(tmp_path: Path) -> None: + candidate_metadata = _metadata() + candidate_metadata["input_sha256"] = "different" + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", _arrays(), candidate_metadata) + + with pytest.raises(ValueError, match="input_sha256"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_compare_artifacts_rejects_different_moe_backends(tmp_path: Path) -> None: + candidate_metadata = _metadata() + candidate_metadata["moe_backend"] = "upstream_triton" + reference = _write_artifact(tmp_path, "reference", _arrays()) + candidate = _write_artifact(tmp_path, "candidate", _arrays(), candidate_metadata) + + with pytest.raises(ValueError, match="moe_backend"): + compare_artifacts(reference, candidate, rtol=0.0, atol=0.0) + + +def test_velocity_trace_snapshots_inputs_and_restores_the_model_instance() -> None: + class _FlowModel: + _use_compile_predict_velocity = True + + def predict_velocity(self, state, prefix_masks, cache, x_t, timestep, **kwargs): + del state, prefix_masks, cache, timestep, kwargs + return torch.ones_like(x_t) + + flow_model = _FlowModel() + capture = TensorCapture() + x_t = torch.zeros(1, 2, 3) + + with trace_predict_velocity(flow_model, capture) as trace: + velocity = flow_model.predict_velocity(None, None, None, x_t, torch.ones(1)) + x_t.add_(velocity) + + assert trace.step == 1 + assert np.array_equal(capture.arrays["initial_noise"], np.zeros((1, 2, 3), dtype=np.float32)) + assert np.array_equal(capture.arrays["x_t_step_00"], np.zeros((1, 2, 3), dtype=np.float32)) + assert np.array_equal(capture.arrays["velocity_step_00"], np.ones((1, 2, 3), dtype=np.float32)) + assert "predict_velocity" not in vars(flow_model) + assert flow_model._use_compile_predict_velocity is True + + +def test_tensor_capture_records_original_bfloat16_dtype() -> None: + capture = TensorCapture() + + capture.add("value", torch.ones(2, dtype=torch.bfloat16)) + + assert capture.arrays["value"].dtype == np.float32 + assert capture.array_metadata["value"] == { + "shape": [2], + "original_dtype": "bfloat16", + "stored_dtype": "float32", + } diff --git a/tests/unit/validation/test_lingbot_vla_v2_benchmark.py b/tests/unit/validation/test_lingbot_vla_v2_benchmark.py new file mode 100644 index 00000000..1a42b45f --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_benchmark.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import argparse + +import pytest +from PIL import Image + +from tools.validation.benchmark_lingbot_vla_v2_service import ( + encode_image, + parse_image_sizes, + percentile, + summarize, +) + + +def test_parse_image_sizes_deduplicates_and_preserves_order() -> None: + assert parse_image_sizes("256x256, 640X480,256x256") == ((256, 256), (640, 480)) + + +@pytest.mark.parametrize("value", ["", "256", "0x256", "axb"]) +def test_parse_image_sizes_rejects_invalid_values(value: str) -> None: + with pytest.raises(argparse.ArgumentTypeError): + parse_image_sizes(value) + + +def test_latency_summary_reports_interpolated_percentiles_and_throughput() -> None: + values = [1.0, 2.0, 3.0, 4.0] + result = summarize(values) + + assert percentile(values, 0.5) == 2.5 + assert result["count"] == 4 + assert result["mean_seconds"] == 2.5 + assert result["p95_seconds"] == pytest.approx(3.85) + assert result["throughput_requests_per_second"] == 0.4 + + +def test_encode_image_reports_decoded_jpeg_size() -> None: + encoded, encoded_bytes = encode_image(Image.new("RGB", (8, 8)), (32, 24), quality=90) + + assert encoded + assert encoded_bytes > 0 + assert len(encoded) >= encoded_bytes diff --git a/tests/unit/validation/test_lingbot_vla_v2_service_faults.py b/tests/unit/validation/test_lingbot_vla_v2_service_faults.py new file mode 100644 index 00000000..cd05f32d --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_service_faults.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from tools.validation import validate_lingbot_vla_v2_service_faults as validator + + +class _Response: + def __init__(self, status_code: int, body: dict[str, Any]) -> None: + self.status_code = status_code + self._body = body + + def json(self) -> dict[str, Any]: + return self._body + + +class _Session: + def __init__(self, responses: list[_Response]) -> None: + self.responses = responses + + def request(self, *args: object, **kwargs: object) -> _Response: + return self.responses.pop(0) + + +def test_expect_invalid_request_accepts_synchronous_rejection() -> None: + result = validator._expect_rejected_or_failed( + _Session([_Response(422, {"detail": "invalid"})]), + "http://127.0.0.1:18080", + {"task": "vla_action"}, + case_name="invalid", + http_timeout_seconds=1.0, + task_timeout_seconds=1.0, + poll_interval_seconds=0.001, + ) + + assert result == {"name": "invalid", "passed": True, "handling": "rejected", "http_status": 422} + + +def test_expect_invalid_request_accepts_asynchronous_failure() -> None: + session = _Session( + [ + _Response(200, {"task_id": "task-1", "task_status": "pending"}), + _Response(200, {"task_id": "task-1", "status": "failed", "error": "bad input"}), + ] + ) + + result = validator._expect_rejected_or_failed( + session, + "http://127.0.0.1:18080", + {"task": "vla_action"}, + case_name="invalid", + http_timeout_seconds=1.0, + task_timeout_seconds=1.0, + poll_interval_seconds=0.001, + ) + + assert result["handling"] == "asynchronous_failure" + assert result["terminal_status"] == "failed" + assert result["error"] == "bad input" + + +def test_select_replica_process_filters_unrelated_and_root_processes() -> None: + rows = "\n".join( + [ + "100, GPU-a", + "101, GPU-a", + "102, GPU-b", + "999, GPU-a", + ] + ) + + selected = validator.select_replica_process( + rows, + service_process_ids={101, 102}, + gpu_uuid_to_index={"GPU-a": "0", "GPU-b": "1"}, + gpu_index="0", + service_pid=100, + ) + + assert selected == 101 + + +def test_select_replica_process_requires_unambiguous_target() -> None: + with pytest.raises(validator.FaultValidationFailure, match="exactly one"): + validator.select_replica_process( + "101, GPU-a\n102, GPU-a", + service_process_ids={101, 102}, + gpu_uuid_to_index={"GPU-a": "0"}, + gpu_index="0", + service_pid=100, + ) + + +def test_gpu_compute_process_ids_ignores_malformed_rows() -> None: + rows = "101, GPU-a\ninvalid, GPU-b\n102, GPU-c\n" + + assert validator.gpu_compute_process_ids(rows) == {101, 102} + + +def test_validate_pool_degradation_requires_one_dead_replica_and_reduced_capacity() -> None: + before = { + "effective_max_concurrent_tasks": 2, + "pool": [{"id": 0, "status": "idle"}, {"id": 1, "status": "idle"}], + } + after = { + "effective_max_concurrent_tasks": 1, + "pool": [{"id": 0, "status": "dead"}, {"id": 1, "status": "idle"}], + } + + result = validator.validate_pool_degradation(before, after) + + assert result["before_capacity"] == 2 + assert result["after_capacity"] == 1 + assert result["dead_replica_ids"] == [0] + assert result["recovery_semantics"] == "graceful_capacity_degradation_without_automatic_restart" + + +def test_validate_pool_degradation_rejects_unchanged_capacity() -> None: + before = { + "effective_max_concurrent_tasks": 2, + "pool": [{"id": 0, "status": "idle"}, {"id": 1, "status": "idle"}], + } + after = { + "effective_max_concurrent_tasks": 2, + "pool": [{"id": 0, "status": "dead"}, {"id": 1, "status": "idle"}], + } + + with pytest.raises(validator.FaultValidationFailure, match="capacity 1"): + validator.validate_pool_degradation(before, after) diff --git a/tests/unit/validation/test_lingbot_vla_v2_structured_service.py b/tests/unit/validation/test_lingbot_vla_v2_structured_service.py new file mode 100644 index 00000000..7b916373 --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_structured_service.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import argparse +import threading +import time +from typing import Any + +import pytest + +from tools.validation import validate_lingbot_vla_v2_structured_service as validator + + +def _action_result(value: float = 0.25) -> dict[str, Any]: + return { + "canonical_normalized_actions": [[value] * 55 for _ in range(50)], + "horizon": 50, + "action_dim": 55, + "checkpoint_variant": "base", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + } + + +def _metadata() -> dict[str, Any]: + parameters = { + "instruction": {"type": "string", "required": True}, + "state": {"type": "array", "required": True}, + "camera_high": {"type": "string", "required": True}, + "camera_left_wrist": {"type": "string", "required": True}, + "camera_right_wrist": {"type": "string", "required": True}, + "seed": {"type": "integer", "required": False}, + } + return { + "declared_pipeline_contract": True, + "supported_tasks": ["vla_action"], + "supported_media_types": ["structured"], + "task_contracts": { + "vla_action": { + "media_type": "structured", + "required_inputs": ["camera_high", "camera_left_wrist", "camera_right_wrist"], + "optional_inputs": [], + "parameters": parameters, + } + }, + } + + +def test_parse_state_json_requires_fourteen_finite_numbers() -> None: + assert validator.parse_state_json("[0,1,2,3,4,5,6,7,8,9,10,11,12,13]") == [float(index) for index in range(14)] + + with pytest.raises(argparse.ArgumentTypeError, match="exactly 14"): + validator.parse_state_json("[0, 1]") + with pytest.raises(argparse.ArgumentTypeError, match="finite numbers"): + validator.parse_state_json("[0,1,2,3,4,5,6,7,8,9,10,11,12,true]") + + +def test_validate_service_metadata_requires_native_structured_contract() -> None: + validator.validate_service_metadata(_metadata()) + + metadata = _metadata() + metadata["task_contracts"]["vla_action"]["media_type"] = "video" + with pytest.raises(validator.ValidationFailure, match="structured task contract"): + validator.validate_service_metadata(metadata) + + metadata = _metadata() + del metadata["task_contracts"]["vla_action"]["parameters"]["seed"] + with pytest.raises(validator.ValidationFailure, match="parameter fields changed"): + validator.validate_service_metadata(metadata) + + +def test_validate_action_result_reports_shape_stats_and_fingerprint() -> None: + summary = validator.validate_action_result(_action_result(), expected_horizon=50, expected_action_dim=55) + + assert summary["shape"] == [50, 55] + assert summary["value_count"] == 2750 + assert summary["minimum"] == 0.25 + assert summary["maximum"] == 0.25 + assert summary["policy_verified"] is False + assert len(summary["sha256_float64_le"]) == 64 + + +def test_validate_action_result_rejects_additive_result_fields() -> None: + result = _action_result() + result["debug"] = "unstable" + + with pytest.raises(validator.ValidationFailure, match="result fields changed"): + validator.validate_action_result(result, expected_horizon=50, expected_action_dim=55) + + +def test_validate_task_status_rejects_sensitive_echo_and_missing_fields() -> None: + status = { + "task_id": "task-1", + "status": "completed", + "inference_time_s": 0.25, + "peak_memory_mb": None, + "result": _action_result(), + } + validator.validate_task_status(status, task_id="task-1") + + leaked = dict(status, camera_high="base64") + with pytest.raises(validator.ValidationFailure, match="sensitive image"): + validator.validate_task_status(leaked, task_id="task-1") + + +def test_compare_windows_reports_first_and_last_measurement_change() -> None: + result = validator.compare_windows([1.0, 1.0, 1.0, 1.2, 1.2, 1.2]) + + assert result is not None + assert result["window_count"] == 1 + assert result["change_percent"] == pytest.approx(20.0) + + +def test_parse_gpu_process_memory_filters_process_tree_and_physical_gpu() -> None: + output = "\n".join( + [ + "101, GPU-a, 1024", + "102, GPU-a, 512", + "101, GPU-b, 2048", + "999, GPU-a, 4096", + "malformed", + ] + ) + + result = validator._parse_gpu_process_memory( + output, + process_ids={101, 102}, + uuid_to_index={"GPU-a": "0", "GPU-b": "1"}, + gpu_indexes={"0"}, + ) + + assert result == {"0": 1536.0} + + +def test_parse_gpu_indexes_requires_integer_indexes() -> None: + assert validator.parse_gpu_indexes("0, 2") == {"0", "2"} + + with pytest.raises(argparse.ArgumentTypeError, match="comma-separated"): + validator.parse_gpu_indexes("0,GPU-a") + + +@pytest.mark.parametrize( + "mutation,match", + [ + (lambda result: result.update(horizon=49), "horizon field"), + (lambda result: result["canonical_normalized_actions"][0].pop(), "row 0"), + (lambda result: result["canonical_normalized_actions"][0].__setitem__(0, float("nan")), "non-finite"), + ], +) +def test_validate_action_result_rejects_invalid_contract(mutation, match: str) -> None: + result = _action_result() + mutation(result) + + with pytest.raises(validator.ValidationFailure, match=match): + validator.validate_action_result(result, expected_horizon=50, expected_action_dim=55) + + +class _Response: + def __init__(self, body: dict[str, Any]) -> None: + self._body = body + self.status_code = 200 + self.text = "" + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, Any]: + return self._body + + +class _Session: + def __init__(self, state: dict[str, Any]) -> None: + self.state = state + self.trust_env = False + + def __enter__(self) -> "_Session": + return self + + def __exit__(self, *args: object) -> None: + return None + + def request(self, method: str, url: str, *, json=None, timeout=None) -> _Response: + if method == "POST": + assert url.endswith("/v1/tasks/structured") + assert json["task"] == "vla_action" + with self.state["lock"]: + self.state["next_id"] += 1 + task_id = f"task-{self.state['next_id']}" + return _Response({"task_id": task_id, "task_status": "pending"}) + task_id = url.rsplit("/", maxsplit=2)[-2] + return _Response( + { + "task_id": task_id, + "status": "completed", + "inference_time_s": 0.25, + "peak_memory_mb": 128.0, + "result": _action_result(), + } + ) + + +def test_run_workload_exercises_concurrent_structured_requests(monkeypatch: pytest.MonkeyPatch) -> None: + state = {"lock": threading.Lock(), "next_id": 0} + monkeypatch.setattr(validator, "_new_session", lambda: _Session(state)) + config = validator.RequestConfig( + base_url="http://127.0.0.1:18080", + payload={"task": "vla_action"}, + http_timeout_seconds=1.0, + task_timeout_seconds=1.0, + poll_interval_seconds=0.001, + expected_horizon=50, + expected_action_dim=55, + ) + + report = validator.run_workload( + config, + request_count=4, + duration_seconds=None, + concurrency=2, + max_records=10, + ) + + assert report["requests"]["total"] == 4 + assert report["requests"]["succeeded"] == 4 + assert report["requests"]["failed"] == 0 + assert report["requests"]["unique_task_ids"] == 4 + assert report["latency_seconds"]["target_inference"]["mean"] == 0.25 + assert len(report["retained_records"]["successful"]) == 4 + + +def test_resource_sampler_bounds_samples_and_reports_trend() -> None: + state = {"value": 100.0} + ready = threading.Event() + + def sample() -> dict[str, Any]: + state["value"] += 10.0 + ready.set() + return {"process_ids": [123], "cpu_rss_mib": state["value"], "gpu_memory_mib": {"0": 2048.0}} + + sampler = validator.ResourceSampler( + 123, + interval_seconds=0.001, + max_samples=4, + sample_function=sample, + ) + sampler.start() + assert ready.wait(1.0) + time.sleep(0.01) + report = sampler.stop() + + assert report["sample_count"] >= 2 + assert len(report["retained_samples"]) <= 4 + assert report["cpu_rss_mib"]["trend"]["last_mean"] > report["cpu_rss_mib"]["trend"]["first_mean"] + assert report["gpu_memory_mib"]["0"]["distribution"]["mean"] == 2048.0 diff --git a/tools/validation/benchmark_lingbot_vla_v2_service.py b/tools/validation/benchmark_lingbot_vla_v2_service.py new file mode 100644 index 00000000..7db06807 --- /dev/null +++ b/tools/validation/benchmark_lingbot_vla_v2_service.py @@ -0,0 +1,362 @@ +"""Benchmark one in-process LingBot-VLA v2 service replica on a single GPU.""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import math +import statistics +import threading +import time +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any + +import psutil +import torch +from PIL import Image + +from telefuser.metrics.runtime import collect_runtime_environment +from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline +from telefuser.pipelines.lingbot_vla_v2.service import ( + LingBotVlaV2ActionRequest, + predict_lingbot_vla_v2_action, +) + +_MIB = 1024**2 + + +class PeakRssSampler: + """Sample process RSS while one benchmark phase is active.""" + + def __init__(self, process: psutil.Process, interval_s: float = 0.01) -> None: + self.process = process + self.interval_s = interval_s + self.peak_bytes = process.memory_info().rss + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + def __enter__(self) -> "PeakRssSampler": + self._thread = threading.Thread(target=self._sample, daemon=True) + self._thread.start() + return self + + def __exit__(self, *args: object) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=max(self.interval_s * 4, 0.1)) + self._record() + + def _record(self) -> None: + try: + self.peak_bytes = max(self.peak_bytes, self.process.memory_info().rss) + except psutil.Error: + pass + + def _sample(self) -> None: + while not self._stop.wait(self.interval_s): + self._record() + + +def parse_image_sizes(value: str) -> tuple[tuple[int, int], ...]: + """Parse a comma-separated WIDTHxHEIGHT list.""" + sizes: list[tuple[int, int]] = [] + for item in value.split(","): + parts = item.strip().lower().split("x", maxsplit=1) + if len(parts) != 2: + raise argparse.ArgumentTypeError(f"invalid image size {item!r}; expected WIDTHxHEIGHT") + try: + width, height = (int(part) for part in parts) + except ValueError as error: + raise argparse.ArgumentTypeError(f"invalid image size {item!r}; expected integers") from error + if width <= 0 or height <= 0: + raise argparse.ArgumentTypeError("image dimensions must be positive") + size = (width, height) + if size not in sizes: + sizes.append(size) + if not sizes: + raise argparse.ArgumentTypeError("at least one image size is required") + return tuple(sizes) + + +def percentile(values: Sequence[float], fraction: float) -> float: + """Return a linearly interpolated percentile for a non-empty sample.""" + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def summarize(values: Sequence[float]) -> dict[str, float | int]: + """Summarize a latency sample in seconds.""" + if not values: + raise ValueError("summary requires at least one value") + total = sum(values) + return { + "count": len(values), + "total_seconds": total, + "mean_seconds": statistics.fmean(values), + "stdev_seconds": statistics.pstdev(values), + "min_seconds": min(values), + "p50_seconds": percentile(values, 0.50), + "p90_seconds": percentile(values, 0.90), + "p95_seconds": percentile(values, 0.95), + "max_seconds": max(values), + "throughput_requests_per_second": len(values) / total, + } + + +def encode_image(source: Image.Image, size: tuple[int, int], *, quality: int) -> tuple[str, int]: + """Resize and JPEG-encode one service input outside measured request time.""" + image = source.resize(size, Image.Resampling.BICUBIC) + buffer = io.BytesIO() + image.save(buffer, format="JPEG", quality=quality, optimize=False) + payload = buffer.getvalue() + return base64.b64encode(payload).decode("ascii"), len(payload) + + +def _cuda_synchronize(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def _memory_snapshot(device: torch.device, process: psutil.Process) -> dict[str, float | None]: + result: dict[str, float | None] = {"cpu_rss_mib": process.memory_info().rss / _MIB} + if device.type != "cuda": + result.update( + gpu_allocated_mib=None, + gpu_reserved_mib=None, + gpu_peak_allocated_mib=None, + gpu_peak_reserved_mib=None, + ) + return result + result.update( + gpu_allocated_mib=torch.cuda.memory_allocated(device) / _MIB, + gpu_reserved_mib=torch.cuda.memory_reserved(device) / _MIB, + gpu_peak_allocated_mib=torch.cuda.max_memory_allocated(device) / _MIB, + gpu_peak_reserved_mib=torch.cuda.max_memory_reserved(device) / _MIB, + ) + return result + + +def measure( + operation: Callable[[], Any], + *, + device: torch.device, + process: psutil.Process, + synchronize_cuda: bool, +) -> tuple[Any, dict[str, Any]]: + """Measure wall time and process/device memory for one operation.""" + if synchronize_cuda: + _cuda_synchronize(device) + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + before = _memory_snapshot(device, process) + with PeakRssSampler(process) as rss_sampler: + started_at = time.perf_counter() + result = operation() + if synchronize_cuda: + _cuda_synchronize(device) + elapsed = time.perf_counter() - started_at + after = _memory_snapshot(device, process) + return result, { + "seconds": elapsed, + "cpu_rss_before_mib": before["cpu_rss_mib"], + "cpu_rss_after_mib": after["cpu_rss_mib"], + "cpu_rss_peak_mib": rss_sampler.peak_bytes / _MIB, + "gpu_allocated_after_mib": after["gpu_allocated_mib"], + "gpu_reserved_after_mib": after["gpu_reserved_mib"], + "gpu_peak_allocated_mib": after["gpu_peak_allocated_mib"], + "gpu_peak_reserved_mib": after["gpu_peak_reserved_mib"], + } + + +def _load_source_image(path: Path | None) -> Image.Image: + if path is not None: + with Image.open(path) as image: + return image.convert("RGB").copy() + return Image.new("RGB", (640, 480), color=(32, 96, 160)) + + +def run_benchmark(args: argparse.Namespace) -> dict[str, Any]: + """Load one replica and return a JSON-serializable benchmark report.""" + if args.warmup < 0 or args.runs < 1: + raise ValueError("--warmup must be non-negative and --runs must be positive") + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("LingBot VLA v2 service benchmarking requires one visible CUDA GPU") + process = psutil.Process() + source = _load_source_image(args.image) + encoded_by_size = {size: encode_image(source, size, quality=args.jpeg_quality) for size in args.image_sizes} + + pipeline, load_metrics = measure( + lambda: get_lingbot_vla_v2_pipeline( + str(args.model_root), + str(args.qwen3vl_root), + device=str(device), + ), + device=device, + process=process, + synchronize_cuda=False, + ) + startup_warmup = None + if args.startup_warmup: + _, startup_warmup = measure( + pipeline.warmup, + device=device, + process=process, + synchronize_cuda=True, + ) + + executor = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="lingbot-vla-v2-benchmark") + if args.execution_mode == "service-thread" + else None + ) + active_phases: dict[str, float] = {} + original_prepare = pipeline.input_processor.prepare + original_predict = pipeline.predict + + def measured_prepare(observation: Any) -> Any: + started_at = time.perf_counter() + result = original_prepare(observation) + active_phases["preprocess_seconds"] = time.perf_counter() - started_at + return result + + def measured_predict(inputs: Any, seed: int | None = None) -> Any: + _cuda_synchronize(device) + started_at = time.perf_counter() + result = original_predict(inputs, seed=seed) + _cuda_synchronize(device) + active_phases["model_seconds"] = time.perf_counter() - started_at + return result + + pipeline.input_processor.prepare = measured_prepare + pipeline.predict = measured_predict + + def request_once(size: tuple[int, int]) -> tuple[dict[str, Any], dict[str, float]]: + active_phases.clear() + encoded, _ = encoded_by_size[size] + payload = { + "task": args.instruction, + "state": [0.0] * 14, + "camera_high": encoded, + "camera_left_wrist": encoded, + "camera_right_wrist": encoded, + "seed": args.seed, + } + request = LingBotVlaV2ActionRequest.model_validate(payload) + + def invoke_request() -> Any: + return predict_lingbot_vla_v2_action( + pipeline, + request, + max_image_bytes=args.max_image_bytes, + ) + + def invoke_service_thread() -> Any: + assert executor is not None + return executor.submit(invoke_request).result() + + operation: Callable[[], Any] = invoke_service_thread if executor is not None else invoke_request + response, metrics = measure( + operation, + device=device, + process=process, + synchronize_cuda=True, + ) + if response.horizon != 50 or response.action_dim != 55: + raise RuntimeError(f"unexpected action shape: {response.horizon}x{response.action_dim}") + if not all(math.isfinite(value) for row in response.canonical_normalized_actions for value in row): + raise RuntimeError("benchmark received non-finite actions") + phases = dict(active_phases) + phases["boundary_seconds"] = max( + metrics["seconds"] - phases.get("preprocess_seconds", 0.0) - phases.get("model_seconds", 0.0), + 0.0, + ) + return metrics, phases + + first_size = args.image_sizes[0] + try: + first_request, first_phases = request_once(first_size) + sizes_report: dict[str, Any] = {} + for size in args.image_sizes: + for _ in range(args.warmup): + request_once(size) + samples: list[dict[str, Any]] = [] + phases: list[dict[str, float]] = [] + for _ in range(args.runs): + sample, phase = request_once(size) + samples.append(sample) + phases.append(phase) + encoded_bytes = encoded_by_size[size][1] + sizes_report[f"{size[0]}x{size[1]}"] = { + "source_image_size": list(size), + "encoded_bytes_per_camera": encoded_bytes, + "total_latency": summarize([sample["seconds"] for sample in samples]), + "preprocess_latency": summarize([phase["preprocess_seconds"] for phase in phases]), + "model_latency": summarize([phase["model_seconds"] for phase in phases]), + "boundary_latency": summarize([phase["boundary_seconds"] for phase in phases]), + "cpu_rss_peak_mib": max(sample["cpu_rss_peak_mib"] for sample in samples), + "gpu_peak_allocated_mib": max(sample["gpu_peak_allocated_mib"] for sample in samples), + "gpu_peak_reserved_mib": max(sample["gpu_peak_reserved_mib"] for sample in samples), + } + report = { + "schema_version": 1, + "benchmark": "lingbot_vla_v2_single_gpu_service", + "model_root": str(args.model_root.resolve()), + "qwen3vl_root": str(args.qwen3vl_root.resolve()), + "device": str(device), + "seed": args.seed, + "instruction": args.instruction, + "internal_model_image_size": [pipeline.input_processor.image_size] * 2, + "warmup_runs_per_size": args.warmup, + "execution_mode": args.execution_mode, + "measured_runs_per_size": args.runs, + "environment": collect_runtime_environment([device], repo_root=Path(__file__).resolve().parents[2]), + "load": load_metrics, + "startup_warmup": startup_warmup, + "first_request": {**first_request, "source_image_size": list(first_size), "phases": first_phases}, + "steady_state_by_source_size": sizes_report, + "memory_after_benchmark": _memory_snapshot(device, process), + } + finally: + if executor is not None: + executor.shutdown(wait=True) + pipeline.close() + return report + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--qwen3vl-root", required=True, type=Path) + parser.add_argument("--image", type=Path, help="Optional source image reused for all three camera inputs.") + parser.add_argument("--image-sizes", type=parse_image_sizes, default=parse_image_sizes("256x256,640x480,1280x720")) + parser.add_argument("--instruction", default="pick up the red block") + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--execution-mode", choices=("service-thread", "direct"), default="service-thread") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--runs", type=int, default=20) + parser.add_argument("--startup-warmup", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--jpeg-quality", type=int, choices=range(1, 101), default=95) + parser.add_argument("--max-image-bytes", type=int, default=10 * 1024 * 1024) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + report = run_benchmark(args) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/capture_lingbot_vla_v2_telefuser.py b/tools/validation/capture_lingbot_vla_v2_telefuser.py new file mode 100644 index 00000000..f440510c --- /dev/null +++ b/tools/validation/capture_lingbot_vla_v2_telefuser.py @@ -0,0 +1,310 @@ +"""Capture a layered TeleFuser LingBot-VLA v2 regression artifact.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, Sequence + +import numpy as np +import torch +import transformers +from transformers import AutoProcessor + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2, resolve_lingbot_vla_v2_shards +from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_CAMERA_KEYS, + LingBotVlaV2Observation, + LingBotVlaV2Pipeline, + LingBotVlaV2PipelineConfig, +) + +ARTIFACT_SCHEMA_VERSION = 1 + + +def _sha256_file(path: Path, digest: Any | None = None) -> str: + result = hashlib.sha256() if digest is None else digest + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + result.update(block) + return result.hexdigest() + + +def _manifest_sha256(paths: Sequence[Path], *, include_contents: bool) -> str: + digest = hashlib.sha256() + for path in sorted((item.resolve() for item in paths), key=lambda item: item.name): + stat = path.stat() + digest.update(path.name.encode("utf-8")) + digest.update(str(stat.st_size).encode("ascii")) + if include_contents: + _sha256_file(path, digest) + return digest.hexdigest() + + +def _processor_files(root: Path) -> list[Path]: + return sorted(path for path in root.iterdir() if path.is_file() and path.suffix != ".safetensors") + + +def _input_sha256(task: str, state: Sequence[float], image_paths: Sequence[Path]) -> str: + digest = hashlib.sha256() + canonical = json.dumps({"task": task, "state": list(state)}, sort_keys=True, separators=(",", ":")) + digest.update(canonical.encode("utf-8")) + for path in image_paths: + digest.update(path.name.encode("utf-8")) + _sha256_file(path, digest) + return digest.hexdigest() + + +def _git_commit() -> str: + repository_root = Path(__file__).resolve().parents[2] + completed = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repository_root, + check=True, + capture_output=True, + text=True, + ) + return completed.stdout.strip() + + +class TensorCapture: + """Own CPU snapshots and their original tensor contracts.""" + + def __init__(self) -> None: + self.arrays: dict[str, np.ndarray] = {} + self.array_metadata: dict[str, dict[str, object]] = {} + + def add(self, key: str, tensor: torch.Tensor) -> None: + if key in self.arrays: + raise ValueError(f"Duplicate capture key: {key}") + snapshot = tensor.detach().cpu().clone() + stored = snapshot.float() if snapshot.is_floating_point() else snapshot + array = stored.numpy() + self.arrays[key] = array + self.array_metadata[key] = { + "shape": list(snapshot.shape), + "original_dtype": str(snapshot.dtype).removeprefix("torch."), + "stored_dtype": str(array.dtype), + } + + +class VelocityTrace: + """Capture the state around each call to the flow-matching velocity model.""" + + def __init__(self, capture: TensorCapture) -> None: + self.capture = capture + self.step = 0 + + def record(self, original: Any, *args: Any, **kwargs: Any) -> torch.Tensor: + if len(args) < 5: + raise RuntimeError("LingBot-VLA v2 predict_velocity trace received an unexpected call signature") + x_t = args[3] + timestep = args[4] + suffix = f"{self.step:02d}" + if self.step == 0: + self.capture.add("initial_noise", x_t) + self.capture.add(f"timestep_step_{suffix}", timestep) + self.capture.add(f"x_t_step_{suffix}", x_t) + velocity = original(*args, **kwargs) + self.capture.add(f"velocity_step_{suffix}", velocity) + self.step += 1 + return velocity + + +@contextmanager +def trace_predict_velocity(flow_model: Any, capture: TensorCapture) -> Iterator[VelocityTrace]: + """Temporarily trace one model instance without changing global classes.""" + original = flow_model.predict_velocity + had_instance_override = "predict_velocity" in vars(flow_model) + previous_override = vars(flow_model).get("predict_velocity") + compile_enabled = bool(getattr(flow_model, "_use_compile_predict_velocity", False)) + trace = VelocityTrace(capture) + + flow_model._use_compile_predict_velocity = False + flow_model.predict_velocity = lambda *args, **kwargs: trace.record(original, *args, **kwargs) + try: + yield trace + finally: + if had_instance_override: + flow_model.predict_velocity = previous_override + else: + del flow_model.predict_velocity + flow_model._use_compile_predict_velocity = compile_enabled + + +def _build_pipeline(model_root: Path, qwen3vl_root: Path, device: str) -> LingBotVlaV2Pipeline: + target_device = torch.device(device) + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + processor = AutoProcessor.from_pretrained(str(qwen3vl_root), local_files_only=True, padding_side="right") + manager = ModuleManager(torch_dtype=dtype, device="cpu") + manager.add_module(processor, "lingbot_vla_v2_processor", path=str(qwen3vl_root)) + load_lingbot_vla_v2(manager, model_root, qwen3vl_root, torch_dtype=dtype) + pipeline = LingBotVlaV2Pipeline(device=device, torch_dtype=dtype) + pipeline.init( + manager, + LingBotVlaV2PipelineConfig( + policy_config=ModelRuntimeConfig( + device_type=target_device.type, + device_id=target_device.index or 0, + torch_dtype=dtype, + ) + ), + ) + return pipeline + + +def capture_artifact( + *, + model_root: Path, + qwen3vl_root: Path, + image_paths: Sequence[Path], + task: str, + state: Sequence[float], + seed: int, + output: Path, + device: str, + full_checkpoint_hash: bool, + deterministic_moe: bool, +) -> tuple[Path, Path]: + if len(image_paths) != len(ROBOTWIN_CAMERA_KEYS): + raise ValueError(f"expected {len(ROBOTWIN_CAMERA_KEYS)} camera paths, got {len(image_paths)}") + output = output.with_suffix(".npz") + metadata_path = output.with_suffix(".json") + output.parent.mkdir(parents=True, exist_ok=True) + + pipeline = _build_pipeline(model_root, qwen3vl_root, device) + if deterministic_moe: + for module in pipeline.policy_stage.policy.modules(): + if hasattr(module, "_use_robby_moe_kernel"): + module._use_robby_moe_kernel = False + capture = TensorCapture() + try: + observation = LingBotVlaV2Observation( + task=task, + state=state, + images=dict(zip(ROBOTWIN_CAMERA_KEYS, image_paths, strict=True)), + ) + inputs = pipeline.input_processor.prepare(observation) + for key in ("images", "img_masks", "image_grid_thw", "lang_tokens", "lang_masks", "state"): + capture.add(key, getattr(inputs, key)) + + flow_model = pipeline.policy_stage.policy.model + with trace_predict_velocity(flow_model, capture) as trace: + chunk = pipeline.predict(inputs, seed=seed) + capture.add("canonical_normalized_actions", chunk.canonical_normalized_actions) + + expected_steps = int(flow_model.config.num_steps) + if trace.step != expected_steps: + raise RuntimeError(f"captured {trace.step} denoising steps, expected {expected_steps}") + + target_device = torch.device(device) + checkpoint_paths = [Path(path) for path in resolve_lingbot_vla_v2_shards(model_root)] + metadata = { + "schema_version": ARTIFACT_SCHEMA_VERSION, + "artifact_kind": "telefuser_regression", + "telefuser_commit": _git_commit(), + "checkpoint_manifest_sha256": _manifest_sha256( + checkpoint_paths, + include_contents=full_checkpoint_hash, + ), + "checkpoint_hash_mode": "full_sha256" if full_checkpoint_hash else "filename_and_size", + "norm_stats_sha256": _sha256_file( + Path(__file__).resolve().parents[2] + / "telefuser/pipelines/lingbot_vla_v2/assets/robotwin_norm_stats.json" + ), + "processor_manifest_sha256": _manifest_sha256( + _processor_files(qwen3vl_root), + include_contents=True, + ), + "input_sha256": _input_sha256(task, state, image_paths), + "seed": seed, + "num_steps": trace.step, + "torch_dtype": str(pipeline.torch_dtype).removeprefix("torch."), + "attention_backend": str(flow_model.config.attention_implementation), + "moe_backend": "deterministic_torch_reference" if deterministic_moe else "upstream_triton", + "device": str(target_device), + "device_name": torch.cuda.get_device_name(target_device) if target_device.type == "cuda" else "cpu", + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + "arrays": capture.array_metadata, + } + np.savez(output, **capture.arrays) + metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + finally: + pipeline.close() + + return output, metadata_path + + +def _parse_state(value: str) -> list[float]: + try: + state = json.loads(value) + except json.JSONDecodeError as error: + raise argparse.ArgumentTypeError("state-json must be valid JSON") from error + if not isinstance(state, list) or len(state) != 14 or any(isinstance(item, bool) for item in state): + raise argparse.ArgumentTypeError("state-json must be a 14-element numeric JSON list") + try: + return [float(item) for item in state] + except (TypeError, ValueError) as error: + raise argparse.ArgumentTypeError("state-json must contain only numeric values") from error + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--qwen3vl-root", required=True, type=Path) + parser.add_argument("--camera-high", required=True, type=Path) + parser.add_argument("--camera-left-wrist", required=True, type=Path) + parser.add_argument("--camera-right-wrist", required=True, type=Path) + parser.add_argument("--task", required=True) + parser.add_argument("--state-json", required=True, type=_parse_state) + parser.add_argument("--seed", required=True, type=int) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--full-checkpoint-hash", + action="store_true", + help="Hash all checkpoint bytes instead of the faster filename-and-size manifest", + ) + parser.add_argument( + "--deterministic-moe", + action="store_true", + help="Disable the upstream atomic Triton MoE kernel for bitwise cross-process parity", + ) + args = parser.parse_args() + + paths = ( + args.model_root, + args.qwen3vl_root, + args.camera_high, + args.camera_left_wrist, + args.camera_right_wrist, + ) + missing = [str(path) for path in paths if not path.exists()] + if missing: + parser.error(f"input paths do not exist: {missing}") + + artifact, metadata = capture_artifact( + model_root=args.model_root, + qwen3vl_root=args.qwen3vl_root, + image_paths=(args.camera_high, args.camera_left_wrist, args.camera_right_wrist), + task=args.task, + state=args.state_json, + seed=args.seed, + output=args.output, + device=args.device, + full_checkpoint_hash=args.full_checkpoint_hash, + deterministic_moe=args.deterministic_moe, + ) + print(f"Saved LingBot-VLA v2 capture: {artifact}") + print(f"Saved LingBot-VLA v2 metadata: {metadata}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/capture_lingbot_vla_v2_upstream.py b/tools/validation/capture_lingbot_vla_v2_upstream.py new file mode 100644 index 00000000..cc724e9e --- /dev/null +++ b/tools/validation/capture_lingbot_vla_v2_upstream.py @@ -0,0 +1,541 @@ +"""Capture a layered artifact from the fixed official LingBot-VLA v2 checkout. + +Run this script with the dedicated upstream uv environment and with the fixed +upstream checkout as ``--upstream-root``. The official code forces +FlashAttention during construction. For reproducible comparison on the local +PyTorch 2.11 stack, this runner intercepts model construction inside this +process only and selects the eager attention implementation used by TeleFuser. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import sys +from contextlib import contextmanager +from pathlib import Path +from types import MethodType, SimpleNamespace +from typing import Any, Iterator, Sequence + +import numpy as np +import torch +import transformers +from PIL import Image +from accelerate import init_empty_weights +from safetensors.torch import load_file +from torchvision.transforms.v2 import Resize +from transformers import AutoConfig, AutoProcessor, PreTrainedModel + +ARTIFACT_SCHEMA_VERSION = 1 +UPSTREAM_COMMIT = "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" +CAMERA_KEYS = ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", +) + + +def _sha256_file(path: Path, digest: Any | None = None) -> str: + result = hashlib.sha256() if digest is None else digest + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + result.update(block) + return result.hexdigest() + + +def _manifest_sha256(paths: Sequence[Path], *, include_contents: bool) -> str: + digest = hashlib.sha256() + for path in sorted((item.resolve() for item in paths), key=lambda item: item.name): + stat = path.stat() + digest.update(path.name.encode("utf-8")) + digest.update(str(stat.st_size).encode("ascii")) + if include_contents: + _sha256_file(path, digest) + return digest.hexdigest() + + +def _processor_files(root: Path) -> list[Path]: + return sorted(path for path in root.iterdir() if path.is_file() and path.suffix != ".safetensors") + + +def _input_sha256(task: str, state: Sequence[float], image_paths: Sequence[Path]) -> str: + digest = hashlib.sha256() + canonical = json.dumps({"task": task, "state": list(state)}, sort_keys=True, separators=(",", ":")) + digest.update(canonical.encode("utf-8")) + for path in image_paths: + digest.update(path.name.encode("utf-8")) + _sha256_file(path, digest) + return digest.hexdigest() + + +def _git_commit(repository: Path) -> str: + completed = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + status = subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=repository, + check=True, + capture_output=True, + text=True, + ) + if status.stdout: + raise RuntimeError(f"Upstream checkout must be clean, got:\n{status.stdout}") + return completed.stdout.strip() + + +def _checkpoint_shards(model_root: Path) -> list[Path]: + index_path = model_root / "model.safetensors.index.json" + payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = payload.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Invalid checkpoint index: {index_path}") + shards = [model_root / name for name in sorted(set(weight_map.values()))] + missing = [str(path) for path in shards if not path.is_file()] + if missing: + raise FileNotFoundError(f"Missing checkpoint shards: {missing}") + return shards + + +class TensorCapture: + def __init__(self) -> None: + self.arrays: dict[str, np.ndarray] = {} + self.array_metadata: dict[str, dict[str, object]] = {} + + def add(self, key: str, tensor: torch.Tensor) -> None: + if key in self.arrays: + raise ValueError(f"Duplicate capture key: {key}") + snapshot = tensor.detach().cpu().clone() + stored = snapshot.float() if snapshot.is_floating_point() else snapshot + array = stored.numpy() + self.arrays[key] = array + self.array_metadata[key] = { + "shape": list(snapshot.shape), + "original_dtype": str(snapshot.dtype).removeprefix("torch."), + "stored_dtype": str(array.dtype), + } + + +class VelocityTrace: + def __init__(self, capture: TensorCapture) -> None: + self.capture = capture + self.step = 0 + + def record(self, original: Any, *args: Any, **kwargs: Any) -> torch.Tensor: + if len(args) < 5: + raise RuntimeError("Official predict_velocity trace received an unexpected call signature") + suffix = f"{self.step:02d}" + x_t, timestep = args[3], args[4] + if self.step == 0: + self.capture.add("initial_noise", x_t) + self.capture.add(f"timestep_step_{suffix}", timestep) + self.capture.add(f"x_t_step_{suffix}", x_t) + velocity = original(*args, **kwargs) + self.capture.add(f"velocity_step_{suffix}", velocity) + self.step += 1 + return velocity + + +@contextmanager +def _trace_predict_velocity(flow_model: Any, capture: TensorCapture) -> Iterator[VelocityTrace]: + original = flow_model.predict_velocity + trace = VelocityTrace(capture) + flow_model.predict_velocity = lambda *args, **kwargs: trace.record(original, *args, **kwargs) + try: + yield trace + finally: + del flow_model.predict_velocity + + +def _force_eager(config: Any) -> None: + for current in (config, getattr(config, "text_config", None), getattr(config, "vision_config", None)): + if current is not None: + current._attn_implementation = "eager" + + +@contextmanager +def _eager_construction() -> Iterator[None]: + """Override the upstream hard-coded FA2 selection in this process only.""" + original = PreTrainedModel._from_config.__func__ + + def from_config(cls: type[PreTrainedModel], config: Any, **kwargs: Any) -> PreTrainedModel: + _force_eager(config) + return original(cls, config, **kwargs) + + PreTrainedModel._from_config = classmethod(from_config) + try: + yield + finally: + PreTrainedModel._from_config = classmethod(original) + + +def _official_model_values(qwen3vl_root: Path) -> dict[str, Any]: + # Base-6B values are documented in the fixed upstream Training_Config.md. + return { + "post_training": False, + "adanorm_time": True, + "moe_implementation": "fused", + "use_robby_moe_kernel": False, + "attention_implementation": "eager", + "vit_attn_implementation": "eager", + "precompute_grid_thw": True, + "vlm_causal": True, + "use_moe": True, + "token_moe_layers": list(range(36)), + "token_num_experts": 32, + "token_top_k": 4, + "token_moe_intermediate_size": 512, + "token_shared_intermediate_size": 704, + "bias_update_speed": 0.0, + "sequence_wise_mode": "per_sequence", + "sequence_wise_loss_coeff": 1e-3, + "router_z_loss_coeff": 1e-4, + "router_activation": "sigmoid", + "routed_scaling_factor": 4.0, + "use_shared_expert_gate": False, + "freeze_vision_encoder": False, + "tokenizer_max_length": 72, + "loss_type": "L1_fm", + "action_dim": 55, + "max_action_dim": 55, + "max_state_dim": 55, + "tokenizer_path": str(qwen3vl_root), + "align_params": { + "mode": "query", + "num_task_tokens": 8, + "depth_loss_weight": 0.004, + "future_depth_loss_weight": 0.004, + "use_future_video": True, + "llm": {"dim_out": 2560, "image_token_size": 8, "image_input_size": 224}, + "depth": { + "model_type": "MoRGBD", + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "token_size": 16, + "dim_out": 1024, + "input_size": 224, + "use_future_depth": True, + "block_future_depth_to_action": True, + "future_depth_head_type": "resampler", + "detach_future_image_feats": True, + }, + "video": { + "attention_mode": "flex_block_causal", + "input_size": 256, + "block_suffix_to_future_video": True, + "share_future_depth_query": True, + "use_shared_future_task_proj": True, + "use_current_shared_task_proj": True, + "num_future_frames": 1, + "use_warmup_frame": True, + "effective_fps": 1.0, + "n_blocks": 1, + "cls_pool": "last", + "detach_image_feats": True, + "num_layers": 1, + "num_heads": 4, + "dim_head": 32, + "ff_mult": 1, + "num_backbone_tokens": 256, + "dim_out": 1024, + "future_video_loss_weight": 0.004, + "use_smooth_l1_loss": False, + "use_mse_loss": True, + "mse_loss_weight": 1.0, + "use_patch_loss": True, + "use_current_patch_loss": True, + "use_cosine_loss": False, + "cosine_loss_weight": 0.2, + "use_cls_loss": False, + "cls_loss_type": "mse", + "cls_loss_weight": 0.2, + }, + }, + } + + +def _build_config(qwen3vl_root: Path) -> Any: + from lingbotvla.models.vla.lingbot_vla.configuration_lingbot_vla import LingbotVLAV2Config + + config = LingbotVLAV2Config(**_official_model_values(qwen3vl_root)) + qwen_config = AutoConfig.from_pretrained(str(qwen3vl_root), local_files_only=True) + for key in ( + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "rms_norm_eps", + "rope_theta", + "vocab_size", + "max_position_embeddings", + "hidden_act", + "tie_word_embeddings", + ): + if hasattr(qwen_config.text_config, key): + setattr(config, key, getattr(qwen_config.text_config, key)) + config.vision_config = qwen_config.vision_config + config.use_cache = True + return config + + +def _load_official_model(model_root: Path, config: Any, device: torch.device) -> Any: + from lingbotvla.models.vla.lingbot_vla.modeling_lingbot_vla_v2 import LingbotVlaV2Policy + from lingbotvla.models.vla.lingbot_vla.qwen3vl_in_vla import apply_lingbot_qwen3_vl_patch + + apply_lingbot_qwen3_vl_patch() + with _eager_construction(), init_empty_weights(): + model = LingbotVlaV2Policy(config, eval=True) + + index = json.loads((model_root / "model.safetensors.index.json").read_text(encoding="utf-8")) + checkpoint_keys = set(index["weight_map"]) + model_keys = set(model.state_dict()) + if checkpoint_keys != model_keys: + raise RuntimeError( + "Official model/checkpoint key mismatch: " + f"missing={sorted(model_keys - checkpoint_keys)[:10]}, " + f"unexpected={sorted(checkpoint_keys - model_keys)[:10]}" + ) + + for shard in _checkpoint_shards(model_root): + model.load_state_dict(load_file(shard, device="cpu"), strict=False, assign=True) + unmaterialized = [name for name, tensor in model.state_dict().items() if tensor.is_meta] + if unmaterialized: + raise RuntimeError(f"Official checkpoint left meta tensors: {unmaterialized[:10]}") + return model.to(device=device, dtype=torch.bfloat16).eval() + + +def _load_rgb(path: Path) -> torch.Tensor: + with Image.open(path) as image: + array = np.asarray(image.convert("RGB")).copy() + return torch.from_numpy(array).permute(2, 0, 1).contiguous() + + +def _prepare_inputs( + upstream_root: Path, + qwen3vl_root: Path, + config: Any, + image_paths: Sequence[Path], + task: str, + state: Sequence[float], +) -> dict[str, torch.Tensor]: + from lingbotvla.data.vla_data.utils import FeatureTransform + + processor = AutoProcessor.from_pretrained(str(qwen3vl_root), local_files_only=True, padding_side="right") + data_config = SimpleNamespace( + joints=["{'arm.position': 14}", "{'end.position': 14}", "{'effector.position': 2}"], + cameras=["camera_top", "camera_wrist_left", "camera_wrist_right"], + norm_type=[ + "{'arm.position': 'bounds_99_woclip'}", + "{'end.position': 'bounds_99_woclip'}", + "{'effector.position': 'bounds_99_woclip'}", + ], + ) + transform = FeatureTransform( + upstream_root / "configs/robot_configs/robotwin.yaml", + data_config, + config, + processor, + chunk_size=config.chunk_size, + norm_stats_path=upstream_root / "assets/norm_stats/robotwin.json", + ) + resize = Resize((256, 256), antialias=True) + item: dict[str, Any] = {"observation.state": torch.tensor(state, dtype=torch.float32), "task": task} + for key, path in zip(CAMERA_KEYS, image_paths, strict=True): + item[key] = resize(_load_rgb(path).to(dtype=torch.float32)) + prepared = transform.apply(item, policy_eval=True) + return { + "images": prepared["images"].unsqueeze(0), + "img_masks": prepared["img_masks"].unsqueeze(0), + "image_grid_thw": prepared["image_grid_thw"].unsqueeze(0), + "lang_tokens": prepared["lang_tokens"].unsqueeze(0), + "lang_masks": prepared["lang_masks"].unsqueeze(0), + "state": prepared["state"].unsqueeze(0), + } + + +def capture_artifact( + *, + upstream_root: Path, + model_root: Path, + qwen3vl_root: Path, + image_paths: Sequence[Path], + task: str, + state: Sequence[float], + seed: int, + output: Path, + device: str, + full_checkpoint_hash: bool, + deterministic_moe: bool, +) -> tuple[Path, Path]: + commit = _git_commit(upstream_root) + if commit != UPSTREAM_COMMIT: + raise RuntimeError(f"Expected upstream commit {UPSTREAM_COMMIT}, got {commit}") + if len(image_paths) != len(CAMERA_KEYS): + raise ValueError(f"expected {len(CAMERA_KEYS)} camera paths, got {len(image_paths)}") + + sys.path.insert(0, str(upstream_root)) + output = output.with_suffix(".npz") + metadata_path = output.with_suffix(".json") + output.parent.mkdir(parents=True, exist_ok=True) + target_device = torch.device(device) + if target_device.type != "cuda": + raise ValueError("Official 6B parity capture currently requires CUDA") + + config = _build_config(qwen3vl_root) + inputs = _prepare_inputs(upstream_root, qwen3vl_root, config, image_paths, task, state) + capture = TensorCapture() + for key, tensor in inputs.items(): + capture.add(key, tensor) + + model = _load_official_model(model_root, config, target_device) + if deterministic_moe: + import lingbotvla.models.vla.lingbot_vla.qwen2_action_expert as qwen2_action_expert + + qwen2_action_expert.robby_moe_forward = None + + def deterministic_forward(experts, module, num_experts, routing_weights, selected_experts, hidden_states): + del module + output = torch.zeros_like(hidden_states) + for expert_id in range(num_experts): + routes = (selected_experts == expert_id).nonzero(as_tuple=False) + if routes.numel() == 0: + continue + token_ids, route_ids = routes[:, 0], routes[:, 1] + expert_input = hidden_states.index_select(0, token_ids) + gate = torch.nn.functional.linear(expert_input, experts.gate_proj[expert_id]) + up = torch.nn.functional.linear(expert_input, experts.up_proj[expert_id]) + intermediate = torch.nn.functional.silu(gate) * up + expert_output = torch.nn.functional.linear(intermediate, experts.down_proj[expert_id]) + weights = routing_weights[token_ids, route_ids].unsqueeze(-1) + output.index_add_(0, token_ids, expert_output * weights) + return output + + for module in model.modules(): + if module.__class__.__name__ == "Qwen2FusedExperts": + module.forward = MethodType(deterministic_forward, module) + tensors = { + "images": inputs["images"].to(device=target_device, dtype=torch.bfloat16), + "img_masks": inputs["img_masks"].to(device=target_device), + "lang_tokens": inputs["lang_tokens"].to(device=target_device), + "lang_masks": inputs["lang_masks"].to(device=target_device), + "state": inputs["state"].to(device=target_device, dtype=torch.bfloat16), + "image_grid_thw": inputs["image_grid_thw"].to(device=target_device, dtype=torch.long), + } + generator = torch.Generator(device=target_device).manual_seed(seed) + noise = torch.randn( + 1, + int(config.n_action_steps), + int(config.max_action_dim), + device=target_device, + dtype=torch.bfloat16, + generator=generator, + ) + with torch.inference_mode(), _trace_predict_velocity(model.model, capture) as trace: + actions = model.sample_actions(**tensors, noise=noise) + capture.add("canonical_normalized_actions", actions.squeeze(0).to(device="cpu", dtype=torch.float32)) + if trace.step != int(config.num_steps): + raise RuntimeError(f"captured {trace.step} denoising steps, expected {config.num_steps}") + + checkpoint_paths = _checkpoint_shards(model_root) + metadata = { + "schema_version": ARTIFACT_SCHEMA_VERSION, + "artifact_kind": "official_upstream_common_eager", + "upstream_commit": commit, + "upstream_attention_override": "process_local_pretrained_model_from_config_intercept", + "checkpoint_manifest_sha256": _manifest_sha256(checkpoint_paths, include_contents=full_checkpoint_hash), + "checkpoint_hash_mode": "full_sha256" if full_checkpoint_hash else "filename_and_size", + "norm_stats_sha256": _sha256_file(upstream_root / "assets/norm_stats/robotwin.json"), + "processor_manifest_sha256": _manifest_sha256(_processor_files(qwen3vl_root), include_contents=True), + "input_sha256": _input_sha256(task, state, image_paths), + "seed": seed, + "num_steps": trace.step, + "torch_dtype": "bfloat16", + "attention_backend": "eager", + "moe_backend": "deterministic_torch_reference" if deterministic_moe else "upstream_triton", + "device": str(target_device), + "device_name": torch.cuda.get_device_name(target_device), + "python_version": sys.version.split()[0], + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + "arrays": capture.array_metadata, + } + np.savez(output, **capture.arrays) + metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return output, metadata_path + + +def _parse_state(value: str) -> list[float]: + try: + state = json.loads(value) + except json.JSONDecodeError as error: + raise argparse.ArgumentTypeError("state-json must be valid JSON") from error + if not isinstance(state, list) or len(state) != 14 or any(isinstance(item, bool) for item in state): + raise argparse.ArgumentTypeError("state-json must be a 14-element numeric JSON list") + try: + return [float(item) for item in state] + except (TypeError, ValueError) as error: + raise argparse.ArgumentTypeError("state-json must contain only numeric values") from error + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream-root", required=True, type=Path) + parser.add_argument("--model-root", required=True, type=Path) + parser.add_argument("--qwen3vl-root", required=True, type=Path) + parser.add_argument("--camera-high", required=True, type=Path) + parser.add_argument("--camera-left-wrist", required=True, type=Path) + parser.add_argument("--camera-right-wrist", required=True, type=Path) + parser.add_argument("--task", required=True) + parser.add_argument("--state-json", required=True, type=_parse_state) + parser.add_argument("--seed", required=True, type=int) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--full-checkpoint-hash", action="store_true") + parser.add_argument( + "--deterministic-moe", + action="store_true", + help="Disable the upstream atomic Triton MoE kernel for bitwise cross-process parity", + ) + args = parser.parse_args() + + paths = ( + args.upstream_root, + args.model_root, + args.qwen3vl_root, + args.camera_high, + args.camera_left_wrist, + args.camera_right_wrist, + ) + missing = [str(path) for path in paths if not path.exists()] + if missing: + parser.error(f"input paths do not exist: {missing}") + + artifact, metadata = capture_artifact( + upstream_root=args.upstream_root.resolve(), + model_root=args.model_root.resolve(), + qwen3vl_root=args.qwen3vl_root.resolve(), + image_paths=(args.camera_high, args.camera_left_wrist, args.camera_right_wrist), + task=args.task, + state=args.state_json, + seed=args.seed, + output=args.output, + device=args.device, + full_checkpoint_hash=args.full_checkpoint_hash, + deterministic_moe=args.deterministic_moe, + ) + print(f"Saved official LingBot-VLA v2 capture: {artifact}") + print(f"Saved official LingBot-VLA v2 metadata: {metadata}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/requirements-lingbot-vla-v2-upstream.txt b/tools/validation/requirements-lingbot-vla-v2-upstream.txt new file mode 100644 index 00000000..fa927ca4 --- /dev/null +++ b/tools/validation/requirements-lingbot-vla-v2-upstream.txt @@ -0,0 +1,17 @@ +# Isolated runtime for the fixed LingBot-VLA v2 upstream parity capture. +# Install this file into a dedicated uv venv; do not install it into TeleFuser's venv. +accelerate==1.7.0 +av==15.0.0 +datasets==3.6.0 +einops==0.8.1 +huggingface-hub==0.34.3 +numpy==2.2.6 +pillow==12.3.0 +psutil==7.0.0 +pydantic==2.13.4 +pyyaml==6.0.3 +safetensors==0.6.2 +torch==2.11.0 +torchdata==0.11.0 +torchvision==0.26.0 +transformers==4.57.3 diff --git a/tools/validation/run_lingbot_vla_v2_parity.py b/tools/validation/run_lingbot_vla_v2_parity.py new file mode 100644 index 00000000..901378d9 --- /dev/null +++ b/tools/validation/run_lingbot_vla_v2_parity.py @@ -0,0 +1,374 @@ +"""Compare layered LingBot-VLA v2 capture artifacts. + +The same comparator is used for local TeleFuser regression artifacts and for +future upstream parity artifacts. Captures stay file based so implementations +with incompatible Python dependencies never need to share a process. +""" + +from __future__ import annotations + +import argparse +import json +import re +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import numpy as np + +UPSTREAM_REPOSITORY = "https://github.com/Robbyant/lingbot-vla-v2" +UPSTREAM_COMMIT = "be27333c9b5f2663b0ec33f069dd7dfd67fa32b5" +ARTIFACT_SCHEMA_VERSION = 1 +PREPROCESSING_KEYS = ( + "images", + "img_masks", + "image_grid_thw", + "lang_tokens", + "lang_masks", + "state", +) +STEP_LAYERS = ("timestep", "x_t", "velocity") +FINAL_ACTION_KEYS = ("canonical_normalized_actions", "actions") +IDENTITY_METADATA_KEYS = ( + "checkpoint_manifest_sha256", + "processor_manifest_sha256", + "norm_stats_sha256", + "input_sha256", + "seed", + "num_steps", + "torch_dtype", + "attention_backend", + "moe_backend", +) +_STEP_KEY = re.compile(r"^(timestep|x_t|velocity)_step_([0-9]+)$") + + +@dataclass(frozen=True) +class ArrayParity: + layer: str + key: str + shape: tuple[int, ...] + expected_dtype: str + actual_dtype: str + max_abs: float + mean_abs: float + mismatch_count: int + rtol: float + atol: float + passed: bool + + +def _load_npz(path: Path) -> dict[str, np.ndarray]: + if not path.is_file(): + raise FileNotFoundError(path) + with np.load(path, allow_pickle=False) as payload: + return {key: payload[key] for key in payload.files} + + +def _metadata_path(artifact_path: Path, explicit_path: Path | None) -> Path: + return explicit_path if explicit_path is not None else artifact_path.with_suffix(".json") + + +def _load_metadata(artifact_path: Path, explicit_path: Path | None = None) -> dict[str, Any]: + path = _metadata_path(artifact_path, explicit_path) + if not path.is_file(): + raise FileNotFoundError(f"Missing LingBot-VLA v2 artifact metadata: {path}") + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise ValueError(f"Artifact metadata must be a JSON object: {path}") + if payload.get("schema_version") != ARTIFACT_SCHEMA_VERSION: + raise ValueError( + f"Unsupported artifact schema in {path}: expected {ARTIFACT_SCHEMA_VERSION}, " + f"got {payload.get('schema_version')!r}" + ) + missing = [key for key in IDENTITY_METADATA_KEYS if key not in payload] + if missing: + raise ValueError(f"Artifact metadata {path} is missing identity fields: {missing}") + return payload + + +def _step_keys(arrays: dict[str, np.ndarray]) -> dict[str, dict[int, str]]: + result: dict[str, dict[int, str]] = {layer: {} for layer in STEP_LAYERS} + for key in arrays: + match = _STEP_KEY.fullmatch(key) + if match is not None: + layer, step_text = match.groups() + step = int(step_text) + if step in result[layer]: + raise ValueError(f"Duplicate {layer} capture for step {step}") + result[layer][step] = key + return result + + +def _validate_contract(arrays: dict[str, np.ndarray], metadata: dict[str, Any], *, side: str) -> None: + required = (*PREPROCESSING_KEYS, "initial_noise") + missing = [key for key in required if key not in arrays] + if missing: + raise ValueError(f"{side} artifact is missing required arrays: {missing}") + + if not any(key in arrays for key in FINAL_ACTION_KEYS): + raise ValueError(f"{side} artifact is missing a final action array: {FINAL_ACTION_KEYS}") + + array_metadata = metadata.get("arrays") + if not isinstance(array_metadata, dict): + raise ValueError(f"{side} metadata must contain an arrays contract") + missing_contracts = sorted(set(arrays) - set(array_metadata)) + if missing_contracts: + raise ValueError(f"{side} metadata is missing array contracts: {missing_contracts}") + for key, array in arrays.items(): + contract = array_metadata[key] + if not isinstance(contract, dict): + raise ValueError(f"{side} metadata contract for {key} must be an object") + expected_contract = { + "shape": list(array.shape), + "stored_dtype": str(array.dtype), + } + mismatches = { + field: {"metadata": contract.get(field), "artifact": value} + for field, value in expected_contract.items() + if contract.get(field) != value + } + if "original_dtype" not in contract: + mismatches["original_dtype"] = {"metadata": None, "artifact": "required"} + if mismatches: + raise ValueError(f"{side} metadata contract for {key} does not match the artifact: {mismatches}") + + num_steps = metadata["num_steps"] + if not isinstance(num_steps, int) or isinstance(num_steps, bool) or num_steps <= 0: + raise ValueError(f"{side} metadata num_steps must be a positive integer, got {num_steps!r}") + expected_steps = set(range(num_steps)) + for layer, keys in _step_keys(arrays).items(): + if set(keys) != expected_steps: + raise ValueError(f"{side} artifact {layer} steps must be {sorted(expected_steps)}, got {sorted(keys)}") + + +def _compare_array( + layer: str, + key: str, + expected: np.ndarray, + actual: np.ndarray, + *, + rtol: float, + atol: float, + expected_original_dtype: str, + actual_original_dtype: str, +) -> ArrayParity: + expected_dtype = expected_original_dtype + actual_dtype = actual_original_dtype + shape_matches = expected.shape == actual.shape + dtype_matches = expected.dtype == actual.dtype and expected_original_dtype == actual_original_dtype + finite = True + if np.issubdtype(expected.dtype, np.number) and not np.isfinite(expected).all(): + finite = False + if np.issubdtype(actual.dtype, np.number) and not np.isfinite(actual).all(): + finite = False + + if not shape_matches or not finite: + max_abs = float("inf") + mean_abs = float("inf") + mismatch_count = max(expected.size, actual.size) + values_match = False + elif expected.dtype == np.bool_ or np.issubdtype(expected.dtype, np.integer): + difference = np.abs(expected.astype(np.int64) - actual.astype(np.int64)) + mismatch_count = int(np.count_nonzero(difference)) + max_abs = float(difference.max()) if difference.size else 0.0 + mean_abs = float(difference.mean()) if difference.size else 0.0 + values_match = mismatch_count == 0 + else: + difference = np.abs(expected.astype(np.float64) - actual.astype(np.float64)) + close = np.isclose(expected, actual, rtol=rtol, atol=atol, equal_nan=False) + mismatch_count = int(np.count_nonzero(~close)) + max_abs = float(difference.max()) if difference.size else 0.0 + mean_abs = float(difference.mean()) if difference.size else 0.0 + values_match = mismatch_count == 0 + + return ArrayParity( + layer=layer, + key=key, + shape=tuple(actual.shape), + expected_dtype=expected_dtype, + actual_dtype=actual_dtype, + max_abs=max_abs, + mean_abs=mean_abs, + mismatch_count=mismatch_count, + rtol=rtol, + atol=atol, + passed=shape_matches and dtype_matches and finite and values_match, + ) + + +def _action_key(arrays: dict[str, np.ndarray]) -> str: + for key in FINAL_ACTION_KEYS: + if key in arrays: + return key + raise ValueError(f"No final action key found: {FINAL_ACTION_KEYS}") + + +def _original_dtype(metadata: dict[str, Any], key: str) -> str: + return str(metadata["arrays"][key]["original_dtype"]) + + +def compare_artifacts( + reference: Path, + candidate: Path, + *, + rtol: float, + atol: float, + reference_metadata: Path | None = None, + candidate_metadata: Path | None = None, +) -> dict[str, object]: + expected = _load_npz(reference) + actual = _load_npz(candidate) + expected_metadata = _load_metadata(reference, reference_metadata) + actual_metadata = _load_metadata(candidate, candidate_metadata) + _validate_contract(expected, expected_metadata, side="reference") + _validate_contract(actual, actual_metadata, side="candidate") + + metadata_mismatches = { + key: {"reference": expected_metadata[key], "candidate": actual_metadata[key]} + for key in IDENTITY_METADATA_KEYS + if expected_metadata[key] != actual_metadata[key] + } + if metadata_mismatches: + raise ValueError(f"Artifact identity metadata does not match: {metadata_mismatches}") + + results: list[ArrayParity] = [] + for key in PREPROCESSING_KEYS: + results.append( + _compare_array( + "preprocessing", + key, + expected[key], + actual[key], + rtol=0.0, + atol=0.0, + expected_original_dtype=_original_dtype(expected_metadata, key), + actual_original_dtype=_original_dtype(actual_metadata, key), + ) + ) + results.append( + _compare_array( + "noise", + "initial_noise", + expected["initial_noise"], + actual["initial_noise"], + rtol=0.0, + atol=0.0, + expected_original_dtype=_original_dtype(expected_metadata, "initial_noise"), + actual_original_dtype=_original_dtype(actual_metadata, "initial_noise"), + ) + ) + + expected_steps = _step_keys(expected) + actual_steps = _step_keys(actual) + for step in range(expected_metadata["num_steps"]): + for layer in STEP_LAYERS: + expected_key = expected_steps[layer][step] + actual_key = actual_steps[layer][step] + layer_rtol = 0.0 if layer == "timestep" else rtol + layer_atol = 0.0 if layer == "timestep" else atol + results.append( + _compare_array( + layer, + expected_key, + expected[expected_key], + actual[actual_key], + rtol=layer_rtol, + atol=layer_atol, + expected_original_dtype=_original_dtype(expected_metadata, expected_key), + actual_original_dtype=_original_dtype(actual_metadata, actual_key), + ) + ) + + expected_action_key = _action_key(expected) + actual_action_key = _action_key(actual) + results.append( + _compare_array( + "action", + expected_action_key, + expected[expected_action_key], + actual[actual_action_key], + rtol=rtol, + atol=atol, + expected_original_dtype=_original_dtype(expected_metadata, expected_action_key), + actual_original_dtype=_original_dtype(actual_metadata, actual_action_key), + ) + ) + + expected_compared_keys = { + *PREPROCESSING_KEYS, + "initial_noise", + expected_action_key, + *(key for layer in expected_steps.values() for key in layer.values()), + } + actual_compared_keys = { + *PREPROCESSING_KEYS, + "initial_noise", + actual_action_key, + *(key for layer in actual_steps.values() for key in layer.values()), + } + unexpected_reference = sorted(set(expected) - expected_compared_keys) + unexpected_candidate = sorted(set(actual) - actual_compared_keys) + first_failed_step = next( + ( + int(match.group(2)) + for item in results + if not item.passed and (match := _STEP_KEY.fullmatch(item.key)) is not None + ), + None, + ) + + return { + "schema_version": ARTIFACT_SCHEMA_VERSION, + "upstream_repository": UPSTREAM_REPOSITORY, + "upstream_commit": UPSTREAM_COMMIT, + "reference": str(reference), + "candidate": str(candidate), + "reference_kind": expected_metadata.get("artifact_kind"), + "candidate_kind": actual_metadata.get("artifact_kind"), + "passed": all(item.passed for item in results), + "first_failed_step": first_failed_step, + "unexpected_reference_keys": unexpected_reference, + "unexpected_candidate_keys": unexpected_candidate, + "results": [asdict(item) for item in results], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--reference-metadata", type=Path, default=None) + parser.add_argument("--candidate-metadata", type=Path, default=None) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--profile", choices=("strict", "portable"), default="strict") + parser.add_argument("--rtol", type=float, default=None) + parser.add_argument("--atol", type=float, default=None) + args = parser.parse_args() + + default_tolerance = 0.0 if args.profile == "strict" else 1e-3 + rtol = default_tolerance if args.rtol is None else args.rtol + atol = default_tolerance if args.atol is None else args.atol + if rtol < 0 or atol < 0: + parser.error("rtol and atol must be non-negative") + + report = compare_artifacts( + args.reference, + args.candidate, + rtol=rtol, + atol=atol, + reference_metadata=args.reference_metadata, + candidate_metadata=args.candidate_metadata, + ) + payload = json.dumps(report, indent=2, sort_keys=True) + if args.output is None: + print(payload) + else: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload + "\n", encoding="utf-8") + if not report["passed"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/validate_lingbot_vla_v2_service_faults.py b/tools/validation/validate_lingbot_vla_v2_service_faults.py new file mode 100644 index 00000000..489b0853 --- /dev/null +++ b/tools/validation/validate_lingbot_vla_v2_service_faults.py @@ -0,0 +1,505 @@ +"""Validate fault handling of a real LingBot-VLA v2 structured service.""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import signal +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import psutil +import requests + +try: + from tools.validation import validate_lingbot_vla_v2_structured_service as structured_validator +except ModuleNotFoundError as error: + if error.name != "tools": + raise + import validate_lingbot_vla_v2_structured_service as structured_validator + +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) + + +class FaultValidationFailure(RuntimeError): + """Raised when the service violates an expected fault-handling behavior.""" + + +def _new_session() -> requests.Session: + session = requests.Session() + session.trust_env = False + return session + + +def _response_body(response: requests.Response) -> dict[str, Any]: + try: + body = response.json() + except ValueError as error: + raise FaultValidationFailure(f"HTTP {response.status_code} response is not JSON") from error + if not isinstance(body, dict): + raise FaultValidationFailure(f"HTTP {response.status_code} response is not a JSON object") + return body + + +def _request_json( + session: requests.Session, + method: str, + url: str, + *, + timeout: float, + payload: dict[str, Any] | None = None, +) -> tuple[int, dict[str, Any]]: + response = session.request(method, url, json=payload, timeout=timeout) + return response.status_code, _response_body(response) + + +def _wait_terminal( + session: requests.Session, + base_url: str, + task_id: str, + *, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> dict[str, Any]: + deadline = time.monotonic() + task_timeout_seconds + while time.monotonic() < deadline: + status_code, body = _request_json( + session, + "GET", + f"{base_url}/v1/tasks/{task_id}/status", + timeout=http_timeout_seconds, + ) + if status_code != 200: + raise FaultValidationFailure(f"task status returned HTTP {status_code}: {body}") + status = body.get("status") or body.get("task_status") + if status in _TERMINAL_STATUSES: + return body + time.sleep(poll_interval_seconds) + raise FaultValidationFailure(f"task {task_id} did not become terminal within {task_timeout_seconds:g}s") + + +def _expect_rejected_or_failed( + session: requests.Session, + base_url: str, + payload: dict[str, Any], + *, + case_name: str, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> dict[str, Any]: + status_code, created = _request_json( + session, + "POST", + f"{base_url}/v1/tasks/structured", + timeout=http_timeout_seconds, + payload=payload, + ) + if 400 <= status_code < 500: + return {"name": case_name, "passed": True, "handling": "rejected", "http_status": status_code} + if status_code != 200: + raise FaultValidationFailure(f"{case_name} returned unexpected HTTP {status_code}: {created}") + task_id = created.get("task_id") + if not isinstance(task_id, str) or not task_id: + raise FaultValidationFailure(f"{case_name} accepted without a task_id") + terminal = _wait_terminal( + session, + base_url, + task_id, + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + if terminal.get("status") != "failed": + raise FaultValidationFailure(f"{case_name} reached unexpected terminal status {terminal.get('status')}") + return { + "name": case_name, + "passed": True, + "handling": "asynchronous_failure", + "http_status": status_code, + "task_id": task_id, + "terminal_status": "failed", + "error": str(terminal.get("error") or "")[:1000], + } + + +def validate_request_faults( + session: requests.Session, + base_url: str, + payload: dict[str, Any], + *, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> list[dict[str, Any]]: + """Validate required fields, payload validation, and request cancellation.""" + cases: list[dict[str, Any]] = [] + + missing_camera = dict(payload) + missing_camera.pop("camera_high", None) + cases.append( + _expect_rejected_or_failed( + session, + base_url, + missing_camera, + case_name="missing_required_camera", + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + ) + + invalid_state = dict(payload) + invalid_state["state"] = [0.0] * 13 + cases.append( + _expect_rejected_or_failed( + session, + base_url, + invalid_state, + case_name="invalid_state_dimension", + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + ) + + invalid_image = dict(payload) + invalid_image["camera_high"] = "not-valid-base64" + cases.append( + _expect_rejected_or_failed( + session, + base_url, + invalid_image, + case_name="invalid_camera_base64", + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + ) + + status_code, created = _request_json( + session, + "POST", + f"{base_url}/v1/tasks/structured", + timeout=http_timeout_seconds, + payload=payload, + ) + if status_code != 200 or not isinstance(created.get("task_id"), str): + raise FaultValidationFailure(f"cancellation case was not accepted: HTTP {status_code}: {created}") + task_id = created["task_id"] + cancel_status, cancellation = _request_json( + session, + "DELETE", + f"{base_url}/v1/tasks/{task_id}", + timeout=http_timeout_seconds, + ) + if cancel_status != 200 or cancellation.get("stop_status") not in {"success", "do_nothing"}: + raise FaultValidationFailure(f"task cancellation failed: HTTP {cancel_status}: {cancellation}") + terminal = _wait_terminal( + session, + base_url, + task_id, + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + ) + terminal_status = terminal.get("status") + if terminal_status not in {"cancelled", "completed"}: + raise FaultValidationFailure(f"cancelled task reached unexpected terminal status {terminal_status}") + cases.append( + { + "name": "client_cancellation", + "passed": True, + "task_id": task_id, + "stop_status": cancellation.get("stop_status"), + "terminal_status": terminal_status, + "race_with_completion": terminal_status == "completed", + } + ) + return cases + + +def _gpu_uuid_to_index() -> dict[str, str]: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return { + fields[1]: fields[0] + for line in result.stdout.splitlines() + if len(fields := [field.strip() for field in line.split(",")]) == 2 + } + + +def select_replica_process( + compute_rows: str, + *, + service_process_ids: set[int], + gpu_uuid_to_index: dict[str, str], + gpu_index: str, + service_pid: int, +) -> int: + """Select exactly one descendant compute process on a physical GPU.""" + candidates: set[int] = set() + for line in compute_rows.splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) < 2: + continue + try: + pid = int(fields[0]) + except ValueError: + continue + observed_index = gpu_uuid_to_index.get(fields[1], fields[1]) + if pid != service_pid and pid in service_process_ids and observed_index == gpu_index: + candidates.add(pid) + if len(candidates) != 1: + raise FaultValidationFailure( + f"expected exactly one service descendant compute process on GPU {gpu_index}, observed {sorted(candidates)}" + ) + return candidates.pop() + + +def discover_replica_process(service_pid: int, gpu_index: str) -> int: + """Discover one replica process without considering unrelated system processes.""" + root = psutil.Process(service_pid) + service_process_ids = {process.pid for process in root.children(recursive=True)} + result = subprocess.run( + ["nvidia-smi", "--query-compute-apps=pid,gpu_uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return select_replica_process( + result.stdout, + service_process_ids=service_process_ids, + gpu_uuid_to_index=_gpu_uuid_to_index(), + gpu_index=gpu_index, + service_pid=service_pid, + ) + + +def gpu_compute_process_ids(compute_rows: str) -> set[int]: + """Parse compute PIDs from nvidia-smi rows.""" + process_ids: set[int] = set() + for line in compute_rows.splitlines(): + try: + process_ids.add(int(line.split(",", 1)[0].strip())) + except (ValueError, IndexError): + continue + return process_ids + + +def wait_for_replica_exit(replica_pid: int, *, timeout_seconds: float = 10.0) -> None: + """Wait until a replica is exited or zombie and no longer owns GPU memory.""" + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + try: + process_exited = psutil.Process(replica_pid).status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + process_exited = True + result = subprocess.run( + ["nvidia-smi", "--query-compute-apps=pid,gpu_uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + if process_exited and replica_pid not in gpu_compute_process_ids(result.stdout): + return + time.sleep(0.1) + raise FaultValidationFailure(f"replica process {replica_pid} did not exit and release GPU memory after SIGTERM") + + +def validate_pool_degradation(before: dict[str, Any], after: dict[str, Any]) -> dict[str, Any]: + """Validate one-replica capacity reduction without requiring automatic restart.""" + before_pool = before.get("pool") + after_pool = after.get("pool") + if not isinstance(before_pool, list) or len(before_pool) < 2: + raise FaultValidationFailure("replica termination requires service status with at least two pool entries") + if not isinstance(after_pool, list) or len(after_pool) != len(before_pool): + raise FaultValidationFailure("pool status disappeared or changed size after replica termination") + before_capacity = before.get("effective_max_concurrent_tasks") + after_capacity = after.get("effective_max_concurrent_tasks") + if not isinstance(before_capacity, int) or not isinstance(after_capacity, int): + raise FaultValidationFailure("service status has no integer effective capacity") + dead = [replica for replica in after_pool if replica.get("status") == "dead"] + live = [replica for replica in after_pool if replica.get("status") != "dead"] + if len(dead) != 1 or not live or after_capacity != before_capacity - 1: + raise FaultValidationFailure( + f"expected one dead replica and capacity {before_capacity - 1}, " + f"observed dead={len(dead)}, capacity={after_capacity}" + ) + return { + "before_capacity": before_capacity, + "after_capacity": after_capacity, + "dead_replica_ids": [replica.get("id") for replica in dead], + "live_replica_ids": [replica.get("id") for replica in live], + "recovery_semantics": "graceful_capacity_degradation_without_automatic_restart", + } + + +def validate_replica_exit( + session: requests.Session, + base_url: str, + payload: dict[str, Any], + *, + service_pid: int, + gpu_index: str, + http_timeout_seconds: float, + task_timeout_seconds: float, + poll_interval_seconds: float, +) -> dict[str, Any]: + """Terminate one explicitly selected replica and validate graceful degradation.""" + status_code, before = _request_json(session, "GET", f"{base_url}/v1/service/status", timeout=http_timeout_seconds) + if status_code != 200 or before.get("execution_mode") != "concurrent_pipeline_pool": + raise FaultValidationFailure("replica termination requires a ready concurrent pipeline pool") + replica_pid = discover_replica_process(service_pid, gpu_index) + os.kill(replica_pid, signal.SIGTERM) + wait_for_replica_exit(replica_pid) + + failed_attempts: list[str] = [] + successful_request: dict[str, Any] | None = None + after: dict[str, Any] | None = None + config = structured_validator.RequestConfig( + base_url=base_url, + payload=payload, + http_timeout_seconds=http_timeout_seconds, + task_timeout_seconds=task_timeout_seconds, + poll_interval_seconds=poll_interval_seconds, + expected_horizon=50, + expected_action_dim=55, + ) + for attempt in range(3): + record = structured_validator.execute_request( + session, + config, + request_index=attempt, + worker_index=0, + run_started_at=time.perf_counter(), + ) + if record.get("outcome") == "succeeded": + successful_request = record + else: + failed_attempts.append(str(record.get("error") or "unknown request failure")) + _, observed = _request_json(session, "GET", f"{base_url}/v1/service/status", timeout=http_timeout_seconds) + if any(replica.get("status") == "dead" for replica in observed.get("pool", [])): + after = observed + if successful_request is not None and after is not None: + break + if successful_request is None or after is None: + raise FaultValidationFailure( + f"service did not degrade cleanly after replica exit; request_errors={failed_attempts}" + ) + degradation = validate_pool_degradation(before, after) + return { + "name": "replica_exit", + "passed": True, + "terminated_pid": replica_pid, + "physical_gpu_index": gpu_index, + "failed_attempts": failed_attempts, + "successful_action": successful_request["action"], + **degradation, + } + + +def _image_base64(path: Path) -> str: + return base64.b64encode(path.read_bytes()).decode("ascii") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080") + parser.add_argument("--image", type=Path, required=True, help="Image reused for all three camera inputs.") + parser.add_argument("--instruction", default="pick up the object") + parser.add_argument("--state", type=structured_validator.parse_state_json, default=[0.0] * 14) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--http-timeout-seconds", type=float, default=10.0) + parser.add_argument("--task-timeout-seconds", type=float, default=120.0) + parser.add_argument("--poll-interval-seconds", type=float, default=0.05) + parser.add_argument("--service-pid", type=int) + parser.add_argument("--kill-replica-gpu-index") + parser.add_argument("--output", type=Path, default=Path("work_dirs/vla_service_fault_validation.json")) + args = parser.parse_args() + if (args.service_pid is None) != (args.kill_replica_gpu_index is None): + parser.error("--service-pid and --kill-replica-gpu-index must be provided together") + if args.http_timeout_seconds <= 0 or args.task_timeout_seconds <= 0 or args.poll_interval_seconds <= 0: + parser.error("timeout and polling values must be positive") + if not args.image.is_file(): + parser.error(f"image does not exist: {args.image}") + return args + + +def main() -> int: + args = parse_args() + encoded_image = _image_base64(args.image) + payload = { + "task": "vla_action", + "instruction": args.instruction, + "state": args.state, + "camera_high": encoded_image, + "camera_left_wrist": encoded_image, + "camera_right_wrist": encoded_image, + "seed": args.seed, + } + report: dict[str, Any] = { + "schema_version": 1, + "validation": "lingbot_vla_v2_structured_service_faults", + "created_at": datetime.now(timezone.utc).isoformat(), + "target": args.base_url.rstrip("/"), + "checks": [], + "passed": False, + } + try: + with _new_session() as session: + structured_validator.inspect_service(report["target"], timeout_seconds=args.http_timeout_seconds) + report["checks"].extend( + validate_request_faults( + session, + report["target"], + payload, + http_timeout_seconds=args.http_timeout_seconds, + task_timeout_seconds=args.task_timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + ) + ) + if args.service_pid is not None: + report["checks"].append( + validate_replica_exit( + session, + report["target"], + payload, + service_pid=args.service_pid, + gpu_index=args.kill_replica_gpu_index, + http_timeout_seconds=args.http_timeout_seconds, + task_timeout_seconds=args.task_timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + ) + ) + report["passed"] = all(check.get("passed") is True for check in report["checks"]) + except ( + FaultValidationFailure, + structured_validator.ValidationFailure, + requests.RequestException, + OSError, + ) as error: + report["error"] = str(error) + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps({"passed": report["passed"], "checks": len(report["checks"]), "output": str(args.output)})) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/validation/validate_lingbot_vla_v2_structured_service.py b/tools/validation/validate_lingbot_vla_v2_structured_service.py new file mode 100644 index 00000000..b35e7e28 --- /dev/null +++ b/tools/validation/validate_lingbot_vla_v2_structured_service.py @@ -0,0 +1,1003 @@ +"""Validate a real LingBot-VLA v2 native structured API service.""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import math +import platform +import statistics +import struct +import subprocess +import sys +import threading +import time +from collections import Counter, deque +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timezone +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +import psutil +import requests + +_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled"}) +_MIB = 1024**2 +_REQUEST_PARAMETER_CONTRACT = { + "instruction": ("string", True), + "state": ("array", True), + "camera_high": ("string", True), + "camera_left_wrist": ("string", True), + "camera_right_wrist": ("string", True), + "seed": ("integer", False), +} +_RESULT_FIELDS = frozenset( + { + "canonical_normalized_actions", + "horizon", + "action_dim", + "checkpoint_variant", + "policy_verified", + "verification_status", + } +) +_SENSITIVE_REQUEST_FIELDS = frozenset({"camera_high", "camera_left_wrist", "camera_right_wrist"}) + + +class ValidationFailure(RuntimeError): + """Raised when the target violates the VLA structured API contract.""" + + +@dataclass(frozen=True) +class RequestConfig: + """Immutable settings shared by validation workers.""" + + base_url: str + payload: dict[str, Any] + http_timeout_seconds: float + task_timeout_seconds: float + poll_interval_seconds: float + expected_horizon: int + expected_action_dim: int + + +def parse_state_json(value: str) -> list[float]: + """Parse and validate a finite 14-dimensional RobotWin state.""" + try: + raw = json.loads(value) + except json.JSONDecodeError as error: + raise argparse.ArgumentTypeError("state must be valid JSON") from error + if not isinstance(raw, list) or len(raw) != 14: + raise argparse.ArgumentTypeError("state must be a JSON array containing exactly 14 values") + state: list[float] = [] + for item in raw: + if isinstance(item, bool) or not isinstance(item, int | float) or not math.isfinite(float(item)): + raise argparse.ArgumentTypeError("state values must be finite numbers") + state.append(float(item)) + return state + + +def percentile(values: Sequence[float], fraction: float) -> float: + """Return a linearly interpolated percentile for a non-empty sample.""" + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def summarize(values: Sequence[float]) -> dict[str, float | int] | None: + """Summarize a possibly empty sample in seconds.""" + if not values: + return None + return { + "count": len(values), + "mean": statistics.fmean(values), + "min": min(values), + "p50": percentile(values, 0.50), + "p90": percentile(values, 0.90), + "p95": percentile(values, 0.95), + "p99": percentile(values, 0.99), + "max": max(values), + } + + +def compare_windows(values: Sequence[float], fraction: float = 0.1) -> dict[str, float | int] | None: + """Compare the first and last windows of an ordered measurement series.""" + if not values: + return None + window_count = max(1, math.ceil(len(values) * fraction)) + first_mean = statistics.fmean(values[:window_count]) + last_mean = statistics.fmean(values[-window_count:]) + delta = last_mean - first_mean + return { + "sample_count": len(values), + "window_count": window_count, + "first_mean": first_mean, + "last_mean": last_mean, + "delta": delta, + "change_percent": delta / first_mean * 100.0 if first_mean else 0.0, + } + + +def validate_service_metadata(metadata: Any) -> None: + """Validate that the target exposes the native VLA structured contract.""" + if not isinstance(metadata, dict): + raise ValidationFailure("service metadata must be a JSON object") + if metadata.get("declared_pipeline_contract") is not True: + raise ValidationFailure("service does not expose a declared pipeline contract") + if "vla_action" not in metadata.get("supported_tasks", []): + raise ValidationFailure("service metadata does not declare the vla_action task") + if "structured" not in metadata.get("supported_media_types", []): + raise ValidationFailure("service metadata does not declare structured output") + task_contract = metadata.get("task_contracts", {}).get("vla_action") + if not isinstance(task_contract, dict) or task_contract.get("media_type") != "structured": + raise ValidationFailure("vla_action does not have a structured task contract") + parameters = task_contract.get("parameters") + if not isinstance(parameters, dict): + raise ValidationFailure("vla_action parameters are missing from service metadata") + if set(parameters) != set(_REQUEST_PARAMETER_CONTRACT): + raise ValidationFailure( + "vla_action parameter fields changed: " + f"expected {sorted(_REQUEST_PARAMETER_CONTRACT)}, observed {sorted(parameters)}" + ) + for name, (expected_type, expected_required) in _REQUEST_PARAMETER_CONTRACT.items(): + parameter = parameters[name] + if not isinstance(parameter, dict): + raise ValidationFailure(f"vla_action parameter contract is invalid: {name}") + if parameter.get("type") != expected_type or parameter.get("required") is not expected_required: + raise ValidationFailure( + f"vla_action parameter {name} changed: expected type={expected_type}, required={expected_required}" + ) + if task_contract.get("required_inputs") != ["camera_high", "camera_left_wrist", "camera_right_wrist"]: + raise ValidationFailure("vla_action required_inputs changed") + if task_contract.get("optional_inputs") != []: + raise ValidationFailure("vla_action optional_inputs changed") + + +def validate_action_result(result: Any, *, expected_horizon: int, expected_action_dim: int) -> dict[str, Any]: + """Validate and summarize one canonical normalized action chunk.""" + if expected_horizon < 1 or expected_action_dim < 1: + raise ValueError("expected action dimensions must be positive") + if not isinstance(result, dict): + raise ValidationFailure("completed task result must be a JSON object") + if set(result) != set(_RESULT_FIELDS): + raise ValidationFailure(f"result fields changed: expected {sorted(_RESULT_FIELDS)}, observed {sorted(result)}") + actions = result.get("canonical_normalized_actions") + if not isinstance(actions, list) or len(actions) != expected_horizon: + observed = len(actions) if isinstance(actions, list) else type(actions).__name__ + raise ValidationFailure(f"expected action horizon {expected_horizon}, observed {observed}") + if result.get("horizon") != expected_horizon: + raise ValidationFailure(f"result horizon field is not {expected_horizon}") + if result.get("action_dim") != expected_action_dim: + raise ValidationFailure(f"result action_dim field is not {expected_action_dim}") + + flat: list[float] = [] + digest = hashlib.sha256() + for row_index, row in enumerate(actions): + if not isinstance(row, list) or len(row) != expected_action_dim: + observed = len(row) if isinstance(row, list) else type(row).__name__ + raise ValidationFailure(f"action row {row_index} has dimension {observed}, expected {expected_action_dim}") + for value in row: + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(float(value)): + raise ValidationFailure("action chunk contains a non-finite or non-numeric value") + number = float(value) + flat.append(number) + digest.update(struct.pack(" None: + """Validate stable terminal task fields without rejecting safe additive metadata.""" + if not isinstance(status, dict): + raise ValidationFailure("task status must be a JSON object") + if status.get("task_id") != task_id: + raise ValidationFailure("task status returned a different task_id") + required = {"status", "inference_time_s", "peak_memory_mb", "result"} + missing = sorted(required.difference(status)) + if missing: + raise ValidationFailure(f"task status is missing fields: {', '.join(missing)}") + leaked = sorted(_SENSITIVE_REQUEST_FIELDS.intersection(status)) + if leaked: + raise ValidationFailure(f"task status echoed sensitive image fields: {', '.join(leaked)}") + + +def _request_json( + session: requests.Session, + method: str, + url: str, + *, + timeout: float, + payload: dict[str, Any] | None = None, +) -> dict[str, Any]: + response = session.request(method, url, json=payload, timeout=timeout) + try: + response.raise_for_status() + except requests.HTTPError as error: + body = response.text[:1000] + raise ValidationFailure(f"{method} {url} returned HTTP {response.status_code}: {body}") from error + try: + body = response.json() + except ValueError as error: + raise ValidationFailure(f"{method} {url} did not return JSON") from error + if not isinstance(body, dict): + raise ValidationFailure(f"{method} {url} did not return a JSON object") + return body + + +def _new_session() -> requests.Session: + session = requests.Session() + session.trust_env = False + return session + + +def inspect_service(base_url: str, *, timeout_seconds: float) -> dict[str, Any]: + """Read and validate native service readiness and metadata.""" + with _new_session() as session: + ready = _request_json(session, "GET", f"{base_url}/v1/service/ready", timeout=timeout_seconds) + if ready.get("ready") is not True: + raise ValidationFailure("service readiness endpoint reports not ready") + metadata = _request_json(session, "GET", f"{base_url}/v1/service/metadata", timeout=timeout_seconds) + validate_service_metadata(metadata) + status = _request_json(session, "GET", f"{base_url}/v1/service/status", timeout=timeout_seconds) + metrics = _request_json(session, "GET", f"{base_url}/v1/service/metrics/json", timeout=timeout_seconds) + return {"ready": ready, "metadata": metadata, "status": status, "metrics": metrics} + + +def execute_request( + session: requests.Session, + config: RequestConfig, + *, + request_index: int, + worker_index: int, + run_started_at: float, +) -> dict[str, Any]: + """Submit, poll, validate, and summarize one real structured request.""" + record: dict[str, Any] = { + "request_index": request_index, + "worker_index": worker_index, + "start_offset_seconds": time.perf_counter() - run_started_at, + } + request_started_at = time.perf_counter() + try: + submit_started_at = time.perf_counter() + created = _request_json( + session, + "POST", + f"{config.base_url}/v1/tasks/structured", + timeout=config.http_timeout_seconds, + payload=config.payload, + ) + accepted_at = time.perf_counter() + record["submit_seconds"] = accepted_at - submit_started_at + task_id = created.get("task_id") + if not isinstance(task_id, str) or not task_id: + raise ValidationFailure("structured task creation response has no task_id") + record["task_id"] = task_id + + deadline = accepted_at + config.task_timeout_seconds + transitions: list[str] = [] + poll_count = 0 + while True: + if time.perf_counter() >= deadline: + raise ValidationFailure(f"task {task_id} exceeded {config.task_timeout_seconds:g}s timeout") + status = _request_json( + session, + "GET", + f"{config.base_url}/v1/tasks/{task_id}/status", + timeout=config.http_timeout_seconds, + ) + poll_count += 1 + task_status = status.get("status") or status.get("task_status") + if not isinstance(task_status, str): + raise ValidationFailure(f"task {task_id} status response has no status") + if not transitions or transitions[-1] != task_status: + transitions.append(task_status) + if task_status in _TERMINAL_STATUSES: + break + time.sleep(config.poll_interval_seconds) + + completed_at = time.perf_counter() + validate_task_status(status, task_id=task_id) + record.update( + end_to_end_seconds=completed_at - request_started_at, + accepted_to_terminal_seconds=completed_at - accepted_at, + poll_count=poll_count, + status_transitions=transitions, + terminal_status=task_status, + ) + inference_time = status.get("inference_time_s") + if inference_time is not None: + if isinstance(inference_time, bool) or not isinstance(inference_time, int | float): + raise ValidationFailure("inference_time_s must be numeric or null") + inference_time = float(inference_time) + if not math.isfinite(inference_time) or inference_time < 0: + raise ValidationFailure("inference_time_s must be finite and non-negative") + record["inference_time_seconds"] = inference_time + peak_memory = status.get("peak_memory_mb") + if peak_memory is not None: + if isinstance(peak_memory, bool) or not isinstance(peak_memory, int | float): + raise ValidationFailure("peak_memory_mb must be numeric or null") + peak_memory = float(peak_memory) + if not math.isfinite(peak_memory) or peak_memory < 0: + raise ValidationFailure("peak_memory_mb must be finite and non-negative") + record["peak_memory_mb"] = peak_memory + + if task_status != "completed": + raise ValidationFailure(f"task {task_id} reached terminal status {task_status}: {status.get('error')}") + record["action"] = validate_action_result( + status.get("result"), + expected_horizon=config.expected_horizon, + expected_action_dim=config.expected_action_dim, + ) + record["outcome"] = "succeeded" + except (requests.RequestException, ValidationFailure, ValueError) as error: + record["outcome"] = "failed" + record["error"] = str(error) + record.setdefault("end_to_end_seconds", time.perf_counter() - request_started_at) + return record + + +def _parse_gpu_process_memory( + output: str, + *, + process_ids: set[int], + uuid_to_index: dict[str, str], + gpu_indexes: set[str] | None, +) -> dict[str, float]: + """Parse nvidia-smi process memory rows for one service process tree.""" + memory_by_gpu: dict[str, float] = {} + for line in output.splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) != 3: + continue + try: + pid = int(fields[0]) + memory_mib = float(fields[2]) + except ValueError: + continue + gpu_index = uuid_to_index.get(fields[1], fields[1]) + if pid not in process_ids or (gpu_indexes is not None and gpu_index not in gpu_indexes): + continue + memory_by_gpu[gpu_index] = memory_by_gpu.get(gpu_index, 0.0) + memory_mib + return memory_by_gpu + + +def _query_gpu_index_map() -> dict[str, str]: + """Resolve stable GPU UUIDs to physical indexes once per validation run.""" + uuid_result = subprocess.run( + ["nvidia-smi", "--query-gpu=index,uuid", "--format=csv,noheader,nounits"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return { + fields[1].strip(): fields[0].strip() + for line in uuid_result.stdout.splitlines() + if len(fields := [field.strip() for field in line.split(",")]) == 2 + } + + +def _query_gpu_process_memory( + process_ids: set[int], + *, + uuid_to_index: dict[str, str], + gpu_indexes: set[str] | None, +) -> dict[str, float]: + """Read GPU memory used by the service process tree through nvidia-smi.""" + process_result = subprocess.run( + [ + "nvidia-smi", + "--query-compute-apps=pid,gpu_uuid,used_gpu_memory", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return _parse_gpu_process_memory( + process_result.stdout, + process_ids=process_ids, + uuid_to_index=uuid_to_index, + gpu_indexes=gpu_indexes, + ) + + +def _sample_local_resources( + root_pid: int, + *, + uuid_to_index: dict[str, str], + gpu_indexes: set[str] | None, +) -> dict[str, Any]: + """Sample process-tree RSS and GPU memory without touching model execution.""" + root = psutil.Process(root_pid) + processes = [root, *root.children(recursive=True)] + process_ids: set[int] = set() + rss_bytes = 0 + for process in processes: + try: + process_ids.add(process.pid) + rss_bytes += process.memory_info().rss + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + return { + "process_ids": sorted(process_ids), + "cpu_rss_mib": rss_bytes / _MIB, + "gpu_memory_mib": _query_gpu_process_memory( + process_ids, + uuid_to_index=uuid_to_index, + gpu_indexes=gpu_indexes, + ), + } + + +class ResourceSampler: + """Periodically sample local service resources in a background thread.""" + + def __init__( + self, + root_pid: int, + *, + interval_seconds: float, + max_samples: int, + gpu_indexes: set[str] | None = None, + sample_function: Callable[[], dict[str, Any]] | None = None, + ) -> None: + if root_pid < 1 or interval_seconds <= 0 or max_samples < 2: + raise ValueError("invalid resource sampler configuration") + self.root_pid = root_pid + self.interval_seconds = interval_seconds + self.max_samples = max_samples + self.gpu_indexes = gpu_indexes + if sample_function is None: + uuid_to_index = _query_gpu_index_map() + + def sample_function() -> dict[str, Any]: + return _sample_local_resources( + root_pid, + uuid_to_index=uuid_to_index, + gpu_indexes=gpu_indexes, + ) + + self.sample_function = sample_function + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._started_at = time.perf_counter() + self._sample_count = 0 + self._cpu_rss: list[float] = [] + self._gpu_memory: dict[str, list[float]] = {} + self._process_ids: set[int] = set() + self._errors: deque[str] = deque(maxlen=100) + self._first_samples: list[dict[str, Any]] = [] + recent_capacity = max(1, max_samples // 2) + self._recent_samples: deque[dict[str, Any]] = deque(maxlen=recent_capacity) + + def start(self) -> None: + """Start sampling and take an initial sample.""" + self._started_at = time.perf_counter() + self._record_once() + self._thread = threading.Thread(target=self._sample_loop, name="vla-resource-sampler", daemon=True) + self._thread.start() + + def stop(self) -> dict[str, Any]: + """Stop sampling, take a final sample, and return the report.""" + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=max(self.interval_seconds * 2, 1.0)) + self._record_once() + return self.report() + + def _sample_loop(self) -> None: + while not self._stop.wait(self.interval_seconds): + self._record_once() + + def _record_once(self) -> None: + offset = time.perf_counter() - self._started_at + try: + sample = self.sample_function() + cpu_rss = float(sample["cpu_rss_mib"]) + if not math.isfinite(cpu_rss) or cpu_rss < 0: + raise ValueError("invalid cpu_rss_mib sample") + gpu_memory = sample.get("gpu_memory_mib", {}) + if not isinstance(gpu_memory, dict): + raise ValueError("invalid gpu_memory_mib sample") + normalized_gpu = { + str(index): float(value) + for index, value in gpu_memory.items() + if math.isfinite(float(value)) and float(value) >= 0 + } + process_ids = {int(pid) for pid in sample.get("process_ids", [])} + self._sample_count += 1 + self._cpu_rss.append(cpu_rss) + self._process_ids.update(process_ids) + for index, value in normalized_gpu.items(): + self._gpu_memory.setdefault(index, []).append(value) + retained = { + "offset_seconds": offset, + "process_ids": sorted(process_ids), + "cpu_rss_mib": cpu_rss, + "gpu_memory_mib": normalized_gpu, + } + first_capacity = self.max_samples - self._recent_samples.maxlen + if len(self._first_samples) < first_capacity: + self._first_samples.append(retained) + else: + self._recent_samples.append(retained) + except Exception as error: # pragma: no cover - hardware errors are environment-dependent + self._errors.append(str(error)) + + def report(self) -> dict[str, Any]: + """Return bounded samples, resource distributions, and first/last trends.""" + retained = self._first_samples + list(self._recent_samples) + return { + "enabled": True, + "root_pid": self.root_pid, + "interval_seconds": self.interval_seconds, + "sample_count": self._sample_count, + "gpu_sample_count": sum(1 for sample in retained if sample["gpu_memory_mib"]), + "observed_process_ids": sorted(self._process_ids), + "errors": list(self._errors), + "cpu_rss_mib": { + "distribution": summarize(self._cpu_rss), + "trend": compare_windows(self._cpu_rss), + }, + "gpu_memory_mib": { + index: { + "distribution": summarize(values), + "trend": compare_windows(values), + } + for index, values in sorted(self._gpu_memory.items()) + }, + "retained_samples": retained, + } + + +class RunAccumulator: + """Collect aggregate measurements while bounding retained request records.""" + + def __init__(self, max_records: int) -> None: + self._lock = threading.Lock() + self.max_records = max_records + self.total = 0 + self.succeeded = 0 + self.failed = 0 + self.task_ids: set[str] = set() + self.duplicate_task_ids: set[str] = set() + self.end_to_end: list[float] = [] + self.submit: list[float] = [] + self.accepted_to_terminal: list[float] = [] + self.inference: list[float] = [] + self.peak_memory: list[float] = [] + self.poll_counts: list[float] = [] + self.terminal_statuses: Counter[str] = Counter() + self.policy_statuses: Counter[str] = Counter() + self.failures: deque[dict[str, Any]] = deque(maxlen=max_records) + self.first_successes: list[dict[str, Any]] = [] + self.recent_successes: deque[dict[str, Any]] = deque(maxlen=max_records // 2) + + def add(self, record: dict[str, Any]) -> None: + with self._lock: + self.total += 1 + task_id = record.get("task_id") + if isinstance(task_id, str): + if task_id in self.task_ids: + self.duplicate_task_ids.add(task_id) + self.task_ids.add(task_id) + terminal_status = record.get("terminal_status") + if isinstance(terminal_status, str): + self.terminal_statuses[terminal_status] += 1 + if record["outcome"] == "failed": + self.failed += 1 + self.failures.append(record) + return + + self.succeeded += 1 + self.end_to_end.append(float(record["end_to_end_seconds"])) + self.submit.append(float(record["submit_seconds"])) + self.accepted_to_terminal.append(float(record["accepted_to_terminal_seconds"])) + self.poll_counts.append(float(record["poll_count"])) + if record.get("inference_time_seconds") is not None: + self.inference.append(float(record["inference_time_seconds"])) + if record.get("peak_memory_mb") is not None: + self.peak_memory.append(float(record["peak_memory_mb"])) + self.policy_statuses[str(record["action"]["verification_status"])] += 1 + first_capacity = self.max_records - self.recent_successes.maxlen + if len(self.first_successes) < first_capacity: + self.first_successes.append(record) + else: + self.recent_successes.append(record) + + def report(self, elapsed_seconds: float) -> dict[str, Any]: + retained_successes = self.first_successes + list(self.recent_successes) + retained_successes.sort(key=lambda record: int(record["request_index"])) + failures = sorted(self.failures, key=lambda record: int(record["request_index"])) + return { + "requests": { + "total": self.total, + "succeeded": self.succeeded, + "failed": self.failed, + "success_rate": self.succeeded / self.total if self.total else 0.0, + "unique_task_ids": len(self.task_ids), + "duplicate_task_ids": sorted(self.duplicate_task_ids), + "terminal_statuses": dict(sorted(self.terminal_statuses.items())), + "policy_statuses": dict(sorted(self.policy_statuses.items())), + }, + "elapsed_seconds": elapsed_seconds, + "throughput_requests_per_second": self.succeeded / elapsed_seconds if elapsed_seconds > 0 else 0.0, + "latency_seconds": { + "end_to_end": summarize(self.end_to_end), + "submission": summarize(self.submit), + "accepted_to_terminal": summarize(self.accepted_to_terminal), + "target_inference": summarize(self.inference), + }, + "latency_trend": { + "end_to_end": compare_windows(self.end_to_end), + "target_inference": compare_windows(self.inference), + }, + "poll_count": summarize(self.poll_counts), + "peak_memory_mb": summarize(self.peak_memory), + "retained_records": { + "limit_per_outcome": self.max_records, + "successful": retained_successes, + "failed": failures, + }, + } + + +def run_workload( + config: RequestConfig, + *, + request_count: int | None, + duration_seconds: float | None, + concurrency: int, + max_records: int, +) -> dict[str, Any]: + """Run a closed-loop fixed-count or duration workload.""" + accumulator = RunAccumulator(max_records) + counter = 0 + counter_lock = threading.Lock() + run_started_at = time.perf_counter() + stop_claiming_at = None if duration_seconds is None else run_started_at + duration_seconds + workers = concurrency if request_count is None else min(concurrency, request_count) + barrier = threading.Barrier(workers) + + def claim_request() -> int | None: + nonlocal counter + with counter_lock: + if request_count is not None and counter >= request_count: + return None + if stop_claiming_at is not None and time.perf_counter() >= stop_claiming_at: + return None + index = counter + counter += 1 + return index + + def worker(worker_index: int) -> None: + with _new_session() as session: + barrier.wait() + while (request_index := claim_request()) is not None: + accumulator.add( + execute_request( + session, + config, + request_index=request_index, + worker_index=worker_index, + run_started_at=run_started_at, + ) + ) + + with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="vla-structured-validator") as executor: + futures = [executor.submit(worker, worker_index) for worker_index in range(workers)] + for future in futures: + future.result() + return accumulator.report(time.perf_counter() - run_started_at) + + +def _encode_image(path: Path) -> str: + if not path.is_file(): + raise ValueError(f"camera image does not exist: {path}") + return base64.b64encode(path.read_bytes()).decode("ascii") + + +def _resolve_camera_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]: + fallback = args.image + paths = tuple(path or fallback for path in (args.camera_high, args.camera_left_wrist, args.camera_right_wrist)) + if any(path is None for path in paths): + raise ValueError("provide --image or all three --camera-* paths") + return paths # type: ignore[return-value] + + +def parse_gpu_indexes(value: str) -> set[str]: + """Parse a comma-separated set of physical GPU indexes.""" + indexes = {item.strip() for item in value.split(",") if item.strip()} + if not indexes or any(not item.isdigit() for item in indexes): + raise argparse.ArgumentTypeError("GPU indexes must be a comma-separated list of integers") + return indexes + + +def _package_version() -> str: + try: + return version("telefuser") + except PackageNotFoundError: + return "source" + + +def _git_commit(repo_root: Path) -> str | None: + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo_root, + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (OSError, subprocess.SubprocessError): + return None + return result.stdout.strip() or None + + +def _metric_delta(before: dict[str, Any], after: dict[str, Any]) -> dict[str, Any]: + delta: dict[str, Any] = {} + for section in ("tasks",): + before_section = before.get(section, {}) + after_section = after.get(section, {}) + if isinstance(before_section, dict) and isinstance(after_section, dict): + delta[section] = { + key: after_section[key] - before_section.get(key, 0) + for key in after_section + if isinstance(after_section[key], int | float) and not isinstance(after_section[key], bool) + } + return delta + + +def run_validation(args: argparse.Namespace) -> dict[str, Any]: + """Validate the service and return a reproducible JSON report.""" + if args.concurrency < 1 or args.warmup < 0 or args.max_records < 2: + raise ValueError("concurrency must be positive, warmup non-negative, and max-records at least 2") + if args.requests is not None and args.requests < 1: + raise ValueError("requests must be positive") + if args.duration_seconds is not None and args.duration_seconds <= 0: + raise ValueError("duration-seconds must be positive") + if args.poll_interval_seconds <= 0 or args.http_timeout_seconds <= 0 or args.task_timeout_seconds <= 0: + raise ValueError("poll interval and HTTP/task timeouts must be positive") + if args.expected_horizon < 1 or args.expected_action_dim < 1: + raise ValueError("expected action dimensions must be positive") + if args.resource_interval_seconds <= 0 or args.max_resource_samples < 2: + raise ValueError("resource interval must be positive and max-resource-samples at least 2") + if args.gpu_indexes is not None and args.service_pid is None: + raise ValueError("--gpu-indexes requires --service-pid") + base_url = args.base_url.rstrip("/") + camera_high, camera_left, camera_right = _resolve_camera_paths(args) + payload = { + "task": "vla_action", + "instruction": args.instruction, + "state": args.state_json, + "camera_high": _encode_image(camera_high), + "camera_left_wrist": _encode_image(camera_left), + "camera_right_wrist": _encode_image(camera_right), + "seed": args.seed, + } + config = RequestConfig( + base_url=base_url, + payload=payload, + http_timeout_seconds=args.http_timeout_seconds, + task_timeout_seconds=args.task_timeout_seconds, + poll_interval_seconds=args.poll_interval_seconds, + expected_horizon=args.expected_horizon, + expected_action_dim=args.expected_action_dim, + ) + before = inspect_service(base_url, timeout_seconds=args.http_timeout_seconds) + warmup_records: list[dict[str, Any]] = [] + warmup_started_at = time.perf_counter() + with _new_session() as session: + for index in range(args.warmup): + warmup_records.append( + execute_request( + session, + config, + request_index=index, + worker_index=0, + run_started_at=warmup_started_at, + ) + ) + if any(record["outcome"] != "succeeded" for record in warmup_records): + raise ValidationFailure("at least one warmup request failed") + + request_count = args.requests + if request_count is None and args.duration_seconds is None: + request_count = 1 + resource_sampler: ResourceSampler | None = None + resource_report: dict[str, Any] = {"enabled": False} + if args.service_pid is not None: + if not psutil.pid_exists(args.service_pid): + raise ValueError(f"service PID does not exist: {args.service_pid}") + resource_sampler = ResourceSampler( + args.service_pid, + interval_seconds=args.resource_interval_seconds, + max_samples=args.max_resource_samples, + gpu_indexes=args.gpu_indexes, + ) + resource_sampler.start() + try: + workload = run_workload( + config, + request_count=request_count, + duration_seconds=args.duration_seconds, + concurrency=args.concurrency, + max_records=args.max_records, + ) + finally: + if resource_sampler is not None: + resource_report = resource_sampler.stop() + after = inspect_service(base_url, timeout_seconds=args.http_timeout_seconds) + requests_report = workload["requests"] + checks = { + "service_ready_before": before["ready"].get("ready") is True, + "service_ready_after": after["ready"].get("ready") is True, + "warmup_succeeded": all(record["outcome"] == "succeeded" for record in warmup_records), + "all_measured_requests_succeeded": requests_report["failed"] == 0 and requests_report["total"] > 0, + "task_ids_unique": not requests_report["duplicate_task_ids"], + "resource_samples_collected": (not resource_report["enabled"] or resource_report["sample_count"] > 0), + "gpu_resource_samples_collected": (not resource_report["enabled"] or resource_report["gpu_sample_count"] > 0), + "queue_drained": ( + after["metrics"].get("queue", {}).get("pending") == 0 + and after["metrics"].get("queue", {}).get("processing") == 0 + ), + } + repo_root = Path(__file__).resolve().parents[2] + return { + "schema_version": 1, + "validation": "lingbot_vla_v2_native_structured_api", + "passed": all(checks.values()), + "checks": checks, + "created_at": datetime.now(timezone.utc).isoformat(), + "environment": { + "python": platform.python_version(), + "python_executable": sys.executable, + "platform": platform.platform(), + "telefuser_version": _package_version(), + "telefuser_commit": _git_commit(repo_root), + }, + "target": { + "base_url": base_url, + "transport": "HTTP native TeleFuser asynchronous structured task API", + "metadata": before["metadata"], + "status_before": before["status"], + "status_after": after["status"], + "health_before": before["ready"], + "health_after": after["ready"], + "metrics_before": before["metrics"], + "metrics_after": after["metrics"], + "metrics_delta": _metric_delta(before["metrics"], after["metrics"]), + }, + "workload": { + "mode": "duration" if args.duration_seconds is not None else "fixed_requests", + "requested_requests": request_count, + "requested_duration_seconds": args.duration_seconds, + "concurrency": args.concurrency, + "warmup_requests": args.warmup, + "instruction": args.instruction, + "state_dimension": len(args.state_json), + "seed": args.seed, + "camera_files": { + "high": str(camera_high.resolve()), + "left_wrist": str(camera_left.resolve()), + "right_wrist": str(camera_right.resolve()), + }, + "expected_action_shape": [args.expected_horizon, args.expected_action_dim], + "poll_interval_seconds": args.poll_interval_seconds, + "task_timeout_seconds": args.task_timeout_seconds, + }, + "warmup_records": warmup_records, + "result": workload, + "resources": resource_report, + "interpretation": ( + "This validates service transport, scheduling, and normalized canonical action structure. " + "It does not establish embodiment-specific robot control semantics." + ), + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080") + camera_group = parser.add_argument_group("camera inputs") + camera_group.add_argument("--image", type=Path, help="Fallback image reused for camera inputs not set explicitly.") + camera_group.add_argument("--camera-high", type=Path) + camera_group.add_argument("--camera-left-wrist", type=Path) + camera_group.add_argument("--camera-right-wrist", type=Path) + parser.add_argument("--instruction", default="pick up the red block") + parser.add_argument( + "--state-json", type=parse_state_json, default=parse_state_json("[0,0,0,0,0,0,0,0,0,0,0,0,0,0]") + ) + parser.add_argument("--seed", type=int, default=7) + workload_group = parser.add_mutually_exclusive_group() + workload_group.add_argument("--requests", type=int) + workload_group.add_argument("--duration-seconds", type=float) + parser.add_argument("--concurrency", type=int, default=1) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--poll-interval-seconds", type=float, default=0.1) + parser.add_argument("--http-timeout-seconds", type=float, default=30.0) + parser.add_argument("--task-timeout-seconds", type=float, default=300.0) + parser.add_argument("--expected-horizon", type=int, default=50) + parser.add_argument("--expected-action-dim", type=int, default=55) + parser.add_argument("--max-records", type=int, default=1000) + parser.add_argument( + "--service-pid", + type=int, + help="Optional local TeleFuser parent PID; enables process-tree RSS and GPU memory sampling.", + ) + parser.add_argument("--gpu-indexes", type=parse_gpu_indexes, help="Optional physical GPU indexes to include.") + parser.add_argument("--resource-interval-seconds", type=float, default=1.0) + parser.add_argument("--max-resource-samples", type=int, default=10000) + parser.add_argument("--output", required=True, type=Path) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + report: dict[str, Any] + exit_code = 0 + try: + report = run_validation(args) + if not report["passed"]: + exit_code = 1 + except Exception as error: + report = { + "schema_version": 1, + "validation": "lingbot_vla_v2_native_structured_api", + "passed": False, + "created_at": datetime.now(timezone.utc).isoformat(), + "fatal_error": str(error), + } + exit_code = 1 + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + summary = { + "passed": report["passed"], + "checks": report.get("checks"), + "requests": report.get("result", {}).get("requests"), + "latency_seconds": report.get("result", {}).get("latency_seconds"), + "fatal_error": report.get("fatal_error"), + "artifact": str(args.output), + } + print(json.dumps(summary, indent=2, sort_keys=True)) + raise SystemExit(exit_code) + + +if __name__ == "__main__": + main()