Yichen WangandDidan Deng
2a9f6aaa8e
Signed-off-by: Yichen Wang <yichen.wang@bytedance.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving.
|
|
|
|
Boots ``vllm serve --omni --deploy-config pi0.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_dreamzero_expansion.py``. Needs a GPU + a
|
|
pi0_base checkpoint; skipped in CI unless explicitly run.
|
|
|
|
The in-process LeRobot parity oracle lives separately in
|
|
``tests/diffusion/models/pi0/test_pi0_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/pi0_base"
|
|
|
|
|
|
test_params = [
|
|
OmniServerParams(
|
|
model=MODEL,
|
|
port=8092,
|
|
server_args=[
|
|
"--deploy-config",
|
|
"vllm_omni/deploy/pi0.yaml",
|
|
"--served-model-name",
|
|
"pi0",
|
|
"--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_pi0_openpi_online(omni_server):
|
|
try:
|
|
pi0_openpi_require_dependencies()
|
|
except ModuleNotFoundError as exc:
|
|
pytest.skip(str(exc))
|
|
|
|
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="pi0-online-e2e",
|
|
num_steps=2,
|
|
)
|
|
|
|
# Asserts every returned chunk is [50, 32] + finite, and the handshake
|
|
# metadata matches pi0.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"
|