sixtyseven 0.1.0__cp312-cp312-macosx_11_0_arm64.whl

sixtyseven-0.1.0.data/purelib/sixtyseven/run.py

@@ -0,0 +1,445 @@
+"""Main Run class for tracking ML experiments."""
+
+from __future__ import annotations
+
+import atexit
+import signal
+import sys
+import threading
+import time
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+from sixtyseven.config import SDKConfig, get_config
+from sixtyseven.exceptions import ValidationError
+from sixtyseven.utils import generate_run_name, get_git_info, get_system_info
+
+if TYPE_CHECKING:
+    from sixtyseven.client import SixtySevenClient
+    from sixtyseven.metrics import MetricsBatcher
+    from sixtyseven.local import LocalWriter, LocalBatcher
+    from sixtyseven.server import ServerManager
+
+
+class Run:
+    """
+    Main class for tracking ML experiments with Sixtyseven.
+
+    A Run represents a single training session or experiment. It tracks:
+    - Configuration/hyperparameters
+    - Metrics over time (loss, accuracy, etc.)
+    - System information
+    - Git information
+
+    Supports two modes:
+    - Local mode (default): Writes to local SQLite files, viewable with `sixtyseven --logdir`
+    - Remote mode: Sends to a Sixtyseven API server
+
+    Usage:
+        # Context manager (recommended)
+        with Run(project="my-project") as run:
+            run.log_config({"lr": 0.001})
+            for epoch in range(100):
+                run.log_metrics({"loss": 0.5}, step=epoch)
+
+        # Explicit management
+        run = Run(project="my-project")
+        run.log_metrics({"loss": 0.5}, step=1)
+        run.complete()  # or run.fail("error message")
+
+    Environment variables:
+        SIXTYSEVEN_LOGDIR: Set to enable local mode with custom log directory
+        SIXTYSEVEN_URL: Set to enable remote mode with API server
+        SIXTYSEVEN_API_KEY: API key for remote mode authentication
+    """
+
+    def __init__(
+        self,
+        project: str,
+        api_key: Optional[str] = None,
+        name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        config: Optional[Dict[str, Any]] = None,
+        # Advanced options
+        mode: Optional[str] = None,  # "local" or "remote"
+        logdir: Optional[str] = None,
+        base_url: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        flush_interval: Optional[float] = None,
+        capture_git: Optional[bool] = None,
+        capture_system: Optional[bool] = None,
+        # Server option (local mode only)
+        start_server: bool = False,
+    ):
+        """
+        Initialize a new run.
+
+        Args:
+            project: Project identifier. In local mode, this is just a project name.
+                     In remote mode, use "team-slug/app-slug" format.
+            api_key: API key for authentication (remote mode only)
+            name: Run name (auto-generated if not provided)
+            tags: List of tags for organizing runs
+            config: Initial configuration/hyperparameters
+            mode: Operating mode ("local" or "remote"). Auto-detected if not specified.
+            logdir: Directory for logs (local mode only)
+            base_url: API base URL (remote mode only)
+            batch_size: Number of metrics to batch before sending/writing
+            flush_interval: Seconds between automatic flushes
+            capture_git: Whether to capture git information
+            capture_system: Whether to capture system information
+            start_server: Automatically start the sixtyseven viewer server and open the
+                browser (local mode only). The server stops when the run ends.
+
+        Raises:
+            ValidationError: If project format is invalid (remote mode)
+            AuthenticationError: If API key is invalid (remote mode)
+            ServerError: If start_server=True but the sixtyseven binary cannot be found
+        """
+        # Get global config
+        global_config = get_config()
+
+        # Build effective config
+        self._config = SDKConfig(
+            mode=mode or global_config.mode,
+            logdir=logdir or global_config.logdir,
+            base_url=base_url or global_config.base_url,
+            api_key=api_key or global_config.api_key,
+            batch_size=batch_size or global_config.batch_size,
+            flush_interval=flush_interval or global_config.flush_interval,
+            capture_git=capture_git
+            if capture_git is not None
+            else global_config.capture_git,
+            capture_system=capture_system
+            if capture_system is not None
+            else global_config.capture_system,
+        )
+
+        self._project = project
+        self._run_name = name or generate_run_name()
+        self._tags = tags or []
+        self._initial_config = config
+
+        # Mode-specific initialization
+        self._local_writer: Optional["LocalWriter"] = None
+        self._local_batcher: Optional["LocalBatcher"] = None
+        self._remote_client: Optional["SixtySevenClient"] = None
+        self._remote_batcher: Optional["MetricsBatcher"] = None
+        self._server_manager: Optional["ServerManager"] = None
+        self._start_server = start_server
+
+        # Capture info before creating run
+        self._git_info = None
+        self._system_info = None
+        if self._config.capture_git:
+            self._git_info = get_git_info()
+        if self._config.capture_system:
+            self._system_info = get_system_info()
+
+        if self._config.mode == "local":
+            self._init_local_mode()
+        else:
+            self._init_remote_mode()
+
+        # Track state
+        self._step = 0
+        self._closed = False
+        self._lock = threading.Lock()
+
+        # Register cleanup handlers
+        atexit.register(self._cleanup)
+        self._setup_signal_handlers()
+
+    def _init_local_mode(self) -> None:
+        """Initialize local file-based storage."""
+        from sixtyseven.local import LocalWriter, LocalBatcher
+
+        self._local_writer = LocalWriter(
+            logdir=self._config.logdir,
+            project=self._project,
+            run_name=self._run_name,
+            tags=self._tags,
+            config=self._initial_config,
+            git_info=self._git_info,
+            system_info=self._system_info,
+        )
+        self._run_id = self._local_writer.run_id
+
+        self._local_batcher = LocalBatcher(
+            writer=self._local_writer,
+            batch_size=self._config.batch_size,
+            flush_interval=self._config.flush_interval,
+        )
+        self._local_batcher.start()
+
+        # Print local mode info
+        print(f"Sixtyseven: Logging to {self._local_writer.run_dir}")
+
+        # Start the viewer server if requested
+        if self._start_server:
+            from sixtyseven.server import ServerManager
+
+            self._server_manager = ServerManager(
+                logdir=self._config.logdir,
+                open_browser=True,
+                project=self._project,
+                run_id=self._run_id,
+            )
+            self._server_manager.start()
+
+    def _init_remote_mode(self) -> None:
+        """Initialize remote API client."""
+        from sixtyseven.client import SixtySevenClient
+        from sixtyseven.metrics import MetricsBatcher
+
+        # Parse project for remote mode
+        self._team_slug, self._app_slug = self._parse_project(self._project)
+
+        self._remote_client = SixtySevenClient(self._config, self._config.api_key)
+
+        # Create run on server
+        self._run_id = self._remote_client.create_run(
+            team_slug=self._team_slug,
+            app_slug=self._app_slug,
+            name=self._run_name,
+            tags=self._tags,
+            config=self._initial_config or {},
+            git_info=self._git_info,
+            system_info=self._system_info,
+        )
+
+        self._remote_batcher = MetricsBatcher(
+            client=self._remote_client,
+            run_id=self._run_id,
+            batch_size=self._config.batch_size,
+            flush_interval=self._config.flush_interval,
+        )
+        self._remote_batcher.start()
+
+    @property
+    def id(self) -> str:
+        """Return the run ID."""
+        return self._run_id
+
+    @property
+    def name(self) -> str:
+        """Return the run name."""
+        return self._run_name
+
+    @property
+    def project(self) -> str:
+        """Return the project identifier."""
+        return self._project
+
+    @property
+    def mode(self) -> str:
+        """Return the operating mode ('local' or 'remote')."""
+        return self._config.mode
+
+    @property
+    def logdir(self) -> Optional[str]:
+        """Return the log directory (local mode only)."""
+        if self._local_writer:
+            return str(self._local_writer.run_dir)
+        return None
+
+    def log_config(self, config: Dict[str, Any]) -> None:
+        """
+        Log configuration/hyperparameters.
+
+        This merges with any config provided at initialization.
+        Can be called multiple times to add more config.
+
+        Args:
+            config: Dictionary of configuration values
+
+        Example:
+            run.log_config({
+                "learning_rate": 0.001,
+                "batch_size": 32,
+                "optimizer": "adam",
+            })
+        """
+        if self._config.mode == "local":
+            self._local_writer.log_config(config)
+        else:
+            self._remote_client.update_run_config(self._run_id, config)
+
+    def log_metrics(
+        self,
+        metrics: Dict[str, float],
+        step: Optional[int] = None,
+        timestamp: Optional[float] = None,
+    ) -> None:
+        """
+        Log metrics for the current step.
+
+        Metrics are automatically batched for efficiency. Use flush()
+        to force immediate sending.
+
+        Args:
+            metrics: Dictionary of metric name -> value
+            step: Step number (auto-incremented if not provided)
+            timestamp: Unix timestamp (current time if not provided)
+
+        Example:
+            run.log_metrics({
+                "train/loss": 0.45,
+                "train/accuracy": 0.82,
+                "val/loss": 0.52,
+            }, step=epoch)
+        """
+        with self._lock:
+            if step is None:
+                step = self._step
+                self._step += 1
+            else:
+                self._step = max(self._step, step + 1)
+
+            ts = timestamp or time.time()
+
+            batcher = (
+                self._local_batcher
+                if self._config.mode == "local"
+                else self._remote_batcher
+            )
+            for name, value in metrics.items():
+                batcher.add(name=name, value=float(value), step=step, timestamp=ts)
+
+    def log(self, name: str, value: float, step: Optional[int] = None) -> None:
+        """
+        Log a single metric.
+
+        Convenience method for logging one metric at a time.
+
+        Args:
+            name: Metric name
+            value: Metric value
+            step: Step number (auto-incremented if not provided)
+        """
+        self.log_metrics({name: value}, step=step)
+
+    def flush(self) -> None:
+        """Force flush all buffered metrics."""
+        if self._config.mode == "local":
+            self._local_batcher.flush()
+        else:
+            self._remote_batcher.flush()
+
+    def add_tags(self, tags: List[str]) -> None:
+        """
+        Add tags to the run.
+
+        Args:
+            tags: List of tags to add
+        """
+        if self._config.mode == "local":
+            # For local mode, update the meta.json
+            meta = self._local_writer._read_meta()
+            existing_tags = meta.get("tags", [])
+            meta["tags"] = list(set(existing_tags + tags))
+            self._local_writer._write_meta(meta)
+        else:
+            self._remote_client.add_run_tags(self._run_id, tags)
+
+    def complete(self) -> None:
+        """Mark the run as completed successfully."""
+        self._finalize("completed")
+
+    def fail(self, error: Optional[str] = None) -> None:
+        """
+        Mark the run as failed.
+
+        Args:
+            error: Optional error message
+        """
+        self._finalize("failed", error=error)
+
+    def abort(self) -> None:
+        """Mark the run as aborted."""
+        self._finalize("aborted")
+
+    def _finalize(self, status: str, error: Optional[str] = None) -> None:
+        """Finalize the run with the given status."""
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+
+        if self._config.mode == "local":
+            # Flush and stop local batcher
+            self._local_batcher.flush()
+            self._local_batcher.stop()
+            # Update status in meta.json
+            self._local_writer.update_status(status, error)
+            self._local_writer.close()
+            # Stop the server if we started it
+            if self._server_manager is not None:
+                self._server_manager.stop()
+        else:
+            # Flush and stop remote batcher
+            self._remote_batcher.flush()
+            self._remote_batcher.stop()
+            # Update status on server
+            self._remote_client.update_run_status(self._run_id, status, error=error)
+
+    def __enter__(self) -> "Run":
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit context manager, marking run as completed or failed."""
+        if exc_type is not None:
+            self.fail(str(exc_val) if exc_val else "Unknown error")
+        else:
+            self.complete()
+
+    def _setup_signal_handlers(self) -> None:
+        """Set up signal handlers for graceful shutdown."""
+        self._original_sigint = signal.getsignal(signal.SIGINT)
+        self._original_sigterm = signal.getsignal(signal.SIGTERM)
+
+        def handle_signal(signum, frame):
+            """Handle termination signals."""
+            # Immediately restore original handlers to prevent re-entry
+            signal.signal(signal.SIGINT, self._original_sigint)
+            signal.signal(signal.SIGTERM, self._original_sigterm)
+
+            # Mark run as canceled
+            if not self._closed:
+                sig_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM"
+                print(f"\nSixtyseven: Run canceled ({sig_name})")
+                self._finalize("canceled", error=f"Interrupted by {sig_name}")
+
+            # Re-raise to exit
+            if signum == signal.SIGINT:
+                raise KeyboardInterrupt
+            else:
+                sys.exit(128 + signum)
+
+        # Only set handlers in main thread
+        if threading.current_thread() is threading.main_thread():
+            signal.signal(signal.SIGINT, handle_signal)
+            signal.signal(signal.SIGTERM, handle_signal)
+
+    def _cleanup(self) -> None:
+        """Cleanup handler for atexit - marks incomplete runs as canceled."""
+        if not self._closed:
+            self._finalize("canceled", error="Process exited unexpectedly")
+
+    @staticmethod
+    def _parse_project(project: str) -> tuple:
+        """Parse project string into team and app slugs (remote mode only)."""
+        parts = project.split("/")
+        if len(parts) != 2:
+            raise ValidationError(
+                f"Invalid project format: '{project}'. Expected 'team-slug/app-slug' format for remote mode."
+            )
+        return parts[0], parts[1]
+
+    def __repr__(self) -> str:
+        mode_info = (
+            f"logdir='{self.logdir}'"
+            if self._config.mode == "local"
+            else f"url='{self._config.base_url}'"
+        )
+        return f"Run(id='{self._run_id}', project='{self.project}', mode='{self.mode}', {mode_info})"