| 365 | return self._stage_registry |
| 366 | |
| 367 | def _validate_pipeline_sequence( |
| 368 | self, |
| 369 | stages: List[StageType], |
| 370 | ) -> None: |
| 371 | if not stages: |
| 372 | raise ValueError("Pipeline stages cannot be empty") |
| 373 | |
| 374 | # Validate pipeline compatibility with input model type |
| 375 | if self._input_model_type == "GraphModule": |
| 376 | # GraphModule input should not run quantization stages |
| 377 | incompatible_stages = {StageType.SOURCE_TRANSFORM, StageType.QUANTIZE} |
| 378 | found_incompatible = set(stages) & incompatible_stages |
| 379 | if found_incompatible: |
| 380 | stage_names = ", ".join(s.name for s in found_incompatible) |
| 381 | raise ValueError( |
| 382 | f"Cannot run {stage_names} stage(s) with GraphModule input. " |
| 383 | f"GraphModule is already quantized. " |
| 384 | f"Remove {stage_names} from pipeline_stages or use nn.Module input." |
| 385 | ) |
| 386 | elif self._input_model_type == "ExportedProgram": |
| 387 | # ExportedProgram input should not run quantization or torch export stages |
| 388 | incompatible_stages = { |
| 389 | StageType.SOURCE_TRANSFORM, |
| 390 | StageType.QUANTIZE, |
| 391 | StageType.TORCH_EXPORT, |
| 392 | } |
| 393 | found_incompatible = set(stages) & incompatible_stages |
| 394 | if found_incompatible: |
| 395 | stage_names = ", ".join(s.name for s in found_incompatible) |
| 396 | raise ValueError( |
| 397 | f"Cannot run {stage_names} stage(s) with ExportedProgram input. " |
| 398 | f"ExportedProgram is already exported. " |
| 399 | f"Remove {stage_names} from pipeline_stages or use nn.Module/GraphModule input." |
| 400 | ) |
| 401 | |
| 402 | # Validate that the first stage can start a pipeline |
| 403 | first_stage = stages[0] |
| 404 | first_stage_instance = self._stage_registry.get(first_stage) |
| 405 | if first_stage_instance is None: |
| 406 | raise ValueError( |
| 407 | f"Stage {first_stage} not found in registry, register it using session.register_stage()" |
| 408 | ) |
| 409 | |
| 410 | if not first_stage_instance.can_start_pipeline: |
| 411 | raise ValueError(f"Stage {first_stage} cannot start a pipeline. ") |
| 412 | |
| 413 | # Validate stage transitions |
| 414 | for i in range(1, len(stages)): |
| 415 | current_stage = stages[i] |
| 416 | previous_stage = stages[i - 1] |
| 417 | |
| 418 | # Get the stage instance to check its valid predecessors |
| 419 | stage_instance = self._stage_registry.get(current_stage) |
| 420 | if stage_instance is None: |
| 421 | raise ValueError( |
| 422 | f"Stage {current_stage} not found in registry, , register it using session.register_stage()" |
| 423 | ) |
| 424 | |