8e738ce21b
Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
372 lines
14 KiB
Python
372 lines
14 KiB
Python
# 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]
|