87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""E2E online serving test for π0.5 OpenPI websocket serving.
|
|
|
|
Boots ``vllm serve --omni --deploy-config pi05.yaml`` and drives the real OpenPI
|
|
websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses
|
|
(handshake metadata → send observation → receive action chunk). Mirrors
|
|
``tests/e2e/online_serving/test_pi0_expansion.py``. Needs one H100 and the full
|
|
π0.5 checkpoint.
|
|
|
|
The ``pi0_openpi_*`` helpers are protocol-level (handshake, send obs, receive
|
|
chunk) and not π0-specific, so π0.5 reuses them; only the deploy config and the
|
|
expected metadata differ. This file is selected explicitly by the nightly H100
|
|
robot-policy job.
|
|
|
|
The in-process LeRobot parity oracle lives separately in
|
|
``tests/diffusion/models/pi05/test_pi05_parity.py``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import (
|
|
OmniServerParams,
|
|
get_open_port,
|
|
pi0_openpi_require_dependencies,
|
|
pi0_openpi_run_policy_session,
|
|
pi0_openpi_validate_session_result,
|
|
)
|
|
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
|
|
|
|
MODEL = "lerobot/pi05_base"
|
|
|
|
|
|
test_params = [
|
|
OmniServerParams(
|
|
model=MODEL,
|
|
port=8093,
|
|
server_args=[
|
|
"--deploy-config",
|
|
"vllm_omni/deploy/pi05.yaml",
|
|
"--enforce-eager",
|
|
"--disable-log-stats",
|
|
],
|
|
env_dict={
|
|
"ATTENTION_BACKEND": "torch",
|
|
"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
|
|
"VLLM_DISABLE_COMPILE_CACHE": "1",
|
|
"MASTER_PORT": str(get_open_port()),
|
|
},
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.full_model
|
|
@pytest.mark.diffusion
|
|
@hardware_test(res={"cuda": "H100"})
|
|
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
|
def test_pi05_openpi_online(omni_server):
|
|
pi0_openpi_require_dependencies()
|
|
|
|
result = pi0_openpi_run_policy_session(
|
|
host=omni_server.host,
|
|
port=omni_server.port,
|
|
prompt="pick up the red block and place it in the bin",
|
|
session_id="pi05-online-e2e",
|
|
num_steps=2,
|
|
num_inference_steps=2,
|
|
)
|
|
|
|
# Asserts every returned chunk is [50, 32] + finite, and the handshake
|
|
# metadata matches pi05.yaml's policy_server_config.
|
|
pi0_openpi_validate_session_result(result)
|
|
|
|
metadata = result["metadata"]
|
|
assert tuple(metadata["image_resolution"]) == (224, 224)
|
|
assert metadata["needs_wrist_camera"] is True
|
|
assert metadata["needs_session_id"] is False
|
|
assert metadata["action_space"] == "joint_position"
|