psystack 0.1.0__py3-none-any.whl

psystack/pipeline/metrics_util.py

@@ -0,0 +1,65 @@
+"""Shared metric computation logic used by both compare and isolation pipelines."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from psystack.core.contracts import AdapterFactory
+from psystack.models import MetricResult
+
+
+def compute_and_filter_metrics(
+    episodes: list[dict[str, Any]],
+    factory: AdapterFactory,
+    world_model: Any = None,
+    env: Any = None,
+    seed: int = 42,
+    skip_metrics: set[str] | None = None,
+) -> list[MetricResult]:
+    """Compute all metrics on episodes, filtering NaN values and optionally skipping metrics.
+
+    Args:
+        episodes: Episode data dicts.
+        factory: Adapter factory providing metric plugins.
+        world_model: Optional world model for prediction_error metric.
+        env: Optional env for prediction_error metric.
+        seed: Random seed.
+        skip_metrics: Set of metric IDs to skip (e.g. {"prediction_error"}).
+
+    Returns:
+        List of MetricResult with NaN values excluded.
+    """
+    skip = skip_metrics or set()
+
+    # Inject world_model/env for prediction_error metric
+    if world_model is not None and env is not None:
+        for ep in episodes:
+            ep["_world_model"] = world_model
+            ep["_env"] = env
+            ep["_seed"] = seed
+
+    results = []
+    for metric in factory.get_metrics():
+        if metric.metric_id() in skip:
+            continue
+        raw = metric.compute(episodes)
+        val = raw["primary_value"]
+        if math.isnan(val):
+            continue
+        results.append(MetricResult(
+            metric_id=metric.metric_id(),
+            value=val,
+            unit=raw.get("unit"),
+            higher_is_better=metric.higher_is_better(),
+            per_episode=raw.get("per_episode", []),
+            breakdown=raw.get("breakdown", {}),
+        ))
+
+    # Clean up injected refs
+    for ep in episodes:
+        ep.pop("_world_model", None)
+        ep.pop("_env", None)
+        ep.pop("_seed", None)
+
+    return results

psystack/pipeline/paired_runner.py

@@ -0,0 +1,185 @@
+"""Paired episode runner — lockstep A/B execution with pair-aware telemetry."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import torch
+
+from psystack.core.contracts import EnvPlugin, PlannerPlugin
+from psystack.pipeline.episodes import _serialize_info
+from psystack.pipeline.live_update import LivePairFrame, LiveStepUpdate
+
+
+class EvalCancelled(Exception):
+    """Raised when the user cancels a running evaluation."""
+
+
+def run_paired_episodes(
+    env_a: EnvPlugin,
+    env_b: EnvPlugin,
+    planner_a: PlannerPlugin,
+    planner_b: PlannerPlugin,
+    num_episodes: int,
+    seed: int = 42,
+    pair_callback: Callable[[LivePairFrame], None] | None = None,
+    max_steps: int = 0,
+    cancel_event: threading.Event | None = None,
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    """Run N episode pairs in lockstep.
+
+    Both envs are stepped once per tick. One LivePairFrame emitted per tick.
+    If one side finishes early, its update freezes.
+
+    Returns (episodes_a, episodes_b) in the same format as run_episodes().
+    """
+    all_episodes_a: list[dict[str, Any]] = []
+    all_episodes_b: list[dict[str, Any]] = []
+
+    for ep_idx in range(num_episodes):
+        ep_seed = seed + ep_idx
+
+        # Deterministic seeding
+        np.random.seed(ep_seed)
+        torch.manual_seed(ep_seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(ep_seed)
+
+        planner_a.reset()
+        planner_b.reset()
+        obs_a = env_a.reset(seed=ep_seed)
+        obs_b = env_b.reset(seed=ep_seed)
+
+        steps_a: list[dict[str, Any]] = []
+        steps_b: list[dict[str, Any]] = []
+        total_reward_a = 0.0
+        total_reward_b = 0.0
+        done_a = False
+        done_b = False
+        termination_a = "max_steps"
+        termination_b = "max_steps"
+        info_a: dict[str, Any] = {}
+        info_b: dict[str, Any] = {}
+
+        # Frozen terminal state for the side that finishes first
+        last_update_a: LiveStepUpdate | None = None
+        last_update_b: LiveStepUpdate | None = None
+        # Track last known progress — env may reset to 0 after done=True
+        last_progress_a = 0.0
+        last_progress_b = 0.0
+
+        tick = 0
+        while not (done_a and done_b):
+            if cancel_event is not None and cancel_event.is_set():
+                raise EvalCancelled("Evaluation cancelled by user")
+            # Step A
+            if not done_a:
+                car_state_a = env_a.get_car_state()
+                progress_a = env_a.get_progress()
+                action_a = planner_a.act(obs_a, car_state=car_state_a)
+                obs_next_a, reward_a, done_a, info_a = env_a.step(action_a)
+                last_progress_a = float(progress_a)
+
+                action_list_a = action_a.tolist() if hasattr(action_a, "tolist") else list(action_a)
+                steps_a.append({
+                    "obs": {"aux": obs_a["aux"].tolist() if hasattr(obs_a["aux"], "tolist") else obs_a["aux"]},
+                    "action": action_list_a,
+                    "reward": float(reward_a),
+                    "done": done_a,
+                    "info": _serialize_info(info_a),
+                    "car_state": car_state_a,
+                    "track_progress": float(progress_a),
+                })
+                total_reward_a += reward_a
+                obs_a = obs_next_a
+
+                if done_a and "termination" in info_a:
+                    termination_a = info_a["termination"]
+
+                last_update_a = LiveStepUpdate(
+                    run_id="a", episode_idx=ep_idx, episode_total=num_episodes,
+                    step=tick, progress=float(progress_a), reward=float(reward_a),
+                    done=done_a,
+                    termination=info_a.get("termination") if done_a else None,
+                    state=car_state_a, action=action_list_a, info=_serialize_info(info_a),
+                )
+
+            if cancel_event is not None and cancel_event.is_set():
+                raise EvalCancelled("Evaluation cancelled by user")
+            # Step B
+            if not done_b:
+                car_state_b = env_b.get_car_state()
+                progress_b = env_b.get_progress()
+                action_b = planner_b.act(obs_b, car_state=car_state_b)
+                obs_next_b, reward_b, done_b, info_b = env_b.step(action_b)
+                last_progress_b = float(progress_b)
+
+                action_list_b = action_b.tolist() if hasattr(action_b, "tolist") else list(action_b)
+                steps_b.append({
+                    "obs": {"aux": obs_b["aux"].tolist() if hasattr(obs_b["aux"], "tolist") else obs_b["aux"]},
+                    "action": action_list_b,
+                    "reward": float(reward_b),
+                    "done": done_b,
+                    "info": _serialize_info(info_b),
+                    "car_state": car_state_b,
+                    "track_progress": float(progress_b),
+                })
+                total_reward_b += reward_b
+                obs_b = obs_next_b
+
+                if done_b and "termination" in info_b:
+                    termination_b = info_b["termination"]
+
+                last_update_b = LiveStepUpdate(
+                    run_id="b", episode_idx=ep_idx, episode_total=num_episodes,
+                    step=tick, progress=float(progress_b), reward=float(reward_b),
+                    done=done_b,
+                    termination=info_b.get("termination") if done_b else None,
+                    state=car_state_b, action=action_list_b, info=_serialize_info(info_b),
+                )
+
+            # Emit pair frame
+            if pair_callback is not None:
+                frame = LivePairFrame(
+                    episode_idx=ep_idx,
+                    episode_total=num_episodes,
+                    tick=tick,
+                    a=last_update_a,
+                    b=last_update_b,
+                    both_done=done_a and done_b,
+                    max_steps=max_steps,
+                )
+                pair_callback(frame)
+
+            tick += 1
+
+        # Build episode records — use last tracked progress, not env.get_progress()
+        # which may return 0 if the env auto-resets on done=True
+        final_progress_a = last_progress_a
+        final_progress_b = last_progress_b
+
+        all_episodes_a.append({
+            "episode_id": f"ep_{ep_idx:04d}",
+            "steps": steps_a,
+            "total_steps": len(steps_a),
+            "final_track_progress": float(final_progress_a),
+            "total_reward": float(total_reward_a),
+            "termination": termination_a,
+            "fastest_lap_time": info_a.get("fastest_lap_time"),
+            "lap_count": info_a.get("lap_count", 0),
+        })
+        all_episodes_b.append({
+            "episode_id": f"ep_{ep_idx:04d}",
+            "steps": steps_b,
+            "total_steps": len(steps_b),
+            "final_track_progress": float(final_progress_b),
+            "total_reward": float(total_reward_b),
+            "termination": termination_b,
+            "fastest_lap_time": info_b.get("fastest_lap_time"),
+            "lap_count": info_b.get("lap_count", 0),
+        })
+
+    return all_episodes_a, all_episodes_b

psystack/pipeline/runner.py

@@ -0,0 +1,107 @@
+"""Stage runner — orchestrates pipeline stages with progress display."""
+
+from __future__ import annotations
+
+import contextlib
+from datetime import datetime, timezone
+from typing import Protocol, runtime_checkable
+
+from psystack.pipeline.context import RunContext
+from psystack.pipeline.stages.base import Stage
+from psystack.pipeline.state import StageResult
+
+
+@runtime_checkable
+class StageObserver(Protocol):
+    def on_stage_start(self, name: str) -> None: ...
+    def on_stage_complete(self, name: str, result: StageResult) -> None: ...
+    def on_stage_fail(self, name: str, error: str) -> None: ...
+    def on_stage_skip(self, name: str, reason: str) -> None: ...
+    def on_stage_reuse(self, name: str) -> None: ...
+
+
+class _NullObserver:
+    """Default no-op observer for CLI usage."""
+    def on_stage_start(self, name: str) -> None: pass
+    def on_stage_complete(self, name: str, result: StageResult) -> None: pass
+    def on_stage_fail(self, name: str, error: str) -> None: pass
+    def on_stage_skip(self, name: str, reason: str) -> None: pass
+    def on_stage_reuse(self, name: str) -> None: pass
+
+
+def utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def run_stages(
+    ctx: RunContext,
+    stages: tuple[Stage, ...],
+    *,
+    observer: StageObserver | None = None,
+) -> None:
+    obs = observer or _NullObserver()
+    all_names = [stage.name for stage in stages]
+
+    # Only use Rich Progress when no observer (CLI mode)
+    if observer is None:
+        from rich.progress import Progress
+        progress_ctx = Progress()
+    else:
+        progress_ctx = contextlib.nullcontext()
+
+    with progress_ctx as progress:
+        if progress is not None:
+            task_id = progress.add_task("PsyStack pipeline", total=len(stages))
+
+        for stage in stages:
+            # 1. Explicit skip
+            if stage.name in ctx.skip:
+                reason = "explicitly skipped"
+                ctx.mark_skipped(stage.name, reason=reason)
+                ctx.save_state()
+                obs.on_stage_skip(stage.name, reason)
+                if progress is not None:
+                    progress.advance(task_id)
+                continue
+
+            # 2. Outside selected range — preserves existing completed state
+            if not ctx.stage_selected(stage.name, all_names):
+                reason = "outside selected range"
+                ctx.mark_skipped(stage.name, reason=reason)
+                ctx.save_state()
+                obs.on_stage_skip(stage.name, reason)
+                if progress is not None:
+                    progress.advance(task_id)
+                continue
+
+            # 3. Dependency check
+            if not ctx.prereqs_satisfied(stage.requires):
+                raise RuntimeError(
+                    f"Stage '{stage.name}' requires completed stages: {stage.requires}"
+                )
+
+            # 4. Resume: reuse if up-to-date — keeps completed status
+            if ctx.resume and stage.is_up_to_date(ctx):
+                ctx.mark_reused(stage.name)
+                ctx.save_state()
+                obs.on_stage_reuse(stage.name)
+                if progress is not None:
+                    progress.advance(task_id)
+                continue
+
+            # 5. Run — fail fast on exception
+            try:
+                ctx.mark_running(stage.name, started_at=utc_now())
+                obs.on_stage_start(stage.name)
+                result = stage.run(ctx)
+                ctx.mark_completed(stage.name, result=result, finished_at=utc_now())
+                obs.on_stage_complete(stage.name, result)
+            except Exception as exc:
+                ctx.mark_failed(stage.name, error=str(exc), finished_at=utc_now())
+                ctx.save_state()
+                obs.on_stage_fail(stage.name, str(exc))
+                raise
+            else:
+                ctx.save_state()
+                if progress is not None:
+                    progress.advance(task_id)

psystack/pipeline/stages/__init__.py

@@ -0,0 +1,22 @@
+from psystack.pipeline.stages.attribute import AttributeStage
+from psystack.pipeline.stages.compare import CompareStage
+from psystack.pipeline.stages.events import EventStage
+from psystack.pipeline.stages.isolate import IsolateStage
+from psystack.pipeline.stages.report import ReportStage
+
+DEFAULT_PIPELINE = (
+    CompareStage(),
+    EventStage(),
+    IsolateStage(),
+    AttributeStage(),
+    ReportStage(),
+)
+
+__all__ = [
+    "CompareStage",
+    "EventStage",
+    "IsolateStage",
+    "AttributeStage",
+    "ReportStage",
+    "DEFAULT_PIPELINE",
+]

psystack/pipeline/stages/attribute.py

@@ -0,0 +1,78 @@
+"""AttributeStage — computes attribution from compare metrics + isolate swap results."""
+
+from __future__ import annotations
+
+import json
+
+from pydantic import TypeAdapter
+
+from psystack.models.comparison import ComparisonReport
+from psystack.models.isolation import AttributionTable, IsolationResultBundle
+from psystack.pipeline.context import RunContext
+from psystack.pipeline.state import StageResult
+
+_ATTR_TABLES = TypeAdapter(list[AttributionTable])
+
+
+def _attribute_output_is_valid(ctx: RunContext) -> bool:
+    path = ctx.output_path("attribute")
+    if not path.exists():
+        return False
+    try:
+        _ATTR_TABLES.validate_json(path.read_text())
+        return True
+    except Exception:
+        return False
+
+
+class AttributeStage:
+    name = "attribute"
+    requires = ("compare", "isolate")
+
+    def is_up_to_date(self, ctx: RunContext) -> bool:
+        return _attribute_output_is_valid(ctx)
+
+    def run(self, ctx: RunContext) -> StageResult:
+        from psystack.pipeline.isolation.attribution import compute_attribution
+
+        report = ComparisonReport.model_validate_json(
+            (ctx.stage_output_dir / "compare_report.json").read_text()
+        )
+        regressions = [m for m in report.metrics if m.status == "regression"]
+
+        if not regressions:
+            output = ctx.output_path("attribute")
+            output.write_text("[]")
+            return StageResult(
+                primary_output=str(output),
+                output_paths=[str(output)],
+                summary="No regressions — no attributions computed",
+                metadata={"num_attributions": 0, "skipped_reason": "no_regressions"},
+            )
+
+        # Read structured isolate bundle
+        bundle = IsolationResultBundle.model_validate_json(
+            ctx.output_path("isolate").read_text()
+        )
+        swap_results = bundle.swap_results
+
+        tables: list[AttributionTable] = []
+        warnings: list[str] = []
+
+        for metric in regressions:
+            table = compute_attribution(metric, swap_results)
+            tables.append(table)
+            if table.decision == "not_attributable":
+                warnings.append(f"{metric.metric_id} not attributable")
+
+        # Typed stage output: list[AttributionTable]
+        output = ctx.output_path("attribute")
+        output.write_text(json.dumps([t.model_dump() for t in tables], indent=2))
+
+        return StageResult(
+            primary_output=str(output),
+            output_paths=[str(output)],
+            summary=f"{len(tables)} attributions computed",
+            warnings=warnings,
+            metadata={"num_attributions": len(tables)},
+        )

psystack/pipeline/stages/base.py

@@ -0,0 +1,18 @@
+"""Stage protocol for the pipeline runner."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+from psystack.pipeline.state import StageResult
+
+if TYPE_CHECKING:
+    from psystack.pipeline.context import RunContext
+
+
+class Stage(Protocol):
+    name: str
+    requires: tuple[str, ...]
+
+    def is_up_to_date(self, ctx: RunContext) -> bool: ...
+    def run(self, ctx: RunContext) -> StageResult: ...

psystack/pipeline/stages/compare.py

@@ -0,0 +1,37 @@
+"""CompareStage — runs baseline vs candidate comparison."""
+
+from __future__ import annotations
+
+from psystack.pipeline.compare import compare_manifests
+from psystack.pipeline.context import RunContext
+from psystack.pipeline.state import StageResult
+
+
+class CompareStage:
+    name = "compare"
+    requires = ()
+
+    def is_up_to_date(self, ctx: RunContext) -> bool:
+        report_path = ctx.stage_output_dir / "compare_report.json"
+        return report_path.exists()
+
+    def run(self, ctx: RunContext) -> StageResult:
+        report = compare_manifests(
+            baseline_manifest=ctx.baseline_manifest,
+            candidate_manifest=ctx.candidate_manifest,
+            workspace=ctx.workspace,
+            factory=ctx.factory,
+            n_resamples=ctx.settings.run.bootstrap_resamples,
+            alpha=ctx.settings.run.alpha,
+        )
+
+        report_path = ctx.stage_output_dir / "compare_report.json"
+        report_path.write_text(report.model_dump_json(indent=2))
+
+        regressions = [m for m in report.metrics if m.status == "regression"]
+        return StageResult(
+            primary_output=str(report_path),
+            output_paths=[str(report_path)],
+            summary=f"{len(report.metrics)} metrics compared, {len(regressions)} regressions",
+            metadata={"num_metrics": len(report.metrics), "num_regressions": len(regressions)},
+        )

psystack/pipeline/stages/events.py

@@ -0,0 +1,53 @@
+"""EventStage — detect divergence events from compare output episode data."""
+
+from __future__ import annotations
+
+import json
+
+from psystack.pipeline.context import RunContext
+from psystack.pipeline.events.config import EventDetectionConfig
+from psystack.pipeline.events.detection import detect_events
+from psystack.pipeline.state import StageResult
+
+
+class EventStage:
+    name = "events"
+    requires = ("compare",)
+
+    def is_up_to_date(self, ctx: RunContext) -> bool:
+        events_path = ctx.stage_output_dir / "events.json"
+        return events_path.exists()
+
+    def run(self, ctx: RunContext) -> StageResult:
+        # Load episode data from compare stage output directory
+        baseline_episodes = self._load_episodes(ctx, "baseline")
+        candidate_episodes = self._load_episodes(ctx, "candidate")
+
+        # Load event config from settings if available
+        config = EventDetectionConfig()
+        settings = ctx.settings
+        if hasattr(settings, "event") and settings.event is not None:
+            config = EventDetectionConfig.model_validate(
+                settings.event.model_dump()
+            )
+
+        events = detect_events(baseline_episodes, candidate_episodes, config)
+
+        events_path = ctx.stage_output_dir / "events.json"
+        events_data = [e.model_dump() for e in events]
+        events_path.write_text(json.dumps(events_data, indent=2))
+
+        event_types = [e.type for e in events]
+        return StageResult(
+            primary_output=str(events_path),
+            output_paths=[str(events_path)],
+            summary=f"{len(events)} events detected: {', '.join(set(event_types)) or 'none'}",
+            metadata={"num_events": len(events), "event_types": list(set(event_types))},
+        )
+
+    def _load_episodes(self, ctx: RunContext, condition: str) -> list[dict]:
+        """Load episode data from the compare stage output."""
+        episodes_path = ctx.workspace / condition / "episodes.json"
+        if episodes_path.exists():
+            return json.loads(episodes_path.read_text())
+        return []

psystack/pipeline/stages/isolate.py

@@ -0,0 +1,88 @@
+"""IsolateStage — runs swap tests from a named isolation design."""
+
+from __future__ import annotations
+
+from psystack.models import SwapTestSpec
+from psystack.models.comparison import ComparisonReport
+from psystack.models.isolation import IsolationResultBundle
+from psystack.pipeline.context import RunContext
+from psystack.pipeline.state import StageResult
+
+
+def _isolate_output_is_valid(ctx: RunContext) -> bool:
+    path = ctx.output_path("isolate")
+    if not path.exists():
+        return False
+    try:
+        IsolationResultBundle.model_validate_json(path.read_text())
+        return True
+    except Exception:
+        return False
+
+
+class IsolateStage:
+    name = "isolate"
+    requires = ("compare",)
+
+    def is_up_to_date(self, ctx: RunContext) -> bool:
+        return _isolate_output_is_valid(ctx)
+
+    def run(self, ctx: RunContext) -> StageResult:
+        from psystack.pipeline.isolation import build_isolation_plan
+        from psystack.pipeline.isolation.executor import execute_swap_test
+
+        report = ComparisonReport.model_validate_json(
+            (ctx.stage_output_dir / "compare_report.json").read_text()
+        )
+        regressions = [m for m in report.metrics if m.status == "regression"]
+
+        plan = build_isolation_plan("screening_v1")
+
+        if not regressions:
+            bundle = IsolationResultBundle(
+                design=plan.design,
+                cases=plan.cases,
+                swap_results=[],
+            )
+            output = ctx.output_path("isolate")
+            output.write_text(bundle.model_dump_json(indent=2))
+            return StageResult(
+                primary_output=str(output),
+                output_paths=[str(output)],
+                summary="No regressions — no swap tests run",
+                metadata={"num_swap_tests": 0, "skipped_reason": "no_regressions"},
+            )
+
+        swap_dir = ctx.workspace / "swap_results"
+        swap_dir.mkdir(parents=True, exist_ok=True)
+
+        swap_results = []
+        warnings: list[str] = []
+        for case in plan.cases:
+            spec = SwapTestSpec(test_id=case.test_id, factors=case.factors)
+            result = execute_swap_test(
+                spec,
+                ctx.baseline_manifest,
+                ctx.candidate_manifest,
+                ctx.factory,
+            )
+            swap_results.append(result)
+            (swap_dir / f"{case.test_id}.json").write_text(result.model_dump_json(indent=2))
+            if result.status == "failed":
+                warnings.append(f"Swap test {case.test_id} failed: {result.error}")
+
+        bundle = IsolationResultBundle(
+            design=plan.design,
+            cases=plan.cases,
+            swap_results=swap_results,
+        )
+        output = ctx.output_path("isolate")
+        output.write_text(bundle.model_dump_json(indent=2))
+
+        return StageResult(
+            primary_output=str(output),
+            output_paths=[str(output)] + [str(swap_dir / f"{c.test_id}.json") for c in plan.cases],
+            summary=f"{len(swap_results)} swap tests completed",
+            warnings=warnings,
+            metadata={"num_swap_tests": len(swap_results)},
+        )