From 8e738ce21b1e2c5389b665ea30a308f7eb528cdd Mon Sep 17 00:00:00 2001 From: Srinivas Krovvidi <194645829+Srinivasoo7@users.noreply.github.com> Date: Sun, 13 Sep 2026 20:26:44 -0500 Subject: [PATCH] [P0][RFC #3747] RL rollout serving: session management + world_model_env (#4770) Signed-off-by: srinivas_oo7 Co-authored-by: srinivas_oo7 Co-authored-by: Zhou Taichang --- .../openai_api/test_api_server_guards.py | 8 + .../openai_api/test_openpi_serving.py | 14 + tests/rl_rollout/__init__.py | 2 + tests/rl_rollout/test_rollout_routes.py | 167 ++++++++ tests/rl_rollout/test_rollout_session.py | 157 ++++++++ tests/rl_rollout/test_serving_rl_rollout.py | 371 ++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 87 ++++ .../entrypoints/openai/protocol/rollout.py | 111 ++++++ .../entrypoints/openai/rollout_session.py | 107 +++++ .../entrypoints/openai/serving_rl_rollout.py | 370 +++++++++++++++++ vllm_omni/entrypoints/openpi/serving.py | 28 +- 11 files changed, 1421 insertions(+), 1 deletion(-) create mode 100644 tests/rl_rollout/__init__.py create mode 100644 tests/rl_rollout/test_rollout_routes.py create mode 100644 tests/rl_rollout/test_rollout_session.py create mode 100644 tests/rl_rollout/test_serving_rl_rollout.py create mode 100644 vllm_omni/entrypoints/openai/protocol/rollout.py create mode 100644 vllm_omni/entrypoints/openai/rollout_session.py create mode 100644 vllm_omni/entrypoints/openai/serving_rl_rollout.py diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index ee141182c..21725183a 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -80,6 +80,11 @@ _EXPECTED_ROUTER_ROUTES = { ("WEBSOCKET", "/v1/realtime"), ("WEBSOCKET", "/v1/realtime/robot/openpi"), ("WEBSOCKET", "/v1/duplex"), + ("POST", "/v1/realtime/sessions"), + ("POST", "/v1/realtime/sessions/{session_id}/step"), + ("POST", "/v1/realtime/sessions/{session_id}/reset"), + ("POST", "/v1/realtime/sessions/{session_id}/close"), + ("GET", "/v1/realtime/sessions/{session_id}/status"), ("GET", "/health"), ("GET", "/v1/models"), ("POST", "/v1/images/generations"), @@ -120,6 +125,7 @@ _DIFFUSION_APP_STATE_KEYS = { "openai_streaming_speech", "openai_streaming_video", "openai_serving_realtime_robot", + "rl_rollout_serving", "enable_server_load_tracking", "server_load_metrics", } @@ -151,12 +157,14 @@ _MULTISTAGE_APP_STATE_KEYS = { "openai_serving_realtime", "openai_serving_video", "openai_serving_realtime_robot", + "rl_rollout_serving", "enable_server_load_tracking", "server_load_metrics", } _MULTISTAGE_MUST_BE_NONE = { "openai_serving_duplex", "openai_serving_realtime_robot", + "rl_rollout_serving", } _MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 82ee6e7dd..46ff0f822 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -440,3 +440,17 @@ def test_extract_actions_does_not_iterate_result_object(): actions = serving._extract_actions(IterableResult()) np.testing.assert_allclose(actions, np.array([[1.0, 2.0, 3.0]], dtype=np.float32)) + + +def test_drop_session_calls_dreamzero_close_hook(): + closed: list[str] = [] + pipeline = SimpleNamespace(close_ar_diffusion_session=lambda session_id: closed.append(session_id)) + engine_client = SimpleNamespace( + model_config={"policy_server_config": TEST_POLICY_SERVER_CONFIG}, + model_runner=SimpleNamespace(pipeline=pipeline), + ) + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine_client) + + serving.drop_session("rollout-1") + + assert closed == ["rollout-1"] diff --git a/tests/rl_rollout/__init__.py b/tests/rl_rollout/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/rl_rollout/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/rl_rollout/test_rollout_routes.py b/tests/rl_rollout/test_rollout_routes.py new file mode 100644 index 000000000..c6a9e2e56 --- /dev/null +++ b/tests/rl_rollout/test_rollout_routes.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Route-level tests for RL rollout serving endpoints.""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm_omni.entrypoints.openai.api_server import router +from vllm_omni.entrypoints.openai.protocol.rollout import ( + ErrorObject, + RolloutStepResponse, + SessionMetadata, +) +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSessionClosedError, + RolloutSessionNotFoundError, +) + + +class _FakeRolloutServing: + def __init__(self, close_error: Exception | None = None, step_error_code: str | None = None) -> None: + self.close_error = close_error + self.step_error_code = step_error_code + self.step_request = None + + async def step(self, session_id, body): + self.step_request = body + error = None + if self.step_error_code is not None: + error = ErrorObject( + code=self.step_error_code, + message=f"{self.step_error_code}: {session_id}", + step_id=body.step_id, + committed_step_id=-1, + ) + return RolloutStepResponse( + step_id=body.step_id, + next_observation=None, + model_metadata=SessionMetadata( + latency_ms=0.0, + steps_generated=0, + context_length=0, + committed_step_id=-1, + ), + error=error, + ) + + async def close_session(self, session_id): + if self.close_error is not None: + raise self.close_error + + +def _client(serving: _FakeRolloutServing | None = None) -> TestClient: + app = FastAPI() + app.include_router(router) + if serving is not None: + app.state.rl_rollout_serving = serving + return TestClient(app) + + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_step_route_rejects_base64_image_payload(): + serving = _FakeRolloutServing() + client = _client(serving) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={ + "step_id": 0, + "observation": {"images": {"front": "base64-frames"}}, + "action": {}, + }, + ) + + assert response.status_code == 422 + assert serving.step_request is None + + +def test_step_route_accepts_nested_image_payload(): + serving = _FakeRolloutServing() + client = _client(serving) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={ + "step_id": 0, + "observation": {"images": {"observation/exterior_image_1_left": [[[0, 0, 0]]]}}, + "action": {}, + }, + ) + + assert response.status_code == 200 + assert serving.step_request.observation.images == {"observation/exterior_image_1_left": [[[0, 0, 0]]]} + + +def test_step_route_requires_action(): + client = _client(_FakeRolloutServing()) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}}, + ) + + assert response.status_code == 422 + + +def test_step_route_maps_missing_session_to_404(): + client = _client(_FakeRolloutServing(step_error_code="session_not_found")) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}, "action": {}}, + ) + + assert response.status_code == 404 + + +def test_step_route_maps_closed_session_to_410(): + client = _client(_FakeRolloutServing(step_error_code="session_closed")) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}, "action": {}}, + ) + + assert response.status_code == 410 + + +def test_step_route_maps_invalid_request_to_400(): + client = _client(_FakeRolloutServing(step_error_code="invalid_request")) + + response = client.post( + "/v1/realtime/sessions/s1/step", + json={"step_id": 0, "observation": {}, "action": {}}, + ) + + assert response.status_code == 400 + + +def test_close_route_maps_missing_session_to_404(): + client = _client(_FakeRolloutServing(close_error=RolloutSessionNotFoundError("s1"))) + + response = client.post("/v1/realtime/sessions/s1/close") + + assert response.status_code == 404 + + +def test_close_route_is_idempotent_for_closed_session(): + client = _client(_FakeRolloutServing(close_error=RolloutSessionClosedError("s1"))) + + response = client.post("/v1/realtime/sessions/s1/close") + + assert response.status_code == 200 + assert response.json()["closed"] is True + + +def test_rollout_routes_return_501_when_serving_is_unavailable(): + client = _client() + + response = client.post("/v1/realtime/sessions/s1/close") + + assert response.status_code == 501 diff --git a/tests/rl_rollout/test_rollout_session.py b/tests/rl_rollout/test_rollout_session.py new file mode 100644 index 000000000..2b3252a1a --- /dev/null +++ b/tests/rl_rollout/test_rollout_session.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for RolloutSessionStore (RFC #3747, P0).""" + +import asyncio + +import pytest + +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSession, + RolloutSessionClosedError, + RolloutSessionNotFoundError, + RolloutSessionStepError, + RolloutSessionStore, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def store() -> RolloutSessionStore: + return RolloutSessionStore() + + +@pytest.mark.asyncio +async def test_create_returns_session(store): + s = await store.create("s1", model="dreamzero", mode="world_model_env") + assert isinstance(s, RolloutSession) + assert s.session_id == "s1" + assert s.committed_step_id == -1 + assert s.context_length == 0 + assert not s.closed + + +@pytest.mark.asyncio +async def test_get_existing(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + s = await store.get("s1") + assert s.session_id == "s1" + + +@pytest.mark.asyncio +async def test_get_missing_raises(store): + with pytest.raises(RolloutSessionNotFoundError): + await store.get("does-not-exist") + + +@pytest.mark.asyncio +async def test_get_closed_raises(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.close("s1") + with pytest.raises(RolloutSessionClosedError): + await store.get("s1") + + +@pytest.mark.asyncio +async def test_close_missing_raises(store): + with pytest.raises(RolloutSessionNotFoundError): + await store.close("does-not-exist") + + +@pytest.mark.asyncio +async def test_close_is_idempotent(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.close("s1") + await store.close("s1") + s = await store.get("s1", include_closed=True) + assert s.closed + + +@pytest.mark.asyncio +async def test_advance_updates_committed_step_id(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + s = await store.get("s1") + assert s.committed_step_id == 0 + assert s.context_length == 1 + + await store.advance("s1", step_id=1) + s = await store.get("s1") + assert s.committed_step_id == 1 + assert s.context_length == 2 + + +@pytest.mark.asyncio +async def test_advance_rejects_duplicate_step_id(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + with pytest.raises(RolloutSessionStepError): + await store.advance("s1", step_id=0) + + s = await store.get("s1") + assert s.committed_step_id == 0 + assert s.context_length == 1 + + +@pytest.mark.asyncio +async def test_advance_rejects_out_of_order_step_id(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + with pytest.raises(RolloutSessionStepError): + await store.advance("s1", step_id=2) + + s = await store.get("s1") + assert s.committed_step_id == -1 + assert s.context_length == 0 + + +@pytest.mark.asyncio +async def test_reset_clears_state(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + await store.advance("s1", step_id=1) + await store.advance("s1", step_id=2) + await store.advance("s1", step_id=3) + await store.reset("s1") + s = await store.get("s1") + assert s.committed_step_id == -1 + assert s.context_length == 0 + + +@pytest.mark.asyncio +async def test_reset_closed_raises(store): + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.close("s1") + with pytest.raises(RolloutSessionClosedError): + await store.reset("s1") + + +@pytest.mark.asyncio +async def test_failed_step_does_not_advance(store): + """committed_step_id must not change when a step raises (RFC section 6.5 invariant).""" + await store.create("s1", model="dreamzero", mode="world_model_env") + await store.advance("s1", step_id=0) + + # Simulate a failed step: advance is NOT called + s = await store.get("s1") + assert s.committed_step_id == 0 # unchanged + + +@pytest.mark.asyncio +async def test_concurrent_steps_are_serialised(store): + """Per-session lock must prevent concurrent context mutations.""" + await store.create("s1", model="dreamzero", mode="world_model_env") + s = await store.get("s1") + + results: list[int] = [] + + async def slow_advance(step_id: int, delay: float) -> None: + await asyncio.sleep(delay) + async with s.lock: + await asyncio.sleep(0.01) + await store.advance("s1", step_id=step_id) + results.append(step_id) + + await asyncio.gather(slow_advance(0, 0.0), slow_advance(1, 0.01)) + # Both steps ran under the per-session lock and committed in order. + assert sorted(results) == [0, 1] diff --git a/tests/rl_rollout/test_serving_rl_rollout.py b/tests/rl_rollout/test_serving_rl_rollout.py new file mode 100644 index 000000000..5474ebc20 --- /dev/null +++ b/tests/rl_rollout/test_serving_rl_rollout.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for ServingRLRollout (RFC #3747, P0). + +The DreamZero engine is mocked. Tests verify session lifecycle, step_id +semantics, committed_step_id invariants, and error paths without a GPU. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import numpy as np +import pytest + +from vllm_omni.entrypoints.openai.protocol.rollout import ( + Action, + CreateSessionRequest, + Observation, + RolloutStepRequest, +) +from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_fake_result(video: Any = None, *, formatter_shaped: bool = True) -> MagicMock: + result = MagicMock() + if video is None: + video = np.zeros((1, 3, 64, 64), dtype=np.float32) + if formatter_shaped: + # Matches output_formatter: video is peeled into images; mm only has actions. + result.images = video + result.multimodal_output = {"actions": np.zeros((1, 8), dtype=np.float32)} + else: + result.images = None + result.multimodal_output = {"video": video} + return result + + +def _make_openpi(video: Any = None, raises: Exception | None = None, *, formatter_shaped: bool = True): + openpi = MagicMock() + openpi.engine_client = MagicMock() + openpi.model_name = None + + async def _generate(*_, **__): + if raises: + raise raises + yield _make_fake_result(video, formatter_shaped=formatter_shaped) + + openpi.engine_client.generate = _generate + openpi.build_request = MagicMock( + return_value=MagicMock( + prompt="", + prompts=[""], + request_id="test-req", + sampling_params=MagicMock(), + ) + ) + return openpi + + +@pytest.fixture +def serving() -> ServingRLRollout: + return ServingRLRollout(_make_openpi()) + + +# --------------------------------------------------------------------------- +# Session lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_session(serving): + resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + assert resp.session_id + assert resp.mode == "world_model_env" + + +@pytest.mark.asyncio +async def test_get_status_after_create(serving): + resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + status = await serving.get_status(resp.session_id) + assert status.committed_step_id == -1 + assert status.context_length == 0 + assert not status.closed + + +@pytest.mark.asyncio +async def test_close_session(serving): + resp = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + await serving.close_session(resp.session_id) + status = await serving.get_status(resp.session_id) + assert status.closed + assert status.committed_step_id == -1 + serving._openpi.drop_session.assert_any_call(resp.session_id) + serving._openpi.drop_session.assert_any_call(f"{resp.session_id}:stateless") + + +# --------------------------------------------------------------------------- +# Step - success path +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_successful_step_advances_committed_step_id(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest( + step_id=0, + observation=Observation(state=[0.1, 0.2]), + action=Action(joint_positions=[0.0] * 7), + ) + resp = await serving.step(sess.session_id, req) + assert resp.error is None + assert resp.model_metadata.committed_step_id == 0 + + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == 0 + assert status.context_length == 1 + + +@pytest.mark.asyncio +async def test_repeated_steps_advance_monotonically(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + for i in range(3): + req = RolloutStepRequest( + step_id=i, + observation=Observation(state=[float(i)]), + action=Action(joint_positions=[0.0] * 7), + ) + resp = await serving.step(sess.session_id, req) + assert resp.error is None + assert resp.model_metadata.committed_step_id == i + + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == 2 + assert status.context_length == 3 + + +@pytest.mark.asyncio +async def test_next_observation_present_on_success(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest( + step_id=0, + observation=Observation(), + action=Action(), + ) + resp = await serving.step(sess.session_id, req) + assert resp.next_observation is not None + assert "video_latent" in resp.next_observation + assert resp.next_observation["dtype"] == "float32" + + +@pytest.mark.asyncio +async def test_duplicate_step_does_not_rerun_or_advance(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + assert (await serving.step(sess.session_id, req)).error is None + + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "step_already_committed" + assert resp.model_metadata.committed_step_id == 0 + assert resp.model_metadata.context_length == 1 + serving._openpi.build_request.assert_called_once() + + +@pytest.mark.asyncio +async def test_out_of_order_step_does_not_rerun_or_advance(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=2, observation=Observation(), action=Action()) + + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "step_out_of_order" + assert resp.model_metadata.committed_step_id == -1 + assert resp.model_metadata.context_length == 0 + serving._openpi.build_request.assert_not_called() + + +@pytest.mark.asyncio +async def test_stateless_step_does_not_commit_context(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest( + step_id=7, + observation=Observation(), + action=Action(), + use_session_context=False, + ) + + resp = await serving.step(sess.session_id, req) + + assert resp.error is None + assert resp.model_metadata.committed_step_id == -1 + assert resp.model_metadata.context_length == 0 + serving._openpi.build_request.assert_called_once() + kwargs = serving._openpi.build_request.call_args.kwargs + assert kwargs["reset"] is True + assert kwargs["session_id"] == f"{sess.session_id}:stateless" + + +# --------------------------------------------------------------------------- +# Step - failure path: committed_step_id must NOT advance +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_failed_step_does_not_advance_committed_step_id(): + serving = ServingRLRollout(_make_openpi(raises=RuntimeError("engine dead"))) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "inference_error" + assert resp.model_metadata.committed_step_id == -1 # must not have advanced + + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == -1 + + +@pytest.mark.asyncio +async def test_output_encoding_failure_does_not_advance_committed_step_id(): + class BadVideo: + def detach(self): + return self + + def float(self): + return self + + def cpu(self): + return self + + def numpy(self): + raise TypeError("cannot encode") + + serving = ServingRLRollout(_make_openpi(video=BadVideo())) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + + resp = await serving.step(sess.session_id, req) + + assert resp.error is not None + assert resp.error.code == "inference_error" + assert resp.model_metadata.committed_step_id == -1 + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == -1 + + +@pytest.mark.asyncio +async def test_failed_step_after_success_preserves_last_committed(): + good_serving = ServingRLRollout(_make_openpi()) + sess = await good_serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + # Step 0 succeeds + await good_serving.step(sess.session_id, RolloutStepRequest(step_id=0, observation=Observation(), action=Action())) + + # Swap engine to failing one - reuse the same store via internal reference + good_serving._openpi = _make_openpi(raises=RuntimeError("oops")) + resp = await good_serving.step( + sess.session_id, RolloutStepRequest(step_id=1, observation=Observation(), action=Action()) + ) + assert resp.error is not None + assert resp.model_metadata.committed_step_id == 0 # still at last good step + assert resp.model_metadata.context_length == 1 + + +# --------------------------------------------------------------------------- +# Step - unknown session +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_step_unknown_session_returns_error(serving): + req = RolloutStepRequest(step_id=0, observation=Observation(), action=Action()) + resp = await serving.step("no-such-session", req) + assert resp.error is not None + assert resp.error.code == "session_not_found" + + +# --------------------------------------------------------------------------- +# Reset +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_reset_clears_committed_step_id(serving): + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + await serving.step(sess.session_id, RolloutStepRequest(step_id=0, observation=Observation(), action=Action())) + await serving.reset_session(sess.session_id) + status = await serving.get_status(sess.session_id) + assert status.committed_step_id == -1 + assert status.context_length == 0 + serving._openpi.drop_session.assert_called_with(sess.session_id) + + +# --------------------------------------------------------------------------- +# action -> obs folding (module assumption) +# --------------------------------------------------------------------------- + + +def test_merge_action_into_obs_emits_dreamzero_state_keys(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + obs = Observation(state=[9.0] * 8, prompt="test") + action = Action(joint_positions=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], gripper_position=0.5) + result = _merge_action_into_obs(obs, action) + + assert result["prompt"] == "test" + np.testing.assert_array_almost_equal(result["observation/joint_position"], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + np.testing.assert_array_almost_equal(result["observation/gripper_position"], [0.5]) + + +def test_merge_action_none_splits_obs_state_for_dreamzero(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + obs = Observation(state=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 0.25]) + result = _merge_action_into_obs(obs, None) + np.testing.assert_array_almost_equal(result["observation/joint_position"], [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]) + np.testing.assert_array_almost_equal(result["observation/gripper_position"], [0.25]) + + +def test_merge_no_state_no_action_omits_state_key(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + obs = Observation() + result = _merge_action_into_obs(obs, None) + assert "observation/joint_position" not in result + assert "observation/gripper_position" not in result + + +def test_merge_rejects_wrong_dreamzero_joint_dimension(): + from vllm_omni.entrypoints.openai.serving_rl_rollout import _merge_action_into_obs + + with pytest.raises(ValueError, match="joint_positions must have 7 values"): + _merge_action_into_obs(Observation(), Action(joint_positions=[1.0, 2.0])) + + +@pytest.mark.asyncio +async def test_step_reads_formatter_images_not_multimodal_video(): + video = np.ones((1, 3, 8, 8), dtype=np.float32) + serving = ServingRLRollout(_make_openpi(video=video, formatter_shaped=True)) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + resp = await serving.step( + sess.session_id, + RolloutStepRequest(step_id=0, observation=Observation(), action=Action()), + ) + assert resp.error is None + assert resp.next_observation is not None + assert "video_latent" in resp.next_observation + assert resp.next_observation["shape"] == [1, 3, 8, 8] + + +@pytest.mark.asyncio +async def test_step_falls_back_to_multimodal_video(): + video = np.ones((2, 4, 4), dtype=np.float32) + serving = ServingRLRollout(_make_openpi(video=video, formatter_shaped=False)) + sess = await serving.create_session(CreateSessionRequest(model="dreamzero", mode="world_model_env")) + resp = await serving.step( + sess.session_id, + RolloutStepRequest(step_id=0, observation=Observation(), action=Action()), + ) + assert resp.error is None + assert resp.next_observation is not None + assert resp.next_observation["shape"] == [2, 4, 4] diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f4f22f95f..8acd1d6fa 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -138,6 +138,10 @@ from vllm_omni.entrypoints.openai.protocol.images import ( ImageGenerationResponse, ResponseFormat, ) +from vllm_omni.entrypoints.openai.protocol.rollout import ( + CreateSessionRequest, + RolloutStepRequest, +) from vllm_omni.entrypoints.openai.protocol.videos import ( VideoDeleteResponse, VideoGenerationRequest, @@ -146,8 +150,14 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoResponse, ) from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSessionCapacityError, + RolloutSessionClosedError, + RolloutSessionNotFoundError, +) from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServingAudioGenerate from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.entrypoints.openai.serving_rl_rollout import ServingRLRollout from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler from vllm_omni.entrypoints.openai.serving_video import ( @@ -604,6 +614,10 @@ async def omni_init_app_state( engine_client=engine_client, model_name=model_name, ) + if state.openai_serving_realtime_robot is not None: + state.rl_rollout_serving = ServingRLRollout(state.openai_serving_realtime_robot) + else: + state.rl_rollout_serving = None state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) state.server_load_metrics = 0 @@ -949,6 +963,7 @@ async def omni_init_app_state( stage_configs=state.stage_configs, ) state.openai_serving_realtime_robot = None + state.rl_rollout_serving = None state.enable_server_load_tracking = args.enable_server_load_tracking state.server_load_metrics = 0 @@ -1529,6 +1544,78 @@ async def duplex_websocket(websocket: WebSocket): await handler.handle_session(websocket) +# RL Rollout serving (RFC #3747, P0) + + +def _rl_rollout_serving(request: Request) -> ServingRLRollout: + serving = getattr(request.app.state, "rl_rollout_serving", None) + if serving is None: + raise HTTPException(status_code=501, detail="RL rollout serving not available for this model.") + return serving + + +@router.post("/v1/realtime/sessions") +async def create_rollout_session(body: CreateSessionRequest, request: Request): + serving = _rl_rollout_serving(request) + try: + return (await serving.create_session(body)).model_dump() + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except RolloutSessionCapacityError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + + +@router.post("/v1/realtime/sessions/{session_id}/step") +async def rollout_step(session_id: str, body: RolloutStepRequest, request: Request): + serving = _rl_rollout_serving(request) + response = await serving.step(session_id, body) + if response.error is not None: + status_code = None + if response.error.code == "session_not_found": + status_code = 404 + elif response.error.code == "session_closed": + status_code = 410 + elif response.error.code in {"invalid_request", "step_already_committed", "step_out_of_order"}: + status_code = 400 + if status_code is not None: + return JSONResponse(content=response.model_dump(), status_code=status_code) + return response.model_dump() + + +@router.post("/v1/realtime/sessions/{session_id}/reset") +async def reset_rollout_session(session_id: str, request: Request): + serving = _rl_rollout_serving(request) + try: + return (await serving.reset_session(session_id)).model_dump() + except RolloutSessionNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.") + except RolloutSessionClosedError: + raise HTTPException(status_code=410, detail=f"Session {session_id!r} is closed.") + + +@router.post("/v1/realtime/sessions/{session_id}/close") +async def close_rollout_session(session_id: str, request: Request): + serving = _rl_rollout_serving(request) + try: + await serving.close_session(session_id) + return {"session_id": session_id, "closed": True} + except RolloutSessionNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.") + except RolloutSessionClosedError: + return {"session_id": session_id, "closed": True} + + +@router.get("/v1/realtime/sessions/{session_id}/status") +async def rollout_session_status(session_id: str, request: Request): + serving = _rl_rollout_serving(request) + try: + return (await serving.get_status(session_id)).model_dump() + except RolloutSessionNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id!r} not found.") + except RolloutSessionClosedError: + raise HTTPException(status_code=410, detail=f"Session {session_id!r} is closed.") + + # Health and Model endpoints for diffusion mode diff --git a/vllm_omni/entrypoints/openai/protocol/rollout.py b/vllm_omni/entrypoints/openai/protocol/rollout.py new file mode 100644 index 000000000..819981d7b --- /dev/null +++ b/vllm_omni/entrypoints/openai/protocol/rollout.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pydantic schemas for RL rollout serving (RFC #3747). + +P0 scope: world_model_env mode only (observation + action -> next observation). + +P0 maps rollout observations/actions onto DreamZero's DROID robot_obs schema: +camera images retain their DreamZero observation/* keys, joint/gripper state is +passed through observation/joint_position and observation/gripper_position, and +the returned video payload is the DreamZero VAE latent unless a later endpoint +adds explicit latent decoding. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +class Observation(BaseModel): + """Single-step robot observation.""" + + images: dict[str, list] | None = Field( + default=None, + description="Named DreamZero camera images as nested numeric lists. Base64 strings are not accepted in P0.", + ) + state: list[float] | None = Field( + default=None, + description="Proprioceptive state vector (joint positions, velocities, etc.).", + ) + prompt: str = Field(default="", description="Optional language conditioning.") + extra: dict[str, Any] = Field( + default_factory=dict, + description="Pass-through fields forwarded verbatim to robot_obs.", + ) + + +class Action(BaseModel): + """Executed action for world_model_env conditioning.""" + + joint_positions: list[float] | None = Field( + default=None, + description="Executed joint positions; forwarded as observation/joint_position for DreamZero.", + ) + gripper_position: float | list[float] | None = Field( + default=None, + description="Executed gripper position; forwarded as observation/gripper_position for DreamZero.", + ) + extra: dict[str, Any] = Field(default_factory=dict) + + +class SessionMetadata(BaseModel): + latency_ms: float + steps_generated: int + context_length: int + committed_step_id: int = Field( + description="Highest step_id whose context has been atomically committed. " + "-1 means no step has been committed yet (fresh or reset session).", + ) + session_memory_bytes: int | None = None + uncertainty: float | None = None + + +class ErrorObject(BaseModel): + code: str + message: str + step_id: int | None = None + committed_step_id: int = -1 + + +class CreateSessionRequest(BaseModel): + model: str + mode: Literal["world_model_env"] = "world_model_env" + + +class CreateSessionResponse(BaseModel): + session_id: str + mode: str + created_at: float + + +class RolloutStepRequest(BaseModel): + step_id: int = Field(ge=0, description="Monotonically increasing per session.") + observation: Observation + action: Action = Field( + description="Required for world_model_env: executed action that produced this observation.", + ) + use_session_context: bool = True + + +class RolloutStepResponse(BaseModel): + step_id: int + next_observation: dict[str, Any] | None = Field( + default=None, + description="Predicted next observation metadata. P0 returns DreamZero video_latent, not RGB frames.", + ) + model_metadata: SessionMetadata + error: ErrorObject | None = None + + +class ResetSessionResponse(BaseModel): + session_id: str + committed_step_id: int + + +class SessionStatusResponse(BaseModel): + session_id: str + committed_step_id: int + context_length: int + closed: bool diff --git a/vllm_omni/entrypoints/openai/rollout_session.py b/vllm_omni/entrypoints/openai/rollout_session.py new file mode 100644 index 000000000..76848ba31 --- /dev/null +++ b/vllm_omni/entrypoints/openai/rollout_session.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""In-memory session store for RL rollout serving (RFC #3747, P0). + +Each RolloutSession holds a per-session asyncio.Lock so that concurrent step +requests for the same session are serialised. committed_step_id advances only +on successful step completion; failed or timed-out steps leave it unchanged, +preserving the atomic-context-commit guarantee from section 6.5 of the RFC. +""" + +from __future__ import annotations + +import asyncio +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import Literal + + +class RolloutSessionNotFoundError(KeyError): + pass + + +class RolloutSessionClosedError(RuntimeError): + pass + + +class RolloutSessionStepError(ValueError): + pass + + +class RolloutSessionCapacityError(RuntimeError): + pass + + +@dataclass +class RolloutSession: + session_id: str + model: str + mode: Literal["world_model_env"] + created_at: float = field(default_factory=time.time) + committed_step_id: int = -1 + context_length: int = 0 + closed: bool = False + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +class RolloutSessionStore: + """Thread-safe (asyncio) store for active rollout sessions.""" + + def __init__(self, max_sessions: int = 64) -> None: + self._sessions: OrderedDict[str, RolloutSession] = OrderedDict() + self._max_sessions = max_sessions + self._store_lock = asyncio.Lock() + + async def create(self, session_id: str, model: str, mode: Literal["world_model_env"]) -> RolloutSession: + async with self._store_lock: + self._evict_closed_sessions_for_create() + if len(self._sessions) >= self._max_sessions: + raise RolloutSessionCapacityError(f"Maximum active rollout sessions reached: {self._max_sessions}.") + session = RolloutSession( + session_id=session_id, + model=model, + mode=mode, + ) + self._sessions[session_id] = session + self._sessions.move_to_end(session_id) + return session + + async def get(self, session_id: str, *, include_closed: bool = False) -> RolloutSession: + session = self._sessions.get(session_id) + if session is None: + raise RolloutSessionNotFoundError(session_id) + self._sessions.move_to_end(session_id) + if session.closed and not include_closed: + raise RolloutSessionClosedError(session_id) + return session + + async def reset(self, session_id: str) -> RolloutSession: + session = await self.get(session_id) + async with session.lock: + session.committed_step_id = -1 + session.context_length = 0 + return session + + async def close(self, session_id: str) -> None: + session = await self.get(session_id, include_closed=True) + async with session.lock: + session.closed = True + + async def advance(self, session_id: str, step_id: int) -> None: + """Commit a successfully completed step. Called only on success.""" + session = await self.get(session_id) + expected_step_id = session.committed_step_id + 1 + if step_id != expected_step_id: + raise RolloutSessionStepError(f"Expected step_id {expected_step_id}, got {step_id}.") + session.committed_step_id = step_id + session.context_length += 1 + + def _evict_closed_sessions_for_create(self) -> None: + while len(self._sessions) >= self._max_sessions: + for session_id, session in self._sessions.items(): + if session.closed: + del self._sessions[session_id] + break + else: + break diff --git a/vllm_omni/entrypoints/openai/serving_rl_rollout.py b/vllm_omni/entrypoints/openai/serving_rl_rollout.py new file mode 100644 index 000000000..38b53d9bd --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_rl_rollout.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Serving layer for RL rollout (RFC #3747, P0). + +P0 scope: world_model_env mode, single-session, DreamZero backbone. + +P0 maps client rollout input onto DreamZero's DROID robot_obs schema and returns +DreamZero video latents explicitly. RGB/frame decoding can be added as a later +contract without pretending the latent is already an image observation. +""" + +from __future__ import annotations + +import binascii +import time +from collections import abc +from typing import Any + +import numpy as np +from vllm.logger import init_logger +from vllm.utils import random_uuid + +from vllm_omni.entrypoints.openai.protocol.rollout import ( + Action, + CreateSessionRequest, + CreateSessionResponse, + ErrorObject, + Observation, + ResetSessionResponse, + RolloutStepRequest, + RolloutStepResponse, + SessionMetadata, + SessionStatusResponse, +) +from vllm_omni.entrypoints.openai.rollout_session import ( + RolloutSession, + RolloutSessionClosedError, + RolloutSessionNotFoundError, + RolloutSessionStore, +) +from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI + +logger = init_logger(__name__) + + +def _merge_action_into_obs(obs: Observation, action: Action | None) -> dict[str, Any]: + """Build DreamZero-compatible robot_obs.""" + robot_obs: dict[str, Any] = {"prompt": obs.prompt} + + if obs.images: + robot_obs.update(obs.images) + + if obs.extra: + for key, value in obs.extra.items(): + if key.startswith("observation/") or key in {"seed", "embodiment_name"}: + robot_obs[key] = value + + joint_positions = action.joint_positions if action is not None else None + gripper_position = action.gripper_position if action is not None else None + + if joint_positions is None and obs.state: + joint_positions = obs.state[:7] + if len(obs.state) > 7 and gripper_position is None: + gripper_position = obs.state[7:8] + + if joint_positions is not None: + joint_arr = np.asarray(joint_positions, dtype=np.float64).flatten() + if joint_arr.size != 7: + raise ValueError(f"DreamZero DROID joint_positions must have 7 values, got {joint_arr.size}.") + robot_obs["observation/joint_position"] = joint_arr + + if gripper_position is not None: + gripper_arr = np.asarray(gripper_position, dtype=np.float64).flatten() + if gripper_arr.size != 1: + raise ValueError(f"DreamZero DROID gripper_position must have 1 value, got {gripper_arr.size}.") + robot_obs["observation/gripper_position"] = gripper_arr + + return robot_obs + + +def _encode_video_output(video: Any) -> dict[str, Any]: + """Encode DreamZero video latent tensor / ndarray to JSON metadata.""" + if video is None: + return {} + if hasattr(video, "detach") and callable(video.detach): + video = video.detach() + if hasattr(video, "float") and callable(video.float): + video = video.float() + if hasattr(video, "cpu") and callable(video.cpu): + video = video.cpu() + if hasattr(video, "numpy"): + video = video.numpy() + if isinstance(video, np.ndarray): + return { + "video_latent": binascii.b2a_base64(video.tobytes(), newline=False).decode(), + "shape": list(video.shape), + "dtype": str(video.dtype), + "encoding": "base64_raw_tensor", + } + return {"video_latent": str(video), "encoding": "string"} + + +class ServingRLRollout: + """HTTP serving layer for RL rollout sessions. + + Wraps ServingRealtimeRobotOpenPI for world_model_env mode. + Session state and committed_step_id tracking live in RolloutSessionStore. + """ + + def __init__(self, openpi_serving: ServingRealtimeRobotOpenPI) -> None: + self._openpi = openpi_serving + self._store = RolloutSessionStore() + + # ------------------------------------------------------------------ # + # Session lifecycle # + # ------------------------------------------------------------------ # + + async def create_session(self, req: CreateSessionRequest) -> CreateSessionResponse: + self._validate_model(req.model) + session_id = random_uuid() + session = await self._store.create( + session_id=session_id, + model=req.model, + mode=req.mode, + ) + logger.info("Created rollout session %s mode=%s", session_id, req.mode) + return CreateSessionResponse( + session_id=session_id, + mode=session.mode, + created_at=session.created_at, + ) + + async def reset_session(self, session_id: str) -> ResetSessionResponse: + session = await self._store.reset(session_id) + # Next infer still passes reset=True; drop GPU KV now if the engine + # session is reachable from this process. + self._drop_engine_session(session_id) + logger.info("Reset rollout session %s", session_id) + return ResetSessionResponse( + session_id=session_id, + committed_step_id=session.committed_step_id, + ) + + async def close_session(self, session_id: str) -> None: + await self._store.close(session_id) + self._drop_engine_session(session_id) + self._drop_engine_session(f"{session_id}:stateless") + logger.info("Closed rollout session %s", session_id) + + async def get_status(self, session_id: str) -> SessionStatusResponse: + session = await self._store.get(session_id, include_closed=True) + return SessionStatusResponse( + session_id=session_id, + committed_step_id=session.committed_step_id, + context_length=session.context_length, + closed=session.closed, + ) + + # ------------------------------------------------------------------ # + # Step # + # ------------------------------------------------------------------ # + + async def step( + self, + session_id: str, + req: RolloutStepRequest, + ) -> RolloutStepResponse: + try: + session = await self._store.get(session_id) + except RolloutSessionNotFoundError: + return self._error_response( + req.step_id, + -1, + 0, + "session_not_found", + f"Session {session_id!r} does not exist.", + ) + except RolloutSessionClosedError: + return self._error_response( + req.step_id, + -1, + 0, + "session_closed", + f"Session {session_id!r} is closed.", + ) + + async with session.lock: + return await self._run_step(session, req) + + async def _run_step( + self, + session: RolloutSession, + req: RolloutStepRequest, + ) -> RolloutStepResponse: + t0 = time.perf_counter() + committed = session.committed_step_id + + if req.use_session_context: + expected_step_id = committed + 1 + if req.step_id != expected_step_id: + if req.step_id <= committed: + code = "step_already_committed" + message = f"Step {req.step_id} is already committed; highest committed step is {committed}." + else: + code = "step_out_of_order" + message = f"Expected step_id {expected_step_id}, got {req.step_id}." + return self._error_response( + req.step_id, + committed, + session.context_length, + code, + message, + ) + + # reset=True on first call after session create/reset + reset = committed == -1 or not req.use_session_context + engine_session_id = session.session_id + if not req.use_session_context: + engine_session_id = f"{session.session_id}:stateless" + try: + robot_obs = _merge_action_into_obs(req.observation, req.action) + video_out = await self._infer_world_model( + obs=robot_obs, + session_id=engine_session_id, + reset=reset, + ) + next_observation = _encode_video_output(video_out) + except ValueError as exc: + logger.info("Invalid rollout step %d for session %s: %s", req.step_id, session.session_id, exc) + return self._error_response( + req.step_id, + committed, + session.context_length, + "invalid_request", + str(exc), + ) + except Exception as exc: + logger.exception("Step %d failed for session %s", req.step_id, session.session_id) + # Do NOT advance committed_step_id on failure (RFC section 6.5). + return self._error_response( + req.step_id, + committed, + session.context_length, + "inference_error", + str(exc), + ) + + if req.use_session_context: + # Only advance committed_step_id on successful contextual steps. + await self._store.advance(session.session_id, req.step_id) + committed = req.step_id + + latency_ms = (time.perf_counter() - t0) * 1000.0 + metadata = SessionMetadata( + latency_ms=round(latency_ms, 2), + steps_generated=1, + context_length=session.context_length, + committed_step_id=committed, + ) + return RolloutStepResponse( + step_id=req.step_id, + next_observation=next_observation, + model_metadata=metadata, + ) + + async def _infer_world_model( + self, + obs: dict[str, Any], + *, + session_id: str, + reset: bool, + ) -> Any: + """Call the engine and return the video (next-observation) output. + + Reuses ServingRealtimeRobotOpenPI.build_request() so request routing, + session_id threading, and OmniDiffusionSamplingParams construction are + identical to the policy_inference path. DreamZero's formatter peels + ``video`` into ``result.images`` and leaves only ``actions`` on + ``multimodal_output``; read images first and fall back to the raw key. + """ + request = self._openpi.build_request(obs, session_id=session_id, reset=reset) + result = None + async for output in self._openpi.engine_client.generate( + prompt=self._request_prompt(request), + request_id=request.request_id, + sampling_params_list=[request.sampling_params], + ): + result = output + + if result is None: + raise RuntimeError("World model request produced no output.") + return self._extract_video_output(result) + + @staticmethod + def _request_prompt(request: Any) -> Any: + prompt = getattr(request, "prompt", None) + if prompt is not None: + return prompt + return request.prompts[0] + + @staticmethod + def _extract_video_output(result: Any) -> Any: + """Prefer formatter-placed ``images``; fall back to multimodal video.""" + images = getattr(result, "images", None) + if isinstance(images, (list, tuple)): + if images: + return images[0] + elif images is not None and not isinstance(images, abc.Mapping): + # Formatter assigns the peeled video tensor directly to images. + return images + + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, abc.Mapping): + video = multimodal_output.get("video") + if video is not None: + return video + + raise RuntimeError( + "No video on result.images or multimodal_output['video']. " + "Confirm DreamZero formatter placement for world_model_env." + ) + + def _drop_engine_session(self, session_id: str) -> None: + drop = getattr(self._openpi, "drop_session", None) + if not callable(drop): + return + try: + drop(session_id) + except Exception: + logger.exception("Failed to drop engine session %s", session_id) + + # ------------------------------------------------------------------ # + # Helpers # + # ------------------------------------------------------------------ # + + def _validate_model(self, requested_model: str) -> None: + served_model = getattr(self._openpi, "model_name", None) + if served_model is None: + return + if requested_model == served_model: + return + if requested_model.lower() == "dreamzero" and "dreamzero" in served_model.lower(): + return + raise ValueError(f"Requested rollout model {requested_model!r} does not match served model {served_model!r}.") + + @staticmethod + def _error_response( + step_id: int, + committed_step_id: int, + context_length: int, + code: str, + message: str, + ) -> RolloutStepResponse: + metadata = SessionMetadata( + latency_ms=0.0, + steps_generated=0, + context_length=context_length, + committed_step_id=committed_step_id, + ) + return RolloutStepResponse( + step_id=step_id, + next_observation=None, + model_metadata=metadata, + error=ErrorObject( + code=code, + message=message, + step_id=step_id, + committed_step_id=committed_step_id, + ), + ) diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 80a3839ab..562ed70d4 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -119,10 +119,32 @@ class ServingRealtimeRobotOpenPI: def reset(self, obs: dict) -> None: """Compatibility hook; per-connection state lives in RobotRealtimeConnection.""" + def drop_session(self, session_id: str) -> None: + """Best-effort release of model-side session state for a closed rollout.""" + drop = getattr(self.engine_client, "drop_session", None) + if callable(drop): + drop(session_id) + return + pipeline = self._pipeline() + for name in ("close_ar_diffusion_session", "drop_session_state"): + close = getattr(pipeline, name, None) + if callable(close): + close(session_id) + return + + def _pipeline(self) -> Any: + engine = self.engine_client + for attr in ("model_runner", "runner", "diffusion_model_runner"): + runner = getattr(engine, attr, None) + pipeline = getattr(runner, "pipeline", None) if runner is not None else None + if pipeline is not None: + return pipeline + return getattr(engine, "pipeline", None) + async def infer(self, obs: dict, *, session_id: str, reset: bool) -> ActionOutput: """raw obs → engine → actions.""" # Build request, run inference through AsyncOmni - request = self._build_request(obs, session_id=session_id, reset=reset) + request = self.build_request(obs, session_id=session_id, reset=reset) result = None # OpenPI policy serving is one request -> one action reply. AsyncOmni # exposes an async iterator, so consume it to completion and use the @@ -141,6 +163,10 @@ class ServingRealtimeRobotOpenPI: def _next_request_id(self, session_id: str) -> str: return f"robot-{session_id}-{next(self._request_counter)}" + def build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: + """Build an engine request from raw robot obs.""" + return self._build_request(obs, session_id=session_id, reset=reset) + def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: """Build engine request from raw robot obs.