synth-ai 0.2.16__py3-none-any.whl → 0.2.19__py3-none-any.whl

synth_ai/streaming/config.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from .types import StreamType
+
+
+@dataclass(slots=True)
+class StreamConfig:
+    """Configuration describing which streams to consume and how to filter them."""
+
+    enabled_streams: set[StreamType] = field(default_factory=lambda: set(StreamType))
+    event_types: set[str] | None = None  # Whitelist: only include these event types
+    event_types_exclude: set[str] | None = None  # Blacklist: exclude these event types
+    event_levels: set[str] | None = None
+    metric_names: set[str] | None = None
+    metric_phases: set[str] | None = None
+    timeline_phases: set[str] | None = None
+    sample_rate: float = 1.0
+    max_events_per_poll: int | None = None
+    deduplicate: bool = True
+
+    @classmethod
+    def default(cls) -> StreamConfig:
+        """Return a configuration representing the default (all streams) view."""
+        return cls(
+            event_types_exclude={
+                # Filter out noisy events that just announce what metrics already show
+                "sft.progress",  # Generic "Training progress" with no data
+                "sft.loss",      # Generic "Loss update" with no data
+                "sft.upstream.status",  # Very verbose status echo events
+            }
+        )
+
+    @classmethod
+    def minimal(cls) -> StreamConfig:
+        """Return a configuration streaming status updates only."""
+        return cls(enabled_streams={StreamType.STATUS})
+
+    @classmethod
+    def verbose(cls) -> StreamConfig:
+        """Return a configuration with all streams and events (no filters)."""
+        return cls()
+
+    @classmethod
+    def progress_only(cls) -> StreamConfig:
+        """Return a configuration tailored to show training progress."""
+        return cls(
+            enabled_streams={StreamType.STATUS, StreamType.EVENTS, StreamType.METRICS},
+            event_types={"sft.progress", "rl.train.step", "sft.validation.summary"},
+            metric_names={"train.loss", "eval.reward_mean"},
+        )
+
+    @classmethod
+    def errors_only(cls) -> StreamConfig:
+        """Return a configuration that focuses on heightened severity signals."""
+        return cls(
+            enabled_streams={StreamType.STATUS, StreamType.EVENTS},
+            event_levels={"error", "warning"},
+        )
+
+    def should_include_event(self, event: dict[str, Any]) -> bool:
+        """Determine whether an event message should be included."""
+        event_type = event.get("type")
+        
+        # Apply blacklist first (takes precedence)
+        if self.event_types_exclude and event_type in self.event_types_exclude:
+            return False
+        
+        # Then apply whitelist
+        if self.event_types and event_type not in self.event_types:
+            return False
+        
+        if self.event_levels:
+            return event.get("level") in self.event_levels
+        return True
+
+    def should_include_metric(self, metric: dict[str, Any]) -> bool:
+        """Determine whether a metric point should be included."""
+        if self.metric_names and metric.get("name") not in self.metric_names:
+            return False
+        if self.metric_phases:
+            return metric.get("phase") in self.metric_phases
+        return True
+
+    def should_include_timeline(self, timeline_entry: dict[str, Any]) -> bool:
+        """Determine whether a timeline entry should be included."""
+        if self.timeline_phases:
+            return timeline_entry.get("phase") in self.timeline_phases
+        return True
+
+
+__all__ = ["StreamConfig"]

synth_ai/streaming/handlers.py

@@ -0,0 +1,518 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import re
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Callable
+
+import click
+
+from .types import StreamMessage, StreamType
+
+
+def _mask_sensitive_urls(text: str) -> str:
+    """Mask S3/Wasabi URLs and sensitive paths in log messages.
+    
+    Replaces full S3/Wasabi URLs with masked versions to prevent leaking
+    bucket names, paths, and infrastructure details in public SDK logs.
+    
+    Examples:
+        s3://synth-artifacts/models/... -> s3://***/***/[masked]
+        Wasabi s3://bucket/path/file.tar.gz -> Wasabi s3://***/***/[masked]
+    """
+    if not text:
+        return text
+    
+    # Pattern matches:
+    # - Optional "Wasabi " prefix
+    # - s3:// or http(s):// scheme
+    # - Any bucket/host
+    # - Any path
+    # - Common model file extensions
+    pattern = r'(Wasabi\s+)?((s3|https?)://[^\s]+\.(tar\.gz|zip|pt|pth|safetensors|ckpt|bin))'
+    
+    def replace_url(match: re.Match) -> str:
+        prefix = match.group(1) or ""  # "Wasabi " or empty
+        url = match.group(2)
+        # Extract just the filename
+        filename = url.split("/")[-1] if "/" in url else "file"
+        return f'{prefix}s3://***/***/[{filename}]'
+    
+    return re.sub(pattern, replace_url, text, flags=re.IGNORECASE)
+
+
+class StreamHandler(ABC):
+    """Base class for log handlers that consume ``StreamMessage`` objects."""
+
+    @abstractmethod
+    def handle(self, message: StreamMessage) -> None:
+        """Process a message produced by the streamer."""
+
+    def should_handle(self, message: StreamMessage) -> bool:  # pragma: no cover - trivial
+        """Predicate allowing handlers to filter messages before processing."""
+        return True
+
+    def flush(self) -> None:  # pragma: no cover - optional
+        """Flush buffered output."""
+        return None
+
+
+class CLIHandler(StreamHandler):
+    """Simple CLI output mirroring current poller behaviour."""
+
+    def __init__(
+        self,
+        *,
+        hidden_event_types: set[str] | None = None,
+        hidden_event_substrings: set[str] | None = None,
+    ) -> None:
+        self._hidden_event_types = set(hidden_event_types or set())
+        self._hidden_event_substrings = {s.lower() for s in (hidden_event_substrings or set())}
+
+    def handle(self, message: StreamMessage) -> None:
+        if not self.should_handle(message):
+            return
+
+        timestamp = datetime.now().strftime("%H:%M:%S")
+        if message.stream_type is StreamType.STATUS:
+            status = str(message.data.get("status") or message.data.get("state") or "unknown")
+            click.echo(f"[{timestamp}] status={status}")
+            return
+
+        if message.stream_type is StreamType.EVENTS:
+            event_type = message.data.get("type", "event")
+            if event_type in self._hidden_event_types:
+                return
+            level = message.data.get("level")
+            msg = message.data.get("message") or ""
+            # Evaluate substring filters against lower-cased concatenated text
+            if self._hidden_event_substrings:
+                blob = " ".join(
+                    [
+                        event_type or "",
+                        str(msg),
+                        json.dumps(message.data.get("data", "")),
+                    ]
+                ).lower()
+                if any(sub in blob for sub in self._hidden_event_substrings):
+                    return
+            prefix = f"[{timestamp}] [{message.seq}] {event_type}"
+            if level:
+                prefix += f" ({level})"
+            # Mask sensitive URLs before displaying
+            sanitized_msg = _mask_sensitive_urls(msg)
+            click.echo(f"{prefix}: {sanitized_msg}".rstrip(": "))
+            return
+
+        if message.stream_type is StreamType.METRICS:
+            name = message.data.get("name")
+            value = message.data.get("value")
+            step = message.data.get("step")
+            data = message.data.get("data", {})
+            
+            # Format metric display
+            metric_str = f"[{timestamp}] [metric] {name}={value:.4f}" if isinstance(value, (int, float)) else f"[{timestamp}] [metric] {name}={value}"
+            if step is not None:
+                metric_str += f" (step={step})"
+            
+            # Add any additional context from data field
+            if isinstance(data, dict):
+                n = data.get("n")
+                if n is not None:
+                    metric_str += f" n={n}"
+            
+            click.echo(metric_str)
+            return
+
+        if message.stream_type is StreamType.TIMELINE:
+            phase = message.data.get("phase", "phase")
+            click.echo(f"[{timestamp}] timeline={phase}")
+
+
+class JSONHandler(StreamHandler):
+    """Emit messages as JSON lines suitable for machine parsing."""
+
+    def __init__(self, output_file: str | None = None, *, indent: int | None = None) -> None:
+        self.output_file = Path(output_file).expanduser() if output_file else None
+        self._indent = indent
+
+    def handle(self, message: StreamMessage) -> None:
+        if not self.should_handle(message):
+            return
+
+        payload: dict[str, Any] = {
+            "stream_type": message.stream_type.name,
+            "timestamp": message.timestamp,
+            "job_id": message.job_id,
+            "data": message.data,
+        }
+        if message.seq is not None:
+            payload["seq"] = message.seq
+        if message.step is not None:
+            payload["step"] = message.step
+        if message.phase is not None:
+            payload["phase"] = message.phase
+
+        line = json.dumps(payload, indent=self._indent)
+        if self.output_file:
+            with self.output_file.open("a", encoding="utf-8") as fh:
+                fh.write(line)
+                if self._indent is None:
+                    fh.write("\n")
+        else:
+            click.echo(line)
+
+    def flush(self) -> None:
+        return None
+
+
+class CallbackHandler(StreamHandler):
+    """Invoke user-provided callbacks for specific stream types."""
+
+    def __init__(
+        self,
+        *,
+        on_status: Callable[[dict[str, Any]], None] | None = None,
+        on_event: Callable[[dict[str, Any]], None] | None = None,
+        on_metric: Callable[[dict[str, Any]], None] | None = None,
+        on_timeline: Callable[[dict[str, Any]], None] | None = None,
+    ) -> None:
+        self._on_status = on_status
+        self._on_event = on_event
+        self._on_metric = on_metric
+        self._on_timeline = on_timeline
+
+    def handle(self, message: StreamMessage) -> None:
+        if not self.should_handle(message):
+            return
+
+        if message.stream_type is StreamType.STATUS and self._on_status:
+            self._on_status(message.data)
+        elif message.stream_type is StreamType.EVENTS and self._on_event:
+            self._on_event(message.data)
+        elif message.stream_type is StreamType.METRICS and self._on_metric:
+            self._on_metric(message.data)
+        elif message.stream_type is StreamType.TIMELINE and self._on_timeline:
+            self._on_timeline(message.data)
+
+
+class BufferedHandler(StreamHandler):
+    """Collect messages and emit them in batches."""
+
+    def __init__(self, *, flush_interval: float = 5.0, max_buffer_size: int = 100) -> None:
+        self.flush_interval = flush_interval
+        self.max_buffer_size = max_buffer_size
+        self._buffer: list[StreamMessage] = []
+        self._last_flush = time.time()
+
+    def handle(self, message: StreamMessage) -> None:
+        if not self.should_handle(message):
+            return
+
+        self._buffer.append(message)
+        now = time.time()
+        if len(self._buffer) >= self.max_buffer_size or now - self._last_flush >= self.flush_interval:
+            self.flush()
+
+    def flush(self) -> None:
+        if not self._buffer:
+            return
+        self.process_batch(self._buffer)
+        self._buffer.clear()
+        self._last_flush = time.time()
+
+    def process_batch(self, messages: list[StreamMessage]) -> None:  # pragma: no cover - abstract
+        """Override to define how buffered messages should be processed."""
+
+
+class IntegrationTestHandler(StreamHandler):
+    """Collect messages for integration tests or programmatic assertions."""
+
+    def __init__(self) -> None:
+        self.messages: list[StreamMessage] = []
+
+    def handle(self, message: StreamMessage) -> None:
+        self.messages.append(message)
+
+    def clear(self) -> None:
+        self.messages.clear()
+
+
+class LossCurveHandler(StreamHandler):
+    """Render a live-updating loss chart inside a fixed Rich panel."""
+
+    def __init__(
+        self,
+        *,
+        metric_name: str = "train.loss",
+        max_points: int = 200,
+        width: int = 60,
+        console: Any | None = None,
+        live: Any | None = None,
+    ) -> None:
+        try:
+            from rich.console import Console
+            from rich.live import Live
+            from rich.panel import Panel
+            from rich.text import Text
+        except ImportError as exc:  # pragma: no cover - optional dependency guard
+            raise RuntimeError(
+                "LossCurveHandler requires the 'rich' package. Install synth-ai[analytics] or rich>=13."
+            ) from exc
+
+        self.metric_name = metric_name
+        self.max_points = max_points
+        self.width = width
+
+        self._console_class = Console
+        self._panel_class = Panel
+        self._text_class = Text
+
+        self._console = console or Console()
+        self._live = live or Live(console=self._console, transient=False, refresh_per_second=8)
+        self._started = False
+
+        self._steps: list[int] = []
+        self._values: list[float] = []
+        self._status = "waiting"
+        self._last_event: str | None = None
+
+    def handle(self, message: StreamMessage) -> None:
+        updated = False
+
+        if message.stream_type is StreamType.STATUS:
+            status = str(message.data.get("status") or message.data.get("state") or "unknown")
+            if status != self._status:
+                self._status = status
+                updated = True
+
+        elif message.stream_type is StreamType.EVENTS:
+            event_type = message.data.get("type", "")
+            msg = message.data.get("message") or ""
+            level = message.data.get("level")
+            summary = f"{event_type}".strip()
+            if level:
+                summary += f" ({level})"
+            if msg:
+                summary += f": {msg}"
+            if summary != self._last_event:
+                self._last_event = summary
+                updated = True
+
+        elif message.stream_type is StreamType.METRICS:
+            if message.data.get("name") != self.metric_name:
+                return
+            value = message.data.get("value")
+            step = message.data.get("step")
+            if not isinstance(value, (int, float)) or not isinstance(step, int):
+                return
+            self._values.append(float(value))
+            self._steps.append(step)
+            if len(self._values) > self.max_points:
+                self._values = self._values[-self.max_points :]
+                self._steps = self._steps[-self.max_points :]
+            updated = True
+
+        elif message.stream_type is StreamType.TIMELINE:
+            phase = message.data.get("phase")
+            if phase:
+                self._status = str(phase)
+                updated = True
+
+        if updated:
+            self._refresh()
+
+    def flush(self) -> None:
+        if self._started:
+            with contextlib.suppress(Exception):
+                self._live.stop()
+            self._started = False
+
+    def _ensure_live(self) -> None:
+        if not self._started:
+            with contextlib.suppress(Exception):
+                self._live.start()
+            self._started = True
+
+    def _refresh(self) -> None:
+        self._ensure_live()
+        body = self._build_body()
+        title = f"{self.metric_name} | status={self._status}"
+        self._live.update(self._panel_class(body, title=title, border_style="cyan"))
+
+    def _build_body(self) -> Any:
+        if not self._values:
+            return self._text_class("Waiting for metrics…", style="yellow")
+
+        chart = self._render_sparkline()
+        last_value = self._values[-1]
+        lines = [
+            chart,
+            f"latest: {last_value:.4f} (step {self._steps[-1]})",
+        ]
+        if self._last_event:
+            lines.append(f"event: {self._last_event}")
+        return "\n".join(lines)
+
+    def _render_sparkline(self) -> str:
+        blocks = "▁▂▃▄▅▆▇█"
+        tail_len = min(self.width, len(self._values))
+        tail = self._values[-tail_len:]
+        minimum = min(tail)
+        maximum = max(tail)
+        if maximum == minimum:
+            level = blocks[0]
+            return f"{minimum:.2f} {level * tail_len} {maximum:.2f}"
+        scale = (len(blocks) - 1) / (maximum - minimum)
+        chars = "".join(blocks[int((v - minimum) * scale + 0.5)] for v in tail)
+        return f"{minimum:.2f} {chars} {maximum:.2f}"
+
+    def __del__(self) -> None:  # pragma: no cover - defensive cleanup
+        with contextlib.suppress(Exception):
+            self.flush()
+
+class RichHandler(StreamHandler):
+    """Rich powered handler with live progress and metrics table."""
+
+    def __init__(
+        self,
+        *,
+        event_log_size: int = 20,
+        console: Any | None = None,
+    ) -> None:
+        try:
+            from rich.console import Console
+            from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn
+            from rich.table import Table
+        except ImportError as exc:  # pragma: no cover - requires optional dependency
+            raise RuntimeError(
+                "RichHandler requires the 'rich' package. Install synth-ai[analytics] or rich>=13."
+            ) from exc
+
+        self._console_class = Console
+        self._progress_class = Progress
+        self._spinner_column = SpinnerColumn
+        self._text_column = TextColumn
+        self._bar_column = BarColumn
+        self._table_class = Table
+
+        self._console = console or Console()
+        self._progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TextColumn("{task.completed}/{task.total}" if console else ""),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            transient=False,
+            console=self._console,
+        )
+        self._task_id: int | None = None
+        self._current_status = "unknown"
+        self._latest_metrics: dict[str, Any] = {}
+        self._event_log: deque[str] = deque(maxlen=event_log_size)
+        self._progress_started = False
+
+    def handle(self, message: StreamMessage) -> None:
+        if not self.should_handle(message):
+            return
+
+        if message.stream_type is StreamType.STATUS:
+            self._current_status = str(message.data.get("status") or message.data.get("state"))
+            self._ensure_progress_started()
+            if self._task_id is not None:
+                description = f"Status: {self._current_status}"
+                self._progress.update(self._task_id, description=description)
+            self._render_summary()
+            return
+
+        if message.stream_type is StreamType.EVENTS:
+            event_type = message.data.get("type", "event")
+            summary = message.data.get("message") or ""
+            level = message.data.get("level")
+            # Mask sensitive URLs before displaying
+            sanitized_summary = _mask_sensitive_urls(summary)
+            formatted = f"[{event_type}] {sanitized_summary}".strip()
+            if level:
+                formatted = f"{formatted} ({level})"
+            self._event_log.append(formatted)
+            data = message.data.get("data") or {}
+            step = data.get("step") or data.get("current_step")
+            total_steps = data.get("total_steps") or data.get("max_steps")
+            if step and total_steps:
+                self._ensure_progress_started(total_steps)
+                if self._task_id is not None:
+                    self._progress.update(self._task_id, completed=int(step), total=int(total_steps))
+            self._render_summary()
+            return
+
+        if message.stream_type is StreamType.METRICS:
+            name = message.data.get("name", "")
+            value = message.data.get("value")
+            if name:
+                self._latest_metrics[name] = value
+            self._render_summary()
+            return
+
+        if message.stream_type is StreamType.TIMELINE:
+            phase = message.data.get("phase", "")
+            if phase and phase.lower() not in {"training", "running"}:
+                self._event_log.append(f"[timeline] {phase}")
+            self._render_summary()
+
+    def flush(self) -> None:
+        if self._progress_started:
+            self._progress.stop()
+            self._progress_started = False
+        self._render_summary(force=True)
+
+    def _ensure_progress_started(self, total: int | float | None = None) -> None:
+        if not self._progress_started:
+            self._progress.start()
+            self._progress_started = True
+        if self._task_id is None:
+            self._task_id = self._progress.add_task(
+                f"Status: {self._current_status}", total=total or 100
+            )
+        elif total is not None and self._task_id is not None:
+            self._progress.update(self._task_id, total=total)
+
+    def _render_summary(self, force: bool = False) -> None:
+        if force and self._progress_started:
+            self._progress.refresh()
+
+        table = self._table_class(title="Latest Metrics")
+        table.add_column("Metric")
+        table.add_column("Value")
+
+        if not self._latest_metrics:
+            table.add_row("—", "—")
+        else:
+            for name, value in sorted(self._latest_metrics.items()):
+                table.add_row(str(name), str(value))
+
+        if self._progress_started:
+            self._progress.console.print(table)
+        else:
+            self._console.print(table)
+
+        if self._event_log:
+            self._console.print("\nRecent events:")
+            for entry in list(self._event_log):
+                self._console.print(f"  • {entry}")
+
+
+__all__ = [
+    "BufferedHandler",
+    "CallbackHandler",
+    "CLIHandler",
+    "JSONHandler",
+    "IntegrationTestHandler",
+    "LossCurveHandler",
+    "RichHandler",
+    "StreamHandler",
+]