forked from genewildish/Mainline
David Gwilliam
828b8489e1
- Add TransformDataSource for filtering/mapping source items - Add MetricsDataSource for rendering live pipeline metrics as ASCII art - Fix display stage registration in StageRegistry - Register sources with both class name and simple name aliases - Fix DisplayStage.init() to pass reuse parameter - Simplify create_default_pipeline to use DataSourceStage wrapper - Set pygame as default display - Remove old pipeline tasks from mise.toml - Add tests for new pipeline architecture
321 lines
9.8 KiB
Python
321 lines
9.8 KiB
Python
"""
|
|
Pipeline controller - DAG-based pipeline execution.
|
|
|
|
The Pipeline class orchestrates stages in dependency order, handling
|
|
the complete render cycle from source to display.
|
|
"""
|
|
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from engine.pipeline.core import PipelineContext, Stage, StageError, StageResult
|
|
from engine.pipeline.params import PipelineParams
|
|
from engine.pipeline.registry import StageRegistry
|
|
|
|
|
|
@dataclass
|
|
class PipelineConfig:
|
|
"""Configuration for a pipeline instance."""
|
|
|
|
source: str = "headlines"
|
|
display: str = "terminal"
|
|
camera: str = "vertical"
|
|
effects: list[str] = field(default_factory=list)
|
|
enable_metrics: bool = True
|
|
|
|
|
|
@dataclass
|
|
class StageMetrics:
|
|
"""Metrics for a single stage execution."""
|
|
|
|
name: str
|
|
duration_ms: float
|
|
chars_in: int = 0
|
|
chars_out: int = 0
|
|
|
|
|
|
@dataclass
|
|
class FrameMetrics:
|
|
"""Metrics for a single frame through the pipeline."""
|
|
|
|
frame_number: int
|
|
total_ms: float
|
|
stages: list[StageMetrics] = field(default_factory=list)
|
|
|
|
|
|
class Pipeline:
|
|
"""Main pipeline orchestrator.
|
|
|
|
Manages the execution of all stages in dependency order,
|
|
handling initialization, processing, and cleanup.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: PipelineConfig | None = None,
|
|
context: PipelineContext | None = None,
|
|
):
|
|
self.config = config or PipelineConfig()
|
|
self.context = context or PipelineContext()
|
|
self._stages: dict[str, Stage] = {}
|
|
self._execution_order: list[str] = []
|
|
self._initialized = False
|
|
|
|
self._metrics_enabled = self.config.enable_metrics
|
|
self._frame_metrics: list[FrameMetrics] = []
|
|
self._max_metrics_frames = 60
|
|
self._current_frame_number = 0
|
|
|
|
def add_stage(self, name: str, stage: Stage) -> "Pipeline":
|
|
"""Add a stage to the pipeline."""
|
|
self._stages[name] = stage
|
|
return self
|
|
|
|
def remove_stage(self, name: str) -> None:
|
|
"""Remove a stage from the pipeline."""
|
|
if name in self._stages:
|
|
del self._stages[name]
|
|
|
|
def get_stage(self, name: str) -> Stage | None:
|
|
"""Get a stage by name."""
|
|
return self._stages.get(name)
|
|
|
|
def build(self) -> "Pipeline":
|
|
"""Build execution order based on dependencies."""
|
|
self._execution_order = self._resolve_dependencies()
|
|
self._initialized = True
|
|
return self
|
|
|
|
def _resolve_dependencies(self) -> list[str]:
|
|
"""Resolve stage execution order using topological sort."""
|
|
ordered = []
|
|
visited = set()
|
|
temp_mark = set()
|
|
|
|
def visit(name: str) -> None:
|
|
if name in temp_mark:
|
|
raise StageError(name, "Circular dependency detected")
|
|
if name in visited:
|
|
return
|
|
|
|
temp_mark.add(name)
|
|
stage = self._stages.get(name)
|
|
if stage:
|
|
for dep in stage.dependencies:
|
|
dep_stage = self._stages.get(dep)
|
|
if dep_stage:
|
|
visit(dep)
|
|
|
|
temp_mark.remove(name)
|
|
visited.add(name)
|
|
ordered.append(name)
|
|
|
|
for name in self._stages:
|
|
if name not in visited:
|
|
visit(name)
|
|
|
|
return ordered
|
|
|
|
def initialize(self) -> bool:
|
|
"""Initialize all stages in execution order."""
|
|
for name in self._execution_order:
|
|
stage = self._stages.get(name)
|
|
if stage and not stage.init(self.context) and not stage.optional:
|
|
return False
|
|
return True
|
|
|
|
def execute(self, data: Any | None = None) -> StageResult:
|
|
"""Execute the pipeline with the given input data."""
|
|
if not self._initialized:
|
|
self.build()
|
|
|
|
if not self._initialized:
|
|
return StageResult(
|
|
success=False,
|
|
data=None,
|
|
error="Pipeline not initialized",
|
|
)
|
|
|
|
current_data = data
|
|
frame_start = time.perf_counter() if self._metrics_enabled else 0
|
|
stage_timings: list[StageMetrics] = []
|
|
|
|
for name in self._execution_order:
|
|
stage = self._stages.get(name)
|
|
if not stage or not stage.is_enabled():
|
|
continue
|
|
|
|
stage_start = time.perf_counter() if self._metrics_enabled else 0
|
|
|
|
try:
|
|
current_data = stage.process(current_data, self.context)
|
|
except Exception as e:
|
|
if not stage.optional:
|
|
return StageResult(
|
|
success=False,
|
|
data=current_data,
|
|
error=str(e),
|
|
stage_name=name,
|
|
)
|
|
continue
|
|
|
|
if self._metrics_enabled:
|
|
stage_duration = (time.perf_counter() - stage_start) * 1000
|
|
chars_in = len(str(data)) if data else 0
|
|
chars_out = len(str(current_data)) if current_data else 0
|
|
stage_timings.append(
|
|
StageMetrics(
|
|
name=name,
|
|
duration_ms=stage_duration,
|
|
chars_in=chars_in,
|
|
chars_out=chars_out,
|
|
)
|
|
)
|
|
|
|
if self._metrics_enabled:
|
|
total_duration = (time.perf_counter() - frame_start) * 1000
|
|
self._frame_metrics.append(
|
|
FrameMetrics(
|
|
frame_number=self._current_frame_number,
|
|
total_ms=total_duration,
|
|
stages=stage_timings,
|
|
)
|
|
)
|
|
if len(self._frame_metrics) > self._max_metrics_frames:
|
|
self._frame_metrics.pop(0)
|
|
self._current_frame_number += 1
|
|
|
|
return StageResult(success=True, data=current_data)
|
|
|
|
def cleanup(self) -> None:
|
|
"""Clean up all stages in reverse order."""
|
|
for name in reversed(self._execution_order):
|
|
stage = self._stages.get(name)
|
|
if stage:
|
|
try:
|
|
stage.cleanup()
|
|
except Exception:
|
|
pass
|
|
self._stages.clear()
|
|
self._initialized = False
|
|
|
|
@property
|
|
def stages(self) -> dict[str, Stage]:
|
|
"""Get all stages."""
|
|
return self._stages.copy()
|
|
|
|
@property
|
|
def execution_order(self) -> list[str]:
|
|
"""Get execution order."""
|
|
return self._execution_order.copy()
|
|
|
|
def get_stage_names(self) -> list[str]:
|
|
"""Get list of stage names."""
|
|
return list(self._stages.keys())
|
|
|
|
def get_metrics_summary(self) -> dict:
|
|
"""Get summary of collected metrics."""
|
|
if not self._frame_metrics:
|
|
return {"error": "No metrics collected"}
|
|
|
|
total_times = [f.total_ms for f in self._frame_metrics]
|
|
avg_total = sum(total_times) / len(total_times)
|
|
min_total = min(total_times)
|
|
max_total = max(total_times)
|
|
|
|
stage_stats: dict[str, dict] = {}
|
|
for frame in self._frame_metrics:
|
|
for stage in frame.stages:
|
|
if stage.name not in stage_stats:
|
|
stage_stats[stage.name] = {"times": [], "total_chars": 0}
|
|
stage_stats[stage.name]["times"].append(stage.duration_ms)
|
|
stage_stats[stage.name]["total_chars"] += stage.chars_out
|
|
|
|
for name, stats in stage_stats.items():
|
|
times = stats["times"]
|
|
stats["avg_ms"] = sum(times) / len(times)
|
|
stats["min_ms"] = min(times)
|
|
stats["max_ms"] = max(times)
|
|
del stats["times"]
|
|
|
|
return {
|
|
"frame_count": len(self._frame_metrics),
|
|
"pipeline": {
|
|
"avg_ms": avg_total,
|
|
"min_ms": min_total,
|
|
"max_ms": max_total,
|
|
},
|
|
"stages": stage_stats,
|
|
}
|
|
|
|
def reset_metrics(self) -> None:
|
|
"""Reset collected metrics."""
|
|
self._frame_metrics.clear()
|
|
self._current_frame_number = 0
|
|
|
|
|
|
class PipelineRunner:
|
|
"""High-level pipeline runner with animation support."""
|
|
|
|
def __init__(
|
|
self,
|
|
pipeline: Pipeline,
|
|
params: PipelineParams | None = None,
|
|
):
|
|
self.pipeline = pipeline
|
|
self.params = params or PipelineParams()
|
|
self._running = False
|
|
|
|
def start(self) -> bool:
|
|
"""Start the pipeline."""
|
|
self._running = True
|
|
return self.pipeline.initialize()
|
|
|
|
def step(self, input_data: Any | None = None) -> Any:
|
|
"""Execute one pipeline step."""
|
|
self.params.frame_number += 1
|
|
self.pipeline.context.params = self.params
|
|
result = self.pipeline.execute(input_data)
|
|
return result.data if result.success else None
|
|
|
|
def stop(self) -> None:
|
|
"""Stop and clean up the pipeline."""
|
|
self._running = False
|
|
self.pipeline.cleanup()
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
"""Check if runner is active."""
|
|
return self._running
|
|
|
|
|
|
def create_pipeline_from_params(params: PipelineParams) -> Pipeline:
|
|
"""Create a pipeline from PipelineParams."""
|
|
config = PipelineConfig(
|
|
source=params.source,
|
|
display=params.display,
|
|
camera=params.camera_mode,
|
|
effects=params.effect_order,
|
|
)
|
|
return Pipeline(config=config)
|
|
|
|
|
|
def create_default_pipeline() -> Pipeline:
|
|
"""Create a default pipeline with all standard components."""
|
|
from engine.pipeline.adapters import DataSourceStage
|
|
from engine.sources_v2 import HeadlinesDataSource
|
|
|
|
pipeline = Pipeline()
|
|
|
|
# Add source stage (wrapped as Stage)
|
|
source = HeadlinesDataSource()
|
|
pipeline.add_stage("source", DataSourceStage(source, name="headlines"))
|
|
|
|
# Add display stage
|
|
display = StageRegistry.create("display", "terminal")
|
|
if display:
|
|
pipeline.add_stage("display", display)
|
|
|
|
return pipeline.build()
|