diff --git a/tests/config/test_config_import_cycle.py b/tests/config/test_config_import_cycle.py index d73f6265e..8a28b7afb 100644 --- a/tests/config/test_config_import_cycle.py +++ b/tests/config/test_config_import_cycle.py @@ -93,3 +93,27 @@ if DiffusionOutput is None or LoRAConfig is None: raise SystemExit("expected LoRAConfig and DiffusionOutput") """ ) + + +def test_pipeline_registry_does_not_import_pi0_runtime() -> None: + """Resolving the pi0 topology must not import its runtime pipeline. + + ``pipeline_pi0`` imports ``diffusion.data``. Importing it while + ``diffusion.data`` is only partially initialized closes the Ascend plugin's + circular-import loop, so the registry must depend only on a lightweight + pipeline-config module. + """ + _assert_isolated_import_succeeds( + """ +import sys + +from vllm_omni.config.pipeline_registry import resolve_pipeline_config + +if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules: + raise SystemExit("pipeline_registry imported the pi0 runtime pipeline") + +pipeline = resolve_pipeline_config("pi0") +if pipeline is None or pipeline.model_type != "pi0": + raise SystemExit("expected the registered pi0 pipeline config") +""" + ) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f498b7bde..fd6d700a9 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -33,7 +33,7 @@ from vllm.logger import init_logger from vllm_omni.config.stage_config import ( PipelineConfig, ) -from vllm_omni.diffusion.models.pi0.pipeline_pi0 import PI0_PIPELINE +from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_S2S_PIPELINE, AUDEX_THINKER_ONLY_PIPELINE, diff --git a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py index 0877df4d6..1eeca2969 100644 --- a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py +++ b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py @@ -22,15 +22,11 @@ import torch from torch import nn from vllm.logger import init_logger -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.models.pi0.config import Pi0Config from vllm_omni.diffusion.models.pi0.modeling_pi0 import Pi0ForActionPrediction from vllm_omni.diffusion.models.pi0.processor_pi0 import build_model_inputs +from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE as PI0_PIPELINE from vllm_omni.diffusion.request import OmniDiffusionRequest logger = init_logger(__name__) @@ -54,32 +50,6 @@ def get_pi0_post_process_func(od_config: OmniDiffusionConfig): return _pi0_post_process -# Single-stage diffusion topology for online OpenPI serving. -# -# π0 is one diffusion stage (robot observation -> action chunk via flow -# matching). It is registered in ``OMNI_PIPELINES`` (rather than relying on the -# default single-stage fallback) because π0 is *online-served* via -# ``vllm serve --deploy-config pi0.yaml``; that path -# (``OmniStagePipelineFactory._create_from_registry``) requires a registered -# pipeline so the deploy yaml's ``policy_server_config`` reaches the OpenPI -# websocket layer. Offline-only models (e.g. ``internvla_a1``) need no topology. -PI0_PIPELINE = PipelineConfig( - model_type="pi0", - model_arch="Pi0Pipeline", - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="diffusion", - execution_type=StageExecutionType.DIFFUSION, - input_sources=(), - final_output=True, - final_output_type="action", - model_arch="Pi0Pipeline", - ), - ), -) - - class Pi0Pipeline(nn.Module): """π0 VLA pipeline: raw robot obs → continuous action chunk. diff --git a/vllm_omni/diffusion/models/pi0_pipeline_config.py b/vllm_omni/diffusion/models/pi0_pipeline_config.py new file mode 100644 index 000000000..d778ca81a --- /dev/null +++ b/vllm_omni/diffusion/models/pi0_pipeline_config.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lightweight pipeline topology for π0 (Pi-Zero).""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +# π0 is one diffusion stage (robot observation -> action chunk via flow +# matching). It is registered in ``OMNI_PIPELINES`` because π0 is online-served +# via ``vllm serve --deploy-config pi0.yaml``. Keep this topology outside the +# ``pi0`` package: importing a package submodule executes ``pi0.__init__``, +# which imports the runtime pipeline and ``diffusion.data``. +PI0_PIPELINE = PipelineConfig( + model_type="pi0", + model_arch="Pi0Pipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Pi0Pipeline", + ), + ), +)