diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 377dcd6e..b094b9fd 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -25,9 +25,10 @@ The parity reference uses [Robbyant/lingbot-vla-v2](https://github.com/Robbyant/ | CUDA Graph | Supported | Dynamic eager prefix with an opt-in fixed-shape action-denoising graph | | Quantization | Partial | Profile-specific release status; see Configuration and Performance | | Native server API | Supported | Asynchronous structured task API and `TFClient` | +| RoboTwin policy protocol | Supported | Standalone persistent MessagePack WebSocket service | | Request replicas | Supported | One complete policy copy per GPU | | Single-policy FSDP, TP, or PP | Unsupported | The integration does not split one policy across GPUs | -| Physical robot action mapping | Unsupported | Output remains in normalized canonical space | +| RoboTwin action mapping | Supported | Unnormalizes canonical output to absolute-position `50 x 14` chunks | ## Requirements @@ -215,6 +216,148 @@ CUDA_VISIBLE_DEVICES=0,1 TF_MODEL_ZOO_PATH=/path/to/model_zoo \ This creates one complete policy per GPU; it does not enable tensor or pipeline parallelism within a policy. +### RoboTwin Policy Server + +The standalone policy server implements the persistent MessagePack WebSocket protocol used by the upstream +`WebsocketClientPolicy`. It is isolated from `telefuser serve`: no TeleFuser API routes, service schemas, or other +model integrations are changed. + +Install the protocol dependency in the TeleFuser inference environment: + +```bash +.venv-vla/bin/python -m pip install -r examples/lingbot_vla_v2/requirements-robotwin.txt +``` + +Start one resident policy process: + +```bash +.venv-vla/bin/python examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py \ + --model-root "$TF_MODEL_ZOO_PATH/lingbot/lingbot-vla-v2-6b" \ + --qwen3vl-root "$TF_MODEL_ZOO_PATH/Qwen3-VL-4B-Instruct" \ + --device cuda:0 --host 0.0.0.0 --port 9330 --use-length 50 +``` + +On NVIDIA H100, this dedicated entrypoint disables cuDNN SDPA before model warmup because the current +PyTorch/cuDNN combination cannot build a valid vision-attention execution plan. Flash, memory-efficient, and math +SDPA remain enabled. The override is process-local and is not applied to other TeleFuser pipelines. + +The server exposes `GET /healthz` and the policy WebSocket at `/`. On connection it sends a MessagePack metadata +frame, including the explicit 16 MiB request limit, then accepts multiple binary MessagePack requests on the same +connection. This matches the upstream client contract: + +```python +from deploy.websocket_client_policy import WebsocketClientPolicy + +policy = WebsocketClientPolicy(host="127.0.0.1", port=9330) +policy.reset("robotwin") +result = policy.infer( + { + "observation.images.cam_high": camera_high, + "observation.images.cam_left_wrist": camera_left_wrist, + "observation.images.cam_right_wrist": camera_right_wrist, + "observation.state": state, + "task": instruction, + } +) +actions = result["action"] # float32 [50, 14] when --use-length=50 +``` + +The initial metadata frame describes protocol version `1.0`, `absolute_qpos` action semantics, `float32` dtype, +horizon, dimension, and the exact dual-arm joint order. Inference requests may include an integer `seed` plus +`request_id` and `episode_id`; the response echoes them and reports decode, lock-wait, pipeline, action-mapping, and +adapter timings. Existing clients may omit all three request fields. + +The endpoint also advertises an additive, latest-wins action scheduler. A client that overlaps simulation and +inference should send a monotonically increasing `sequence_id` within each `episode_id`, plus a positive +`request_ttl_ms`. The server has one GPU worker, retains at most one pending request per connection/episode, and +accepts new observations while inference is running. A newer observation replaces queued work; because an in-flight +CUDA call cannot be cancelled, its result is discarded after completion when it has become stale. Successful +responses use `scheduler_status="completed"`. Responses with `superseded`, `expired`, `stale_sequence`, or +`overloaded` contain `action=None` and a structured `error`; clients must never execute those responses. + +`request_ttl_ms` starts when the H100 server receives the request. Do not compare monotonic timestamps between the +H100 and RTX machines. The RTX client should separately enforce its round-trip deadline and hold the current joint +positions when no fresh action is available. + +Validate this direct endpoint before a simulator is available. This sends reset and repeated inference requests to +the resident model, validates the returned `[H, 14]` action contract, and optionally verifies exact fixed-seed replay: + +```bash +.venv-vla/bin/python -m tools.validation.validate_lingbot_vla_v2_robotwin_ws \ + --host 127.0.0.1 --port 9330 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --max-image-edge 640 \ + --task "pick up the object" --seed 7 --requests 10 \ + --output work_dirs/robotwin_ws_validation/smoke.json +``` + +The validator preserves aspect ratio and downsizes only images whose longest edge exceeds `--max-image-edge`, then +checks the encoded MessagePack request against the limit advertised by the server before sending it. This keeps the +large repository sample representative of normal RoboTwin camera payloads. Add `--require-exact-replay` only when +validating a runtime profile that promises bitwise determinism; BF16 H100 inference is validated with numerical +tolerances rather than identical action hashes. + +Exercise overlapping submissions and stale-action rejection without a simulator: + +```bash +.venv-vla/bin/python -m tools.validation.validate_lingbot_vla_v2_robotwin_ws \ + --host 127.0.0.1 --port 9330 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --task "pick up the object" --seed 7 --requests 3 \ + --request-ttl-ms 5000 --overlap-requests \ + --output work_dirs/robotwin_ws_validation/overlap.json +``` + +This mode sends all observations before receiving responses, requires the newest request to return an action, and +requires at least one older request to be reported as `superseded`. + +Each request runs the existing pipeline, converts normalized canonical `50 x 55` output through the bundled RoboTwin +profile, and returns absolute-position actions in raw RoboTwin order. `--use-length` may truncate the returned chunk; +start with 50 for upstream-equivalent open-loop execution. The adapter accepts episode reset messages but deliberately +rejects runtime checkpoint switching. + +For split-machine deployment, run the model endpoint and the repository-owned XPolicyLab proxy on the H100 inference +host. The proxy does not load a second model; it translates XPolicyLab observations to the direct TeleFuser protocol: + +```bash +cd /data/RoboTwin +bash XPolicyLab/policy/TeleFuser_LingBot_VLA/setup_eval_policy_server.sh \ + RoboTwin lift_pot remote_base arx_x5 joint 0 0 \ + /data/RoboTwin/.venv 19000 0.0.0.0 \ + 127.0.0.1 9330 +``` + +On the remote RTX/Vulkan workstation, use the standard RoboTwin evaluation client and point it at the proxy. No +TeleFuser files or model weights are required on that workstation: + +```bash +cd /data/RoboTwin +bash scripts/eval_policy.sh \ + --bench_name RoboTwin \ + --task_name lift_pot \ + --env_cfg_type arx_x5 \ + --policy_name TeleFuser_LingBot_VLA \ + --host INFERENCE_HOST --port 19000 --protocol ws \ + --eval_batch false --root_dir /data/RoboTwin --device_id 0 \ + --additional_info ckpt_name=remote_base,action_type=joint \ + --seed 0 --task_config demo_clean --test_num 1 +``` + +The current XPolicyLab proxy calls `infer()` synchronously, so it remains compatible but does not yet overlap action +execution with inference. Full overlap requires an incremental RTX-side change: execute chunk N while submitting a +newer observation for chunk N+1, keep only the newest completed chunk in an atomic action buffer, and apply the same +sequence/deadline checks before execution. That simulator-side change is outside this repository and is not required +for the no-simulation server validation above. + +Keep ports `9330` and `19000` on a trusted private network or an SSH/VPN tunnel. These WebSocket endpoints do not +provide authentication or transport encryption. The direct validator covers preprocessing, inference, mapping, and +the inner WebSocket contract; only the RTX smoke episode can additionally establish XPolicyLab translation and one +real SAPIEN simulation step. + +The base checkpoint remains marked `unverified_official_6b_base`. This endpoint establishes preprocessing, inference, +action mapping, transport, and simulator execution continuity; it does not establish RoboTwin task success without +an embodiment-validated checkpoint. + ## Validation The repository includes strict upstream parity, runtime, quantization, structured-service, fault, and AIPerf diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py index e358863d..9095ad17 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py @@ -106,6 +106,7 @@ def run_structured( camera_left_wrist: str, camera_right_wrist: str, seed: int | None = None, + stop_event: Any | None = None, **_: Any, ) -> dict[str, Any]: """Return one JSON-serializable canonical normalized action chunk.""" @@ -122,5 +123,6 @@ def run_structured( request, max_image_bytes=int(PPL_CONFIG["max_image_bytes"]), max_image_pixels=int(PPL_CONFIG["max_image_pixels"]), + stop_event=stop_event, ) return response.model_dump(mode="json") diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py new file mode 100644 index 00000000..62905647 --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py @@ -0,0 +1,415 @@ +"""Serve LingBot-VLA v2 through the upstream RoboTwin policy protocol.""" + +from __future__ import annotations + +import asyncio +import contextlib +import operator +import threading +import time +from contextlib import asynccontextmanager +from typing import Any, Mapping, Protocol + +import click +import msgpack +import numpy as np +import torch +import uvicorn +from fastapi import FastAPI, WebSocket, WebSocketDisconnect + +from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_CAMERA_KEYS, + LingBotVlaV2Observation, + RobotWinProfile, +) +from telefuser.pipelines.lingbot_vla_v2.action_scheduler import ActionChunkScheduler +from telefuser.pipelines.lingbot_vla_v2.runtime import ( + LINGBOT_VLA_V2_QUANTIZATION_CHOICES, + get_lingbot_vla_v2_pipeline, +) +from telefuser.utils.logging import logger + +ROBOTWIN_PROTOCOL_VERSION = "1.0" +ROBOTWIN_ACTION_TYPE = "absolute_qpos" +ROBOTWIN_ACTION_DTYPE = "float32" +ROBOTWIN_MAX_REQUEST_BYTES = 16 * 1024 * 1024 +ROBOTWIN_ACTION_ORDER = ( + "left_arm_joint_0", + "left_arm_joint_1", + "left_arm_joint_2", + "left_arm_joint_3", + "left_arm_joint_4", + "left_arm_joint_5", + "left_gripper", + "right_arm_joint_0", + "right_arm_joint_1", + "right_arm_joint_2", + "right_arm_joint_3", + "right_arm_joint_4", + "right_arm_joint_5", + "right_gripper", +) +_TRACE_ID_FIELDS = ("request_id", "episode_id") + + +class _Pipeline(Protocol): + config: Any + + def __call__(self, observation: LingBotVlaV2Observation, seed: int | None = None) -> Any: ... + + def close(self) -> None: ... + + +def _pack_numpy(value: Any) -> Any: + """Encode NumPy values using the upstream msgpack_numpy wire format.""" + if isinstance(value, (np.ndarray, np.generic)) and value.dtype.kind in ("V", "O", "c"): + raise ValueError(f"unsupported NumPy dtype: {value.dtype}") + if isinstance(value, np.ndarray): + return { + b"__ndarray__": True, + b"data": value.tobytes(), + b"dtype": value.dtype.str, + b"shape": value.shape, + } + if isinstance(value, np.generic): + return { + b"__npgeneric__": True, + b"data": value.item(), + b"dtype": value.dtype.str, + } + raise TypeError(f"cannot encode value of type {type(value)!r}") + + +def _unpack_numpy(value: dict[Any, Any]) -> Any: + """Decode NumPy values produced by the upstream msgpack_numpy helper.""" + if b"__ndarray__" in value: + return np.ndarray( + buffer=value[b"data"], + dtype=np.dtype(value[b"dtype"]), + shape=value[b"shape"], + ) + if b"__npgeneric__" in value: + return np.dtype(value[b"dtype"]).type(value[b"data"]) + return value + + +def pack_message(payload: Mapping[str, Any]) -> bytes: + """Pack one RoboTwin policy protocol message.""" + return msgpack.packb(dict(payload), default=_pack_numpy) + + +def unpack_message(payload: bytes) -> dict[str, Any]: + """Unpack and validate one RoboTwin policy protocol message.""" + decoded = msgpack.unpackb(payload, object_hook=_unpack_numpy, raw=False) + if not isinstance(decoded, dict): + raise ValueError("RoboTwin request must be a MessagePack object") + return decoded + + +def _optional_seed(request: Mapping[str, Any]) -> int | None: + value = request.get("seed") + if value is None: + return None + if isinstance(value, bool): + raise ValueError("seed must be an integer") + try: + return operator.index(value) + except TypeError as error: + raise ValueError("seed must be an integer") from error + + +def _trace_fields(request: Mapping[str, Any]) -> dict[str, str | int]: + fields: dict[str, str | int] = {} + for name in _TRACE_ID_FIELDS: + value = request.get(name) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, str | int) or isinstance(value, str) and not value: + raise ValueError(f"{name} must be a non-empty string or integer") + fields[name] = value + return fields + + +class RobotWinPolicyAdapter: + """Translate upstream RoboTwin observations to the TeleFuser VLA SDK.""" + + def __init__( + self, + pipeline: _Pipeline, + *, + profile: RobotWinProfile | None = None, + use_length: int = 50, + ) -> None: + if not 1 <= use_length <= 50: + raise ValueError(f"use_length must be in [1, 50], got {use_length}") + self.pipeline = pipeline + self.profile = profile or pipeline.config.robot_profile + self.use_length = use_length + self._lock = threading.Lock() + + @property + def metadata(self) -> dict[str, Any]: + """Describe the action contract sent when a client connects.""" + return { + "protocol_version": ROBOTWIN_PROTOCOL_VERSION, + "robot_profile": self.profile.name, + "action_type": ROBOTWIN_ACTION_TYPE, + "action_horizon": self.use_length, + "action_dim": self.profile.raw_state_dim, + "action_dtype": ROBOTWIN_ACTION_DTYPE, + "action_order": list(ROBOTWIN_ACTION_ORDER), + "max_request_bytes": ROBOTWIN_MAX_REQUEST_BYTES, + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + } + + def infer(self, request: Mapping[str, Any]) -> dict[str, Any]: + """Return one absolute-position RoboTwin action chunk.""" + trace_fields = _trace_fields(request) + if request.get("reset", False): + return {**self._reset(request), **trace_fields} + + missing = [key for key in (*ROBOTWIN_CAMERA_KEYS, "observation.state", "task") if key not in request] + if missing: + raise ValueError(f"RoboTwin observation is missing fields: {missing}") + + observation = LingBotVlaV2Observation( + task=request["task"], + state=request["observation.state"], + images={key: request[key] for key in ROBOTWIN_CAMERA_KEYS}, + ) + seed = _optional_seed(request) + adapter_started_at = time.monotonic() + with self._lock: + lock_wait_ms = (time.monotonic() - adapter_started_at) * 1000.0 + pipeline_started_at = time.monotonic() + canonical_chunk = self.pipeline(observation, seed=seed) + pipeline_ms = (time.monotonic() - pipeline_started_at) * 1000.0 + mapping_started_at = time.monotonic() + action_chunk = self.profile.structure_actions( + canonical_chunk.canonical_normalized_actions, + ) + action_mapping_ms = (time.monotonic() - mapping_started_at) * 1000.0 + if action_chunk.horizon < self.use_length: + raise RuntimeError( + f"policy returned horizon {action_chunk.horizon}, shorter than use_length={self.use_length}" + ) + actions = np.ascontiguousarray(action_chunk.raw_actions[: self.use_length].numpy(), dtype=np.float32) + expected_shape = (self.use_length, self.profile.raw_state_dim) + if actions.shape != expected_shape: + raise RuntimeError(f"mapped actions must have shape {expected_shape}, got {actions.shape}") + if not np.isfinite(actions).all(): + raise RuntimeError("mapped actions must contain only finite values") + response: dict[str, Any] = { + "action": actions, + "policy_verified": canonical_chunk.policy_verified, + "verification_status": canonical_chunk.verification_status, + "server_timing": { + "lock_wait_ms": lock_wait_ms, + "pipeline_ms": pipeline_ms, + "action_mapping_ms": action_mapping_ms, + "adapter_total_ms": (time.monotonic() - adapter_started_at) * 1000.0, + }, + **trace_fields, + } + if seed is not None: + response["seed"] = seed + return response + + def _reset(self, request: Mapping[str, Any]) -> dict[str, Any]: + robot_name = request.get("robo_name", self.profile.name) + if robot_name != self.profile.name: + raise ValueError(f"unsupported robot profile: {robot_name!r}") + if request.get("path_to_pi_model") not in (None, ""): + raise ValueError("runtime checkpoint switching is not supported") + return {"action": None} + + def close(self) -> None: + """Release resources owned by the resident policy.""" + self.pipeline.close() + + +def create_robotwin_app( + adapter: RobotWinPolicyAdapter, + *, + max_pending_sessions: int = 32, +) -> FastAPI: + """Create a standalone app compatible with upstream WebsocketClientPolicy.""" + scheduler = ActionChunkScheduler(adapter.infer, max_pending_sessions=max_pending_sessions) + + @asynccontextmanager + async def lifespan(_app: FastAPI): + await scheduler.start() + try: + yield + finally: + await scheduler.close() + + app = FastAPI(title="LingBot-VLA v2 RoboTwin Policy", lifespan=lifespan) + + @app.get("/healthz") + async def healthz() -> dict[str, str]: + return {"status": "ok"} + + @app.websocket("/") + async def policy_socket(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_bytes(pack_message({**adapter.metadata, **scheduler.metadata})) + connection_key = str(id(websocket)) + session_keys: set[str] = set() + delivery_tasks: set[asyncio.Task[None]] = set() + send_lock = asyncio.Lock() + previous_total_ms: float | None = None + + async def deliver( + response_future: asyncio.Future[dict[str, Any]], + *, + decode_ms: float, + round_started_at: float, + ) -> None: + nonlocal previous_total_ms + try: + response = dict(await response_future) + server_timing = dict(response.get("server_timing", {})) + server_timing["decode_ms"] = decode_ms + response["server_timing"] = server_timing + async with send_lock: + if previous_total_ms is not None: + server_timing["prev_total_ms"] = previous_total_ms + await websocket.send_bytes(pack_message(response)) + previous_total_ms = (time.monotonic() - round_started_at) * 1000.0 + except asyncio.CancelledError: + raise + except WebSocketDisconnect: + return + except Exception as error: + logger.exception("LingBot-VLA v2 RoboTwin request failed") + async with send_lock: + with contextlib.suppress(WebSocketDisconnect, RuntimeError): + await websocket.send_text(f"{type(error).__name__}: {error}") + await websocket.close(code=1011) + + try: + while True: + message = await websocket.receive() + if message["type"] == "websocket.disconnect": + break + payload = message.get("bytes") + if payload is None: + raise ValueError("RoboTwin requests must use binary MessagePack frames") + + round_started_at = time.monotonic() + decode_started_at = time.monotonic() + request = unpack_message(payload) + decode_ms = (time.monotonic() - decode_started_at) * 1000.0 + trace_fields = _trace_fields(request) + episode_id = trace_fields.get("episode_id", "default") + session_key = f"{connection_key}:{episode_id!r}" + if request.get("reset", False): + for previous_session_key in session_keys - {session_key}: + scheduler.release_session(previous_session_key) + session_keys.intersection_update({session_key}) + session_keys.add(session_key) + response_future = scheduler.submit(request, session_key=session_key) + task = asyncio.create_task( + deliver( + response_future, + decode_ms=decode_ms, + round_started_at=round_started_at, + ) + ) + delivery_tasks.add(task) + task.add_done_callback(delivery_tasks.discard) + except WebSocketDisconnect: + return + except Exception as error: + logger.exception("LingBot-VLA v2 RoboTwin request failed") + with contextlib.suppress(WebSocketDisconnect, RuntimeError): + await websocket.send_text(f"{type(error).__name__}: {error}") + await websocket.close(code=1011) + finally: + for session_key in session_keys: + scheduler.release_session(session_key) + for task in delivery_tasks: + task.cancel() + if delivery_tasks: + await asyncio.gather(*delivery_tasks, return_exceptions=True) + + return app + + +def _configure_h100_sdpa_backends(device: str) -> None: + """Avoid unsupported cuDNN SDPA plans in the isolated H100 policy process.""" + resolved_device = torch.device(device) + if resolved_device.type != "cuda" or not torch.cuda.is_available(): + return + if "H100" not in torch.cuda.get_device_name(resolved_device): + return + + if hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) + if hasattr(torch.backends.cuda, "enable_flash_sdp"): + torch.backends.cuda.enable_flash_sdp(True) + if hasattr(torch.backends.cuda, "enable_math_sdp"): + torch.backends.cuda.enable_math_sdp(True) + if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"): + torch.backends.cuda.enable_mem_efficient_sdp(True) + logger.info("Disabled cuDNN SDPA for the LingBot-VLA v2 H100 policy process") + + +@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("--host", default="0.0.0.0", show_default=True) +@click.option("--port", default=9330, show_default=True, type=click.IntRange(1, 65535)) +@click.option("--device", default="cuda:0", show_default=True) +@click.option("--use-length", default=50, show_default=True, type=click.IntRange(1, 50)) +@click.option( + "--max-pending-sessions", + default=32, + show_default=True, + type=click.IntRange(1), + help="Bound the number of sessions waiting behind the GPU worker", +) +@click.option("--cuda-graph", is_flag=True, help="Enable fixed-shape CUDA Graph inference") +@click.option( + "--quantization", + type=click.Choice(LINGBOT_VLA_V2_QUANTIZATION_CHOICES), + default=None, +) +def main( + model_root: str, + qwen3vl_root: str, + host: str, + port: int, + device: str, + use_length: int, + max_pending_sessions: int, + cuda_graph: bool, + quantization: str | None, +) -> None: + """Start one resident LingBot-VLA v2 policy for a RoboTwin client.""" + _configure_h100_sdpa_backends(device) + pipeline = get_lingbot_vla_v2_pipeline( + model_root, + qwen3vl_root, + device=device, + warmup=True, + quantization=quantization, + cuda_graph=cuda_graph, + ) + adapter = RobotWinPolicyAdapter(pipeline, use_length=use_length) + try: + uvicorn.run( + create_robotwin_app(adapter, max_pending_sessions=max_pending_sessions), + host=host, + port=port, + workers=1, + ws_max_size=ROBOTWIN_MAX_REQUEST_BYTES, + ) + finally: + adapter.close() + + +if __name__ == "__main__": + main() diff --git a/examples/lingbot_vla_v2/requirements-robotwin.txt b/examples/lingbot_vla_v2/requirements-robotwin.txt new file mode 100644 index 00000000..7beb7e9c --- /dev/null +++ b/examples/lingbot_vla_v2/requirements-robotwin.txt @@ -0,0 +1,2 @@ +msgpack>=1.0,<2.0 +websockets>=12.0,<16.0 diff --git a/telefuser/models/lingbot_vla_v2.py b/telefuser/models/lingbot_vla_v2.py index 5525cf7c..b6c129ff 100644 --- a/telefuser/models/lingbot_vla_v2.py +++ b/telefuser/models/lingbot_vla_v2.py @@ -1391,8 +1391,11 @@ def sample_actions( state, noise=None, image_grid_thw=None, + stop_event=None, ) -> Tensor: """Do a full Qwen3-VL inference forward and compute the action.""" + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") bsize = state.shape[0] device = state.device dtype = state.dtype @@ -1414,7 +1417,7 @@ def sample_actions( raise RuntimeError("LingBot-VLA v2 CUDA Graph and torch.compile cannot be enabled together") if self._cuda_graph_runner is None: self._cuda_graph_runner = LingBotVlaV2CudaGraphs(self) - return self._cuda_graph_runner.run( + result = self._cuda_graph_runner.run( images, img_masks, lang_tokens, @@ -1424,6 +1427,9 @@ def sample_actions( image_grid_thw, ) + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") + return result prefix_pad_masks, prefix_position_ids, past_key_values = self.build_prefix_cache( images, img_masks, @@ -1433,6 +1439,8 @@ def sample_actions( ) dt = torch.tensor(-1.0 / self.config.num_steps, dtype=dtype, device=device) + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") x_t = noise time = torch.tensor(1.0, dtype=dtype, device=device) predict_velocity_fn = self.predict_velocity @@ -1448,6 +1456,8 @@ def sample_actions( self._compiled_predict_velocity = predict_velocity_fn for _ in range(int(self.config.num_steps)): + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") expanded_time = time.expand(bsize) v_t = predict_velocity_fn( state, diff --git a/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py b/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py new file mode 100644 index 00000000..73439d4f --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py @@ -0,0 +1,271 @@ +"""Bounded latest-wins scheduling for LingBot-VLA v2 action chunks.""" + +from __future__ import annotations + +import asyncio +import math +import operator +import time +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any, Callable, Mapping + + +@dataclass(slots=True) +class _ScheduledAction: + request: Mapping[str, Any] + session_key: str + generation: int + received_at: float + deadline_at: float | None + future: asyncio.Future[dict[str, Any]] + + +class ActionChunkScheduler: + """Serialize GPU inference while retaining only the newest pending chunk per session.""" + + def __init__( + self, + infer: Callable[[Mapping[str, Any]], dict[str, Any]], + *, + max_pending_sessions: int = 32, + ) -> None: + if max_pending_sessions < 1: + raise ValueError("max_pending_sessions must be positive") + self._infer = infer + self._max_pending_sessions = max_pending_sessions + self._pending: OrderedDict[str, _ScheduledAction] = OrderedDict() + self._latest_generation: dict[str, int] = {} + self._latest_sequence: dict[str, int] = {} + self._wake = asyncio.Event() + self._worker: asyncio.Task[None] | None = None + self._closed = False + + @property + def metadata(self) -> dict[str, Any]: + """Describe the additive scheduling controls accepted by the endpoint.""" + return { + "scheduling": { + "mode": "latest_wins", + "max_pending_per_session": 1, + "max_pending_sessions": self._max_pending_sessions, + "sequence_field": "sequence_id", + "ttl_field": "request_ttl_ms", + "inflight_cancellation": False, + } + } + + async def start(self) -> None: + """Start the single inference worker on the current event loop.""" + if self._worker is not None: + return + if self._closed: + raise RuntimeError("action scheduler is closed") + self._worker = asyncio.create_task(self._run(), name="lingbot-vla-v2-action-scheduler") + + def submit(self, request: Mapping[str, Any], *, session_key: str) -> asyncio.Future[dict[str, Any]]: + """Admit one request and return a future without waiting for inference.""" + if self._worker is None or self._closed: + raise RuntimeError("action scheduler is not running") + if not session_key: + raise ValueError("session_key must be non-empty") + + loop = asyncio.get_running_loop() + future: asyncio.Future[dict[str, Any]] = loop.create_future() + received_at = time.monotonic() + ttl_ms = self._optional_ttl_ms(request) + sequence_id = self._optional_sequence_id(request) + previous = self._pending.get(session_key) + if previous is None and len(self._pending) >= self._max_pending_sessions: + future.set_result( + self._discarded_response( + request, + status="overloaded", + message="the action scheduler has no free pending-session slot", + ) + ) + return future + if sequence_id is not None: + latest_sequence = self._latest_sequence.get(session_key) + if latest_sequence is not None and sequence_id <= latest_sequence: + future.set_result( + self._discarded_response( + request, + status="stale_sequence", + message=f"sequence_id={sequence_id} is not newer than {latest_sequence}", + ) + ) + return future + self._latest_sequence[session_key] = sequence_id + + generation = self._latest_generation.get(session_key, 0) + 1 + self._latest_generation[session_key] = generation + + deadline_at = None if ttl_ms is None else received_at + ttl_ms / 1000.0 + job = _ScheduledAction( + request=dict(request), + session_key=session_key, + generation=generation, + received_at=received_at, + deadline_at=deadline_at, + future=future, + ) + if previous is not None: + self._resolve( + previous, + self._discarded_response( + previous.request, + status="superseded", + message="a newer observation replaced this pending action request", + ), + ) + self._pending[session_key] = job + self._wake.set() + return future + + def release_session(self, session_key: str) -> None: + """Discard pending work and sequence state for a disconnected session.""" + pending = self._pending.pop(session_key, None) + if pending is not None: + self._resolve( + pending, + self._discarded_response( + pending.request, + status="session_closed", + message="the client session closed before inference", + ), + ) + self._latest_generation.pop(session_key, None) + self._latest_sequence.pop(session_key, None) + + async def close(self) -> None: + """Drain an in-flight call and reject work that has not started.""" + if self._closed: + return + self._closed = True + for job in self._pending.values(): + self._resolve( + job, + self._discarded_response( + job.request, + status="server_stopping", + message="the action scheduler is stopping", + ), + ) + self._pending.clear() + self._wake.set() + if self._worker is not None: + await self._worker + self._worker = None + + async def _run(self) -> None: + while True: + await self._wake.wait() + if self._closed and not self._pending: + return + if not self._pending: + self._wake.clear() + continue + + _, job = self._pending.popitem(last=False) + if not self._pending: + self._wake.clear() + if job.future.done(): + continue + discard = self._discard_reason(job) + if discard is not None: + self._resolve(job, discard) + continue + + inference_started_at = time.monotonic() + try: + response = await asyncio.to_thread(self._infer, job.request) + except Exception as error: + if not job.future.done(): + job.future.set_exception(error) + continue + inference_ms = (time.monotonic() - inference_started_at) * 1000.0 + + discard = self._discard_reason(job) + if discard is not None: + self._resolve(job, discard) + continue + completed_at = time.monotonic() + response = dict(response) + server_timing = dict(response.get("server_timing", {})) + server_timing.update( + queue_wait_ms=(inference_started_at - job.received_at) * 1000.0, + infer_ms=inference_ms, + scheduler_total_ms=(completed_at - job.received_at) * 1000.0, + ) + response.update(scheduler_status="completed", server_timing=server_timing) + for field in ("request_id", "episode_id", "sequence_id"): + if field in job.request: + response.setdefault(field, job.request[field]) + if job.deadline_at is not None: + response["request_ttl_ms"] = (job.deadline_at - job.received_at) * 1000.0 + self._resolve(job, response) + + def _discard_reason(self, job: _ScheduledAction) -> dict[str, Any] | None: + if job.deadline_at is not None and time.monotonic() >= job.deadline_at: + return self._discarded_response( + job.request, + status="expired", + message="the action request exceeded request_ttl_ms", + ) + if self._latest_generation.get(job.session_key) != job.generation: + return self._discarded_response( + job.request, + status="superseded", + message="a newer observation superseded this action result", + ) + return None + + @staticmethod + def _resolve(job: _ScheduledAction, response: dict[str, Any]) -> None: + if not job.future.done(): + job.future.set_result(response) + + @staticmethod + def _optional_sequence_id(request: Mapping[str, Any]) -> int | None: + value = request.get("sequence_id") + if value is None: + return None + if isinstance(value, bool): + raise ValueError("sequence_id must be a non-negative integer") + try: + sequence_id = operator.index(value) + except TypeError as error: + raise ValueError("sequence_id must be a non-negative integer") from error + if sequence_id < 0: + raise ValueError("sequence_id must be a non-negative integer") + return sequence_id + + @staticmethod + def _optional_ttl_ms(request: Mapping[str, Any]) -> float | None: + value = request.get("request_ttl_ms") + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError("request_ttl_ms must be a positive finite number") + ttl_ms = float(value) + if not math.isfinite(ttl_ms) or ttl_ms <= 0: + raise ValueError("request_ttl_ms must be a positive finite number") + return ttl_ms + + @staticmethod + def _discarded_response( + request: Mapping[str, Any], + *, + status: str, + message: str, + ) -> dict[str, Any]: + response: dict[str, Any] = { + "action": None, + "scheduler_status": status, + "error": {"code": status, "message": message}, + } + for field in ("request_id", "episode_id", "sequence_id"): + if field in request: + response[field] = request[field] + return response diff --git a/telefuser/pipelines/lingbot_vla_v2/pipeline.py b/telefuser/pipelines/lingbot_vla_v2/pipeline.py index fc3cc646..ffc6c004 100644 --- a/telefuser/pipelines/lingbot_vla_v2/pipeline.py +++ b/telefuser/pipelines/lingbot_vla_v2/pipeline.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import Any import torch @@ -76,9 +77,10 @@ def predict( self, inputs: LingBotVlaV2Inputs, seed: int | None = None, + stop_event: Any | None = None, ) -> LingBotVlaV2CanonicalActionChunk: """Run prepared tensors and return normalized canonical actions.""" - actions = self.policy_stage.process(inputs, seed=seed) + actions = self.policy_stage.process(inputs, seed=seed, stop_event=stop_event) if actions.shape[0] != 1: raise RuntimeError(f"LingBot-VLA v2 pipeline expects batch size 1, got {actions.shape[0]}") canonical_actions = actions[0] @@ -97,9 +99,15 @@ def __call__( self, observation: LingBotVlaV2Observation, seed: int | None = None, + stop_event: Any | None = None, ) -> LingBotVlaV2CanonicalActionChunk: """Predict one normalized canonical action chunk.""" - return self.predict(self.input_processor.prepare(observation), seed=seed) + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") + inputs = self.input_processor.prepare(observation) + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") + return self.predict(inputs, seed=seed, stop_event=stop_event) def prepare_for_inference(self) -> None: """Move the policy to its target device before the service becomes ready.""" diff --git a/telefuser/pipelines/lingbot_vla_v2/policy.py b/telefuser/pipelines/lingbot_vla_v2/policy.py index 51a5b09e..948b5ddf 100644 --- a/telefuser/pipelines/lingbot_vla_v2/policy.py +++ b/telefuser/pipelines/lingbot_vla_v2/policy.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Any + import torch from telefuser.core.base_stage import BaseStage, with_model_offload @@ -31,8 +33,15 @@ def _validate_parallelism(self) -> None: @with_model_offload(["policy"]) @torch.inference_mode() @with_metrics - def process(self, inputs: LingBotVlaV2Inputs, seed: int | None = None) -> torch.Tensor: + def process( + self, + inputs: LingBotVlaV2Inputs, + seed: int | None = None, + stop_event: Any | None = None, + ) -> torch.Tensor: """Return a CPU float32 normalized action chunk with shape ``[1, H, 55]``.""" + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") device = self.device dtype = self.torch_dtype tensors = { @@ -55,7 +64,9 @@ def process(self, inputs: LingBotVlaV2Inputs, seed: int | None = None) -> torch. dtype=dtype, generator=generator, ) - actions = self.policy.sample_actions(**tensors, noise=noise) + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") + actions = self.policy.sample_actions(**tensors, noise=noise, stop_event=stop_event) 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 diff --git a/telefuser/pipelines/lingbot_vla_v2/service.py b/telefuser/pipelines/lingbot_vla_v2/service.py index 79efb291..98c6da38 100644 --- a/telefuser/pipelines/lingbot_vla_v2/service.py +++ b/telefuser/pipelines/lingbot_vla_v2/service.py @@ -6,7 +6,7 @@ import binascii import io import math -from typing import Protocol +from typing import Any, Protocol from PIL import Image, UnidentifiedImageError from pydantic import BaseModel, ConfigDict, Field, field_validator @@ -16,6 +16,7 @@ from .robot_profile import ROBOTWIN_CAMERA_KEYS DEFAULT_MAX_IMAGE_PIXELS = 16 * 1024 * 1024 +DEFAULT_MAX_IMAGE_BYTES = 10 * 1024 * 1024 class LingBotVlaV2ActionRequest(BaseModel): @@ -64,6 +65,7 @@ def __call__( self, observation: LingBotVlaV2Observation, seed: int | None = None, + stop_event: Any | None = None, ) -> LingBotVlaV2CanonicalActionChunk: ... @@ -109,19 +111,21 @@ def predict_lingbot_vla_v2_action( *, max_image_bytes: int, max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS, + stop_event: Any | None = None, ) -> 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, - max_image_pixels=max_image_pixels, - ) - for key, value in zip(ROBOTWIN_CAMERA_KEYS, encoded_images, strict=True) - } + images = _decode_request_images( + request, + max_image_bytes=max_image_bytes, + max_image_pixels=max_image_pixels, + ) + if stop_event is not None and stop_event.is_set(): + raise RuntimeError("LingBot-VLA v2 inference cancelled") observation = LingBotVlaV2Observation(task=request.task, state=request.state, images=images) - chunk = pipeline(observation, seed=request.seed) + if stop_event is None: + chunk = pipeline(observation, seed=request.seed) + else: + chunk = pipeline(observation, seed=request.seed, stop_event=stop_event) return LingBotVlaV2ActionResponse( canonical_normalized_actions=chunk.canonical_normalized_actions.tolist(), horizon=chunk.horizon, @@ -130,3 +134,34 @@ def predict_lingbot_vla_v2_action( policy_verified=chunk.policy_verified, verification_status=chunk.verification_status, ) + + +def _decode_request_images( + request: LingBotVlaV2ActionRequest, + *, + max_image_bytes: int, + max_image_pixels: int, +) -> dict[str, Image.Image]: + encoded_images = (request.camera_high, request.camera_left_wrist, request.camera_right_wrist) + return { + key: _decode_image( + value, + max_image_bytes=max_image_bytes, + max_image_pixels=max_image_pixels, + ) + for key, value in zip(ROBOTWIN_CAMERA_KEYS, encoded_images, strict=True) + } + + +def validate_lingbot_vla_v2_action_request( + request: LingBotVlaV2ActionRequest, + *, + max_image_bytes: int = DEFAULT_MAX_IMAGE_BYTES, + max_image_pixels: int = DEFAULT_MAX_IMAGE_PIXELS, +) -> None: + """Validate all encoded camera images without invoking the policy.""" + _decode_request_images( + request, + max_image_bytes=max_image_bytes, + max_image_pixels=max_image_pixels, + ) diff --git a/telefuser/service/api/routers/tasks.py b/telefuser/service/api/routers/tasks.py index e48beb62..fd5c063e 100644 --- a/telefuser/service/api/routers/tasks.py +++ b/telefuser/service/api/routers/tasks.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio import json from pathlib import Path from typing import TYPE_CHECKING, Any @@ -110,6 +111,7 @@ async def create_structured_task(self, message: StructuredTaskRequest) -> Struct return await self.api.task_app_service.submit_structured( message, explicit_fields=set(getattr(message, "model_fields_set", set())), + validate_inputs=self._validate_structured_message_inputs, ) except HTTPException: raise @@ -404,6 +406,41 @@ def _validate_message_inputs(self, message: TaskRequest, contract: dict[str, Any ref_video_path=message.ref_video_path, ) + async def _validate_structured_message_inputs( + self, + message: StructuredTaskRequest, + contract: dict[str, Any] | None, + ) -> None: + if message.task != "vla_action": + return + + from telefuser.pipelines.lingbot_vla_v2.service import ( + DEFAULT_MAX_IMAGE_BYTES, + DEFAULT_MAX_IMAGE_PIXELS, + LingBotVlaV2ActionRequest, + validate_lingbot_vla_v2_action_request, + ) + + try: + request = LingBotVlaV2ActionRequest( + task=getattr(message, "instruction", None), + state=getattr(message, "state", None), + camera_high=getattr(message, "camera_high", None), + camera_left_wrist=getattr(message, "camera_left_wrist", None), + camera_right_wrist=getattr(message, "camera_right_wrist", None), + seed=getattr(message, "seed", None), + ) + await asyncio.to_thread( + validate_lingbot_vla_v2_action_request, + request, + max_image_bytes=DEFAULT_MAX_IMAGE_BYTES, + max_image_pixels=DEFAULT_MAX_IMAGE_PIXELS, + ) + except ValidationError as error: + raise HTTPException(status_code=422, detail=error.errors(include_url=False)) from error + except ValueError as error: + raise HTTPException(status_code=422, detail=str(error)) from error + def _collect_available_inputs( self, *, diff --git a/telefuser/service/api/task_application_service.py b/telefuser/service/api/task_application_service.py index 70b16d08..fd24ecc6 100644 --- a/telefuser/service/api/task_application_service.py +++ b/telefuser/service/api/task_application_service.py @@ -5,7 +5,7 @@ import asyncio import time import uuid -from collections.abc import Callable +from collections.abc import Awaitable, Callable from pathlib import Path from typing import TYPE_CHECKING, Any @@ -61,6 +61,7 @@ async def submit_structured( message: StructuredTaskRequest, *, explicit_fields: set[str], + validate_inputs: Callable[[StructuredTaskRequest, dict[str, Any] | None], Awaitable[None] | None] | None = None, ensure_processing: bool = True, ) -> StructuredTaskResponse: """Validate and enqueue a task whose result is returned as JSON.""" @@ -75,6 +76,10 @@ async def submit_structured( ) apply_task_contract_defaults(message, task_contract=contract, explicit_fields=explicit_fields) validate_required_task_parameters(message, task_contract=contract) + if validate_inputs is not None: + validation = validate_inputs(message, contract) + if validation is not None: + await validation task_id = self.api.task_manager.create_task(message) message.task_id = task_id diff --git a/telefuser/service/core/pipeline_runner.py b/telefuser/service/core/pipeline_runner.py index 09d5fddf..21f39e51 100644 --- a/telefuser/service/core/pipeline_runner.py +++ b/telefuser/service/core/pipeline_runner.py @@ -16,6 +16,7 @@ import asyncio import inspect import os +import threading from dataclasses import dataclass from pathlib import Path from types import ModuleType @@ -52,6 +53,7 @@ def _select_kwargs( *, task_data: dict[str, Any], module: ModuleType | None, + stop_event: Any | None = None, ) -> dict[str, Any]: """Build kwargs for calling run_with_file based on signature inspection. @@ -65,8 +67,12 @@ def _select_kwargs( params = list(sig.parameters.values()) accepts_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params) + call_data = dict(task_data) + if stop_event is not None and "stop_event" in sig.parameters: + call_data["stop_event"] = stop_event + if accepts_var_kw: - return dict(task_data) + return call_data kwargs: dict[str, Any] = {} for p in params: @@ -74,8 +80,8 @@ def _select_kwargs( continue name = p.name - if name in task_data: - kwargs[name] = task_data[name] + if name in call_data: + kwargs[name] = call_data[name] continue # Common aliases used in examples/orchestrator runners. @@ -117,6 +123,7 @@ def __init__( self._started = False self._start_lock = asyncio.Lock() + self._run_lock = asyncio.Lock() async def ensure_started(self) -> None: """Best-effort pipeline startup (once).""" @@ -140,14 +147,24 @@ async def shutdown(self) -> None: if not self._started: return - if hasattr(self._pipeline, "astop"): - 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) + async with self._run_lock: + if hasattr(self._pipeline, "astop"): + 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 + self._started = False + + def _release_run_lock(self, task: asyncio.Task[Any]) -> None: + """Release a timed-out invocation only after its underlying work exits.""" + try: + task.exception() + except asyncio.CancelledError: + pass + if self._run_lock.locked(): + self._run_lock.release() async def run( self, @@ -173,7 +190,13 @@ async def run( if output_root and self._output_root_env: os.environ[self._output_root_env] = str(output_root) - kwargs = _select_kwargs(self._run_with_file, task_data=task_data, module=self._module) + inference_stop_event = stop_event if stop_event is not None else threading.Event() + kwargs = _select_kwargs( + self._run_with_file, + task_data=task_data, + module=self._module, + stop_event=inference_stop_event, + ) async def _invoke() -> Any: if inspect.iscoroutinefunction(self._run_with_file): @@ -184,8 +207,36 @@ async def _invoke() -> Any: return await result return result + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s if timeout_s else None + + def remaining_timeout() -> float | None: + if deadline is None: + return None + return max(0.0, deadline - loop.time()) + + try: + if deadline is None: + await self._run_lock.acquire() + else: + await asyncio.wait_for(self._run_lock.acquire(), timeout=remaining_timeout()) + except asyncio.TimeoutError: + inference_stop_event.set() + return PipelineRunResult( + status=PipelineRunStatus.ERROR, output_path=None, message="Task processing timeout" + ) + + if inference_stop_event.is_set(): + self._run_lock.release() + return PipelineRunResult(status=PipelineRunStatus.CANCELLED, message="Task cancelled before inference") + + invocation = asyncio.create_task(_invoke()) + release_on_completion = False try: - raw = await asyncio.wait_for(_invoke(), timeout=timeout_s) if timeout_s else await _invoke() + if deadline is None: + raw = await asyncio.shield(invocation) + else: + raw = await asyncio.wait_for(asyncio.shield(invocation), timeout=remaining_timeout()) output_path = None if isinstance(raw, dict): @@ -199,9 +250,20 @@ async def _invoke() -> Any: ) except asyncio.TimeoutError: + inference_stop_event.set() + release_on_completion = True + invocation.add_done_callback(self._release_run_lock) return PipelineRunResult( status=PipelineRunStatus.ERROR, output_path=None, message="Task processing timeout" ) + except asyncio.CancelledError: + inference_stop_event.set() + release_on_completion = True + invocation.add_done_callback(self._release_run_lock) + raise except Exception as e: logger.exception(f"Pipeline run failed: {e}") return PipelineRunResult(status=PipelineRunStatus.ERROR, output_path=None, message=str(e)) + finally: + if not release_on_completion: + self._run_lock.release() diff --git a/telefuser/service/core/task_manager.py b/telefuser/service/core/task_manager.py index 3258efe7..fa7f7b16 100644 --- a/telefuser/service/core/task_manager.py +++ b/telefuser/service/core/task_manager.py @@ -24,6 +24,7 @@ class TaskInfo: task_id: str status: TaskStatus message: Any + request_metadata: dict[str, Any] = field(default_factory=dict) start_time: datetime = field(default_factory=datetime.now) end_time: datetime | None = None error: str | None = None @@ -85,6 +86,7 @@ def create_task(self, message: Any) -> str: task_id=task_id, status=TaskStatus.PENDING, message=message, + request_metadata=self._serialize_task_message(message), output_path=getattr(message, "output_path", None), ) @@ -154,6 +156,7 @@ def complete_task( get_service_metrics().record_task_completed(duration) task.peak_memory_mb = peak_memory_mb task.result = result + task.message = None def fail_task(self, task_id: str, error: str) -> None: """Mark task as failed with metrics.""" @@ -170,6 +173,7 @@ def fail_task(self, task_id: str, error: str) -> None: task.status = TaskStatus.FAILED task.end_time = datetime.now() task.error = error + task.message = None self.failed_tasks += 1 get_service_metrics().record_task_failed() @@ -187,10 +191,13 @@ def cancel_task(self, task_id: str) -> bool: if task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED]: return False + previous_status = task.status task.stop_event.set() task.status = TaskStatus.CANCELLED task.end_time = datetime.now() task.error = "Task cancelled by user" + if previous_status == TaskStatus.PENDING: + task.message = None get_service_metrics().record_task_cancelled() @@ -232,7 +239,7 @@ def get_task_status(self, task_id: str) -> dict[str, Any] | None: } if task.result is not None: status["result"] = task.result - status.update(self._serialize_task_message(task.message)) + status.update(task.request_metadata) return status def get_all_tasks(self) -> dict[str, dict[str, Any] | None]: @@ -335,6 +342,9 @@ def release_processing_slot(self, task_id: str) -> None: """ with self._lock: self._current_processing_tasks.pop(task_id, None) + task = self._tasks.get(task_id) + if task is not None and task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED): + task.message = None def get_service_status(self) -> dict[str, Any]: """Get overall service status.""" @@ -388,6 +398,7 @@ def _fail_all_pending(self, error: str) -> None: task.status = TaskStatus.FAILED task.end_time = datetime.now() task.error = error + task.message = None self.failed_tasks += 1 get_service_metrics().record_task_failed() logger.warning(f"Task {task_id} failed: {error}") diff --git a/telefuser/service/core/task_processor.py b/telefuser/service/core/task_processor.py index 18577478..8e5b5192 100644 --- a/telefuser/service/core/task_processor.py +++ b/telefuser/service/core/task_processor.py @@ -133,6 +133,7 @@ async def _process_task(self, task_id: str) -> None: if not task_info: logger.warning(f"Task {task_id} disappeared before processing") return + message = task_info.message logger.info(f"Processing task {task_id}") @@ -145,16 +146,12 @@ async def _process_task(self, task_id: str) -> None: return try: - if isinstance(task_info.message, StructuredTaskRequest): + if isinstance(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 - ) + result = await self.structured_service.execute_with_stop_event(message, task_info.stop_event) else: - result = await self.media_service.generate_media_with_stop_event( - task_info.message, task_info.stop_event - ) + result = await self.media_service.generate_media_with_stop_event(message, task_info.stop_event) if result: self.task_manager.complete_task( @@ -166,14 +163,14 @@ async def _process_task(self, task_id: str) -> None: ) logger.info(f"Task {task_id} completed successfully") else: - if task_info.stop_event.is_set(): + if task_info.status == TaskStatus.CANCELLED: logger.info(f"Task {task_id} cancelled during processing") else: self.task_manager.fail_task(task_id, "Generation failed") logger.error(f"Task {task_id} generation failed") except Exception as e: - if task_info.stop_event.is_set(): + if task_info.status == TaskStatus.CANCELLED: logger.info(f"Task {task_id} exited after cancellation") return logger.exception(f"Task {task_id} processing failed") diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py b/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py new file mode 100644 index 00000000..cae1ede1 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import asyncio +import threading +import time +from typing import Any, Mapping + +import pytest + +from telefuser.pipelines.lingbot_vla_v2.action_scheduler import ActionChunkScheduler + + +def test_scheduler_discards_inflight_and_pending_work_when_newer_observation_arrives() -> None: + started = threading.Event() + release = threading.Event() + + def infer(request: Mapping[str, Any]) -> dict[str, Any]: + if request["sequence_id"] == 0: + started.set() + assert release.wait(timeout=2) + return {"action": request["sequence_id"]} + + async def scenario() -> None: + scheduler = ActionChunkScheduler(infer) + await scheduler.start() + first = scheduler.submit({"request_id": "first", "sequence_id": 0}, session_key="session") + assert await asyncio.to_thread(started.wait, 2) + second = scheduler.submit({"request_id": "second", "sequence_id": 1}, session_key="session") + third = scheduler.submit({"request_id": "third", "sequence_id": 2}, session_key="session") + + assert (await second)["scheduler_status"] == "superseded" + release.set() + assert (await first)["scheduler_status"] == "superseded" + completed = await third + assert completed["action"] == 2 + assert completed["sequence_id"] == 2 + assert completed["scheduler_status"] == "completed" + assert completed["server_timing"]["queue_wait_ms"] >= 0 + await scheduler.close() + + asyncio.run(scenario()) + + +def test_scheduler_expires_result_after_non_cancellable_inference() -> None: + def infer(_request: Mapping[str, Any]) -> dict[str, Any]: + time.sleep(0.02) + return {"action": "too-late"} + + async def scenario() -> None: + scheduler = ActionChunkScheduler(infer) + await scheduler.start() + response = await scheduler.submit( + {"request_id": "expired", "request_ttl_ms": 1}, + session_key="session", + ) + assert response["action"] is None + assert response["scheduler_status"] == "expired" + assert response["error"]["code"] == "expired" + await scheduler.close() + + asyncio.run(scenario()) + + +def test_scheduler_rejects_stale_sequence_without_running_inference() -> None: + calls: list[int] = [] + + def infer(request: Mapping[str, Any]) -> dict[str, Any]: + calls.append(request["sequence_id"]) + return {"action": request["sequence_id"]} + + async def scenario() -> None: + scheduler = ActionChunkScheduler(infer) + await scheduler.start() + assert (await scheduler.submit({"sequence_id": 4}, session_key="session"))["action"] == 4 + stale = await scheduler.submit({"sequence_id": 4}, session_key="session") + assert stale["scheduler_status"] == "stale_sequence" + assert calls == [4] + await scheduler.close() + + asyncio.run(scenario()) + + +def test_scheduler_bounds_pending_sessions() -> None: + started = threading.Event() + release = threading.Event() + + def infer(request: Mapping[str, Any]) -> dict[str, Any]: + if request["request_id"] == "active": + started.set() + assert release.wait(timeout=2) + return {"action": request["request_id"]} + + async def scenario() -> None: + scheduler = ActionChunkScheduler(infer, max_pending_sessions=1) + await scheduler.start() + active = scheduler.submit({"request_id": "active"}, session_key="active") + assert await asyncio.to_thread(started.wait, 2) + pending = scheduler.submit({"request_id": "pending"}, session_key="pending") + overloaded = await scheduler.submit({"request_id": "overloaded"}, session_key="overloaded") + assert overloaded["scheduler_status"] == "overloaded" + release.set() + assert (await active)["scheduler_status"] == "completed" + assert (await pending)["scheduler_status"] == "completed" + await scheduler.close() + + asyncio.run(scenario()) + + +@pytest.mark.parametrize( + ("payload", "match"), + [ + ({"sequence_id": -1}, "sequence_id"), + ({"sequence_id": True}, "sequence_id"), + ({"request_ttl_ms": 0}, "request_ttl_ms"), + ({"request_ttl_ms": float("inf")}, "request_ttl_ms"), + ], +) +def test_scheduler_rejects_invalid_controls(payload: dict[str, Any], match: str) -> None: + async def scenario() -> None: + scheduler = ActionChunkScheduler(lambda _request: {"action": 1}) + await scheduler.start() + with pytest.raises(ValueError, match=match): + scheduler.submit(payload, session_key="session") + await scheduler.close() + + asyncio.run(scenario()) diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py new file mode 100644 index 00000000..22c79d07 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +import threading +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch +from click.testing import CliRunner +from fastapi.testclient import TestClient + +from examples.lingbot_vla_v2 import lingbot_vla_v2_robotwin_server as server +from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_CAMERA_KEYS, + LingBotVlaV2CanonicalActionChunk, + RobotWinProfile, +) + + +def _profile() -> RobotWinProfile: + return RobotWinProfile( + { + "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}, + } + ) + + +class _Pipeline: + def __init__(self, profile: RobotWinProfile) -> None: + self.config = SimpleNamespace(robot_profile=profile) + self.observations: list[Any] = [] + self.seeds: list[int | None] = [] + self.closed = False + + def __call__(self, observation, seed: int | None = None) -> LingBotVlaV2CanonicalActionChunk: + self.observations.append(observation) + self.seeds.append(seed) + return LingBotVlaV2CanonicalActionChunk( + canonical_normalized_actions=torch.zeros(50, 55), + horizon=50, + action_dim=55, + ) + + def close(self) -> None: + self.closed = True + + +def _observation() -> dict[str, Any]: + images = {key: np.zeros((8, 8, 3), dtype=np.uint8) for key in ROBOTWIN_CAMERA_KEYS} + return { + **images, + "observation.state": np.zeros(14, dtype=np.float32), + "task": "pick up the block", + } + + +def test_message_codec_matches_upstream_numpy_contract() -> None: + payload = { + "image": np.arange(24, dtype=np.uint8).reshape(2, 4, 3), + "state": np.float32(1.5), + } + + decoded = server.unpack_message(server.pack_message(payload)) + + assert np.array_equal(decoded["image"], payload["image"]) + assert decoded["image"].dtype == np.uint8 + assert decoded["state"] == np.float32(1.5) + + +def test_message_codec_rejects_unsafe_object_arrays() -> None: + with pytest.raises(ValueError, match="unsupported NumPy dtype"): + server.pack_message({"value": np.asarray([object()], dtype=object)}) + + +def test_adapter_maps_observation_and_returns_robotwin_chunk() -> None: + profile = _profile() + pipeline = _Pipeline(profile) + adapter = server.RobotWinPolicyAdapter(pipeline, use_length=8) + + request = {**_observation(), "seed": 7, "request_id": "request-1", "episode_id": "episode-1"} + result = adapter.infer(request) + + assert result["action"].shape == (8, 14) + assert result["action"].dtype == np.float32 + assert np.allclose(result["action"][:, [6, 13]], 0.0, atol=1e-6) + assert np.allclose(np.delete(result["action"], [6, 13], axis=1), 1.0000005) + assert result["policy_verified"] is False + assert result["seed"] == 7 + assert result["request_id"] == "request-1" + assert result["episode_id"] == "episode-1" + assert result["server_timing"]["lock_wait_ms"] >= 0 + assert result["server_timing"]["pipeline_ms"] >= 0 + assert result["server_timing"]["action_mapping_ms"] >= 0 + assert result["server_timing"]["adapter_total_ms"] >= 0 + assert len(pipeline.observations) == 1 + assert pipeline.seeds == [7] + observation = pipeline.observations[0] + assert observation.task == "pick up the block" + assert set(observation.images) == set(ROBOTWIN_CAMERA_KEYS) + + +def test_adapter_reset_does_not_run_or_reload_policy() -> None: + pipeline = _Pipeline(_profile()) + adapter = server.RobotWinPolicyAdapter(pipeline) + + assert adapter.infer({"reset": True, "robo_name": "robotwin"}) == {"action": None} + assert pipeline.observations == [] + with pytest.raises(ValueError, match="runtime checkpoint switching"): + adapter.infer({"reset": True, "path_to_pi_model": "/different/checkpoint"}) + + +@pytest.mark.parametrize("field", ["request_id", "episode_id"]) +def test_adapter_rejects_invalid_trace_fields(field: str) -> None: + adapter = server.RobotWinPolicyAdapter(_Pipeline(_profile())) + + with pytest.raises(ValueError, match=field): + adapter.infer({**_observation(), field: ""}) + + +@pytest.mark.parametrize("seed", [True, 1.5, "7"]) +def test_adapter_rejects_non_integer_seed(seed: Any) -> None: + adapter = server.RobotWinPolicyAdapter(_Pipeline(_profile())) + + with pytest.raises(ValueError, match="seed must be an integer"): + adapter.infer({**_observation(), "seed": seed}) + + +def test_adapter_rejects_missing_observation_fields() -> None: + adapter = server.RobotWinPolicyAdapter(_Pipeline(_profile())) + + with pytest.raises(ValueError, match="missing fields"): + adapter.infer({"task": "pick up the block"}) + + +def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None: + pipeline = _Pipeline(_profile()) + adapter = server.RobotWinPolicyAdapter(pipeline, use_length=3) + + with TestClient(server.create_robotwin_app(adapter)) as client: + assert client.get("/healthz").json() == {"status": "ok"} + with client.websocket_connect("/") as websocket: + metadata = server.unpack_message(websocket.receive_bytes()) + assert metadata["robot_profile"] == "robotwin" + assert metadata["protocol_version"] == "1.0" + assert metadata["action_type"] == "absolute_qpos" + assert metadata["action_horizon"] == 3 + assert metadata["action_dim"] == 14 + assert metadata["action_dtype"] == "float32" + assert metadata["action_order"] == list(server.ROBOTWIN_ACTION_ORDER) + assert metadata["max_request_bytes"] == server.ROBOTWIN_MAX_REQUEST_BYTES + assert metadata["scheduling"]["mode"] == "latest_wins" + assert metadata["scheduling"]["max_pending_per_session"] == 1 + assert metadata["scheduling"]["inflight_cancellation"] is False + + websocket.send_bytes( + server.pack_message( + { + **_observation(), + "seed": 11, + "request_id": "request-1", + "episode_id": "episode-1", + "sequence_id": 0, + "request_ttl_ms": 5_000, + } + ) + ) + response = server.unpack_message(websocket.receive_bytes()) + assert response["action"].shape == (3, 14) + assert response["seed"] == 11 + assert response["request_id"] == "request-1" + assert response["episode_id"] == "episode-1" + assert response["sequence_id"] == 0 + assert response["scheduler_status"] == "completed" + assert response["server_timing"]["decode_ms"] >= 0 + assert response["server_timing"]["infer_ms"] >= 0 + assert response["server_timing"]["queue_wait_ms"] >= 0 + assert response["server_timing"]["scheduler_total_ms"] >= 0 + assert response["server_timing"]["pipeline_ms"] >= 0 + assert response["server_timing"]["action_mapping_ms"] >= 0 + + websocket.send_bytes( + server.pack_message( + { + "reset": True, + "robo_name": "robotwin", + "request_id": "reset-1", + "episode_id": "episode-1", + } + ) + ) + reset_response = server.unpack_message(websocket.receive_bytes()) + assert reset_response["action"] is None + assert reset_response["request_id"] == "reset-1" + assert reset_response["episode_id"] == "episode-1" + assert reset_response["server_timing"]["prev_total_ms"] >= 0 + + +def test_websocket_accepts_overlapping_chunks_and_discards_superseded_actions() -> None: + first_started = threading.Event() + release_first = threading.Event() + call_count = 0 + + class BlockingPipeline(_Pipeline): + def __call__(self, observation, seed: int | None = None): + nonlocal call_count + call_count += 1 + if call_count == 1: + first_started.set() + assert release_first.wait(timeout=2) + return super().__call__(observation, seed=seed) + + pipeline = BlockingPipeline(_profile()) + adapter = server.RobotWinPolicyAdapter(pipeline, use_length=3) + + with TestClient(server.create_robotwin_app(adapter)) as client: + with client.websocket_connect("/") as websocket: + server.unpack_message(websocket.receive_bytes()) + for sequence_id in range(3): + websocket.send_bytes( + server.pack_message( + { + **_observation(), + "request_id": f"request-{sequence_id}", + "episode_id": "episode", + "sequence_id": sequence_id, + } + ) + ) + if sequence_id == 0: + assert first_started.wait(timeout=2) + release_first.set() + responses = { + response["request_id"]: response + for response in (server.unpack_message(websocket.receive_bytes()) for _ in range(3)) + } + + assert responses["request-0"]["scheduler_status"] == "superseded" + assert responses["request-0"]["action"] is None + assert responses["request-1"]["scheduler_status"] == "superseded" + assert responses["request-1"]["action"] is None + assert responses["request-2"]["scheduler_status"] == "completed" + assert responses["request-2"]["action"].shape == (3, 14) + assert call_count == 2 + + +def test_cli_exposes_isolated_robotwin_server_options() -> None: + result = CliRunner().invoke(server.main, ["--help"]) + + assert result.exit_code == 0 + assert "--model-root" in result.output + assert "--qwen3vl-root" in result.output + assert "--use-length" in result.output + assert "--max-pending-sessions" in result.output + assert "--cuda-graph" in result.output + + +def test_h100_policy_process_disables_only_cudnn_sdpa(monkeypatch) -> None: + calls: list[tuple[str, bool]] = [] + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda _device: "NVIDIA H100 80GB HBM3") + monkeypatch.setattr(torch.backends.cuda, "enable_cudnn_sdp", lambda enabled: calls.append(("cudnn", enabled))) + monkeypatch.setattr(torch.backends.cuda, "enable_flash_sdp", lambda enabled: calls.append(("flash", enabled))) + monkeypatch.setattr(torch.backends.cuda, "enable_math_sdp", lambda enabled: calls.append(("math", enabled))) + monkeypatch.setattr( + torch.backends.cuda, + "enable_mem_efficient_sdp", + lambda enabled: calls.append(("mem_efficient", enabled)), + ) + + server._configure_h100_sdpa_backends("cuda:0") + + assert calls == [("cudnn", False), ("flash", True), ("math", True), ("mem_efficient", True)] + + +def test_non_h100_policy_process_preserves_sdpa_backends(monkeypatch) -> None: + calls: list[bool] = [] + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda _device: "NVIDIA RTX 4090") + monkeypatch.setattr(torch.backends.cuda, "enable_cudnn_sdp", calls.append) + + server._configure_h100_sdpa_backends("cuda:0") + + assert calls == [] diff --git a/tests/unit/service/test_structured_tasks.py b/tests/unit/service/test_structured_tasks.py index a2004672..e1214079 100644 --- a/tests/unit/service/test_structured_tasks.py +++ b/tests/unit/service/test_structured_tasks.py @@ -255,3 +255,69 @@ async def scenario() -> None: asyncio.run(scenario()) assert pipeline.closed is True + + +@pytest.mark.parametrize( + ("field", "value", "expected_detail"), + ( + ("state", [0.0] * 13, "state"), + ("camera_high", "not-base64", "valid base64"), + ), +) +def test_structured_route_rejects_invalid_vla_payload_before_enqueue( + tmp_path: Path, + field: str, + value, + expected_detail: str, +) -> None: + task_manager = TaskManager() + server = ApiServer(task_manager=task_manager, enable_openai_api=False) + server.initialize_services(tmp_path, _StructuredPipelineService()) + payload = _payload() + payload[field] = value + + with TestClient(server.get_app()) as client: + response = client.post("/v1/tasks/structured", json=payload) + + assert response.status_code == 422 + assert expected_detail in str(response.json()["detail"]) + assert task_manager.get_all_tasks() == {} + + +def test_pipeline_runner_timeout_signals_stop_and_isolates_late_inference() -> None: + from telefuser.service.core.pipeline_runner import PipelineRunner + from telefuser.service_types import PipelineRunStatus + + first_started = threading.Event() + release_first = threading.Event() + calls: list[str] = [] + received_stop_events: list[threading.Event] = [] + + def run_structured(pipeline, request: str, stop_event=None): + calls.append(request) + received_stop_events.append(stop_event) + if request == "first": + first_started.set() + assert release_first.wait(timeout=2.0) + return {"request": request} + + async def scenario() -> None: + runner = PipelineRunner(pipeline=object(), run_with_file=run_structured) + first = await runner.run(task_data={"request": "first"}, timeout_s=0.05) + + assert first.status == PipelineRunStatus.ERROR + assert first.message == "Task processing timeout" + assert first_started.is_set() + assert received_stop_events[0].is_set() + + second_task = asyncio.create_task(runner.run(task_data={"request": "second"}, timeout_s=1.0)) + await asyncio.sleep(0.05) + assert calls == ["first"] + + release_first.set() + second = await asyncio.wait_for(second_task, timeout=1.0) + assert second.status == PipelineRunStatus.SUCCESS + assert second.raw == {"request": "second"} + assert calls == ["first", "second"] + + asyncio.run(scenario()) diff --git a/tests/unit/service/test_task_runtime.py b/tests/unit/service/test_task_runtime.py index 5c181d08..e340cd07 100644 --- a/tests/unit/service/test_task_runtime.py +++ b/tests/unit/service/test_task_runtime.py @@ -4,7 +4,7 @@ from datetime import timedelta from types import SimpleNamespace -from telefuser.service.api.schema import TaskRequest, TaskResponse +from telefuser.service.api.schema import StructuredTaskRequest, TaskRequest, TaskResponse from telefuser.service.core.task_manager import TaskManager, TaskStatus from telefuser.service.core.task_processor import AsyncTaskProcessor @@ -198,3 +198,75 @@ def worker() -> None: # Every task claimed at most once; all 20 eventually claimed. assert len(claimed) == len(set(claimed)) assert set(claimed) == set(task_manager._tasks.keys()) + + +def test_terminal_tasks_release_original_structured_request_payloads() -> None: + task_manager = TaskManager(max_queue_size=10) + + def request() -> StructuredTaskRequest: + return StructuredTaskRequest( + task="vla_action", + instruction="pick up the block", + state=[0.0] * 14, + camera_high="encoded-camera-payload", + camera_left_wrist="encoded-camera-payload", + camera_right_wrist="encoded-camera-payload", + ) + + completed_id = task_manager.create_task(request()) + task_manager.complete_task(completed_id, result={"horizon": 1}) + completed = task_manager.get_task(completed_id) + assert completed is not None + assert completed.message is None + completed_status = task_manager.get_task_status(completed_id) + assert completed_status is not None + assert completed_status["task"] == "vla_action" + assert completed_status["media_type"] == "structured" + assert "camera_high" not in completed_status + + failed_id = task_manager.create_task(request()) + task_manager.fail_task(failed_id, "inference failed") + failed = task_manager.get_task(failed_id) + assert failed is not None + assert failed.message is None + + pending_cancel_id = task_manager.create_task(request()) + assert task_manager.cancel_task(pending_cancel_id) is True + pending_cancelled = task_manager.get_task(pending_cancel_id) + assert pending_cancelled is not None + assert pending_cancelled.message is None + + processing_cancel_id = task_manager.create_task(request()) + assert task_manager.claim_next_pending_task() == processing_cancel_id + processing_cancelled = task_manager.get_task(processing_cancel_id) + assert processing_cancelled is not None + assert task_manager.cancel_task(processing_cancel_id) is True + assert processing_cancelled.message is not None + task_manager.release_processing_slot(processing_cancel_id) + assert processing_cancelled.message is None + + +def test_timeout_stop_signal_fails_task_instead_of_leaving_it_processing() -> None: + class TimedOutMediaService: + async def generate_media_with_stop_event(self, message: TaskRequest, stop_event): + stop_event.set() + return None + + async def scenario() -> None: + task_manager = TaskManager(max_queue_size=10) + processor = AsyncTaskProcessor( + task_manager=task_manager, + media_service=TimedOutMediaService(), + max_concurrent=1, + ) + task_id = task_manager.create_task(TaskRequest(task="t2i")) + assert task_manager.claim_next_pending_task() == task_id + + await processor._process_task(task_id) + + status = task_manager.get_task_status(task_id) + assert status is not None + assert status["status"] == TaskStatus.FAILED.value + assert status["error"] == "Generation failed" + + asyncio.run(scenario()) diff --git a/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py new file mode 100644 index 00000000..8014901b --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import argparse + +import numpy as np +import pytest +from PIL import Image + +from tools.validation import validate_lingbot_vla_v2_robotwin_ws as validator + + +def _metadata() -> dict: + return { + "protocol_version": validator.PROTOCOL_VERSION, + "robot_profile": "robotwin", + "action_type": validator.ACTION_TYPE, + "action_horizon": 3, + "action_dim": len(validator.ACTION_ORDER), + "action_dtype": validator.ACTION_DTYPE, + "action_order": list(validator.ACTION_ORDER), + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + "max_request_bytes": 16 * 1024 * 1024, + "scheduling": { + "mode": "latest_wins", + "max_pending_per_session": 1, + "max_pending_sessions": 32, + "sequence_field": "sequence_id", + "ttl_field": "request_ttl_ms", + "inflight_cancellation": False, + }, + } + + +def _response() -> dict: + return { + "action": np.arange(42, dtype=np.float32).reshape(3, 14), + "seed": 7, + "request_id": "request-1", + "episode_id": "episode-1", + "sequence_id": 3, + "scheduler_status": "completed", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + "server_timing": { + "decode_ms": 0.1, + "infer_ms": 2.0, + "queue_wait_ms": 0.1, + "scheduler_total_ms": 2.1, + "lock_wait_ms": 0.05, + "pipeline_ms": 1.5, + "action_mapping_ms": 0.2, + "adapter_total_ms": 1.8, + }, + } + + +def test_validate_metadata_accepts_contract_and_additive_fields() -> None: + metadata = {**_metadata(), "future_field": "ignored"} + + summary = validator.validate_metadata(metadata) + + assert summary["action_shape"] == [3, 14] + assert summary["action_type"] == "absolute_qpos" + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("protocol_version", "2.0"), + ("action_type", "delta_qpos"), + ("action_dtype", "float64"), + ("action_order", ["unknown"] * 14), + ("action_horizon", 0), + ], +) +def test_validate_metadata_rejects_contract_mismatch(field: str, value) -> None: + metadata = {**_metadata(), field: value} + + with pytest.raises(validator.ValidationFailure, match=field): + validator.validate_metadata(metadata) + + +def test_validate_action_response_returns_stable_float32_digest() -> None: + summary = validator.validate_action_response( + _response(), + expected_horizon=3, + request_id="request-1", + episode_id="episode-1", + sequence_id=3, + seed=7, + ) + + assert summary["shape"] == [3, 14] + assert summary["dtype"] == "float32" + assert len(summary["sha256_float32_le"]) == 64 + assert summary["server_timing_ms"]["pipeline_ms"] == 1.5 + + +@pytest.mark.parametrize( + "actions", + [ + np.zeros((2, 14), dtype=np.float32), + np.zeros((3, 14), dtype=np.float64), + np.full((3, 14), np.nan, dtype=np.float32), + ], +) +def test_validate_action_response_rejects_invalid_action(actions: np.ndarray) -> None: + response = {**_response(), "action": actions} + + with pytest.raises(validator.ValidationFailure, match="action"): + validator.validate_action_response( + response, + expected_horizon=3, + request_id="request-1", + episode_id="episode-1", + sequence_id=3, + seed=7, + ) + + +def test_validate_reset_response_checks_trace_and_timings() -> None: + summary = validator.validate_reset_response( + { + "action": None, + "request_id": "reset-1", + "episode_id": "episode-1", + "server_timing": {"decode_ms": 0.1, "infer_ms": 0.2}, + }, + request_id="reset-1", + episode_id="episode-1", + ) + + assert summary["server_timing_ms"] == {"decode_ms": 0.1, "infer_ms": 0.2} + + +def test_validate_discarded_response_rejects_action_execution() -> None: + summary = validator.validate_discarded_response( + { + "action": None, + "request_id": "request-1", + "episode_id": "episode-1", + "sequence_id": 3, + "scheduler_status": "superseded", + "error": {"code": "superseded", "message": "newer observation received"}, + "server_timing": {"decode_ms": 0.1}, + }, + request_id="request-1", + episode_id="episode-1", + sequence_id=3, + ) + + assert summary["scheduler_status"] == "superseded" + + +def test_require_exact_replay_rejects_divergent_digests() -> None: + records = [ + {"action": {"sha256_float32_le": "a"}}, + {"action": {"sha256_float32_le": "b"}}, + ] + + with pytest.raises(validator.ValidationFailure, match="replay diverged"): + validator.require_exact_replay(records) + + +def test_parse_state_json_rejects_non_finite_state() -> None: + with pytest.raises(argparse.ArgumentTypeError, match="finite numbers"): + validator.parse_state_json("[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1e999]") + + +def test_load_validation_image_bounds_longest_edge_without_upscaling(tmp_path) -> None: + large_path = tmp_path / "large.png" + small_path = tmp_path / "small.png" + Image.new("RGB", (800, 400)).save(large_path) + Image.new("RGB", (32, 16)).save(small_path) + + large, source_shape = validator.load_validation_image(large_path, max_image_edge=640) + small, _ = validator.load_validation_image(small_path, max_image_edge=640) + + assert source_shape == [400, 800, 3] + assert large.shape == (320, 640, 3) + assert small.shape == (16, 32, 3) diff --git a/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py b/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py new file mode 100644 index 00000000..738874a1 --- /dev/null +++ b/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py @@ -0,0 +1,537 @@ +"""Validate a running LingBot-VLA v2 RoboTwin WebSocket endpoint without simulation.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import statistics +import time +from pathlib import Path +from typing import Any, Mapping + +import numpy as np +from PIL import Image + +PROTOCOL_VERSION = "1.0" +ACTION_TYPE = "absolute_qpos" +ACTION_DTYPE = "float32" +DEFAULT_MAX_IMAGE_EDGE = 640 +ACTION_ORDER = ( + "left_arm_joint_0", + "left_arm_joint_1", + "left_arm_joint_2", + "left_arm_joint_3", + "left_arm_joint_4", + "left_arm_joint_5", + "left_gripper", + "right_arm_joint_0", + "right_arm_joint_1", + "right_arm_joint_2", + "right_arm_joint_3", + "right_arm_joint_4", + "right_arm_joint_5", + "right_gripper", +) +CAMERA_KEYS = ( + "observation.images.cam_high", + "observation.images.cam_left_wrist", + "observation.images.cam_right_wrist", +) +_ACTION_TIMING_FIELDS = ( + "decode_ms", + "infer_ms", + "queue_wait_ms", + "scheduler_total_ms", + "lock_wait_ms", + "pipeline_ms", + "action_mapping_ms", + "adapter_total_ms", +) + + +class ValidationFailure(RuntimeError): + """Raised when the endpoint violates the RoboTwin action contract.""" + + +def parse_state_json(value: str) -> list[float]: + """Parse 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 validate_metadata(metadata: Any) -> dict[str, Any]: + """Validate and summarize the server's advertised action contract.""" + if not isinstance(metadata, dict): + raise ValidationFailure("metadata frame must be a MessagePack object") + expected = { + "protocol_version": PROTOCOL_VERSION, + "robot_profile": "robotwin", + "action_type": ACTION_TYPE, + "action_dim": len(ACTION_ORDER), + "action_dtype": ACTION_DTYPE, + "action_order": list(ACTION_ORDER), + } + for field, expected_value in expected.items(): + if metadata.get(field) != expected_value: + raise ValidationFailure( + f"metadata {field} mismatch: expected {expected_value!r}, observed {metadata.get(field)!r}" + ) + horizon = metadata.get("action_horizon") + if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: + raise ValidationFailure("metadata action_horizon must be a positive integer") + policy_verified = metadata.get("policy_verified") + verification_status = metadata.get("verification_status") + if not isinstance(policy_verified, bool): + raise ValidationFailure("metadata policy_verified must be boolean") + if not isinstance(verification_status, str) or not verification_status: + raise ValidationFailure("metadata verification_status must be a non-empty string") + max_request_bytes = metadata.get("max_request_bytes") + if isinstance(max_request_bytes, bool) or not isinstance(max_request_bytes, int) or max_request_bytes < 1: + raise ValidationFailure("metadata max_request_bytes must be a positive integer") + scheduling = metadata.get("scheduling") + if not isinstance(scheduling, dict): + raise ValidationFailure("metadata scheduling must be an object") + expected_scheduling = { + "mode": "latest_wins", + "max_pending_per_session": 1, + "sequence_field": "sequence_id", + "ttl_field": "request_ttl_ms", + "inflight_cancellation": False, + } + for field, expected_value in expected_scheduling.items(): + if scheduling.get(field) != expected_value: + raise ValidationFailure( + f"metadata scheduling.{field} mismatch: expected {expected_value!r}, observed {scheduling.get(field)!r}" + ) + max_pending_sessions = scheduling.get("max_pending_sessions") + if isinstance(max_pending_sessions, bool) or not isinstance(max_pending_sessions, int) or max_pending_sessions < 1: + raise ValidationFailure("metadata scheduling.max_pending_sessions must be a positive integer") + return { + "protocol_version": PROTOCOL_VERSION, + "action_shape": [horizon, len(ACTION_ORDER)], + "action_type": ACTION_TYPE, + "action_dtype": ACTION_DTYPE, + "policy_verified": policy_verified, + "verification_status": verification_status, + "max_request_bytes": max_request_bytes, + "scheduling": dict(scheduling), + } + + +def _validate_trace_fields(response: Mapping[str, Any], *, request_id: str, episode_id: str) -> None: + if response.get("request_id") != request_id: + raise ValidationFailure("response did not echo the request_id") + if response.get("episode_id") != episode_id: + raise ValidationFailure("response did not echo the episode_id") + + +def _validate_timings(response: Mapping[str, Any], required_fields: tuple[str, ...]) -> dict[str, float]: + timings = response.get("server_timing") + if not isinstance(timings, dict): + raise ValidationFailure("response server_timing must be an object") + validated: dict[str, float] = {} + for field in required_fields: + value = timings.get(field) + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(float(value)): + raise ValidationFailure(f"response server_timing.{field} must be a finite number") + if float(value) < 0: + raise ValidationFailure(f"response server_timing.{field} must be non-negative") + validated[field] = float(value) + if "prev_total_ms" in timings: + value = timings["prev_total_ms"] + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(float(value)): + raise ValidationFailure("response server_timing.prev_total_ms must be a finite number") + if float(value) < 0: + raise ValidationFailure("response server_timing.prev_total_ms must be non-negative") + validated["prev_total_ms"] = float(value) + return validated + + +def validate_reset_response(response: Any, *, request_id: str, episode_id: str) -> dict[str, Any]: + """Validate one reset acknowledgement.""" + if not isinstance(response, dict): + raise ValidationFailure("reset response must be a MessagePack object") + if response.get("action", object()) is not None: + raise ValidationFailure("reset response action must be null") + _validate_trace_fields(response, request_id=request_id, episode_id=episode_id) + return {"server_timing_ms": _validate_timings(response, ("decode_ms", "infer_ms"))} + + +def validate_action_response( + response: Any, + *, + expected_horizon: int, + request_id: str, + episode_id: str, + sequence_id: int, + seed: int, +) -> dict[str, Any]: + """Validate and summarize one raw RoboTwin action response.""" + if not isinstance(response, dict): + raise ValidationFailure("action response must be a MessagePack object") + _validate_trace_fields(response, request_id=request_id, episode_id=episode_id) + if response.get("sequence_id") != sequence_id: + raise ValidationFailure("response did not echo the sequence_id") + if response.get("scheduler_status") != "completed": + raise ValidationFailure("action response scheduler_status must be 'completed'") + if response.get("seed") != seed: + raise ValidationFailure("response did not echo the inference seed") + actions = response.get("action") + if not isinstance(actions, np.ndarray): + raise ValidationFailure("response action must be a NumPy array") + if actions.shape != (expected_horizon, len(ACTION_ORDER)): + raise ValidationFailure( + f"response action shape must be {(expected_horizon, len(ACTION_ORDER))}, got {actions.shape}" + ) + if actions.dtype != np.dtype(np.float32): + raise ValidationFailure(f"response action dtype must be float32, got {actions.dtype}") + if not np.isfinite(actions).all(): + raise ValidationFailure("response action contains non-finite values") + policy_verified = response.get("policy_verified") + verification_status = response.get("verification_status") + if not isinstance(policy_verified, bool): + raise ValidationFailure("response policy_verified must be boolean") + if not isinstance(verification_status, str) or not verification_status: + raise ValidationFailure("response verification_status must be a non-empty string") + + contiguous = np.ascontiguousarray(actions, dtype=" dict[str, Any]: + """Validate a structured scheduler response that must not be executed.""" + if not isinstance(response, dict): + raise ValidationFailure("discarded response must be a MessagePack object") + _validate_trace_fields(response, request_id=request_id, episode_id=episode_id) + if response.get("sequence_id") != sequence_id: + raise ValidationFailure("discarded response did not echo the sequence_id") + if response.get("action", object()) is not None: + raise ValidationFailure("discarded response action must be null") + status = response.get("scheduler_status") + if status not in {"superseded", "expired", "stale_sequence", "overloaded"}: + raise ValidationFailure(f"unexpected scheduler_status: {status!r}") + error = response.get("error") + if not isinstance(error, dict) or error.get("code") != status or not isinstance(error.get("message"), str): + raise ValidationFailure("discarded response must include a matching structured error") + return { + "scheduler_status": status, + "server_timing_ms": _validate_timings(response, ("decode_ms",)), + } + + +def require_exact_replay(records: list[dict[str, Any]]) -> None: + """Require every recorded action digest to match the first response.""" + if len(records) < 2: + raise ValidationFailure("exact replay validation requires at least two requests") + digests = {record["action"]["sha256_float32_le"] for record in records} + if len(digests) != 1: + raise ValidationFailure(f"fixed-seed action replay diverged across {len(digests)} digests") + + +def _receive_binary(connection: Any, *, timeout_seconds: float) -> bytes: + payload = connection.recv(timeout=timeout_seconds) + if not isinstance(payload, bytes): + raise ValidationFailure(f"server returned a text error frame: {payload}") + return payload + + +def load_validation_image(image_path: Path, *, max_image_edge: int) -> tuple[np.ndarray, list[int]]: + """Load an RGB image and bound its encoded request size while preserving aspect ratio.""" + with Image.open(image_path) as opened: + image = opened.convert("RGB") + source_shape = [image.height, image.width, 3] + if max(image.size) > max_image_edge: + image.thumbnail((max_image_edge, max_image_edge), Image.Resampling.LANCZOS) + return np.asarray(image, dtype=np.uint8), source_shape + + +def run_validation( + *, + host: str, + port: int, + image_path: Path, + task: str, + state: list[float], + seed: int, + request_count: int, + timeout_seconds: float, + exact_replay: bool, + request_ttl_ms: float = 5_000.0, + overlap_requests: bool = False, + max_image_edge: int = DEFAULT_MAX_IMAGE_EDGE, +) -> dict[str, Any]: + """Connect to a resident policy and exercise reset plus fixed-seed inference.""" + from websockets.sync.client import connect + + from examples.lingbot_vla_v2.lingbot_vla_v2_robotwin_server import pack_message, unpack_message + + if request_count < 1: + raise ValueError("request_count must be positive") + if exact_replay and request_count < 2: + raise ValueError("exact replay validation requires request_count >= 2") + if overlap_requests and request_count < 2: + raise ValueError("overlap validation requires request_count >= 2") + if overlap_requests and exact_replay: + raise ValueError("overlap validation cannot require exact replay") + if not math.isfinite(request_ttl_ms) or request_ttl_ms <= 0: + raise ValueError("request_ttl_ms must be a positive finite number") + image, source_image_shape = load_validation_image(image_path, max_image_edge=max_image_edge) + + episode_id = f"validator-{time.time_ns()}" + uri = f"ws://{host}:{port}/" + records: list[dict[str, Any]] = [] + started_at = time.perf_counter() + with connect( + uri, + open_timeout=timeout_seconds, + close_timeout=timeout_seconds, + max_size=None, + compression=None, + ) as connection: + metadata = unpack_message(_receive_binary(connection, timeout_seconds=timeout_seconds)) + metadata_summary = validate_metadata(metadata) + expected_horizon = metadata["action_horizon"] + + reset_request_id = "validator-reset" + connection.send( + pack_message( + { + "reset": True, + "robo_name": "robotwin", + "request_id": reset_request_id, + "episode_id": episode_id, + } + ) + ) + reset = unpack_message(_receive_binary(connection, timeout_seconds=timeout_seconds)) + reset_summary = validate_reset_response(reset, request_id=reset_request_id, episode_id=episode_id) + + base_request = { + **{key: image for key in CAMERA_KEYS}, + "observation.state": np.asarray(state, dtype=np.float32), + "task": task, + "episode_id": episode_id, + "seed": seed, + "request_ttl_ms": request_ttl_ms, + } + request_payload_bytes = len( + pack_message( + { + **base_request, + "request_id": "validator-size-check", + "sequence_id": request_count - 1, + } + ) + ) + if request_payload_bytes > metadata["max_request_bytes"]: + raise ValidationFailure( + f"encoded request uses {request_payload_bytes} bytes, exceeding the server limit " + f"of {metadata['max_request_bytes']} bytes" + ) + request_started_at: dict[str, float] = {} + for index in range(request_count): + request_id = f"validator-{index}" + request_started_at[request_id] = time.perf_counter() + connection.send( + pack_message( + { + **base_request, + "request_id": request_id, + "sequence_id": index, + } + ) + ) + if not overlap_requests: + response = unpack_message(_receive_binary(connection, timeout_seconds=timeout_seconds)) + action_summary = validate_action_response( + response, + expected_horizon=expected_horizon, + request_id=request_id, + episode_id=episode_id, + sequence_id=index, + seed=seed, + ) + records.append( + { + "request_id": request_id, + "sequence_id": index, + "round_trip_ms": (time.perf_counter() - request_started_at[request_id]) * 1000.0, + "scheduler_status": "completed", + "action": action_summary, + } + ) + + if overlap_requests: + for _ in range(request_count): + response = unpack_message(_receive_binary(connection, timeout_seconds=timeout_seconds)) + request_id = response.get("request_id") + if not isinstance(request_id, str) or not request_id.startswith("validator-"): + raise ValidationFailure("overlap response has an unknown request_id") + try: + sequence_id = int(request_id.removeprefix("validator-")) + except ValueError as error: + raise ValidationFailure("overlap response request_id has an invalid sequence") from error + if request_id not in request_started_at or not 0 <= sequence_id < request_count: + raise ValidationFailure("overlap response request_id was not sent by this validator") + if any(record["request_id"] == request_id for record in records): + raise ValidationFailure("overlap endpoint returned a duplicate response") + if response.get("scheduler_status") == "completed": + action_summary = validate_action_response( + response, + expected_horizon=expected_horizon, + request_id=request_id, + episode_id=episode_id, + sequence_id=sequence_id, + seed=seed, + ) + record = {"scheduler_status": "completed", "action": action_summary} + else: + discard_summary = validate_discarded_response( + response, + request_id=request_id, + episode_id=episode_id, + sequence_id=sequence_id, + ) + record = {"scheduler_status": discard_summary["scheduler_status"], "discard": discard_summary} + records.append( + { + "request_id": request_id, + "sequence_id": sequence_id, + "round_trip_ms": (time.perf_counter() - request_started_at[request_id]) * 1000.0, + **record, + } + ) + statuses = {record["request_id"]: record["scheduler_status"] for record in records} + if statuses.get(f"validator-{request_count - 1}") != "completed": + raise ValidationFailure("the newest overlap request did not produce an action") + if "superseded" not in statuses.values(): + raise ValidationFailure("overlap validation did not observe a superseded action") + + if exact_replay: + require_exact_replay(records) + return { + "passed": True, + "endpoint": uri, + "elapsed_seconds": time.perf_counter() - started_at, + "configuration": { + "image": str(image_path), + "task": task, + "seed": seed, + "request_count": request_count, + "exact_replay": exact_replay, + "overlap_requests": overlap_requests, + "request_ttl_ms": request_ttl_ms, + "max_image_edge": max_image_edge, + "source_image_shape": source_image_shape, + "transmitted_image_shape": list(image.shape), + "request_payload_bytes": request_payload_bytes, + }, + "metadata": metadata, + "metadata_summary": metadata_summary, + "reset": reset_summary, + "records": records, + } + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def _positive_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError("value must be a positive finite number") + return parsed + + +def build_parser() -> argparse.ArgumentParser: + """Build the command-line parser.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=9330) + parser.add_argument( + "--image", + type=Path, + default=Path("examples/data/lingbot_world_fast/image.jpg"), + ) + parser.add_argument("--task", default="pick up the object") + parser.add_argument("--state-json", type=parse_state_json, default=[0.0] * 14) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--requests", type=_positive_int, default=2) + parser.add_argument("--timeout-seconds", type=_positive_float, default=120.0) + parser.add_argument( + "--request-ttl-ms", + type=_positive_float, + default=5_000.0, + help="Server-local lifetime for each queued action request", + ) + parser.add_argument( + "--overlap-requests", + action="store_true", + help="Send all requests before receiving to verify latest-wins scheduling", + ) + parser.add_argument("--max-image-edge", type=_positive_int, default=DEFAULT_MAX_IMAGE_EDGE) + parser.add_argument("--require-exact-replay", action="store_true") + parser.add_argument("--output", type=Path) + return parser + + +def main() -> None: + """Run the no-simulation RoboTwin WebSocket validation.""" + args = build_parser().parse_args() + report = run_validation( + host=args.host, + port=args.port, + image_path=args.image, + task=args.task, + state=args.state_json, + seed=args.seed, + request_count=args.requests, + timeout_seconds=args.timeout_seconds, + exact_replay=args.require_exact_replay, + request_ttl_ms=args.request_ttl_ms, + overlap_requests=args.overlap_requests, + max_image_edge=args.max_image_edge, + ) + rendered = json.dumps(report, indent=2, sort_keys=True) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(f"{rendered}\n", encoding="utf-8") + print(rendered) + + +if __name__ == "__main__": + main()