tsagentkit 1.0.2__py3-none-any.whl

tsagentkit/serving/packaging.py

@@ -0,0 +1,73 @@
+"""Run artifact packaging.
+
+Bundles all outputs from a forecasting run into a comprehensive artifact.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from tsagentkit.contracts import RunArtifact
+
+if TYPE_CHECKING:
+    from tsagentkit.backtest import BacktestReport
+    from tsagentkit.contracts import ForecastResult, ModelArtifact, Provenance
+    from tsagentkit.qa import QAReport
+    from tsagentkit.router import PlanSpec
+
+
+def package_run(
+    forecast: ForecastResult,
+    plan: PlanSpec | dict[str, Any],
+    task_spec: Any | None = None,
+    plan_spec: dict[str, Any] | None = None,
+    validation_report: dict[str, Any] | None = None,
+    backtest_report: BacktestReport | None = None,
+    qa_report: QAReport | None = None,
+    model_artifact: ModelArtifact | None = None,
+    provenance: Provenance | None = None,
+    calibration_artifact: dict[str, Any] | None = None,
+    anomaly_report: dict[str, Any] | None = None,
+    metadata: dict[str, Any] | None = None,
+) -> RunArtifact:
+    """Package all run outputs into a comprehensive artifact.
+
+    Args:
+        forecast: ForecastResult with predictions + provenance
+        plan: Execution plan (PlanSpec or dict)
+        backtest_report: Optional backtest results
+        qa_report: Optional QA report
+        model_artifact: Optional fitted model
+        provenance: Optional provenance information (overrides forecast provenance)
+        calibration_artifact: Optional calibration artifact
+        anomaly_report: Optional anomaly report
+        metadata: Optional metadata
+
+    Returns:
+        RunArtifact containing all run outputs
+    """
+    if hasattr(plan, "to_dict"):
+        plan_dict = plan.to_dict()
+    elif hasattr(plan, "model_dump"):
+        plan_dict = plan.model_dump()
+    else:
+        plan_dict = plan
+    if plan_spec is None:
+        plan_spec = plan_dict
+    backtest_dict = backtest_report.to_dict() if backtest_report else None
+    qa_dict = qa_report.to_dict() if qa_report and hasattr(qa_report, "to_dict") else None
+
+    return RunArtifact(
+        forecast=forecast,
+        plan=plan_dict,
+        task_spec=task_spec,
+        plan_spec=plan_spec,
+        validation_report=validation_report,
+        backtest_report=backtest_dict,
+        qa_report=qa_dict,
+        model_artifact=model_artifact,
+        provenance=provenance or forecast.provenance,
+        calibration_artifact=calibration_artifact,
+        anomaly_report=anomaly_report,
+        metadata=metadata or {},
+    )

tsagentkit/serving/provenance.py

@@ -0,0 +1,317 @@
+"""Provenance tracking for forecasting runs.
+
+Provides utilities for tracking data lineage and reproducibility.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    pass
+
+    from tsagentkit.contracts import Provenance, TaskSpec
+    from tsagentkit.router import PlanSpec
+
+
+from datetime import UTC
+
+from tsagentkit.utils import compute_config_signature, compute_data_signature
+
+
+def create_provenance(
+    data: pd.DataFrame,
+    task_spec: TaskSpec,
+    plan: PlanSpec,
+    model_config: dict[str, Any] | None = None,
+    qa_repairs: list[Any] | None = None,
+    fallbacks_triggered: list[dict[str, Any]] | None = None,
+    feature_matrix: Any | None = None,
+    drift_report: Any | None = None,
+    column_map: dict[str, str] | None = None,
+    original_panel_contract: dict[str, Any] | None = None,
+    route_decision: Any | None = None,
+) -> Provenance:
+    """Create a provenance record for a forecasting run.
+
+    Args:
+        data: Input data
+        task_spec: Task specification
+        plan: Execution plan
+        model_config: Model configuration
+        qa_repairs: List of QA repairs applied (RepairReport objects)
+        fallbacks_triggered: List of fallback events
+        feature_matrix: Optional FeatureMatrix for feature signature (v0.2)
+        drift_report: Optional DriftReport for drift info (v0.2)
+        route_decision: Optional RouteDecision for routing audit trail (v1.0)
+
+    Returns:
+        Provenance object with signatures and metadata
+    """
+    from datetime import datetime
+
+    from tsagentkit.contracts import Provenance
+
+    metadata: dict[str, Any] = {}
+
+    # v0.2: Add feature signature if available
+    if feature_matrix is not None:
+        metadata["feature_signature"] = feature_matrix.signature
+        metadata["feature_config_hash"] = feature_matrix.config_hash
+        metadata["n_features"] = len(feature_matrix.feature_cols)
+
+    # v0.2: Add drift info if available
+    if drift_report is not None:
+        metadata["drift_detected"] = drift_report.drift_detected
+        metadata["drift_score"] = drift_report.overall_drift_score
+        metadata["drift_threshold"] = drift_report.threshold_used
+        if drift_report.drift_detected:
+            metadata["drifting_features"] = drift_report.get_drifting_features()
+
+    if column_map:
+        metadata["column_map"] = column_map
+    if original_panel_contract:
+        metadata["original_panel_contract"] = original_panel_contract
+
+    # v1.0: Add route decision for audit trail
+    if route_decision is not None:
+        metadata["route_decision"] = {
+            "buckets": route_decision.buckets,
+            "reasons": route_decision.reasons,
+            "stats": route_decision.stats,
+        }
+
+    from tsagentkit.router import compute_plan_signature
+
+    # Convert RepairReport objects to dicts for serialization
+    repairs_serialized: list[dict[str, Any]] = []
+    for repair in (qa_repairs or []):
+        if hasattr(repair, "to_dict"):
+            repairs_serialized.append(repair.to_dict())
+        else:
+            repairs_serialized.append(repair)
+
+    return Provenance(
+        run_id=str(uuid4()),
+        timestamp=datetime.now(UTC).isoformat(),
+        data_signature=compute_data_signature(data),
+        task_signature=task_spec.model_hash(),
+        plan_signature=compute_plan_signature(plan),
+        model_signature=compute_config_signature(model_config or {}),
+        qa_repairs=repairs_serialized,
+        fallbacks_triggered=fallbacks_triggered or [],
+        metadata=metadata,
+    )
+
+
+def log_event(
+    step_name: str,
+    status: str,
+    duration_ms: float,
+    error_code: str | None = None,
+    artifacts_generated: list[str] | None = None,
+    context: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Log a structured event.
+
+    Creates an event dictionary with all required fields per PRD section 6.2:
+    - step_name: Pipeline step name
+    - status: Execution status
+    - duration_ms: Execution duration
+    - error_code: Error code if applicable
+    - artifacts_generated: List of generated artifacts
+    - timestamp: ISO 8601 timestamp
+    - context: Additional context
+
+    Args:
+        step_name: Name of the pipeline step
+        status: Status (e.g., "success", "failed")
+        duration_ms: Duration in milliseconds
+        error_code: Error code if failed
+        artifacts_generated: List of artifact names generated
+        context: Additional context dictionary
+
+    Returns:
+        Event dictionary with all structured logging fields
+    """
+    from datetime import datetime
+
+    event = {
+        "step_name": step_name,
+        "status": status,
+        "duration_ms": round(duration_ms, 3),
+        "timestamp": datetime.now(UTC).isoformat(),
+        "error_code": error_code,
+        "artifacts_generated": artifacts_generated or [],
+    }
+
+    if context:
+        event["context"] = context
+
+    return event
+
+
+def format_event_json(event: dict[str, Any]) -> str:
+    """Format an event as JSON string for structured logging.
+
+    Args:
+        event: Event dictionary from log_event()
+
+    Returns:
+        JSON string representation
+    """
+    return json.dumps(event, sort_keys=True, separators=(",", ":"))
+
+
+class StructuredLogger:
+    """Structured logger for tsagentkit pipeline events.
+
+    Provides consistent JSON-formatted logging with all required fields
+    per PRD section 6.2 (Observability & Error Codes).
+
+    Example:
+        >>> logger = StructuredLogger()
+        >>> logger.start_step("fit")
+        >>> # ... do work ...
+        >>> event = logger.end_step("fit", status="success")
+        >>> print(logger.to_json())
+
+    Attributes:
+        events: List of logged events
+    """
+
+    def __init__(self) -> None:
+        """Initialize the structured logger."""
+        self.events: list[dict[str, Any]] = []
+        self._start_times: dict[str, float] = {}
+
+    def start_step(self, step_name: str) -> None:
+        """Record the start time for a step.
+
+        Args:
+            step_name: Name of the step
+        """
+        import time
+
+        self._start_times[step_name] = time.time()
+
+    def end_step(
+        self,
+        step_name: str,
+        status: str = "success",
+        error_code: str | None = None,
+        artifacts_generated: list[str] | None = None,
+        context: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """End a step and log the event.
+
+        Args:
+            step_name: Name of the step
+            status: Execution status
+            error_code: Error code if failed
+            artifacts_generated: List of artifacts generated
+            context: Additional context
+
+        Returns:
+            The logged event
+        """
+        import time
+
+        start_time = self._start_times.get(step_name)
+        if start_time is not None:
+            duration_ms = (time.time() - start_time) * 1000
+            del self._start_times[step_name]
+        else:
+            duration_ms = 0.0
+
+        event = log_event(
+            step_name=step_name,
+            status=status,
+            duration_ms=duration_ms,
+            error_code=error_code,
+            artifacts_generated=artifacts_generated,
+            context=context,
+        )
+
+        self.events.append(event)
+        return event
+
+    def log(
+        self,
+        step_name: str,
+        status: str,
+        duration_ms: float,
+        error_code: str | None = None,
+        artifacts_generated: list[str] | None = None,
+        context: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Log an event directly.
+
+        Args:
+            step_name: Name of the step
+            status: Execution status
+            duration_ms: Duration in milliseconds
+            error_code: Error code if failed
+            artifacts_generated: List of artifacts generated
+            context: Additional context
+
+        Returns:
+            The logged event
+        """
+        event = log_event(
+            step_name=step_name,
+            status=status,
+            duration_ms=duration_ms,
+            error_code=error_code,
+            artifacts_generated=artifacts_generated,
+            context=context,
+        )
+
+        self.events.append(event)
+        return event
+
+    def to_json(self) -> str:
+        """Export all events as JSON.
+
+        Returns:
+            JSON array string
+        """
+        return json.dumps(self.events, sort_keys=True, separators=(",", ":"))
+
+    def get_events(self) -> list[dict[str, Any]]:
+        """Get all logged events.
+
+        Returns:
+            List of event dictionaries (copy)
+        """
+        return self.events.copy()
+
+    def to_dict(self) -> list[dict[str, Any]]:
+        """Export all events as list of dictionaries.
+
+        Returns:
+            List of event dictionaries
+        """
+        return self.events.copy()
+
+    def get_summary(self) -> dict[str, Any]:
+        """Get summary statistics of logged events.
+
+        Returns:
+            Summary dictionary with counts and timing
+        """
+        total_duration = sum(e.get("duration_ms", 0) for e in self.events)
+        success_count = sum(1 for e in self.events if e.get("status") == "success")
+        failed_count = sum(1 for e in self.events if e.get("status") == "failed")
+
+        return {
+            "total_events": len(self.events),
+            "success_count": success_count,
+            "failed_count": failed_count,
+            "total_duration_ms": round(total_duration, 3),
+            "steps": [e.get("step_name") for e in self.events],
+        }

tsagentkit/serving/tsfm_cache.py

@@ -0,0 +1,214 @@
+"""TSFM (Time-Series Foundation Model) caching for serving.
+
+Provides model caching and lazy loading for TSFM adapters to enable
+ efficient inference in production serving environments.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from tsagentkit.models.adapters import TSFMAdapter
+
+
+class TSFMModelCache:
+    """Thread-safe cache for TSFM model instances.
+
+    Implements a singleton cache pattern for TSFM adapters to avoid
+    reloading large foundation models on each request. Uses weak
+    references to allow garbage collection when memory is constrained.
+
+    Attributes:
+        _cache: Dictionary mapping model names to cached instances
+        _lock: Threading lock for concurrent access
+        _metadata: Cache metadata (load time, access count, etc.)
+
+    Example:
+        >>> cache = TSFMModelCache()
+        >>> model = cache.get_model("chronos", pipeline="large")
+        >>> # Model is loaded once and cached for subsequent calls
+    """
+
+    _instance: TSFMModelCache | None = None
+    _instance_lock = threading.Lock()
+
+    def __new__(cls) -> TSFMModelCache:
+        """Ensure singleton pattern for global cache."""
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self) -> None:
+        """Initialize the cache (only runs once due to singleton)."""
+        if getattr(self, "_initialized", False):
+            return
+
+        self._cache: dict[str, Any] = {}
+        self._lock = threading.RLock()
+        self._metadata: dict[str, dict[str, Any]] = {}
+        self._initialized = True
+
+    def get_model(
+        self,
+        model_name: str,
+        **model_kwargs,
+    ) -> TSFMAdapter:
+        """Get a TSFM model from cache or load it.
+
+        Args:
+            model_name: Name of the TSFM (e.g., "chronos", "moirai", "timesfm")
+            **model_kwargs: Arguments passed to model initialization
+                (e.g., pipeline="large", device="cuda")
+
+        Returns:
+            Cached or newly loaded TSFMAdapter instance
+
+        Raises:
+            EModelLoadFailed: If model loading fails
+            EAdapterNotAvailable: If adapter is not installed
+        """
+        cache_key = self._make_cache_key(model_name, **model_kwargs)
+
+        with self._lock:
+            # Check cache
+            if cache_key in self._cache:
+                model = self._cache[cache_key]
+                self._metadata[cache_key]["access_count"] += 1
+                return model
+
+            # Load model
+            model = self._load_model(model_name, **model_kwargs)
+
+            # Cache model
+            self._cache[cache_key] = model
+            self._metadata[cache_key] = {
+                "model_name": model_name,
+                "load_time": self._get_timestamp(),
+                "access_count": 1,
+                "kwargs": model_kwargs,
+            }
+
+            return model
+
+    def clear_cache(self, model_name: str | None = None) -> None:
+        """Clear cached models.
+
+        Args:
+            model_name: If specified, only clear this model. Otherwise clear all.
+        """
+        with self._lock:
+            if model_name is None:
+                self._cache.clear()
+                self._metadata.clear()
+            else:
+                keys_to_remove = [
+                    k for k, v in self._metadata.items()
+                    if v.get("model_name") == model_name
+                ]
+                for key in keys_to_remove:
+                    del self._cache[key]
+                    del self._metadata[key]
+
+    def get_cache_stats(self) -> dict[str, Any]:
+        """Get cache statistics.
+
+        Returns:
+            Dictionary with cache statistics:
+            - num_models: Number of models in cache
+            - models: List of cached model names
+            - total_accesses: Total access count across all models
+        """
+        with self._lock:
+            return {
+                "num_models": len(self._cache),
+                "models": [
+                    v["model_name"] for v in self._metadata.values()
+                ],
+                "total_accesses": sum(
+                    v["access_count"] for v in self._metadata.values()
+                ),
+                "details": {
+                    k: {
+                        "model_name": v["model_name"],
+                        "access_count": v["access_count"],
+                        "load_time": v["load_time"],
+                    }
+                    for k, v in self._metadata.items()
+                },
+            }
+
+    def _make_cache_key(self, model_name: str, **kwargs) -> str:
+        """Create a unique cache key from model name and kwargs."""
+        # Sort kwargs for consistent keys
+        kv_pairs = sorted(kwargs.items())
+        kv_str = ",".join(f"{k}={v}" for k, v in kv_pairs)
+        return f"{model_name}:{kv_str}"
+
+    def _load_model(self, model_name: str, **kwargs) -> TSFMAdapter:
+        """Load a TSFM model via the adapter registry."""
+        from tsagentkit.models.adapters import AdapterConfig, AdapterRegistry
+
+        try:
+            adapter_class = AdapterRegistry.get(model_name)
+        except ValueError as exc:
+            from tsagentkit.contracts import EAdapterNotAvailable
+
+            raise EAdapterNotAvailable(
+                f"TSFM adapter '{model_name}' not found. "
+                f"Ensure the required package is installed.",
+                context={"adapter_name": model_name, "error": str(exc)},
+            ) from exc
+
+        try:
+            config = AdapterConfig(model_name=model_name, **kwargs)
+            return adapter_class(config)
+        except Exception as exc:
+            from tsagentkit.contracts import EModelLoadFailed
+
+            raise EModelLoadFailed(
+                f"Failed to load TSFM model '{model_name}': {exc}",
+                context={"adapter_name": model_name, "error": str(exc)},
+            ) from exc
+
+    def _get_timestamp(self) -> float:
+        """Get current timestamp."""
+        import time
+
+        return time.time()
+
+
+def get_tsfm_model(model_name: str, **kwargs) -> TSFMAdapter:
+    """Convenience function to get a cached TSFM model.
+
+    Args:
+        model_name: Name of the TSFM (e.g., "chronos", "moirai", "timesfm")
+        **kwargs: Model initialization arguments
+
+    Returns:
+        TSFMAdapter instance (cached or newly loaded)
+
+    Example:
+        >>> model = get_tsfm_model("chronos", pipeline="large")
+        >>> forecast = model.predict(series, horizon=7)
+    """
+    cache = TSFMModelCache()
+    return cache.get_model(model_name, **kwargs)
+
+
+def clear_tsfm_cache(model_name: str | None = None) -> None:
+    """Clear TSFM model cache.
+
+    Args:
+        model_name: If specified, only clear this model. Otherwise clear all.
+
+    Example:
+        >>> clear_tsfm_cache("chronos")  # Clear only Chronos
+        >>> clear_tsfm_cache()  # Clear all cached models
+    """
+    cache = TSFMModelCache()
+    cache.clear_cache(model_name)