[Refactor] Remove dead legacy pipeline.yaml loader and duplicate diff… (#4023)
Signed-off-by: dph97 <nenbaying@163.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: zhuhonghan <z00806815@china.huawei.com>
This commit is contained in:
6 files changed
+7
-685
No files matched your search
@@ -92,7 +92,6 @@ Core scheduling and caching components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.config.stage_config.ModelPipeline][]
|
||||
- [vllm_omni.config.stage_config.StageConfig][]
|
||||
- [vllm_omni.config.stage_config.StageConfigFactory][]
|
||||
- [vllm_omni.config.stage_config.StageType][]
|
||||
|
||||
@@ -532,7 +532,10 @@ async def test_pure_diffusion_scenario(tp_size: int):
|
||||
try:
|
||||
await engine.sleep(level=1)
|
||||
await engine.wake_up()
|
||||
async for _ in engine.generate("test", sampling_params=SamplingParams()):
|
||||
async for _ in engine.generate(
|
||||
"test",
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
|
||||
):
|
||||
pass
|
||||
logger.info("Pure diffusion OK (TP=%s)", tp_size)
|
||||
finally:
|
||||
|
||||
@@ -17,7 +17,6 @@ from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config.stage_config import (
|
||||
_PIPELINE_REGISTRY,
|
||||
DeployConfig,
|
||||
ModelPipeline,
|
||||
PipelineConfig,
|
||||
StageConfig,
|
||||
StageConfigFactory,
|
||||
@@ -351,93 +350,6 @@ class TestStageConfig:
|
||||
assert "parallel_config" not in omega_config.engine_args
|
||||
|
||||
|
||||
class TestModelPipeline:
|
||||
"""Tests for ModelPipeline class."""
|
||||
|
||||
def test_valid_linear_dag(self):
|
||||
"""Test validation of a valid linear DAG."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
|
||||
StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected errors: {errors}"
|
||||
|
||||
def test_valid_branching_dag(self):
|
||||
"""Test validation of a valid branching DAG."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
|
||||
StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected errors: {errors}"
|
||||
|
||||
def test_missing_entry_point(self):
|
||||
"""Test that missing entry point is detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
|
||||
StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("entry point" in e.lower() for e in errors)
|
||||
|
||||
def test_missing_dependency(self):
|
||||
"""Test that missing stage reference is detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="output", input_sources=[99]), # Invalid
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("non-existent" in e.lower() for e in errors)
|
||||
|
||||
def test_duplicate_stage_ids(self):
|
||||
"""Test that duplicate stage IDs are detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
|
||||
StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]), # Duplicate
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("duplicate" in e.lower() for e in errors)
|
||||
|
||||
def test_self_reference(self):
|
||||
"""Test that self-references are detected."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]), # Self
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("itself" in e.lower() for e in errors)
|
||||
|
||||
def test_get_stage_by_id(self):
|
||||
"""Test getting stage by ID."""
|
||||
stages = [
|
||||
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
|
||||
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
|
||||
]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
|
||||
stage = pipeline.get_stage(1)
|
||||
assert stage is not None
|
||||
assert stage.model_stage == "talker"
|
||||
|
||||
missing = pipeline.get_stage(99)
|
||||
assert missing is None
|
||||
|
||||
def test_empty_pipeline(self):
|
||||
"""Test validation of empty pipeline."""
|
||||
pipeline = ModelPipeline(model_type="test", stages=[])
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert any("no stages" in e.lower() for e in errors)
|
||||
|
||||
|
||||
class TestStageConfigFactory:
|
||||
"""Tests for StageConfigFactory class."""
|
||||
|
||||
@@ -598,304 +510,6 @@ class TestStageResolutionHelpers:
|
||||
assert overridden == []
|
||||
|
||||
|
||||
class TestPipelineYamlParsing:
|
||||
"""Tests for pipeline YAML file parsing (@ZJY0516)."""
|
||||
|
||||
def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
|
||||
"""Test parsing the qwen3_omni_moe pipeline YAML."""
|
||||
yaml_content = """\
|
||||
model_type: qwen3_omni_moe
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: thinker
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
hf_config_name: thinker_config
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
|
||||
- stage_id: 1
|
||||
model_stage: talker
|
||||
stage_type: llm
|
||||
input_sources: [0]
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
hf_config_name: talker_config
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
|
||||
- stage_id: 2
|
||||
model_stage: code2wav
|
||||
stage_type: llm
|
||||
input_sources: [1]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
hf_config_name: thinker_config
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
"""
|
||||
yaml_file = tmp_path / "qwen3_omni_moe.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
|
||||
|
||||
assert pipeline.model_type == "qwen3_omni_moe"
|
||||
assert len(pipeline.stages) == 3
|
||||
|
||||
# Stage 0: thinker
|
||||
s0 = pipeline.stages[0]
|
||||
assert s0.stage_id == 0
|
||||
assert s0.model_stage == "thinker"
|
||||
assert s0.stage_type == StageType.LLM
|
||||
assert s0.input_sources == []
|
||||
assert s0.worker_type == "ar"
|
||||
assert s0.final_output is True
|
||||
assert s0.final_output_type == "text"
|
||||
assert s0.is_comprehension is True
|
||||
|
||||
# Stage 1: talker
|
||||
s1 = pipeline.stages[1]
|
||||
assert s1.stage_id == 1
|
||||
assert s1.input_sources == [0]
|
||||
assert s1.custom_process_input_func == (
|
||||
"vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
|
||||
)
|
||||
assert s1.final_output is False
|
||||
|
||||
# Stage 2: code2wav
|
||||
s2 = pipeline.stages[2]
|
||||
assert s2.stage_id == 2
|
||||
assert s2.input_sources == [1]
|
||||
assert s2.worker_type == "generation"
|
||||
assert s2.final_output_type == "audio"
|
||||
|
||||
def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
|
||||
"""Test backward compatibility with engine_input_source field."""
|
||||
yaml_content = """\
|
||||
model_type: legacy_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
- stage_id: 1
|
||||
model_stage: downstream
|
||||
stage_type: llm
|
||||
engine_input_source: [0]
|
||||
"""
|
||||
yaml_file = tmp_path / "legacy.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
|
||||
assert pipeline.stages[1].input_sources == [0]
|
||||
|
||||
def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
|
||||
"""Test parsing pipeline with optional connectors and edges."""
|
||||
yaml_content = """\
|
||||
model_type: test_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
|
||||
connectors:
|
||||
type: ray
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
"""
|
||||
yaml_file = tmp_path / "with_connectors.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
|
||||
assert pipeline.connectors == {"type": "ray"}
|
||||
assert pipeline.edges == [{"from": 0, "to": 1}]
|
||||
|
||||
def test_parsed_pipeline_passes_validation(self, tmp_path):
|
||||
"""Test that a well-formed YAML produces a valid pipeline."""
|
||||
yaml_content = """\
|
||||
model_type: valid_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
- stage_id: 1
|
||||
model_stage: next
|
||||
stage_type: llm
|
||||
input_sources: [0]
|
||||
"""
|
||||
yaml_file = tmp_path / "valid.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected validation errors: {errors}"
|
||||
|
||||
def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
|
||||
"""Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
|
||||
yaml_content = """\
|
||||
model_type: legacy_batch
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 16
|
||||
engine_args:
|
||||
model_arch: SomeModel
|
||||
"""
|
||||
yaml_file = tmp_path / "legacy_batch.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
|
||||
s0 = pipeline.stages[0]
|
||||
assert "max_batch_size" not in s0.yaml_runtime
|
||||
assert s0.yaml_engine_args.get("max_num_seqs") == 16
|
||||
|
||||
def test_parse_diffusion_stage_type(self, tmp_path):
|
||||
"""Test parsing a diffusion stage type from YAML."""
|
||||
yaml_content = """\
|
||||
model_type: diff_model
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: dit
|
||||
stage_type: diffusion
|
||||
input_sources: []
|
||||
final_output: true
|
||||
final_output_type: image
|
||||
"""
|
||||
yaml_file = tmp_path / "diffusion.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
|
||||
assert pipeline.stages[0].stage_type == StageType.DIFFUSION
|
||||
|
||||
def test_parse_mimo_audio_yaml(self, tmp_path):
|
||||
"""Test parsing the MiMo Audio pipeline YAML."""
|
||||
yaml_content = """\
|
||||
model_type: mimo_audio
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: fused_thinker_talker
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
|
||||
- stage_id: 1
|
||||
model_stage: code2wav
|
||||
stage_type: llm
|
||||
input_sources: [0]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
"""
|
||||
yaml_file = tmp_path / "mimo_audio.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
|
||||
|
||||
assert pipeline.model_type == "mimo_audio"
|
||||
assert len(pipeline.stages) == 2
|
||||
|
||||
# Stage 0: fused_thinker_talker
|
||||
s0 = pipeline.stages[0]
|
||||
assert s0.stage_id == 0
|
||||
assert s0.model_stage == "fused_thinker_talker"
|
||||
assert s0.input_sources == []
|
||||
assert s0.worker_type == "ar"
|
||||
assert s0.is_comprehension is True
|
||||
assert s0.final_output is True
|
||||
assert s0.final_output_type == "text"
|
||||
|
||||
# Stage 1: code2wav
|
||||
s1 = pipeline.stages[1]
|
||||
assert s1.stage_id == 1
|
||||
assert s1.input_sources == [0]
|
||||
assert s1.worker_type == "generation"
|
||||
assert s1.custom_process_input_func == (
|
||||
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
|
||||
)
|
||||
assert s1.final_output is True
|
||||
assert s1.final_output_type == "audio"
|
||||
|
||||
# Pipeline validation
|
||||
errors = pipeline.validate_pipeline()
|
||||
assert errors == [], f"Unexpected errors: {errors}"
|
||||
|
||||
|
||||
class TestAsyncChunk:
|
||||
"""Tests for async_chunk pipeline flag."""
|
||||
|
||||
def test_model_pipeline_async_chunk_default(self):
|
||||
"""Test ModelPipeline defaults async_chunk to False."""
|
||||
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages)
|
||||
assert pipeline.async_chunk is False
|
||||
|
||||
def test_model_pipeline_async_chunk_set(self):
|
||||
"""Test ModelPipeline with async_chunk=True."""
|
||||
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
|
||||
pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
|
||||
assert pipeline.async_chunk is True
|
||||
|
||||
def test_parse_async_chunk_from_yaml(self, tmp_path):
|
||||
"""Test that async_chunk is parsed from pipeline YAML."""
|
||||
yaml_content = """\
|
||||
model_type: test_async
|
||||
async_chunk: true
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
"""
|
||||
yaml_file = tmp_path / "async.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
|
||||
assert pipeline.async_chunk is True
|
||||
|
||||
def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
|
||||
"""Test that missing async_chunk defaults to False."""
|
||||
yaml_content = """\
|
||||
model_type: test_no_async
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
model_stage: entry
|
||||
stage_type: llm
|
||||
input_sources: []
|
||||
"""
|
||||
yaml_file = tmp_path / "no_async.yaml"
|
||||
yaml_file.write_text(yaml_content)
|
||||
|
||||
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
|
||||
assert pipeline.async_chunk is False
|
||||
|
||||
|
||||
class TestPipelineDiscovery:
|
||||
"""Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
|
||||
|
||||
|
||||
@@ -6,7 +6,6 @@ from vllm_omni.config.lora import LoRAConfig
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
from vllm_omni.config.stage_config import (
|
||||
DeployConfig,
|
||||
ModelPipeline,
|
||||
PipelineConfig,
|
||||
StageConfig,
|
||||
StageConfigFactory,
|
||||
@@ -30,7 +29,6 @@ __all__ = [
|
||||
"LoRAConfig",
|
||||
"StageConfig",
|
||||
"StageConfigFactory",
|
||||
"ModelPipeline",
|
||||
"StageType",
|
||||
"StageExecutionType",
|
||||
"StagePipelineConfig",
|
||||
|
||||
@@ -20,13 +20,6 @@ from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models"
|
||||
|
||||
|
||||
def get_pipeline_path(model_dir: str, filename: str) -> Path:
|
||||
return _MODELS_DIR / model_dir / filename
|
||||
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@@ -1061,77 +1054,6 @@ class StageConfig:
|
||||
return create_config(config_dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelPipeline:
|
||||
"""Complete pipeline definition for a multi-stage model (legacy).
|
||||
|
||||
TODO(@lishunyang12): remove once all models migrate to PipelineConfig.
|
||||
"""
|
||||
|
||||
model_type: str
|
||||
stages: list[StageConfig]
|
||||
|
||||
# Pipeline-wide behavior flags
|
||||
async_chunk: bool = False
|
||||
|
||||
# Optional distributed configuration
|
||||
connectors: dict[str, Any] | None = None
|
||||
edges: list[dict[str, Any]] | None = None
|
||||
|
||||
def get_stage(self, stage_id: int) -> StageConfig | None:
|
||||
"""Look up a stage by its ID.
|
||||
|
||||
Args:
|
||||
stage_id: The stage ID to search for.
|
||||
|
||||
Returns:
|
||||
The matching StageConfig, or None if not found.
|
||||
"""
|
||||
for stage in self.stages:
|
||||
if stage.stage_id == stage_id:
|
||||
return stage
|
||||
return None
|
||||
|
||||
def validate_pipeline(self) -> list[str]:
|
||||
"""Validate pipeline topology at model integration time (not runtime).
|
||||
|
||||
Checks:
|
||||
- All stage IDs are unique
|
||||
- All input_sources reference valid stage IDs
|
||||
- At least one entry point (stage with empty input_sources)
|
||||
|
||||
Returns:
|
||||
List of validation error messages. Empty list if valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not self.stages:
|
||||
errors.append("Topology has no stages defined")
|
||||
return errors
|
||||
|
||||
# Check for unique stage IDs
|
||||
stage_ids = [s.stage_id for s in self.stages]
|
||||
if len(stage_ids) != len(set(stage_ids)):
|
||||
errors.append("Duplicate stage IDs found")
|
||||
|
||||
stage_id_set = set(stage_ids)
|
||||
|
||||
# Check input_sources reference valid stages
|
||||
for stage in self.stages:
|
||||
for source_id in stage.input_sources:
|
||||
if source_id not in stage_id_set:
|
||||
errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}")
|
||||
if source_id == stage.stage_id:
|
||||
errors.append(f"Stage {stage.stage_id} references itself as input source")
|
||||
|
||||
# Check for at least one entry point
|
||||
entry_points = [s for s in self.stages if not s.input_sources]
|
||||
if not entry_points:
|
||||
errors.append("No entry point found (stage with empty input_sources)")
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
class StageConfigFactory:
|
||||
"""Factory that loads pipeline YAML and merges CLI overrides.
|
||||
|
||||
@@ -1206,33 +1128,9 @@ class StageConfigFactory:
|
||||
continue
|
||||
return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path)
|
||||
|
||||
# --- Legacy path: load from pipeline YAML ---
|
||||
pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code)
|
||||
|
||||
if pipeline is None:
|
||||
return None
|
||||
|
||||
errors = pipeline.validate_pipeline()
|
||||
if errors:
|
||||
logger.warning(f"Pipeline validation warnings for {model}: {errors}")
|
||||
|
||||
# Materialize the resolved pipeline-wide async_chunk value into every
|
||||
# stage so build_engine_args_dict() can inject the stage connector
|
||||
# spec and explicit False overrides are preserved.
|
||||
resolved_async_chunk = cli_overrides.get("async_chunk")
|
||||
if resolved_async_chunk is None:
|
||||
resolved_async_chunk = bool(pipeline.async_chunk)
|
||||
for stage in pipeline.stages:
|
||||
stage.yaml_engine_args["async_chunk"] = bool(resolved_async_chunk)
|
||||
|
||||
# Apply CLI overrides
|
||||
result: list[StageConfig] = []
|
||||
for stage in pipeline.stages:
|
||||
# Merge global CLI overrides
|
||||
stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides)
|
||||
result.append(stage)
|
||||
|
||||
return result
|
||||
# Not in the pipeline registry — let the caller fall back to the
|
||||
# legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _create_from_registry(
|
||||
@@ -1355,193 +1253,6 @@ class StageConfigFactory:
|
||||
|
||||
return [config_dict]
|
||||
|
||||
@classmethod
|
||||
def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None:
|
||||
"""Load a legacy ``pipeline.yaml`` for the model.
|
||||
|
||||
Searches ``model_executor/models/<dir>/pipeline.yaml`` by trying
|
||||
(a) the raw ``model_type`` as the directory name, then
|
||||
(b) ``model_type`` with hyphens replaced by underscores,
|
||||
and finally (c) scanning every ``pipeline.yaml`` for one that
|
||||
declares a matching ``model_type`` or ``hf_architectures``.
|
||||
|
||||
Returns None if no pipeline.yaml is found — caller handles the
|
||||
``resolve_model_config_path`` fallback via stage_configs/ YAMLs.
|
||||
"""
|
||||
model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
|
||||
if model_type is None:
|
||||
return None
|
||||
|
||||
# Direct lookups by convention
|
||||
candidates = [model_type, model_type.replace("-", "_")]
|
||||
for dir_name in candidates:
|
||||
pipeline_path = get_pipeline_path(dir_name, "pipeline.yaml")
|
||||
if pipeline_path.exists():
|
||||
return cls._parse_pipeline_yaml(pipeline_path, model_type)
|
||||
|
||||
# Scan fallback: read every pipeline.yaml and match on declared fields
|
||||
hf_archs = set(getattr(hf_config, "architectures", []) or []) if hf_config else set()
|
||||
if _MODELS_DIR.exists():
|
||||
for subdir in sorted(_MODELS_DIR.iterdir()):
|
||||
if not subdir.is_dir():
|
||||
continue
|
||||
pipeline_path = subdir / "pipeline.yaml"
|
||||
if not pipeline_path.exists():
|
||||
continue
|
||||
try:
|
||||
cfg = load_yaml_config(pipeline_path)
|
||||
except Exception as exc:
|
||||
logger.debug("Skip %s: %s", pipeline_path, exc)
|
||||
continue
|
||||
declared_type = getattr(cfg, "model_type", None)
|
||||
declared_archs = set(getattr(cfg, "hf_architectures", None) or [])
|
||||
if declared_type == model_type or (hf_archs and hf_archs.intersection(declared_archs)):
|
||||
return cls._parse_pipeline_yaml(pipeline_path, declared_type or model_type)
|
||||
|
||||
logger.debug("No pipeline.yaml found for model_type %s (archs=%s)", model_type, sorted(hf_archs))
|
||||
return None
|
||||
|
||||
# Keys consumed as explicit StageConfig fields — everything else is
|
||||
# passed through via yaml_extras.
|
||||
_KNOWN_STAGE_KEYS: set[str] = {
|
||||
"stage_id",
|
||||
"model_stage",
|
||||
"stage_type",
|
||||
"input_sources",
|
||||
"engine_input_source",
|
||||
"custom_process_input_func",
|
||||
"final_output",
|
||||
"final_output_type",
|
||||
"worker_type",
|
||||
"scheduler_cls",
|
||||
"hf_config_name",
|
||||
"is_comprehension",
|
||||
"engine_args",
|
||||
"runtime",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
|
||||
"""Parse a pipeline YAML file.
|
||||
|
||||
Args:
|
||||
path: Path to the YAML file.
|
||||
model_type: Model type identifier.
|
||||
|
||||
Returns:
|
||||
ModelPipeline object.
|
||||
"""
|
||||
config_data = load_yaml_config(path)
|
||||
|
||||
stages: list[StageConfig] = []
|
||||
for stage_data in config_data.stages:
|
||||
# Use .get() for optional fields — idiomatic for OmegaConf DictConfig
|
||||
stage_type_str = stage_data.get("stage_type", "llm")
|
||||
stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM
|
||||
|
||||
# Handle both 'input_sources' (new) and 'engine_input_source' (legacy)
|
||||
input_sources = stage_data.get("input_sources", None)
|
||||
if input_sources is None:
|
||||
input_sources = stage_data.get("engine_input_source", [])
|
||||
if input_sources is None:
|
||||
input_sources = []
|
||||
input_sources = list(input_sources)
|
||||
|
||||
# Extract per-stage engine_args and runtime dicts
|
||||
raw_ea = stage_data.get("engine_args", None)
|
||||
yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {}
|
||||
raw_rt = stage_data.get("runtime", None)
|
||||
yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
|
||||
|
||||
# Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
|
||||
if "max_batch_size" in yaml_runtime:
|
||||
mbs = yaml_runtime.pop("max_batch_size")
|
||||
yaml_engine_args.setdefault("max_num_seqs", int(mbs))
|
||||
logger.debug(
|
||||
"Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
|
||||
stage_data.get("stage_id", "?"),
|
||||
mbs,
|
||||
)
|
||||
|
||||
# Topology-level fields that also live inside engine_args in legacy
|
||||
# YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
|
||||
worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
|
||||
scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None)
|
||||
if scheduler_cls:
|
||||
async_sched = yaml_engine_args.get("async_scheduling")
|
||||
if async_sched is not None:
|
||||
logger.warning(
|
||||
"Stage %s: async_scheduling=%r and scheduler_cls=%r "
|
||||
"should not be set together. scheduler_cls will take "
|
||||
"precedence for which scheduler is used.",
|
||||
stage_data.stage_id,
|
||||
async_sched,
|
||||
scheduler_cls,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Stage %s: scheduler_cls=%r is deprecated. Use async_scheduling instead.",
|
||||
stage_data.stage_id,
|
||||
scheduler_cls,
|
||||
)
|
||||
hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None)
|
||||
model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None)
|
||||
|
||||
# Collect pass-through fields (default_sampling_params,
|
||||
# output_connectors, input_connectors, tts_args, etc.)
|
||||
yaml_extras: dict[str, Any] = {}
|
||||
for key in stage_data:
|
||||
if key not in cls._KNOWN_STAGE_KEYS:
|
||||
val = stage_data[key]
|
||||
try:
|
||||
yaml_extras[key] = to_dict(val)
|
||||
except ValueError:
|
||||
yaml_extras[key] = val
|
||||
|
||||
stage = StageConfig(
|
||||
stage_id=stage_data.stage_id,
|
||||
model_stage=model_stage or "",
|
||||
stage_type=stage_type,
|
||||
input_sources=input_sources,
|
||||
custom_process_input_func=stage_data.get("custom_process_input_func", None),
|
||||
final_output=stage_data.get("final_output", False),
|
||||
final_output_type=stage_data.get("final_output_type", None),
|
||||
worker_type=worker_type,
|
||||
scheduler_cls=scheduler_cls,
|
||||
hf_config_name=hf_config_name,
|
||||
is_comprehension=stage_data.get("is_comprehension", False),
|
||||
yaml_engine_args=yaml_engine_args,
|
||||
yaml_runtime=yaml_runtime,
|
||||
yaml_extras=yaml_extras,
|
||||
)
|
||||
stages.append(stage)
|
||||
|
||||
# Get pipeline-wide flags
|
||||
async_chunk = config_data.get("async_chunk", False)
|
||||
|
||||
# Get optional connector config — check both top-level and nested
|
||||
# under ``runtime`` (legacy stage_configs format).
|
||||
connectors = None
|
||||
edges = None
|
||||
if hasattr(config_data, "connectors"):
|
||||
connectors = to_dict(config_data.connectors)
|
||||
if hasattr(config_data, "edges"):
|
||||
edges = to_dict(config_data.edges)
|
||||
if hasattr(config_data, "runtime") and config_data.runtime is not None:
|
||||
top_runtime = config_data.runtime
|
||||
if connectors is None and hasattr(top_runtime, "connectors"):
|
||||
connectors = to_dict(top_runtime.connectors)
|
||||
if edges is None and hasattr(top_runtime, "edges"):
|
||||
edges = to_dict(top_runtime.edges)
|
||||
|
||||
return ModelPipeline(
|
||||
model_type=getattr(config_data, "model_type", model_type),
|
||||
stages=stages,
|
||||
async_chunk=async_chunk,
|
||||
connectors=connectors,
|
||||
edges=edges,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
|
||||
"""Auto-detect model_type from model directory.
|
||||
|
||||
@@ -610,9 +610,6 @@ class OmniDiffusionConfig:
|
||||
# Step mode settings
|
||||
step_execution: bool = False
|
||||
|
||||
# sleep mode
|
||||
enable_sleep_mode: bool = False
|
||||
|
||||
# Maximum number of sequences to generate in a batch
|
||||
max_num_seqs: int = 1
|
||||
|
||||
|
||||
Reference in new issue
Block a user