From e1bfd15db09225ce6f8b00f3d59c07673e7abd29 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 31 Aug 2026 03:57:43 +0000 Subject: [PATCH 1/5] fix(vla): harden task lifecycle and validation --- .../lingbot_vla_v2_native_service.py | 2 + telefuser/models/lingbot_vla_v2.py | 12 ++- .../pipelines/lingbot_vla_v2/pipeline.py | 12 ++- telefuser/pipelines/lingbot_vla_v2/policy.py | 15 +++- telefuser/pipelines/lingbot_vla_v2/service.py | 57 +++++++++--- telefuser/service/api/routers/tasks.py | 37 ++++++++ .../service/api/task_application_service.py | 7 +- telefuser/service/core/pipeline_runner.py | 86 ++++++++++++++++--- telefuser/service/core/task_manager.py | 13 ++- telefuser/service/core/task_processor.py | 15 ++-- tests/unit/service/test_structured_tasks.py | 66 ++++++++++++++ tests/unit/service/test_task_runtime.py | 74 +++++++++++++++- 12 files changed, 356 insertions(+), 40 deletions(-) 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/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/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/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()) From d2b21d7393847bf5b58e46f201e3158d8b0dfe76 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Thu, 3 Sep 2026 03:35:59 +0000 Subject: [PATCH 2/5] feat(vla): add RoboTwin action server adapter Add an isolated persistent MessagePack WebSocket endpoint for LingBot VLA v2 that accepts RoboTwin observations and returns absolute-position action chunks through the existing RobotWin profile. Document the standalone deployment flow and protocol dependency, and cover message encoding, observation mapping, reset handling, persistent connections, timing metadata, and CLI options. Verification: 18 focused LingBot VLA v2 tests passed; ruff check and git diff --check passed. --- examples/lingbot_vla_v2/README.md | 73 +++++- .../lingbot_vla_v2_robotwin_server.py | 238 ++++++++++++++++++ .../lingbot_vla_v2/requirements-robotwin.txt | 1 + .../lingbot_vla_v2/test_robotwin_server.py | 141 +++++++++++ 4 files changed, 452 insertions(+), 1 deletion(-) create mode 100644 examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py create mode 100644 examples/lingbot_vla_v2/requirements-robotwin.txt create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 377dcd6e..1643a8d1 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,76 @@ 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 +``` + +The server exposes `GET /healthz` and the policy WebSocket at `/`. On connection it sends a MessagePack metadata +frame, 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 +``` + +On the simulation side, place the upstream `experiment/robotwin/eval_policy_client_lingbotvla.py` at +`/script/eval_policy_client_lingbotvla.py`, together with its `script/deploy/websocket_client_policy.py` and +`msgpack_numpy.py` helpers, then run one smoke episode configuration against the same port: + +```bash +cd /path/to/RoboTwin +python -u script/eval_policy_client_lingbotvla.py --config policy/ACT/deploy_policy.yml \ + --overrides \ + --task_name lift_pot \ + --task_config demo_clean \ + --train_config_name 0 \ + --seed 0 \ + --policy_name ACT \ + --port 9330 \ + --robo_name robotwin \ + --eval_video_log False \ + --output_dir ./eval_results +``` + +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. + +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_robotwin_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py new file mode 100644 index 00000000..0ea91f60 --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py @@ -0,0 +1,238 @@ +"""Serve LingBot-VLA v2 through the upstream RoboTwin policy protocol.""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading +import time +from typing import Any, Mapping, Protocol + +import click +import msgpack +import numpy as np +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.runtime import ( + LINGBOT_VLA_V2_QUANTIZATION_CHOICES, + get_lingbot_vla_v2_pipeline, +) +from telefuser.utils.logging import logger + + +class _Pipeline(Protocol): + config: Any + + def __call__(self, observation: LingBotVlaV2Observation) -> 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 + + +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 { + "robot_profile": self.profile.name, + "action_horizon": self.use_length, + "action_dim": self.profile.raw_state_dim, + "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.""" + if request.get("reset", False): + return self._reset(request) + + 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}, + ) + with self._lock: + canonical_chunk = self.pipeline(observation) + action_chunk = self.profile.structure_actions( + canonical_chunk.canonical_normalized_actions, + ) + if action_chunk.horizon < self.use_length: + raise RuntimeError( + f"policy returned horizon {action_chunk.horizon}, shorter than use_length={self.use_length}" + ) + actions = action_chunk.raw_actions[: self.use_length].numpy() + return { + "action": actions, + "policy_verified": canonical_chunk.policy_verified, + "verification_status": canonical_chunk.verification_status, + } + + 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) -> FastAPI: + """Create a standalone app compatible with upstream WebsocketClientPolicy.""" + app = FastAPI(title="LingBot-VLA v2 RoboTwin Policy") + + @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)) + previous_total_ms: float | None = None + 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() + request = unpack_message(payload) + inference_started_at = time.monotonic() + response = await asyncio.to_thread(adapter.infer, request) + inference_ms = (time.monotonic() - inference_started_at) * 1000.0 + response = dict(response) + response["server_timing"] = {"infer_ms": inference_ms} + if previous_total_ms is not None: + response["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 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) + + return app + + +@click.command() +@click.option("--model-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--qwen3vl-root", required=True, type=click.Path(exists=True, file_okay=False)) +@click.option("--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("--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, + cuda_graph: bool, + quantization: str | None, +) -> None: + """Start one resident LingBot-VLA v2 policy for a RoboTwin client.""" + 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), host=host, port=port, workers=1) + 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..1ebb3bda --- /dev/null +++ b/examples/lingbot_vla_v2/requirements-robotwin.txt @@ -0,0 +1 @@ +msgpack>=1.0,<2.0 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..1036fe46 --- /dev/null +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +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.closed = False + + def __call__(self, observation) -> LingBotVlaV2CanonicalActionChunk: + self.observations.append(observation) + 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) + + result = adapter.infer(_observation()) + + 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 len(pipeline.observations) == 1 + 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"}) + + +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["action_horizon"] == 3 + assert metadata["action_dim"] == 14 + + websocket.send_bytes(server.pack_message(_observation())) + response = server.unpack_message(websocket.receive_bytes()) + assert response["action"].shape == (3, 14) + assert response["server_timing"]["infer_ms"] >= 0 + + websocket.send_bytes(server.pack_message({"reset": True, "robo_name": "robotwin"})) + reset_response = server.unpack_message(websocket.receive_bytes()) + assert reset_response["action"] is None + assert reset_response["server_timing"]["prev_total_ms"] >= 0 + + +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 "--cuda-graph" in result.output From 65c48a4395baa4f5be24be78e7ff8974e8975825 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 7 Sep 2026 07:57:58 +0000 Subject: [PATCH 3/5] feat(vla): harden RoboTwin remote inference contract Advertise the raw action contract, propagate deterministic seeds and trace identifiers, validate mapped action chunks, and expose phase timings without changing shared TeleFuser APIs. Add a no-simulation WebSocket validator, focused unit coverage, dedicated client dependency, and the H100-to-RTX XPolicyLab deployment guide. Verification: ruff check; 46 LingBot-VLA v2 unit tests; 13 validator unit tests; CLI help; git diff --check. --- examples/lingbot_vla_v2/README.md | 60 ++- .../lingbot_vla_v2_robotwin_server.py | 91 ++++- .../lingbot_vla_v2/requirements-robotwin.txt | 1 + .../lingbot_vla_v2/test_robotwin_server.py | 64 +++- .../test_lingbot_vla_v2_robotwin_ws.py | 133 +++++++ .../validate_lingbot_vla_v2_robotwin_ws.py | 361 ++++++++++++++++++ 6 files changed, 682 insertions(+), 28 deletions(-) create mode 100644 tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py create mode 100644 tools/validation/validate_lingbot_vla_v2_robotwin_ws.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 1643a8d1..72e23f1e 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -258,23 +258,21 @@ result = policy.infer( actions = result["action"] # float32 [50, 14] when --use-length=50 ``` -On the simulation side, place the upstream `experiment/robotwin/eval_policy_client_lingbotvla.py` at -`/script/eval_policy_client_lingbotvla.py`, together with its `script/deploy/websocket_client_policy.py` and -`msgpack_numpy.py` helpers, then run one smoke episode configuration against the same port: +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. + +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 -cd /path/to/RoboTwin -python -u script/eval_policy_client_lingbotvla.py --config policy/ACT/deploy_policy.yml \ - --overrides \ - --task_name lift_pot \ - --task_config demo_clean \ - --train_config_name 0 \ - --seed 0 \ - --policy_name ACT \ - --port 9330 \ - --robo_name robotwin \ - --eval_video_log False \ - --output_dir ./eval_results +.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_robotwin_ws.py \ + --host 127.0.0.1 --port 9330 \ + --image examples/data/lingbot_world_fast/image.jpg \ + --task "pick up the object" --seed 7 --requests 2 \ + --require-exact-replay \ + --output work_dirs/robotwin_ws_validation/smoke.json ``` Each request runs the existing pipeline, converts normalized canonical `50 x 55` output through the bundled RoboTwin @@ -282,6 +280,38 @@ profile, and returns absolute-position actions in raw RoboTwin order. `--use-len 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 +``` + +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. diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py index 0ea91f60..f953a899 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py @@ -4,6 +4,7 @@ import asyncio import contextlib +import operator import threading import time from typing import Any, Mapping, Protocol @@ -25,11 +26,32 @@ ) from telefuser.utils.logging import logger +ROBOTWIN_PROTOCOL_VERSION = "1.0" +ROBOTWIN_ACTION_TYPE = "absolute_qpos" +ROBOTWIN_ACTION_DTYPE = "float32" +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) -> Any: ... + def __call__(self, observation: LingBotVlaV2Observation, seed: int | None = None) -> Any: ... def close(self) -> None: ... @@ -80,6 +102,30 @@ def unpack_message(payload: bytes) -> dict[str, Any]: 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.""" @@ -101,17 +147,22 @@ def __init__( 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), "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) + return {**self._reset(request), **trace_fields} missing = [key for key in (*ROBOTWIN_CAMERA_KEYS, "observation.state", "task") if key not in request] if missing: @@ -122,21 +173,43 @@ def infer(self, request: Mapping[str, Any]) -> dict[str, Any]: 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: - canonical_chunk = self.pipeline(observation) + 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 = action_chunk.raw_actions[: self.use_length].numpy() - return { + 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) @@ -174,14 +247,18 @@ async def policy_socket(websocket: WebSocket) -> 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 inference_started_at = time.monotonic() response = await asyncio.to_thread(adapter.infer, request) inference_ms = (time.monotonic() - inference_started_at) * 1000.0 response = dict(response) - response["server_timing"] = {"infer_ms": inference_ms} + server_timing = dict(response.get("server_timing", {})) + server_timing.update(decode_ms=decode_ms, infer_ms=inference_ms) if previous_total_ms is not None: - response["server_timing"]["prev_total_ms"] = previous_total_ms + server_timing["prev_total_ms"] = previous_total_ms + response["server_timing"] = server_timing await websocket.send_bytes(pack_message(response)) previous_total_ms = (time.monotonic() - round_started_at) * 1000.0 except WebSocketDisconnect: diff --git a/examples/lingbot_vla_v2/requirements-robotwin.txt b/examples/lingbot_vla_v2/requirements-robotwin.txt index 1ebb3bda..7beb7e9c 100644 --- a/examples/lingbot_vla_v2/requirements-robotwin.txt +++ b/examples/lingbot_vla_v2/requirements-robotwin.txt @@ -1 +1,2 @@ msgpack>=1.0,<2.0 +websockets>=12.0,<16.0 diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py index 1036fe46..e5e997ee 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py @@ -32,10 +32,12 @@ 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) -> LingBotVlaV2CanonicalActionChunk: + 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, @@ -78,14 +80,23 @@ def test_adapter_maps_observation_and_returns_robotwin_chunk() -> None: pipeline = _Pipeline(profile) adapter = server.RobotWinPolicyAdapter(pipeline, use_length=8) - result = adapter.infer(_observation()) + 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) @@ -101,6 +112,22 @@ def test_adapter_reset_does_not_run_or_reload_policy() -> None: 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())) @@ -117,17 +144,42 @@ def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None: 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 - - websocket.send_bytes(server.pack_message(_observation())) + assert metadata["action_dtype"] == "float32" + assert metadata["action_order"] == list(server.ROBOTWIN_ACTION_ORDER) + + websocket.send_bytes( + server.pack_message( + {**_observation(), "seed": 11, "request_id": "request-1", "episode_id": "episode-1"} + ) + ) 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["server_timing"]["decode_ms"] >= 0 assert response["server_timing"]["infer_ms"] >= 0 - - websocket.send_bytes(server.pack_message({"reset": True, "robo_name": "robotwin"})) + 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 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..d8892e55 --- /dev/null +++ b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import argparse + +import numpy as np +import pytest + +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", + } + + +def _response() -> dict: + return { + "action": np.arange(42, dtype=np.float32).reshape(3, 14), + "seed": 7, + "request_id": "request-1", + "episode_id": "episode-1", + "policy_verified": False, + "verification_status": "unverified_official_6b_base", + "server_timing": { + "decode_ms": 0.1, + "infer_ms": 2.0, + "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", + 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", + 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_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]") 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..e83998d8 --- /dev/null +++ b/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py @@ -0,0 +1,361 @@ +"""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" +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", + "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") + 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, + } + + +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, + 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("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=" 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 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, +) -> 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") + with Image.open(image_path) as opened: + image = np.asarray(opened.convert("RGB"), dtype=np.uint8) + + 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, + } + for index in range(request_count): + request_id = f"validator-{index}" + request_started_at = time.perf_counter() + connection.send(pack_message({**base_request, "request_id": request_id})) + response = unpack_message(_receive_binary(connection, timeout_seconds=timeout_seconds)) + round_trip_ms = (time.perf_counter() - request_started_at) * 1000.0 + action_summary = validate_action_response( + response, + expected_horizon=expected_horizon, + request_id=request_id, + episode_id=episode_id, + seed=seed, + ) + records.append({"request_id": request_id, "round_trip_ms": round_trip_ms, "action": action_summary}) + + 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, + }, + "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("--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, + ) + 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() From 9377f0429c5b5fe9602e4927f3cfa54fe67aa37d Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Tue, 8 Sep 2026 04:20:50 +0000 Subject: [PATCH 4/5] fix(vla): harden H100 RoboTwin validation Disable unsupported cuDNN SDPA plans only in the dedicated H100 RoboTwin policy process. Advertise and enforce the WebSocket request limit, resize oversized validation images, and document the module-based validator invocation and BF16 replay semantics. Verification: ruff check; ruff format --check; 141 LingBot VLA unit tests; real H100 direct WebSocket and XPolicyLab proxy smoke tests. --- examples/lingbot_vla_v2/README.md | 20 +++++++++--- .../lingbot_vla_v2_robotwin_server.py | 31 ++++++++++++++++++- .../lingbot_vla_v2/test_robotwin_server.py | 30 ++++++++++++++++++ .../test_lingbot_vla_v2_robotwin_ws.py | 16 ++++++++++ .../validate_lingbot_vla_v2_robotwin_ws.py | 31 +++++++++++++++++-- 5 files changed, 120 insertions(+), 8 deletions(-) diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 72e23f1e..f14b807b 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -237,9 +237,13 @@ Start one resident policy process: --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, then accepts multiple binary MessagePack requests on the same connection. This matches the upstream client -contract: +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 @@ -267,14 +271,20 @@ Validate this direct endpoint before a simulator is available. This sends reset the resident model, validates the returned `[H, 14]` action contract, and optionally verifies exact fixed-seed replay: ```bash -.venv-vla/bin/python tools/validation/validate_lingbot_vla_v2_robotwin_ws.py \ +.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 2 \ - --require-exact-replay \ + --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. + 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 diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py index f953a899..67849dad 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py @@ -12,6 +12,7 @@ import click import msgpack import numpy as np +import torch import uvicorn from fastapi import FastAPI, WebSocket, WebSocketDisconnect @@ -29,6 +30,7 @@ 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", @@ -154,6 +156,7 @@ def metadata(self) -> dict[str, Any]: "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", } @@ -272,6 +275,25 @@ async def policy_socket(websocket: WebSocket) -> None: 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)) @@ -296,6 +318,7 @@ def main( 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, @@ -306,7 +329,13 @@ def main( ) adapter = RobotWinPolicyAdapter(pipeline, use_length=use_length) try: - uvicorn.run(create_robotwin_app(adapter), host=host, port=port, workers=1) + uvicorn.run( + create_robotwin_app(adapter), + host=host, + port=port, + workers=1, + ws_max_size=ROBOTWIN_MAX_REQUEST_BYTES, + ) finally: adapter.close() diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py index e5e997ee..f06a50ff 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py @@ -150,6 +150,7 @@ def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None: 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 websocket.send_bytes( server.pack_message( @@ -191,3 +192,32 @@ def test_cli_exposes_isolated_robotwin_server_options() -> None: assert "--qwen3vl-root" in result.output assert "--use-length" 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/validation/test_lingbot_vla_v2_robotwin_ws.py b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py index d8892e55..5f73c07d 100644 --- a/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py +++ b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py @@ -4,6 +4,7 @@ import numpy as np import pytest +from PIL import Image from tools.validation import validate_lingbot_vla_v2_robotwin_ws as validator @@ -19,6 +20,7 @@ def _metadata() -> dict: "action_order": list(validator.ACTION_ORDER), "policy_verified": False, "verification_status": "unverified_official_6b_base", + "max_request_bytes": 16 * 1024 * 1024, } @@ -131,3 +133,17 @@ def test_require_exact_replay_rejects_divergent_digests() -> None: 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 index e83998d8..4675b061 100644 --- a/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py +++ b/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py @@ -17,6 +17,7 @@ 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", @@ -94,6 +95,9 @@ def validate_metadata(metadata: Any) -> dict[str, Any]: 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") return { "protocol_version": PROTOCOL_VERSION, "action_shape": [horizon, len(ACTION_ORDER)], @@ -101,6 +105,7 @@ def validate_metadata(metadata: Any) -> dict[str, Any]: "action_dtype": ACTION_DTYPE, "policy_verified": policy_verified, "verification_status": verification_status, + "max_request_bytes": max_request_bytes, } @@ -207,6 +212,16 @@ def _receive_binary(connection: Any, *, timeout_seconds: float) -> bytes: 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, @@ -218,6 +233,7 @@ def run_validation( request_count: int, timeout_seconds: float, exact_replay: bool, + 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 @@ -228,8 +244,7 @@ def run_validation( raise ValueError("request_count must be positive") if exact_replay and request_count < 2: raise ValueError("exact replay validation requires request_count >= 2") - with Image.open(image_path) as opened: - image = np.asarray(opened.convert("RGB"), dtype=np.uint8) + 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}/" @@ -267,6 +282,12 @@ def run_validation( "episode_id": episode_id, "seed": seed, } + request_payload_bytes = len(pack_message({**base_request, "request_id": "validator-size-check"})) + 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" + ) for index in range(request_count): request_id = f"validator-{index}" request_started_at = time.perf_counter() @@ -294,6 +315,10 @@ def run_validation( "seed": seed, "request_count": request_count, "exact_replay": exact_replay, + "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, @@ -331,6 +356,7 @@ def build_parser() -> argparse.ArgumentParser: 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("--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 @@ -349,6 +375,7 @@ def main() -> None: request_count=args.requests, timeout_seconds=args.timeout_seconds, exact_replay=args.require_exact_replay, + max_image_edge=args.max_image_edge, ) rendered = json.dumps(report, indent=2, sort_keys=True) if args.output is not None: From ef8e40f2d8a04c421c7ca444a41fb7bd32a4afd8 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Tue, 8 Sep 2026 07:30:00 +0000 Subject: [PATCH 5/5] feat(vla): schedule fresh RoboTwin action chunks Add a bounded LingBot VLA v2 scheduler that serializes GPU inference while retaining only the newest pending request per session. Support sequence IDs, server-local request TTLs, structured stale-action responses, and explicit scheduling metadata without changing the shared service or model interfaces. Decouple WebSocket receive and response delivery so clients can submit observations while inference is active. Extend the no-simulation validator with sequential and overlap modes, and document the remaining RTX-side execution-buffer work. Verification: ruff check and format checks passed; focused scheduler/WebSocket tests passed (38); all LingBot VLA v2 CPU tests passed (151); real H100 sequential and three-request overlap validation passed. --- examples/lingbot_vla_v2/README.md | 32 +++ .../lingbot_vla_v2_robotwin_server.py | 101 ++++++- .../lingbot_vla_v2/action_scheduler.py | 271 ++++++++++++++++++ .../lingbot_vla_v2/test_action_scheduler.py | 126 ++++++++ .../lingbot_vla_v2/test_robotwin_server.py | 66 ++++- .../test_lingbot_vla_v2_robotwin_ws.py | 33 +++ .../validate_lingbot_vla_v2_robotwin_ws.py | 173 ++++++++++- 7 files changed, 774 insertions(+), 28 deletions(-) create mode 100644 telefuser/pipelines/lingbot_vla_v2/action_scheduler.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index f14b807b..b094b9fd 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -267,6 +267,18 @@ horizon, dimension, and the exact dual-arm joint order. Inference requests may i `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: @@ -285,6 +297,20 @@ large repository sample representative of normal RoboTwin camera payloads. Add ` 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 @@ -317,6 +343,12 @@ bash scripts/eval_policy.sh \ --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 diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py index 67849dad..62905647 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py @@ -7,6 +7,7 @@ import operator import threading import time +from contextlib import asynccontextmanager from typing import Any, Mapping, Protocol import click @@ -21,6 +22,7 @@ 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, @@ -227,9 +229,23 @@ def close(self) -> None: self.pipeline.close() -def create_robotwin_app(adapter: RobotWinPolicyAdapter) -> FastAPI: +def create_robotwin_app( + adapter: RobotWinPolicyAdapter, + *, + max_pending_sessions: int = 32, +) -> FastAPI: """Create a standalone app compatible with upstream WebsocketClientPolicy.""" - app = FastAPI(title="LingBot-VLA v2 RoboTwin Policy") + 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]: @@ -238,8 +254,41 @@ async def healthz() -> dict[str, str]: @app.websocket("/") async def policy_socket(websocket: WebSocket) -> None: await websocket.accept() - await websocket.send_bytes(pack_message(adapter.metadata)) + 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() @@ -253,17 +302,24 @@ async def policy_socket(websocket: WebSocket) -> None: decode_started_at = time.monotonic() request = unpack_message(payload) decode_ms = (time.monotonic() - decode_started_at) * 1000.0 - inference_started_at = time.monotonic() - response = await asyncio.to_thread(adapter.infer, request) - inference_ms = (time.monotonic() - inference_started_at) * 1000.0 - response = dict(response) - server_timing = dict(response.get("server_timing", {})) - server_timing.update(decode_ms=decode_ms, infer_ms=inference_ms) - if previous_total_ms is not None: - server_timing["prev_total_ms"] = previous_total_ms - response["server_timing"] = server_timing - await websocket.send_bytes(pack_message(response)) - previous_total_ms = (time.monotonic() - round_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: @@ -271,6 +327,13 @@ async def policy_socket(websocket: WebSocket) -> None: 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 @@ -301,6 +364,13 @@ def _configure_h100_sdpa_backends(device: str) -> None: @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", @@ -314,6 +384,7 @@ def main( port: int, device: str, use_length: int, + max_pending_sessions: int, cuda_graph: bool, quantization: str | None, ) -> None: @@ -330,7 +401,7 @@ def main( adapter = RobotWinPolicyAdapter(pipeline, use_length=use_length) try: uvicorn.run( - create_robotwin_app(adapter), + create_robotwin_app(adapter, max_pending_sessions=max_pending_sessions), host=host, port=port, workers=1, 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/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 index f06a50ff..22c79d07 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py @@ -1,5 +1,6 @@ from __future__ import annotations +import threading from types import SimpleNamespace from typing import Any @@ -151,10 +152,20 @@ def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None: 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"} + { + **_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()) @@ -162,8 +173,12 @@ def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None: 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 @@ -184,6 +199,54 @@ def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None: 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"]) @@ -191,6 +254,7 @@ def test_cli_exposes_isolated_robotwin_server_options() -> None: 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 diff --git a/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py index 5f73c07d..8014901b 100644 --- a/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py +++ b/tests/unit/validation/test_lingbot_vla_v2_robotwin_ws.py @@ -21,6 +21,14 @@ def _metadata() -> dict: "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, + }, } @@ -30,11 +38,15 @@ def _response() -> dict: "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, @@ -75,6 +87,7 @@ def test_validate_action_response_returns_stable_float32_digest() -> None: expected_horizon=3, request_id="request-1", episode_id="episode-1", + sequence_id=3, seed=7, ) @@ -101,6 +114,7 @@ def test_validate_action_response_rejects_invalid_action(actions: np.ndarray) -> expected_horizon=3, request_id="request-1", episode_id="episode-1", + sequence_id=3, seed=7, ) @@ -120,6 +134,25 @@ def test_validate_reset_response_checks_trace_and_timings() -> None: 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"}}, diff --git a/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py b/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py index 4675b061..738874a1 100644 --- a/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py +++ b/tools/validation/validate_lingbot_vla_v2_robotwin_ws.py @@ -42,6 +42,8 @@ _ACTION_TIMING_FIELDS = ( "decode_ms", "infer_ms", + "queue_wait_ms", + "scheduler_total_ms", "lock_wait_ms", "pipeline_ms", "action_mapping_ms", @@ -98,6 +100,24 @@ def validate_metadata(metadata: Any) -> dict[str, Any]: 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)], @@ -106,6 +126,7 @@ def validate_metadata(metadata: Any) -> dict[str, Any]: "policy_verified": policy_verified, "verification_status": verification_status, "max_request_bytes": max_request_bytes, + "scheduling": dict(scheduling), } @@ -154,12 +175,17 @@ def validate_action_response( 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") @@ -196,6 +222,33 @@ def validate_action_response( } +def validate_discarded_response( + response: Any, + *, + request_id: str, + episode_id: str, + sequence_id: int, +) -> 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: @@ -233,6 +286,8 @@ def run_validation( 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.""" @@ -244,6 +299,12 @@ def run_validation( 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()}" @@ -281,27 +342,100 @@ def run_validation( "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"})) + 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 = time.perf_counter() - connection.send(pack_message({**base_request, "request_id": request_id})) - response = unpack_message(_receive_binary(connection, timeout_seconds=timeout_seconds)) - round_trip_ms = (time.perf_counter() - request_started_at) * 1000.0 - action_summary = validate_action_response( - response, - expected_horizon=expected_horizon, - request_id=request_id, - episode_id=episode_id, - seed=seed, + request_started_at[request_id] = time.perf_counter() + connection.send( + pack_message( + { + **base_request, + "request_id": request_id, + "sequence_id": index, + } + ) ) - records.append({"request_id": request_id, "round_trip_ms": round_trip_ms, "action": action_summary}) + 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) @@ -315,6 +449,8 @@ def run_validation( "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), @@ -356,6 +492,17 @@ def build_parser() -> argparse.ArgumentParser: 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) @@ -375,6 +522,8 @@ def main() -> None: 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)