diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index f81e5e887..6ef24bcf4 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -50,9 +50,10 @@ mode is selected per request: **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way (`domain_name=droid_lerobot`). -- **DROID OpenPI policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and - connect an OpenPI-compatible websocket client to `/v1/realtime/robot/openpi`. - This path returns action chunks directly instead of an mp4. +- **DROID policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and connect + RoboLab or another OpenPI-compatible client to + `/v1/realtime/robot/openpi`. This endpoint returns action chunks directly + instead of an mp4. Action requests can use `input_reference` or `video_reference` for video input. `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics` @@ -314,40 +315,71 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \ curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}' curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4 -# DROID OpenPI policy server (websocket action serving). -# Requires cosmos_framework on PYTHONPATH because the pipeline reuses the -# reference RoboLab action transforms. If your checkpoint config already -# includes policy_server_config, omit the stage_overrides file and flag. -cat > cosmos3_droid_openpi_stage_overrides.json <<'JSON' -{ - "0": { - "model_config": { - "policy_server_config": { - "image_resolution": [540, 640], - "n_external_cameras": 2, - "needs_wrist_camera": true, - "needs_stereo_camera": false, - "needs_session_id": true, - "action_space": "joint_position" - } - } - } -} -JSON +# DROID websocket policy server. Use the tested cosmos-framework revision +# directly from source; installing the package can introduce dependency conflicts. +export COSMOS_FRAMEWORK_ROOT=/path/to/cosmos-framework +git -C "$COSMOS_FRAMEWORK_ROOT" checkout c14617c2bc93dacbf69674fb964eec93182933d9 +export PYTHONPATH="$COSMOS_FRAMEWORK_ROOT" +# Validate every cosmos-framework symbol used by the action-policy pipeline +# before allocating the model. +python - <<'PY' +from vllm_omni.diffusion.models.cosmos3.utils import ( + get_robolab_domain_id, + preflight_cosmos3_action_framework_imports, +) + +preflight_cosmos3_action_framework_imports() +print("cosmos-framework action imports OK; DROID domain:", get_robolab_domain_id("droid_lerobot")) +PY + +# The bundled deploy config (vllm_omni/deploy/cosmos3_policy_droid.yaml) +# selects the registered cosmos3_policy pipeline and carries the DROID +# checkpoint's serving defaults: JSON prompt formatting (a property of this +# checkpoint's training recipe, not a generic server default), the +# model-specific OpenPI binary handshake metadata (policy_server_config), and +# guardrails off (policy serving emits robot actions; the guardrail stack is +# not part of it, so --no-guardrails is implied). +# Policy checkpoints cannot be auto-detected — they share their HF metadata +# with the T2I/video Cosmos3 checkpoints — so --deploy-config is required. +# Per-checkpoint model_config tweaks can be layered on top with +# --stage-overrides; dict overrides deep-merge with the deploy yaml. +export VLLM_OMNI_ROOT=/absolute/path/to/vllm-omni vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ --omni \ --host 0.0.0.0 --port 8000 \ - --model-class-name Cosmos3OmniDiffusersPipeline \ - --no-guardrails \ - --stage-overrides "$(cat cosmos3_droid_openpi_stage_overrides.json)" + --deploy-config "$VLLM_OMNI_ROOT/vllm_omni/deploy/cosmos3_policy_droid.yaml" \ + --robot-openpi-idle-timeout 0 -# Point an OpenPI websocket client at: -# ws://localhost:8000/v1/realtime/robot/openpi -# The first server message is policy_server_config. Each infer request sends a -# msgpack-numpy observation dict and receives a writable float32 action array. +# From a RoboLab checkout: +python policies/cosmos3/run.py \ + --remote-uri ws://localhost:8000/v1/realtime/robot/openpi \ + --task BananaInBowlTask + +# For a server started with "--api-key $POLICY_API_KEY", pass the same value as +# a Bearer token. --remote-token takes precedence over COSMOS3_API_TOKEN. +COSMOS3_API_TOKEN="$POLICY_API_KEY" python policies/cosmos3/run.py \ + --remote-uri wss://policy.example/v1/realtime/robot/openpi \ + --task BananaInBowlTask ``` +The endpoint sends `policy_server_config` as its initial binary MsgPack +message and returns the action array directly. RoboLab accepts that direct +array as well as proxy responses shaped as `{"action": ...}` or +`{"actions": ...}`, and turns structured server errors into client exceptions. + +RoboLab supplies `session_id=robolab-episode--env-`, so +parallel environments have independent policy state even though they share +one WebSocket. The server tracks interleaved session IDs independently. +The default receive-idle timeout is 30 seconds; set +`--robot-openpi-idle-timeout 0` when simulator steps between replans can take +longer, or set another non-negative timeout in seconds. + +If vLLM-Omni is started with `--api-key` or `VLLM_API_KEY`, the standard +OpenPI route uses the normal API authentication middleware. RoboLab's +`--remote-token` and `COSMOS3_API_TOKEN` send the required +`Authorization: Bearer ...` header. + #### Notes - **Measured latency (1x B300, bf16, guardrails off):** diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 51c7d549e..af898ae50 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -197,6 +197,25 @@ class TestStageConfig: omega_config = config.to_omegaconf() assert omega_config.engine_args.max_num_seqs == 32 + def test_to_omegaconf_dict_override_deep_merges_yaml_dict(self): + """A partial dict CLI override (e.g. --no-guardrails riding on + model_config) must layer onto the YAML dict, not clobber siblings.""" + config = StageConfig( + stage_id=0, + model_stage="diffusion", + stage_type=StageType.DIFFUSION, + yaml_engine_args={ + "model_config": { + "guardrails": True, + "policy_server_config": {"action_space": "joint_position"}, + } + }, + runtime_overrides={"model_config": {"guardrails": False}}, + ) + omega_config = config.to_omegaconf() + assert omega_config.engine_args.model_config.guardrails is False + assert omega_config.engine_args.model_config.policy_server_config.action_space == "joint_position" + def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self): config = StageConfig( stage_id=1, @@ -623,6 +642,58 @@ class TestPipelineDiscovery: assert p.hf_architectures == ("SomeCollidingArch",) +class TestCosmos3PolicyPipeline: + """Cosmos3 policy serving resolves via deploy yaml selection (π0 precedent).""" + + def test_registered_without_capturing_base_cosmos3_checkpoints(self): + assert "cosmos3_policy" in OMNI_PIPELINES + # T2I/video Cosmos3 checkpoints report model_type=cosmos3_omni and the + # same model_index.json _class_name as policy checkpoints; they must + # keep resolving through the single-stage diffusion fallback, so the + # policy pipeline must not be reachable by auto-detection. + assert "cosmos3_omni" not in OMNI_PIPELINES + pipeline = OMNI_PIPELINES["cosmos3_policy"] + assert pipeline.hf_architectures == () + assert pipeline.diffusers_class_name is None + + def test_droid_deploy_yaml_carries_policy_server_config(self): + deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) + assert deploy.pipeline == "cosmos3_policy" + + stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_policy"], deploy) + assert len(stages) == 1 + stage = stages[0].to_omegaconf() + + assert stage.stage_type == "diffusion" + assert stage.final_output_type == "action" + assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline" + # OpenPI websocket handshake metadata must reach the serving layer via + # engine_args.model_config (ServingRealtimeRobotOpenPI reads it there). + psc = stage.engine_args.model_config.policy_server_config + assert list(psc.image_resolution) == [540, 640] + assert psc.n_external_cameras == 2 + assert psc.needs_wrist_camera is True + assert psc.needs_stereo_camera is False + assert psc.needs_session_id is True + assert psc.action_space == "joint_position" + # The DROID checkpoint's training recipe expects JSON-formatted prompts. + assert stage.default_sampling_params.extra_args.format_prompt_as_json is True + + def test_no_guardrails_cli_flag_keeps_policy_server_config(self): + """--no-guardrails becomes a partial model_config CLI override; it must + deep-merge with the deploy yaml's model_config, not clobber the + policy_server_config the OpenPI handshake depends on.""" + deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) + stages, _ = StageConfigFactory._create_legacy_from_registry( + OMNI_PIPELINES["cosmos3_policy"], + {"model_config": {"guardrails": False}}, + user_deploy_config=deploy, + ) + stage = stages[0].to_omegaconf() + assert stage.engine_args.model_config.guardrails is False + assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position" + + class TestStagePipelineConfig: def test_frozen(self): s = StagePipelineConfig(stage_id=0, model_stage="a") diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py b/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py new file mode 100644 index 000000000..25f06b7a9 --- /dev/null +++ b/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from importlib.util import find_spec + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_FRAMEWORK_AVAILABLE = find_spec("cosmos_framework") is not None +requires_cosmos_framework = pytest.mark.skipif( + not _FRAMEWORK_AVAILABLE, + reason="cosmos_framework source is not present on PYTHONPATH", +) + + +def test_framework_preflight_checks_every_policy_symbol(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils + + imported = [] + + def record_import(module_name: str, symbol_name: str, error_message: str): + imported.append((module_name, symbol_name, error_message)) + return object() + + monkeypatch.setattr(cosmos3_utils, "lazy_import", record_import) + + cosmos3_utils.preflight_cosmos3_action_framework_imports() + + expected_symbols = { + "ActionTransformPipeline", + "FlowUniPCMultistepScheduler", + "build_abs_pose_from_components", + "convert_rotation", + "get_domain_id", + "pose_abs_to_rel", + "pose_rel_to_abs", + } + assert {symbol_name for _, symbol_name, _ in imported} == expected_symbols + assert [(module_name, symbol_name) for module_name, symbol_name, _ in imported] == list( + cosmos3_utils._ROBOLAB_FRAMEWORK_IMPORTS.values() + ) + + +def test_lazy_import_reports_framework_symbol_skew(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils + + monkeypatch.setattr(cosmos3_utils, "import_module", lambda _module_name: object()) + + with pytest.raises(ImportError, match="MissingSymbol") as exc_info: + cosmos3_utils.lazy_import( + "cosmos_framework.fake_module", + "MissingSymbol", + "Cosmos Framework compatibility error.", + ) + + assert isinstance(exc_info.value.__cause__, AttributeError) + + +@requires_cosmos_framework +def test_action_domain_table_matches_current_framework_and_legacy_aliases() -> None: + from cosmos_framework.data.generator.action.utils.domain_utils import ( + EMBODIMENT_TO_DOMAIN_ID as FRAMEWORK_DOMAIN_IDS, + ) + + from vllm_omni.diffusion.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID as VLLM_DOMAIN_IDS + from vllm_omni.diffusion.models.cosmos3.utils import _LEGACY_EMBODIMENT_ALIASES + + canonical_vllm_domain_ids = { + name: domain_id for name, domain_id in VLLM_DOMAIN_IDS.items() if name not in _LEGACY_EMBODIMENT_ALIASES + } + + assert canonical_vllm_domain_ids == FRAMEWORK_DOMAIN_IDS + assert {alias: VLLM_DOMAIN_IDS[alias] for alias in _LEGACY_EMBODIMENT_ALIASES} == { + alias: FRAMEWORK_DOMAIN_IDS[canonical_name] for alias, canonical_name in _LEGACY_EMBODIMENT_ALIASES.items() + } + + +@requires_cosmos_framework +def test_current_cosmos_framework_policy_modules_import_and_instantiate() -> None: + from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id + from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + convert_rotation, + pose_abs_to_rel, + pose_rel_to_abs, + ) + from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline + from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import ( + FlowUniPCMultistepScheduler, + ) + + transform = ActionTransformPipeline( + max_action_dim=64, + cfg_dropout_rate=0.0, + format_prompt_as_json=True, + ) + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, + shift=1.0, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(4, device=torch.device("cpu"), shift=5.0) + + assert transform.prompt_json_formatter is not None + assert len(scheduler.timesteps) == 4 + assert get_domain_id("droid_lerobot") == 8 + assert convert_rotation([0.0, 0.0, 0.0, 1.0], "quat_xyzw", "rot6d").shape == (6,) + assert callable(pose_abs_to_rel) + assert callable(pose_rel_to_abs) + assert callable(build_abs_pose_from_components) + + +@requires_cosmos_framework +def test_real_action_transform_pipeline_matches_droid_policy_contract() -> None: + from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline + + transform = ActionTransformPipeline( + max_action_dim=64, + cfg_dropout_rate=0.0, + format_prompt_as_json=True, + ) + sample = { + "ai_caption": "Pick up the cube.", + "video": torch.zeros((3, 3, 192, 320), dtype=torch.uint8), + "action": torch.zeros((3, 8), dtype=torch.float32), + "conditioning_fps": torch.tensor(15, dtype=torch.long), + "mode": "wam", + "domain_id": torch.tensor(8, dtype=torch.long), + "viewpoint": "concat_view", + "additional_view_description": "Wrist view above two exterior views.", + } + + transformed = transform(sample, "256") + + assert isinstance(transformed["ai_caption"], dict) + assert transformed["action"].shape == (3, 64) + assert int(transformed["raw_action_dim"].item()) == 8 + assert transformed["sequence_plan"].has_action is True + assert transformed["mode"] == "wam" + assert int(transformed["domain_id"].item()) == 8 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 72e4ab962..302409aeb 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -3,6 +3,7 @@ from __future__ import annotations +import json import sys import types from dataclasses import dataclass @@ -21,6 +22,20 @@ from vllm_omni.experimental.world_models.session_state import SessionStateManage pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.mark.parametrize( + ("alias", "canonical_name"), + [ + ("galbot", "embodiment_b"), + ("agibot_gear_gripper", "embodiment_c_gripper"), + ("agibot_gear_gripper_ext", "embodiment_c_gripper_ext"), + ], +) +def test_action_domain_table_preserves_legacy_aliases(alias: str, canonical_name: str) -> None: + from vllm_omni.diffusion.models.cosmos3.action import resolve_domain_id + + assert resolve_domain_id(domain_name=alias) == resolve_domain_id(domain_name=canonical_name) + + def test_pipeline_declares_layerwise_offload_components() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline @@ -479,6 +494,69 @@ def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None: assert captured["session_id"] == "robolab-request-7" +@pytest.mark.parametrize("format_prompt_as_json", [False, True]) +def test_robolab_input_builder_threads_prompt_format_and_uses_wam( + make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, + format_prompt_as_json: bool, +) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + pipeline.transformer = StubCosmos3Transformer(action_gen=True, action_dim=64) + captured: dict[str, Any] = {} + + def fake_transform(sample, resolution): + captured["sample_mode"] = sample["mode"] + captured["resolution"] = resolution + sample["sequence_plan"] = SimpleNamespace( + condition_frame_indexes_action=[0], + action_start_frame_offset=1, + ) + sample["raw_action_dim"] = torch.tensor(8) + sample["image_size"] = torch.tensor([16, 16, 16, 16]) + if format_prompt_as_json: + sample["ai_caption"] = {"actions": {"instruction": sample["ai_caption"]}} + return sample + + def fake_get_transform(*, format_prompt_as_json: bool): + captured["format_prompt_as_json"] = format_prompt_as_json + return fake_transform + + pipeline._get_robolab_transform = fake_get_transform + monkeypatch.setattr(pipeline_cosmos3, "get_robolab_domain_id", lambda name: 8) + obs = { + "prompt": "Pick up the cube.", + "observation/image": np.zeros((16, 16, 3), dtype=np.uint8), + "observation/joint_position": np.zeros(7, dtype=np.float32), + "observation/gripper_position": np.zeros(1, dtype=np.float32), + } + sampling_params = make_sampling_params( + extra_args={ + "robot_obs": obs, + "action_chunk_size": 2, + "image_height": 16, + "image_width": 16, + "format_prompt_as_json": format_prompt_as_json, + } + ) + + inputs = pipeline._build_robolab_policy_inputs(sampling_params, request_id="request-1") + + assert inputs is not None + assert captured == { + "sample_mode": "wam", + "resolution": "480", + "format_prompt_as_json": format_prompt_as_json, + } + assert inputs.domain_id == 8 + assert inputs.raw_action_dim == 8 + if format_prompt_as_json: + assert json.loads(inputs.prompt) == {"actions": {"instruction": "Pick up the cube."}} + else: + assert inputs.prompt == "Pick up the cube." + + @pytest.mark.parametrize( ("prompt", "sampling_params", "message"), [ diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 19ec23b6d..a6f12f661 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -310,6 +310,29 @@ def test_handle_connection_closes_websocket_on_idle_timeout(monkeypatch): serving.infer.assert_not_called() +def test_handle_connection_can_disable_idle_timeout(monkeypatch): + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + + async def fail_wait_for(*_args, **_kwargs): + pytest.fail("asyncio.wait_for should not be used when the idle timeout is disabled") + + monkeypatch.setattr(openpi_connection.asyncio, "wait_for", fail_wait_for) + websocket = FakeWebSocket([{"type": "websocket.disconnect"}]) + serving = _serving_mock() + + asyncio.run( + openpi_connection.RobotRealtimeConnection( + websocket, + serving, + idle_timeout=None, + ).handle_connection() + ) + + assert websocket.accepted is True + assert websocket.closed is False + serving.infer.assert_not_called() + + def test_handle_connection_keeps_session_state_per_websocket(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) requests = { @@ -343,6 +366,73 @@ def test_handle_connection_keeps_session_state_per_websocket(monkeypatch): assert calls[2].kwargs == {"session_id": "session-b", "reset": True} +def test_handle_connection_keeps_interleaved_parallel_session_state(monkeypatch): + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + requests = { + b"a1": {"prompt": "first-a", "session_id": "session-a"}, + b"b1": {"prompt": "first-b", "session_id": "session-b"}, + b"a2": {"prompt": "second-a", "session_id": "session-a"}, + b"b2": {"prompt": "second-b", "session_id": "session-b"}, + } + monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data])) + serving = _serving_mock() + websocket = FakeWebSocket( + [{"type": "websocket.receive", "bytes": key} for key in (b"a1", b"b1", b"a2", b"b2")] + + [{"type": "websocket.disconnect"}] + ) + + asyncio.run(openpi_connection.RobotRealtimeConnection(websocket, serving).handle_connection()) + + assert [call.kwargs for call in serving.infer.await_args_list] == [ + {"session_id": "session-a", "reset": True}, + {"session_id": "session-b", "reset": True}, + {"session_id": "session-a", "reset": False}, + {"session_id": "session-b", "reset": False}, + ] + + +def test_handle_connection_bounds_tracked_sessions_without_resetting_live_ones(monkeypatch): + """A client may hold one socket across many episodes, minting a fresh + session id per episode and never sending ``reset`` (RoboLab does exactly + this). Session tracking must stay bounded, and evicting finished episodes + must not disturb the sessions that are still running. + """ + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + monkeypatch.setattr(openpi_connection, "MAX_TRACKED_SESSIONS", 8) + + # One long-lived session interleaved with far more short episodes than the cap. + episodes = 40 + keys = [] + requests = {} + for episode in range(episodes): + for session_id in (f"episode-{episode}", "live-env"): + key = f"{session_id}@{episode}".encode() + keys.append(key) + requests[key] = {"prompt": "pick", "session_id": session_id} + + monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data])) + serving = _serving_mock() + websocket = FakeWebSocket( + [{"type": "websocket.receive", "bytes": key} for key in keys] + [{"type": "websocket.disconnect"}] + ) + connection = openpi_connection.RobotRealtimeConnection(websocket, serving) + + asyncio.run(connection.handle_connection()) + + assert len(connection._seen_sessions) <= 8 + + resets_by_session: dict[str, list[bool]] = {} + for call in serving.infer.await_args_list: + resets_by_session.setdefault(call.kwargs["session_id"], []).append(call.kwargs["reset"]) + + # Each new episode resets exactly once, on its first observation. + for episode in range(episodes): + assert resets_by_session[f"episode-{episode}"] == [True] + # The continuously active session is never evicted, so it resets only once + # despite ``episodes`` other sessions churning past the cap. + assert resets_by_session["live-env"] == [True] + [False] * (episodes - 1) + + def test_handle_connection_reset_endpoint_resets_next_infer(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) requests = { diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 072e734e4..82ee6e7dd 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -15,6 +15,7 @@ from starlette.testclient import TestClient from vllm_omni.entrypoints.openpi import connection as openpi_connection from vllm_omni.entrypoints.openpi import serving as openpi_serving +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -177,6 +178,22 @@ def test_create_policy_server_returns_none_without_policy_config(): assert serving is None +def test_legacy_policy_request_defaults_key_no_longer_controls_serving(): + od_config = SimpleNamespace( + model_config={ + "policy_server_config": {}, + "policy_request_defaults": "malformed legacy value", + } + ) + + serving = openpi_serving.ServingRealtimeRobotOpenPI.create_policy_server( + engine_client=SimpleNamespace(get_diffusion_od_config=lambda: od_config), + model_name="generic-policy", + ) + + assert serving is not None + + def test_policy_server_config_allows_explicit_empty_config(): serving = openpi_serving.ServingRealtimeRobotOpenPI( engine_client=_engine_with_policy_config(policy_config={}), @@ -221,6 +238,44 @@ def test_build_request_uses_unique_engine_request_id_per_inference(): assert request_a.request_id != request_b.request_id +def test_build_request_clones_stage_defaults_before_protocol_fields(): + default_params = OmniDiffusionSamplingParams( + extra_args={ + "format_prompt_as_json": True, + "session_id": "configured-session-must-not-win", + "reset": False, + "nested": {"values": []}, + } + ) + od_config = SimpleNamespace(model_config={"policy_server_config": {}}) + engine_client = SimpleNamespace( + get_diffusion_od_config=lambda: od_config, + default_sampling_params_list=[default_params], + ) + serving = openpi_serving.ServingRealtimeRobotOpenPI( + engine_client=engine_client, + model_name="custom-policy", + ) + + request_a = serving._build_request( + {"prompt": "pick up the object"}, + session_id="wire-session", + reset=True, + ) + request_b = serving._build_request( + {"prompt": "pick up the object"}, + session_id="wire-session", + reset=False, + ) + + assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True + assert request_a.sampling_params.extra_args["session_id"] == "wire-session" + assert request_a.sampling_params.extra_args["reset"] is True + request_a.sampling_params.extra_args["nested"]["values"].append("request-a") + assert request_b.sampling_params.extra_args["nested"] == {"values": []} + assert default_params.extra_args["nested"] == {"values": []} + + def test_build_request_forwards_seed_to_sampling_params(): """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds.""" serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) @@ -257,6 +312,29 @@ def test_infer_keeps_session_state_but_uses_unique_engine_request_ids(): assert sampling_params_b.extra_args["reset"] is False +def test_infer_yields_event_loop_while_engine_is_running(): + class DelayedEngine(RecordingEngine): + def generate(self, *, prompt, request_id, sampling_params_list): + async def _generate(): + await asyncio.sleep(0.03) + yield SimpleNamespace(multimodal_output={"actions": [0.0]}) + + return _generate() + + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=DelayedEngine()) + + async def run_test(): + infer_task = asyncio.create_task(serving.infer({"prompt": "pick"}, session_id="session-a", reset=True)) + ticks = 0 + while not infer_task.done(): + await asyncio.sleep(0.005) + ticks += 1 + await infer_task + return ticks + + assert asyncio.run(run_test()) >= 2 + + def test_two_websocket_clients_without_session_id_do_not_conflict(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", _json_pack) monkeypatch.setattr(openpi_connection, "_unpack", _json_unpack) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 0089330b2..083f96d9f 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -112,6 +112,26 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: assert args.cfg_parallel_size == 4 +def test_serve_parser_accepts_robot_openpi_idle_timeout() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "0"]) + + assert args.robot_openpi_idle_timeout == 0 + assert args.get_explicit_kwargs_dict()["robot_openpi_idle_timeout"] == 0 + + +def test_serve_parser_rejects_negative_robot_openpi_idle_timeout() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + with pytest.raises(SystemExit): + parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "-1"]) + + def test_serve_parser_accepts_ulysses_a2a_permute() -> None: parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="subcommand") diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 6bf823339..e55691e10 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -177,6 +177,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na "resolution", "image_size", "use_state", + "format_prompt_as_json", "observation", "robot_obs", "deterministic_seed", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index b711f0cfb..a69fb08c6 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -46,6 +46,7 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_SINGLE_STAGE_PIPELINE, BAGEL_THINK_PIPELINE, ) +from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE @@ -149,6 +150,11 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "lingbot_world": LINGBOT_WORLD_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, + # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3 + # checkpoints (which stay on the single-stage diffusion fallback), so this + # entry is only reachable through a deploy yaml's ``pipeline:`` key + # (see deploy/cosmos3_policy_droid.yaml). + "cosmos3_policy": COSMOS3_POLICY_PIPELINE, "gepard": GEPARD_PIPELINE, "glm_image": GLM_IMAGE_PIPELINE, "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f7d714890..0e962b4cc 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1130,10 +1130,22 @@ class StageConfig: _apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides) reconcile_diffusion_attention_overrides(engine_args, runtime_overrides) - # CLI overrides take precedence over YAML defaults + # CLI overrides take precedence over YAML defaults. Most dict-valued + # overrides are deep-merged so a partial CLI dict (e.g. --no-guardrails + # riding on ``model_config``) layers onto the deploy YAML instead of + # clobbering sibling keys such as ``policy_server_config`` — the same + # rationale as the platform-overlay deep-merge. Legacy atomic mappings + # are handled explicitly below. for key, value in runtime_overrides.items(): if value is not None and key not in ("devices", "max_batch_size", "num_replicas"): - engine_args[key] = value + existing = engine_args.get(key) + # ``omni_kv_config`` is an atomic legacy override: callers use + # a partial mapping to replace the topology-provided transfer + # role, rather than to add fields to it. + if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, dict): + engine_args[key] = _get_recursively_merged_dict(existing, value) + else: + engine_args[key] = value # Build runtime config from YAML defaults + CLI overrides runtime: dict[str, Any] = dict(self.yaml_runtime) diff --git a/vllm_omni/deploy/cosmos3_policy_droid.yaml b/vllm_omni/deploy/cosmos3_policy_droid.yaml new file mode 100644 index 000000000..cbcbcb42c --- /dev/null +++ b/vllm_omni/deploy/cosmos3_policy_droid.yaml @@ -0,0 +1,49 @@ +# Cosmos3 policy (DROID) deploy: single diffusion stage serving the OpenPI +# realtime API (/v1/realtime/robot/openpi). +# +# Topology (COSMOS3_POLICY_PIPELINE) is declared in +# vllm_omni/model_executor/models/cosmos3/pipeline.py. Policy checkpoints share +# their HF metadata (model_type=cosmos3_omni, model_index.json _class_name= +# Cosmos3OmniDiffusersPipeline) with the T2I/video Cosmos3 checkpoints, which +# must keep using the default single-stage diffusion fallback — so this +# pipeline cannot be auto-detected and must be selected explicitly: +# +# vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni \ +# --deploy-config /absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml +# +# Per-checkpoint model_config tweaks (e.g. policy_server_config values) can +# still be layered on top with --stage-overrides; dict overrides deep-merge +# with this file. + +pipeline: cosmos3_policy +# async_chunk defaults to True, which the deploy validator rejects for a +# single-stage pipeline (no next-stage processor). Must be false here. +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + # OpenPI policy serving is one observation -> one action reply. + max_num_seqs: 1 + model_class_name: Cosmos3OmniDiffusersPipeline + # JSON prompt formatting is a property of this DROID checkpoint's training + # recipe, not a generic server default. + default_sampling_params: + extra_args: + format_prompt_as_json: true + model_config: + # Policy serving emits robot actions, not text/video; the Cosmos3 + # guardrail stack (gated nvidia/Cosmos-1.0-Guardrail repo) is not part + # of it. Equivalent to launching with --no-guardrails. + guardrails: false + # OpenPI handshake metadata advertised to the robot client. Sent verbatim + # (the pipeline never reads these); they describe the camera inputs and + # action space the client must produce for this checkpoint. + policy_server_config: + image_resolution: [540, 640] + n_external_cameras: 2 + needs_wrist_camera: true + needs_stereo_camera: false + needs_session_id: true + action_space: joint_position diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index dc755b39c..affba4195 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -38,13 +38,24 @@ EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { "bridge_orig_lerobot": 7, "droid_lerobot": 8, "robomind-franka": 8, + "embodiment_b": 9, + # Backward-compatible alias for the pre-generator Cosmos Framework name. "galbot": 9, "robomind-franka-dual": 12, "robomind-ur": 13, "agibotworld": 15, + "embodiment_c_gripper": 15, + "embodiment_c_gripper_ext": 15, + # Backward-compatible aliases for the pre-generator Cosmos Framework names. "agibot_gear_gripper": 15, "agibot_gear_gripper_ext": 15, + "xdof_yam": 16, + "molmoact2_yam": 16, + "abc_yam": 16, "fractal": 20, + "drawanything": 21, + "behavior1k_lerobot": 22, + "maniparena": 23, } diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index c60817afb..241eae805 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -138,6 +138,7 @@ from .utils import ( ensure_gripper_array, extract_robolab_image, extract_robolab_prompt_image, + get_robolab_domain_id, lazy_action_transform_pipeline, make_robolab_action_postprocess_inputs, next_robolab_seed, @@ -1075,7 +1076,7 @@ class Cosmos3OmniDiffusersPipeline( self._guidance_scale = None self._num_timesteps = None self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None - self._robolab_transform = None + self._robolab_transforms: dict[bool, Any] = {} # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on # this pipeline. Tells the sequential-CFG loop to keep paired @@ -1394,11 +1395,18 @@ class Cosmos3OmniDiffusersPipeline( return val return default - def _get_robolab_transform(self): - if self._robolab_transform is None: + def _get_robolab_transform(self, *, format_prompt_as_json: bool = False): + transforms = getattr(self, "_robolab_transforms", None) + if transforms is None: + transforms = {} + self._robolab_transforms = transforms + if format_prompt_as_json not in transforms: action_dim = int(getattr(self.transformer, "action_dim", 64)) - self._robolab_transform = lazy_action_transform_pipeline(action_dim) - return self._robolab_transform + transforms[format_prompt_as_json] = lazy_action_transform_pipeline( + action_dim, + format_prompt_as_json=format_prompt_as_json, + ) + return transforms[format_prompt_as_json] def _build_robolab_policy_inputs( self, @@ -1443,7 +1451,8 @@ class Cosmos3OmniDiffusersPipeline( resolution = str(extra_param("resolution", ROBOLAB_DEFAULT_RESOLUTION)) fps = float(extra_param("conditioning_fps", ROBOLAB_DEFAULT_CONDITIONING_FPS)) domain_name = str(extra_param("domain_name", ROBOLAB_DEFAULT_DOMAIN_NAME)) - domain_id = resolve_domain_id(domain_name=domain_name, require_explicit=True) + domain_id = get_robolab_domain_id(domain_name) + format_prompt_as_json = self._truthy(extra_param("format_prompt_as_json", False)) if use_state and history_length < 1: raise ValueError("RoboLab history_length must be >= 1 when use_state is true.") @@ -1506,7 +1515,9 @@ class Cosmos3OmniDiffusersPipeline( "action": action, # Cosmos Framework consumes this as an integer conditioning bucket. "conditioning_fps": torch.tensor(fps, dtype=torch.long), - "mode": ACTION_MODE_POLICY, + # Cosmos Framework 1.2 renamed the internal policy transform mode to + # ``wam``. The public vLLM action_mode remains ``policy``. + "mode": "wam", "domain_id": torch.tensor(domain_id, dtype=torch.long), "viewpoint": "concat_view", "additional_view_description": ROBOLAB_CONCAT_VIEW_DESCRIPTION, @@ -1514,7 +1525,9 @@ class Cosmos3OmniDiffusersPipeline( if history_action is not None: sample["history_action"] = history_action - sample = self._get_robolab_transform()(sample, resolution) + sample = self._get_robolab_transform(format_prompt_as_json=format_prompt_as_json)(sample, resolution) + if isinstance(sample.get("ai_caption"), dict): + sample["ai_caption"] = json.dumps(sample["ai_caption"]) sequence_plan = sample["sequence_plan"] video_tensor = sample["video"].float() / 127.5 - 1.0 raw_action_dim_tensor = sample.get("raw_action_dim") diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py index 76df64b3b..8d104e987 100644 --- a/vllm_omni/diffusion/models/cosmos3/utils.py +++ b/vllm_omni/diffusion/models/cosmos3/utils.py @@ -78,6 +78,12 @@ ROBOLAB_CONCAT_VIEW_DESCRIPTION = ( "sides, with the robot visible." ) +_LEGACY_EMBODIMENT_ALIASES = { + "galbot": "embodiment_b", + "agibot_gear_gripper": "embodiment_c_gripper", + "agibot_gear_gripper_ext": "embodiment_c_gripper_ext", +} + @dataclass(frozen=True) class RoboLabPolicyInputs: @@ -252,33 +258,83 @@ def extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None: return ensure_rgb_uint8_image(image, "multi_modal_data.image") +_ROBOLAB_FRAMEWORK_IMPORT_ERROR = ( + "Cosmos3 RoboLab policy serving requires a compatible cosmos_framework source on PYTHONPATH." +) +_ROBOLAB_FRAMEWORK_IMPORTS: dict[str, tuple[str, str]] = { + "ActionTransformPipeline": ( + "cosmos_framework.data.generator.action.utils.transforms", + "ActionTransformPipeline", + ), + "convert_rotation": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "convert_rotation", + ), + "pose_abs_to_rel": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "pose_abs_to_rel", + ), + "pose_rel_to_abs": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "pose_rel_to_abs", + ), + "build_abs_pose_from_components": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "build_abs_pose_from_components", + ), + "FlowUniPCMultistepScheduler": ( + "cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc", + "FlowUniPCMultistepScheduler", + ), + "get_domain_id": ( + "cosmos_framework.data.generator.action.utils.domain_utils", + "get_domain_id", + ), +} + + def lazy_import(module_name: str, symbol_name: str, error_message: str): try: module = import_module(module_name) - except ModuleNotFoundError as exc: - raise ModuleNotFoundError(error_message) from exc - return getattr(module, symbol_name) + return getattr(module, symbol_name) + except (ImportError, AttributeError) as exc: + raise ImportError( + f"{error_message} Could not import required symbol {symbol_name!r} from {module_name!r}." + ) from exc -def lazy_action_transform_pipeline(max_action_dim: int): - ActionTransformPipeline = lazy_import( - "cosmos_framework.data.vfm.action.transforms", - "ActionTransformPipeline", - "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the " - "golden ActionTransformPipeline can be reused.", +def _load_robolab_framework_symbol(symbol_name: str): + module_name, exported_name = _ROBOLAB_FRAMEWORK_IMPORTS[symbol_name] + return lazy_import( + module_name, + exported_name, + _ROBOLAB_FRAMEWORK_IMPORT_ERROR, ) - return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0) + + +def preflight_cosmos3_action_framework_imports() -> None: + """Validate every Cosmos Framework symbol used by action policy serving.""" + for symbol_name in _ROBOLAB_FRAMEWORK_IMPORTS: + _load_robolab_framework_symbol(symbol_name) + + +def lazy_action_transform_pipeline(max_action_dim: int, *, format_prompt_as_json: bool = False): + ActionTransformPipeline = _load_robolab_framework_symbol("ActionTransformPipeline") + return ActionTransformPipeline( + max_action_dim=max_action_dim, + cfg_dropout_rate=0.0, + format_prompt_as_json=format_prompt_as_json, + ) + + +def get_robolab_domain_id(domain_name: str) -> int: + get_domain_id = _load_robolab_framework_symbol("get_domain_id") + canonical_name = _LEGACY_EMBODIMENT_ALIASES.get(domain_name.lower().strip(), domain_name) + return int(get_domain_id(canonical_name)) def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device): - FlowUniPCMultistepScheduler = lazy_import( - "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc", - "FlowUniPCMultistepScheduler", - ( - "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the " - "golden FlowUniPCMultistepScheduler can be reused." - ), - ) + FlowUniPCMultistepScheduler = _load_robolab_framework_symbol("FlowUniPCMultistepScheduler") scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, @@ -290,38 +346,22 @@ def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.de def convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray: - convert_rotation = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "convert_rotation", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + convert_rotation = _load_robolab_framework_symbol("convert_rotation") return convert_rotation(value, src, dst) def pose_abs_to_rel(*args, **kwargs) -> np.ndarray: - pose_abs_to_rel_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "pose_abs_to_rel", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + pose_abs_to_rel_func = _load_robolab_framework_symbol("pose_abs_to_rel") return pose_abs_to_rel_func(*args, **kwargs) def pose_rel_to_abs(*args, **kwargs) -> np.ndarray: - pose_rel_to_abs_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "pose_rel_to_abs", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + pose_rel_to_abs_func = _load_robolab_framework_symbol("pose_rel_to_abs") return pose_rel_to_abs_func(*args, **kwargs) def build_abs_pose_from_components(*args, **kwargs) -> np.ndarray: - build_abs_pose_from_components_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "build_abs_pose_from_components", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + build_abs_pose_from_components_func = _load_robolab_framework_symbol("build_abs_pose_from_components") return build_abs_pose_from_components_func(*args, **kwargs) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 348e41316..598cbe893 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -930,6 +930,15 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Disable Cosmos3 text/video safety guardrails for this server.", ) + omni_config_group.add_argument( + "--robot-openpi-idle-timeout", + type=_nonneg_finite_float, + default=30.0, + help=( + "Seconds the /v1/realtime/robot/openpi endpoint waits for the next request " + "before closing an idle WebSocket (default: 30). Set to 0 to disable the timeout." + ), + ) # Enable diffusion pipeline profiling omni_config_group.add_argument( @@ -1011,6 +1020,8 @@ def run_headless(args: TrackingNamespace) -> None: # Filter down to a dict of things explicitly requested by the user args_dict = args.get_explicit_kwargs_dict() + # This controls only the API-process WebSocket and is not a stage-engine option. + args_dict.pop("robot_openpi_idle_timeout", None) # This CLI-only negative alias is consumed below when selecting the # launcher log_stats value; it is not a per-stage config override. args_dict.pop("disable_log_stats", None) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index adb7d6e9d..565362312 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -485,6 +485,8 @@ async def build_async_omni_from_stage_config( try: kwargs = args.get_explicit_kwargs_dict() + # This controls only the API-process WebSocket and is not an AsyncOmni option. + kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) async_omni = AsyncOmni(model=model, **kwargs) @@ -1494,11 +1496,24 @@ async def realtime_robot_openpi(websocket: WebSocket): serving = getattr(websocket.app.state, "openai_serving_realtime_robot", None) if serving is None: + logger.warning( + "Rejecting robot OpenPI WebSocket: policy serving is disabled. " + "The diffusion stage must provide model_config.policy_server_config." + ) await websocket.accept() - await websocket.send_json({"type": "error", "error": "Robot policy not available", "code": "unsupported"}) + await websocket.send_json( + { + "type": "error", + "error": "Robot policy not available", + "code": "unsupported", + } + ) await websocket.close() return - connection = RobotRealtimeConnection(websocket, serving) + state_args = getattr(websocket.app.state, "args", None) + configured_timeout = getattr(state_args, "robot_openpi_idle_timeout", 30.0) + idle_timeout = None if configured_timeout == 0 else configured_timeout + connection = RobotRealtimeConnection(websocket, serving, idle_timeout=idle_timeout) await connection.handle_connection() diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index 9b8bfe0d8..112254fbe 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -29,6 +29,7 @@ mappings. from __future__ import annotations import asyncio +from collections import OrderedDict from typing import Any import msgspec @@ -45,10 +46,12 @@ logger = init_logger(__name__) _DEFAULT_IDLE_TIMEOUT = 30.0 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024 _MISSING = object() +# Upper bound on the per-connection set of seen session ids. +MAX_TRACKED_SESSIONS = 1024 def _pack_numpy(obj: Any) -> Any: - if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"): + if isinstance(obj, np.ndarray | np.generic) and obj.dtype.kind in ("V", "O", "c"): raise ValueError(f"Unsupported dtype: {obj.dtype}") if isinstance(obj, np.ndarray): if not obj.flags.c_contiguous: @@ -165,17 +168,38 @@ class RobotRealtimeConnection: self, websocket: WebSocket, serving: ServingRealtimeRobotOpenPI, - idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + idle_timeout: float | None = _DEFAULT_IDLE_TIMEOUT, ) -> None: self.websocket = websocket self.serving = serving self._idle_timeout = idle_timeout self._current_session_id: str | None = None - self._call_count = 0 + # Session ids seen on this connection, most-recently-used last. + self._seen_sessions: OrderedDict[str, None] = OrderedDict() def reset(self) -> None: self._current_session_id = None - self._call_count = 0 + self._seen_sessions.clear() + + def _mark_session_seen(self, session_id: str) -> bool: + """Record ``session_id``; return True when it is the first sighting. + + A session's first observation is the one that carries ``reset`` to the + policy, so this needs per-session memory. + """ + if session_id in self._seen_sessions: + self._seen_sessions.move_to_end(session_id) + return False + + self._seen_sessions[session_id] = None + if len(self._seen_sessions) > MAX_TRACKED_SESSIONS: + evicted, _ = self._seen_sessions.popitem(last=False) + logger.debug( + "Robot OpenPI connection is tracking more than %d sessions; dropped least-recently-used %s", + MAX_TRACKED_SESSIONS, + evicted, + ) + return True async def _send_error(self, message: str) -> None: await self.websocket.send_bytes(_pack({"type": "error", "message": message})) @@ -199,18 +223,22 @@ class RobotRealtimeConnection: await self.websocket.send_bytes(_pack(metadata)) while True: - try: - msg = await asyncio.wait_for( - self.websocket.receive(), - timeout=self._idle_timeout, - ) - except asyncio.TimeoutError: - logger.info("Robot OpenPI connection idle timeout after %.1f seconds", self._idle_timeout) + idle_timeout = self._idle_timeout + if idle_timeout is None: + msg = await self.websocket.receive() + else: try: - await self.websocket.close() - except Exception: - logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True) - return + msg = await asyncio.wait_for( + self.websocket.receive(), + timeout=idle_timeout, + ) + except asyncio.TimeoutError: + logger.info("Robot OpenPI connection idle timeout after %.1f seconds", idle_timeout) + try: + await self.websocket.close() + except Exception: + logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True) + return if msg.get("type") == "websocket.disconnect": break @@ -245,13 +273,12 @@ class RobotRealtimeConnection: session_id, ) self._current_session_id = session_id - self._call_count = 0 - self._call_count += 1 + reset = self._mark_session_seen(session_id) actions = await self.serving.infer( obs, session_id=session_id, - reset=self._call_count <= 1, + reset=reset, ) await self.websocket.send_bytes(_pack(actions)) except Exception: diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index f2448dfa1..80a3839ab 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -148,16 +148,32 @@ class ServingRealtimeRobotOpenPI: `AsyncOmni.generate()` and routed to the diffusion stage. """ from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.entrypoints.openai.stage_params import ( + clone_sampling_params, + get_default_sampling_params_list, + ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams - # An optional integer ``seed`` in the inference message becomes the engine request's - # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one. + # The engine applies stage default_sampling_params only to requests + # that carry no explicit params; this endpoint always passes explicit + # params, so start from a clone of the diffusion stage's defaults + # (e.g. a policy deploy yaml's ``extra_args``) and layer the OpenPI + # protocol fields on top. seed = obs.pop("seed", None) - extra_args = { - "reset": reset, - "session_id": session_id, - "robot_obs": obs, - } + sampling_params = OmniDiffusionSamplingParams() + for default_params in get_default_sampling_params_list(self.engine_client): + if isinstance(default_params, OmniDiffusionSamplingParams): + sampling_params = clone_sampling_params(default_params) + break + + extra_args = sampling_params.extra_args or {} + extra_args.update( + { + "reset": reset, + "session_id": session_id, + "robot_obs": obs, + } + ) prompt = obs.get("prompt", "") sampling_params = OmniDiffusionSamplingParams( diff --git a/vllm_omni/model_executor/models/cosmos3/__init__.py b/vllm_omni/model_executor/models/cosmos3/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/model_executor/models/cosmos3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/cosmos3/pipeline.py b/vllm_omni/model_executor/models/cosmos3/pipeline.py new file mode 100644 index 000000000..d26b4f46c --- /dev/null +++ b/vllm_omni/model_executor/models/cosmos3/pipeline.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 policy single-stage topology for online OpenPI serving. + +Cosmos3 policy checkpoints (e.g. ``nvidia/Cosmos3-Nano-Policy-DROID``) are one +diffusion stage: robot observation -> action chunk. They are +registered in ``OMNI_PIPELINES`` because online OpenPI serving +(``/v1/realtime/robot/openpi``) requires a registered pipeline so the deploy +yaml's ``model_config.policy_server_config`` reaches the websocket handshake. + +This pipeline declares neither ``hf_architectures`` nor +``diffusers_class_name``: policy checkpoints share their HF metadata +(``model_type=cosmos3_omni``, ``model_index.json`` ``_class_name= +Cosmos3OmniDiffusersPipeline``) with the T2I/video Cosmos3 checkpoints, which +must keep resolving through the default single-stage diffusion fallback. +Select this pipeline explicitly via a deploy yaml ``pipeline:`` key, e.g. +``vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni --deploy-config +/absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml``. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +COSMOS3_POLICY_PIPELINE = PipelineConfig( + model_type="cosmos3_policy", + model_arch="Cosmos3OmniDiffusersPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Cosmos3OmniDiffusersPipeline", + ), + ), +) diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py index 01cecf94f..452946b08 100644 --- a/vllm_omni/model_extras/cosmos3.py +++ b/vllm_omni/model_extras/cosmos3.py @@ -35,6 +35,7 @@ COSMOS3_EXTRA_BODY_PARAMS = frozenset( "resolution", "image_size", "use_state", + "format_prompt_as_json", "observation", "robot_obs", "deterministic_seed",