wafer-core 0.1.37__py3-none-any.whl → 0.1.39__py3-none-any.whl

wafer_core/lib/trace_compare/fusion_analyzer.py

@@ -737,6 +737,8 @@ def analyze_fusion_differences(
         "metadata": {
             "trace1_gpu": trace1_gpu,
             "trace2_gpu": trace2_gpu,
+            "trace1_platform": trace1_platform,
+            "trace2_platform": trace2_platform,
             "trace1_total_kernels": len(trace1_kernels),
             "trace2_total_kernels": len(trace2_kernels),
             "trace1_correlation_groups": len(trace1_large),

wafer_core/rollouts/_logging/__init__.py

@@ -9,11 +9,15 @@ Provides standardized logging configuration with:
 
 from .._logging.color_formatter import ColorFormatter, Colors
 from .._logging.json_formatter import JSONFormatter
-from .._logging.logging_config import setup_logging
+from .._logging.logging_config import EvalLoggingContext, setup_eval_logging, setup_logging
+from .._logging.sample_handler import SampleRoutingHandler
 
 __all__ = [
     "setup_logging",
+    "setup_eval_logging",
+    "EvalLoggingContext",
     "ColorFormatter",
     "Colors",
     "JSONFormatter",
+    "SampleRoutingHandler",
 ]

wafer_core/rollouts/_logging/logging_config.py

@@ -3,10 +3,17 @@
 Tiger Style: Explicit configuration, bounded resources, fail-fast.
 """
 
+import atexit
+import logging
 import logging.config
+import logging.handlers
 import os
+from pathlib import Path
 from typing import Any
 
+from .json_formatter import JSONFormatter
+from .sample_handler import SampleRoutingHandler
+
 
 def setup_logging(
     level: str | None = None,
@@ -120,7 +127,7 @@ def setup_logging(
                 "class": "logging.StreamHandler",
                 "level": console_level,
                 "formatter": console_formatter,
-                "stream": "ext://sys.stdout",
+                "stream": "ext://sys.stderr",
             }
         }
 
@@ -179,6 +186,91 @@ def setup_logging(
         if queue_handler is not None and hasattr(queue_handler, "listener"):
             queue_handler.listener.start()
             # Register cleanup on exit (mCoding pattern)
-            import atexit
-
             atexit.register(queue_handler.listener.stop)
+
+
+class EvalLoggingContext:
+    """Manages eval-specific logging handlers for the duration of an eval run.
+
+    Returned by setup_eval_logging(). Call teardown() when the eval is done
+    to remove handlers and close file handles.
+    """
+
+    def __init__(
+        self,
+        logger: logging.Logger,
+        queue_handler: logging.handlers.QueueHandler,
+        listener: logging.handlers.QueueListener,
+        sample_handler: SampleRoutingHandler,
+    ) -> None:
+        self.logger = logger
+        self.queue_handler = queue_handler
+        self.listener = listener
+        self.sample_handler = sample_handler
+
+    def teardown(self) -> None:
+        """Stop listener, remove handler, close files."""
+        self.listener.stop()
+        self.logger.removeHandler(self.queue_handler)
+        self.sample_handler.close()
+
+
+def setup_eval_logging(
+    output_dir: Path,
+    logger_name: str = "wafer.eval.events",
+    max_log_bytes: int = 100_000_000,
+    backup_count: int = 5,
+) -> EvalLoggingContext:
+    """Configure eval-specific logging: overview JSONL + per-sample JSONL.
+
+    Adds handlers to the named logger (not root) for the duration of an
+    eval run. Returns an EvalLoggingContext — call .teardown() when done.
+
+    Handlers:
+      1. RotatingFileHandler -> {output_dir}/events.jsonl (INFO+ overview)
+      2. SampleRoutingHandler -> {output_dir}/samples/{sample_id}.jsonl (all levels)
+
+    Both are wrapped in a QueueHandler for non-blocking writes.
+    """
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    formatter = JSONFormatter()
+
+    # 1. Overview: events.jsonl (INFO+ only)
+    events_handler = logging.handlers.RotatingFileHandler(
+        str(output_dir / "events.jsonl"),
+        mode="a",
+        maxBytes=max_log_bytes,
+        backupCount=backup_count,
+    )
+    events_handler.setFormatter(formatter)
+    events_handler.setLevel(logging.INFO)
+
+    # 2. Per-sample: samples/{sample_id}.jsonl (all levels including DEBUG)
+    sample_handler = SampleRoutingHandler(output_dir)
+    sample_handler.setFormatter(formatter)
+    sample_handler.setLevel(logging.DEBUG)
+
+    # Wrap both in QueueHandler for non-blocking writes.
+    # QueueHandler puts records onto a queue; QueueListener consumes
+    # from that queue and dispatches to the real handlers in a background thread.
+    import queue as queue_mod
+
+    q: queue_mod.Queue[logging.LogRecord] = queue_mod.Queue()
+    queue_handler = logging.handlers.QueueHandler(q)
+    listener = logging.handlers.QueueListener(
+        q, events_handler, sample_handler, respect_handler_level=True
+    )
+    listener.start()
+
+    logger = logging.getLogger(logger_name)
+    logger.addHandler(queue_handler)
+    logger.setLevel(logging.DEBUG)  # Let handlers decide what to filter
+    logger.propagate = False  # Events go to files only, not root
+
+    return EvalLoggingContext(
+        logger=logger,
+        queue_handler=queue_handler,
+        listener=listener,
+        sample_handler=sample_handler,
+    )

wafer_core/rollouts/_logging/sample_handler.py

@@ -0,0 +1,66 @@
+"""Per-sample JSONL routing handler for eval logging.
+
+Routes log records to per-sample files based on extra["sample_id"].
+Eagerly closes files on sample_end to bound open FDs to max_concurrent.
+"""
+
+import logging
+from pathlib import Path
+from typing import TextIO
+
+
+class SampleRoutingHandler(logging.Handler):
+    """Routes log records to per-sample JSONL files.
+
+    Each record with a sample_id in its extra dict gets written to
+    {output_dir}/samples/{sample_id}.jsonl. Records without sample_id
+    are silently dropped.
+
+    File handles are opened lazily on first write and closed eagerly
+    when a sample_end message is seen. Peak open FDs = max concurrent
+    samples, not total samples.
+    """
+
+    def __init__(self, output_dir: Path) -> None:
+        super().__init__()
+        self.samples_dir = output_dir / "samples"
+        self.samples_dir.mkdir(parents=True, exist_ok=True)
+        self._files: dict[str, TextIO] = {}
+
+    def emit(self, record: logging.LogRecord) -> None:
+        sample_id = getattr(record, "sample_id", None)
+        if sample_id is None:
+            return
+
+        try:
+            f = self._get_or_open(sample_id)
+            f.write(self.format(record) + "\n")
+            f.flush()
+
+            # Eagerly close on sample_end — QueueHandler serializes
+            # writes so this is safe.
+            if record.getMessage() == "sample_end":
+                self._close_sample(sample_id)
+        except Exception:
+            self.handleError(record)
+
+    def _get_or_open(self, sample_id: str) -> TextIO:
+        f = self._files.get(sample_id)
+        if f is not None:
+            return f
+        path = self.samples_dir / f"{sample_id}.jsonl"
+        f = open(path, "a")
+        self._files[sample_id] = f
+        return f
+
+    def _close_sample(self, sample_id: str) -> None:
+        f = self._files.pop(sample_id, None)
+        if f is not None:
+            f.close()
+
+    def close(self) -> None:
+        """Close all open file handles on shutdown."""
+        for f in self._files.values():
+            f.close()
+        self._files.clear()
+        super().close()

wafer_core/rollouts/_pytui/__init__.py

@@ -0,0 +1,114 @@
+"""pytui - Elm-style Python TUI framework.
+
+Zero deps. Differential rendering. Pure functions.
+
+Quick start:
+    from dataclasses import dataclass, replace
+    from pytui import App, Cmd, KeyPress
+
+    @dataclass(frozen=True)
+    class Model:
+        count: int = 0
+
+    def update(model, msg):
+        match msg:
+            case KeyPress(key="q"):   return model, Cmd.quit()
+            case KeyPress(key="j"):   return replace(model, count=model.count + 1), Cmd.none()
+        return model, Cmd.none()
+
+    def view(model, width, height):
+        return [f"Count: {model.count}", "", "j: up  q: quit"]
+
+    App(init=(Model(), Cmd.none()), update=update, view=view).run()
+"""
+
+# Elm core
+from .app import App, Cmd, FocusEvent, KeyPress, MouseEvent, PasteEvent, Resize, Sub
+
+# Console (unified output with spinner/logging coordination)
+from .console import Console, ConsoleHandler, SpinnerHandle
+
+# Renderer
+from .renderer import RenderState, diff_render
+
+# Spinner
+from .spinner import Spinner
+
+# Terminal
+from .terminal import Terminal, TerminalProtocol
+
+# Text utilities
+from .text import (
+    AnsiCode,
+    AnsiCodeTracker,
+    apply_background_to_line,
+    extract_ansi_code,
+    slice_ansi,
+    strip_terminal_control_sequences,
+    truncate_to_width,
+    visible_width,
+    wrap_text_with_ansi,
+)
+
+# Theme
+from .theme import (
+    DARK_THEME,
+    MINIMAL_THEME,
+    RESET,
+    ROUNDED_THEME,
+    SOFT_DARK_THEME,
+    MinimalTheme,
+    RoundedTheme,
+    Theme,
+    ToolDisplayMode,
+    hex_to_bg,
+    hex_to_fg,
+    hex_to_rgb,
+)
+
+__all__ = [
+    # Elm core
+    "App",
+    "Cmd",
+    "Sub",
+    "KeyPress",
+    "MouseEvent",
+    "PasteEvent",
+    "FocusEvent",
+    "Resize",
+    # Terminal
+    "Terminal",
+    "TerminalProtocol",
+    # Renderer
+    "RenderState",
+    "diff_render",
+    # Text
+    "visible_width",
+    "truncate_to_width",
+    "slice_ansi",
+    "wrap_text_with_ansi",
+    "apply_background_to_line",
+    "extract_ansi_code",
+    "strip_terminal_control_sequences",
+    "AnsiCode",
+    "AnsiCodeTracker",
+    # Spinner
+    "Spinner",
+    # Console
+    "Console",
+    "ConsoleHandler",
+    "SpinnerHandle",
+    # Theme
+    "Theme",
+    "MinimalTheme",
+    "RoundedTheme",
+    "DARK_THEME",
+    "SOFT_DARK_THEME",
+    "MINIMAL_THEME",
+    "ROUNDED_THEME",
+    "RESET",
+    "ToolDisplayMode",
+    "hex_to_fg",
+    "hex_to_bg",
+    "hex_to_rgb",
+]