gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/run.py

@@ -0,0 +1,1781 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""Gateway Runs - Flexible metric tracking for experiments.
+
+This module provides wandb-style experiment tracking with automatic crash detection
+via heartbeats, signal handling, and context managers.
+
+Example usage::
+
+    import gatewaysdk
+
+    # Recommended: Context manager (handles crashes automatically)
+    with gatewaysdk.run("experiment-v1", config={"lr": 0.01}) as run:
+        for epoch in range(100):
+            loss = train()
+            run.log({"loss": loss, "accuracy": acc}, step=epoch)
+
+    # Alternative: Manual lifecycle (must call finish!)
+    run = gatewaysdk.init_run("experiment-v1")
+    run.log({"metric": value})
+    run.finish()
+"""
+
+from __future__ import annotations
+
+import asyncio
+import atexit
+import functools
+import hashlib
+import json
+import logging
+import os
+import signal
+import subprocess
+import sys
+import threading
+import time
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, TypeVar, Union
+
+if TYPE_CHECKING:
+    from gatewaysdk.store.listener import StorageListener
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+# Type variable for decorators
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+# =============================================================================
+# Reproducibility Helpers
+# =============================================================================
+
+
+def _compute_config_hash(config: Optional[Dict[str, Any]]) -> Optional[str]:
+    """SHA256 hash of config for detecting identical configs.
+
+    Returns first 16 characters of the hash for brevity.
+    """
+    if not config:
+        return None
+    try:
+        config_str = json.dumps(config, sort_keys=True, default=str)
+        return hashlib.sha256(config_str.encode()).hexdigest()[:16]
+    except Exception:
+        return None
+
+
+def _get_git_info() -> Dict[str, Any]:
+    """Capture git commit, branch, and dirty state.
+
+    Returns empty dict if not in a git repo or git is not available.
+    """
+    try:
+        # Get commit SHA (first 12 chars)
+        sha = (
+            subprocess.check_output(
+                ["git", "rev-parse", "HEAD"],
+                stderr=subprocess.DEVNULL,
+                timeout=5,
+            )
+            .decode()
+            .strip()[:12]
+        )
+
+        # Get branch name
+        branch = (
+            subprocess.check_output(
+                ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+                stderr=subprocess.DEVNULL,
+                timeout=5,
+            )
+            .decode()
+            .strip()
+        )
+
+        # Check if working directory is dirty
+        dirty = (
+            subprocess.call(
+                ["git", "diff", "--quiet"],
+                stderr=subprocess.DEVNULL,
+                timeout=5,
+            )
+            != 0
+        )
+
+        return {"gitSha": sha, "gitBranch": branch, "gitDirty": dirty}
+    except Exception:
+        return {}
+
+
+# =============================================================================
+# Types
+# =============================================================================
+
+RunState = Literal["pending", "running", "completed", "failed", "cancelled", "crashed"]
+
+
+class GatewayRunError(RuntimeError):
+    """Run creation or completion could not satisfy its durability contract."""
+
+
+# =============================================================================
+# Constants
+# =============================================================================
+
+MAX_RUN_NAME_LENGTH = 256
+MAX_METRIC_BUFFER_SIZE = 10000
+
+
+@dataclass
+class Metric:
+    """A self-describing metric with metadata."""
+
+    name: str
+    value: float
+    step: Optional[int] = None
+    unit: str = ""
+    metric_type: str = "gauge"
+    labels: Dict[str, str] = field(default_factory=dict)
+
+
+# =============================================================================
+# Run and Rollout Context Management
+# =============================================================================
+
+_current_run: ContextVar[Optional["Run"]] = ContextVar("gateway_current_run", default=None)
+_current_rollout: ContextVar[Optional["Rollout"]] = ContextVar("gateway_current_rollout", default=None)
+
+
+def get_current_run() -> Optional["Run"]:
+    """Get the currently active run, if any."""
+    return _current_run.get()
+
+
+def _set_current_run(run: Optional["Run"]) -> None:
+    """Set the current active run."""
+    _current_run.set(run)
+
+
+def get_current_rollout() -> Optional["Rollout"]:
+    """Get the currently active rollout (episode), if any."""
+    return _current_rollout.get()
+
+
+def _set_current_rollout(rollout: Optional["Rollout"]) -> None:
+    """Set the current active rollout."""
+    _current_rollout.set(rollout)
+
+
+# =============================================================================
+# Rollout Class
+# =============================================================================
+
+
+class Rollout:
+    """A Rollout/Episode within a Gateway Run.
+
+    Represents a single episode (multi-turn interaction) within a larger
+    experiment (Run). Each rollout gets its own session_id that is injected
+    into traces created within its context.
+
+    The trace_id is captured automatically by TracingContextSpanProcessor
+    when the first span starts within this rollout. This allows score()
+    to link scores to the correct trace even after instrumented spans end.
+
+    Usage::
+
+        with gatewaysdk.run("gepa-v1") as run:
+            for candidate in population:
+                with run.rollout(f"candidate-{candidate.id}") as ep:
+                    result = agent.execute(task)
+                    ep.score("rubric", grade)
+                    # Traces get: run_id="run-xxx", session_id="ep-yyy"
+    """
+
+    def __init__(
+        self,
+        run: "Run",
+        name: Optional[str] = None,
+    ):
+        """Initialize a Rollout.
+
+        Args:
+            run: The parent Run instance.
+            name: Optional human-readable name for this rollout.
+        """
+        self._run = run
+        self._name = name
+        self._session_id = f"ep-{uuid.uuid4().hex[:8]}"
+        self._trace_id: Optional[str] = None
+        self._started = False
+        self._finished = False
+
+    @property
+    def session_id(self) -> str:
+        """Get the session ID for this rollout."""
+        return self._session_id
+
+    @property
+    def trace_id(self) -> Optional[str]:
+        """Get the trace ID captured from the first span in this rollout."""
+        return self._trace_id
+
+    @property
+    def name(self) -> Optional[str]:
+        """Get the rollout name."""
+        return self._name
+
+    @property
+    def run(self) -> "Run":
+        """Get the parent run."""
+        return self._run
+
+    def _capture_trace_id(self, trace_id: str) -> None:
+        """Store the trace_id from the first span in this rollout.
+
+        Called by TracingContextSpanProcessor.on_start() when a span
+        starts within this rollout's context.
+        """
+        if self._trace_id is None:
+            self._trace_id = trace_id
+
+    def __enter__(self) -> "Rollout":
+        """Context manager entry - set this as the current rollout."""
+        self._started = True
+        _set_current_rollout(self)
+        logger.debug(f"Started rollout: {self._name or self._session_id}")
+        return self
+
+    def score(
+        self,
+        name: str,
+        value: Union[float, str],
+        **kwargs,
+    ) -> None:
+        """Post a score linked to this rollout's session and trace.
+
+        Convenience shorthand that auto-fills session_id and trace_id
+        from this rollout. See :meth:`Run.score` for full documentation.
+
+        Args:
+            name: Score name (e.g., "rubric", "accuracy").
+            value: Score value.
+            **kwargs: Additional arguments passed to Run.score().
+
+        Example::
+
+            with run.rollout("task-1") as rollout:
+                result = await run_agent(task)
+                rollout.score("rubric", 0.85)
+        """
+        # Use the trace_id captured from spans within this rollout
+        if "trace_id" not in kwargs and self._trace_id:
+            kwargs["trace_id"] = self._trace_id
+        self._run.score(name, value, session_id=self._session_id, **kwargs)
+
+    def success(
+        self,
+        name: str,
+        value: Union[bool, float] = True,
+        **kwargs,
+    ) -> None:
+        """Emit a success signal linked to this rollout's session and trace.
+
+        Convenience shorthand; see :meth:`Run.success` for full documentation.
+        """
+        if "trace_id" not in kwargs and self._trace_id:
+            kwargs["trace_id"] = self._trace_id
+        self._run.success(name, value, session_id=self._session_id, **kwargs)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - clear the current rollout."""
+        _set_current_rollout(None)
+        self._finished = True
+        if exc_type is not None:
+            logger.debug(f"Rollout {self._session_id} exited with exception: {exc_val}")
+        else:
+            logger.debug(f"Rollout {self._session_id} completed")
+        return False  # Don't suppress exceptions
+
+
+# =============================================================================
+# Run Class
+# =============================================================================
+
+
+class Run:
+    """A Gateway Run for tracking metrics with automatic heartbeat and crash detection.
+
+    Features:
+    - Automatic heartbeat to detect crashed runs
+    - Signal handlers for SIGINT/SIGTERM (Ctrl+C)
+    - Context manager support with exception handling
+    - atexit handler for scripts that exit without calling finish()
+    - Automatic trace linkage via session_id
+    """
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        *,
+        project_id: Optional[str] = None,
+        heartbeat_interval: int = 30,
+        source: str = "sdk",
+        labels: Optional[Dict[str, Any]] = None,
+        config: Optional[Dict[str, Any]] = None,
+        tags: Optional[List[str]] = None,
+        api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        listener: Optional["StorageListener"] = None,
+    ):
+        """Initialize a Run.
+
+        Args:
+            name: Human-readable name for the run.
+            project_id: Project ID (from env GATEWAY_PROJECT_ID if not provided).
+            heartbeat_interval: Seconds between heartbeats (default: 30).
+            source: Source identifier (e.g., "gepa", "trainer", "notebook").
+            labels: Arbitrary key-value labels for filtering.
+            config: Hyperparameters/config dict (logged for reproducibility).
+            tags: Simple string tags for categorization.
+            api_url: API URL (from env GATEWAY_HOST if not provided).
+            api_key: API key (from env GATEWAY_SECRET_KEY if not provided).
+            listener: Optional StorageListener for GEPA/optimization callbacks.
+
+        Note:
+            Runs can have MANY sessions (episodes/rollouts). Use run.rollout()
+            to create individual episodes within a run.
+        """
+        # Validate and truncate name if too long
+        if name and len(name) > MAX_RUN_NAME_LENGTH:
+            logger.warning(f"Run name truncated from {len(name)} to {MAX_RUN_NAME_LENGTH} characters")
+            name = name[:MAX_RUN_NAME_LENGTH]
+        self._name = name or f"run-{uuid.uuid4().hex[:8]}"
+        self._project_id = project_id or os.environ.get("GATEWAY_PROJECT_ID")
+        self._heartbeat_interval = heartbeat_interval
+        self._source = source
+        self._labels = labels or {}
+        self._config = config
+        self._tags = tags or []
+
+        # API configuration
+        self._api_url = api_url or os.environ.get("GATEWAY_HOST", "http://localhost:3000")
+        self._api_key = api_key or os.environ.get("GATEWAY_SECRET_KEY")
+        # Store public key for auth (read from env since it's set by GatewayLightningStore)
+        self._public_key = os.environ.get("GATEWAY_PUBLIC_KEY", "")
+
+        # Internal state
+        self._run_id: Optional[str] = None
+        self._run_number: int = 0
+        self._state: RunState = "pending"
+        self._heartbeat_thread: Optional[threading.Thread] = None
+        self._stop_heartbeat = threading.Event()
+        self._flush_ticker_thread: Optional[threading.Thread] = None
+        self._stop_flush_ticker = threading.Event()
+        self._finished = False
+        self._lock = threading.Lock()
+        self._score_threads: List[threading.Thread] = []
+        self._score_delivery_failures: List[str] = []
+        self._last_finish_error: Optional[str] = None
+        self._last_finish_exception: Optional[BaseException] = None
+
+        # Signal handling state
+        self._original_sigint: Optional[Callable] = None
+        self._original_sigterm: Optional[Callable] = None
+
+        # Metric buffer for batching
+        self._metric_buffer: List[Dict[str, Any]] = []
+        self._metric_buffer_lock = threading.Lock()
+        self._metric_flush_lock = threading.Lock()
+        self._dropped_metric_count = 0
+        self._last_flush_time = time.time()
+        self._flush_interval = 1.0  # Flush every 1 second for live updates
+
+        # Optional listener for GEPA/optimization callbacks
+        self._listener: Optional["StorageListener"] = listener
+
+        # Cancellation callback -- set by Trainer/algorithm to propagate
+        # remote cancellation (from UI "Cancel Run") to the algorithm.
+        self._on_cancel: Optional[Callable[[], None]] = None
+
+    @property
+    def id(self) -> Optional[str]:
+        """Get the run ID."""
+        return self._run_id
+
+    @property
+    def run_number(self) -> int:
+        """Get the auto-incrementing run number for this project."""
+        return self._run_number
+
+    @property
+    def name(self) -> str:
+        """Get the run name."""
+        return self._name
+
+    @property
+    def state(self) -> RunState:
+        """Get the current run state."""
+        return self._state
+
+    @property
+    def project_id(self) -> Optional[str]:
+        """Get the project ID."""
+        return self._project_id
+
+    @property
+    def listener(self) -> Optional["StorageListener"]:
+        """Get the attached listener, if any."""
+        return self._listener
+
+    def _notify_listener(self, method_name: str, *args: Any, **kwargs: Any) -> None:
+        """Notify the listener of an event (synchronous wrapper for async methods).
+
+        This method safely calls listener methods, handling both sync and async cases.
+        Exceptions are logged but not raised to avoid disrupting the main flow.
+        """
+        if self._listener is None:
+            return
+        try:
+            method = getattr(self._listener, method_name, None)
+            if method is not None:
+                import asyncio
+
+                if asyncio.iscoroutinefunction(method):
+                    # Run async method in a new event loop if needed
+                    try:
+                        loop = asyncio.get_running_loop()
+                        # Already in an async context, schedule it
+                        asyncio.create_task(method(*args, **kwargs))
+                    except RuntimeError:
+                        # No running loop, run synchronously
+                        asyncio.run(method(*args, **kwargs))
+                else:
+                    method(*args, **kwargs)
+        except Exception as e:
+            logger.warning(f"Listener {method_name} failed: {e}")
+
+    # =========================================================================
+    # Rollout/Episode Methods
+    # =========================================================================
+
+    def rollout(self, name: Optional[str] = None) -> Rollout:
+        """Create a new rollout (episode) within this run.
+
+        A rollout represents a single episode or interaction sequence.
+        Traces created within the rollout context get both run_id and
+        session_id attributes for proper grouping.
+
+        Usage::
+
+            with gatewaysdk.run("gepa-v1") as run:
+                for candidate in population:
+                    with run.rollout(f"candidate-{candidate.id}") as ep:
+                        result = agent.execute(task)
+                        # Traces get: run_id="run-xxx", session_id="ep-yyy"
+
+        Args:
+            name: Optional human-readable name for this rollout.
+
+        Returns:
+            A Rollout context manager.
+        """
+        return Rollout(run=self, name=name)
+
+    # =========================================================================
+    # Lifecycle Methods
+    # =========================================================================
+
+    def _auto_init_tracing(self) -> None:
+        """Auto-initialize tracing if not already done and Gateway env vars are set.
+
+        This enables the "Runs standalone" experience: when a user creates a Run,
+        tracing is automatically set up so all LLM calls within the run context
+        are traced and exported to Gateway, linked to this run via run_id.
+
+        Only initializes if:
+        - Tracing is not already initialized
+        - GATEWAY_HOST or GATEWAY_PUBLIC_KEY env vars are configured
+        """
+        try:
+            from gatewaysdk.tracing.init import _initialized as tracing_initialized
+            from gatewaysdk.tracing.init import init as tracing_init
+
+            if tracing_initialized:
+                logger.debug("Tracing already initialized, skipping auto-init")
+                return
+
+            # Only auto-init if Gateway env vars are present (somewhere to send spans)
+            has_gateway_host = bool(os.environ.get("GATEWAY_HOST"))
+            has_gateway_keys = bool(os.environ.get("GATEWAY_PUBLIC_KEY") and os.environ.get("GATEWAY_SECRET_KEY"))
+
+            if has_gateway_host or has_gateway_keys:
+                logger.info("Auto-initializing tracing from msk.run() (Gateway env vars detected)")
+                tracing_init(service_name=self._source or "gatewaysdk")
+            else:
+                logger.debug("Skipping tracing auto-init: no Gateway env vars configured")
+
+        except Exception as e:
+            logger.warning(f"Failed to auto-init tracing: {e}")
+
+    def _start(self) -> "Run":
+        """Initialize the run on the backend and start heartbeat.
+
+        Automatically initializes tracing if not already done and Gateway
+        environment variables are configured. This ensures that traces created
+        within this run are exported to Gateway and linked via run_id.
+        """
+        if not self._project_id:
+            raise ValueError(
+                "project_id is required. Set GATEWAY_PROJECT_ID environment variable "
+                "or pass project_id to init_run()."
+            )
+
+        # Create run via API (skip if run_id is pre-set, e.g., platform runs)
+        if self._run_id is None:
+            run_id, run_number = self._create_run_on_backend()
+            self._run_id = run_id
+            self._run_number = run_number
+        else:
+            logger.info(f"Reusing existing run {self._run_id} (pre-set)")
+
+        # Backend identity is the activation commit point. Do not initialize
+        # tracing or install lifecycle side effects before it is acquired.
+        self._auto_init_tracing()
+        self._state = "running"
+
+        # Set as current run (for trace context)
+        _set_current_run(self)
+
+        # Start heartbeat thread
+        self._start_heartbeat_loop()
+
+        # Start flush ticker for live metric updates (every 1s)
+        self._start_flush_ticker()
+
+        # Register signal handlers
+        self._register_signal_handlers()
+
+        # Register atexit handler
+        atexit.register(self._atexit_handler)
+
+        # Print run URL for convenience
+        run_url = f"{self._api_url}/project/{self._project_id}/runs/{self._run_id}"
+        logger.info(f"Gateway Run #{self._run_number} started: {run_url}")
+
+        return self
+
+    def _create_run_on_backend(self) -> tuple[str, int]:
+        """Create run record on backend, return (run_id, run_number)."""
+        if not self._public_key:
+            raise GatewayRunError("Run create failed: GATEWAY_PUBLIC_KEY is missing or empty")
+
+        if not self._api_key:
+            raise GatewayRunError("Run create failed: GATEWAY_SECRET_KEY (API secret key) is missing or empty")
+
+        try:
+            # Ensure config is always a dict (API rejects null)
+            config = self._config if self._config is not None else {}
+
+            # Auto-capture reproducibility info
+            config_hash = _compute_config_hash(config) or ""
+            git_info = _get_git_info()
+
+            url = f"{self._api_url}/api/public/runs"
+            headers = self._get_headers()
+            payload = {
+                "name": self._name,
+                "heartbeatInterval": self._heartbeat_interval,
+                "source": self._source,
+                "labels": self._labels or {},
+                "config": config,
+                "tags": self._tags or [],
+                # Reproducibility tracking (API requires strings, not null)
+                "configHash": config_hash,
+                "gitSha": git_info.get("gitSha") or "",
+                "gitBranch": git_info.get("gitBranch") or "",
+                "gitDirty": git_info.get("gitDirty", False),
+                # SDK metadata
+                "sdkVersion": self._get_sdk_version(),
+                "clientId": str(uuid.uuid4()),
+                "hostname": self._get_hostname(),
+            }
+
+            logger.debug(f"Creating run at {url}")
+            response = requests.post(url, json=payload, headers=headers, timeout=30)
+
+            if not response.ok:
+                try:
+                    error_body = response.json()
+                except Exception:
+                    error_body = response.text
+                raise GatewayRunError(f"Run create failed with HTTP status {response.status_code}: {error_body}")
+
+            try:
+                data = response.json()
+            except Exception as error:
+                raise GatewayRunError(f"Run create failed: response JSON could not be decoded ({error})") from error
+
+            if not isinstance(data, dict):
+                raise GatewayRunError("Run create failed: response must be a JSON object")
+
+            run_id = data.get("runId")
+            if not isinstance(run_id, str) or not run_id:
+                raise GatewayRunError("Run create failed: response runId must be a non-empty string")
+
+            run_number = data.get("runNumber")
+            if isinstance(run_number, bool) or not isinstance(run_number, int) or run_number < 0:
+                raise GatewayRunError("Run create failed: response runNumber must be a non-negative integer")
+
+            logger.info(f"Created run {run_id} (#{run_number}) on backend")
+            return run_id, run_number
+
+        except GatewayRunError:
+            raise
+        except requests.exceptions.ConnectionError as error:
+            raise GatewayRunError(f"Run create failed: connection to {self._api_url} failed ({error})") from error
+        except requests.exceptions.Timeout as error:
+            raise GatewayRunError(f"Run create failed: request to {self._api_url} timed out ({error})") from error
+        except Exception as error:
+            raise GatewayRunError(f"Run create failed unexpectedly: {type(error).__name__}: {error}") from error
+
+    def _get_headers(self) -> Dict[str, str]:
+        """Get HTTP headers for API requests."""
+        import base64
+
+        headers = {"Content-Type": "application/json"}
+        if self._api_key and self._public_key:
+            # Use Basic auth with public_key:secret_key
+            credentials = base64.b64encode(f"{self._public_key}:{self._api_key}".encode()).decode()
+            headers["Authorization"] = f"Basic {credentials}"
+        elif self._api_key:
+            # Fallback: try reading public key from env (in case it was set after __init__)
+            public_key = os.environ.get("GATEWAY_PUBLIC_KEY", "")
+            if public_key:
+                credentials = base64.b64encode(f"{public_key}:{self._api_key}".encode()).decode()
+                headers["Authorization"] = f"Basic {credentials}"
+            else:
+                logger.warning("GATEWAY_PUBLIC_KEY not set - API requests will fail auth")
+        return headers
+
+    def _get_sdk_version(self) -> str:
+        """Get the gatewaysdk version."""
+        try:
+            import gatewaysdk
+
+            return getattr(gatewaysdk, "__version__", "unknown")
+        except Exception:
+            return "unknown"
+
+    def _get_hostname(self) -> str:
+        """Get the hostname."""
+        import socket
+
+        try:
+            return socket.gethostname()
+        except Exception:
+            return "unknown"
+
+    # =========================================================================
+    # Heartbeat
+    # =========================================================================
+
+    def _start_heartbeat_loop(self):
+        """Start background heartbeat thread."""
+        self._stop_heartbeat.clear()
+
+        def heartbeat_loop():
+            while not self._stop_heartbeat.wait(timeout=self._heartbeat_interval):
+                try:
+                    self._send_heartbeat()
+                    # Also flush metrics periodically
+                    self._maybe_flush_metrics()
+                except Exception as e:
+                    logger.warning(f"Heartbeat failed: {e}")
+
+        self._heartbeat_thread = threading.Thread(
+            target=heartbeat_loop,
+            name=f"gateway-heartbeat-{self._run_id}",
+            daemon=True,
+        )
+        self._heartbeat_thread.start()
+
+    def _send_heartbeat(self):
+        """Send heartbeat to backend and check for remote cancellation."""
+        if not self._run_id or self._run_id.startswith("local-"):
+            return
+
+        try:
+            response = requests.post(
+                f"{self._api_url}/api/public/runs/{self._run_id}/heartbeat",
+                headers=self._get_headers(),
+                timeout=10,
+            )
+            response.raise_for_status()
+
+            # Check if the run was cancelled via the UI.
+            try:
+                data = response.json()
+                if data.get("state") == "cancelled":
+                    logger.info(f"Run {self._run_id} cancelled via UI")
+                    self._fire_cancel_callback()
+                    # Force-exit the entire process.  The UI already set
+                    # state=cancelled in Postgres so no finish request needed.
+                    # os._exit() skips Python cleanup but that's intentional --
+                    # we need to kill forked subprocesses (shm strategy) and
+                    # can't use SIGINT from the heartbeat thread (deadlock).
+                    import os as _os
+
+                    _os._exit(130)
+            except (ValueError, KeyError):
+                pass  # Response wasn't JSON or missing fields -- ignore
+        except Exception as e:
+            logger.warning(f"Heartbeat request failed: {e}")
+
+    # =========================================================================
+    # Remote cancellation
+    # =========================================================================
+
+    def set_on_cancel(self, callback: Callable[[], None]) -> None:
+        """Register a callback invoked when the run is cancelled remotely.
+
+        Typically set to ``algorithm.request_cancel`` so that a "Cancel Run"
+        click in the UI propagates to the running algorithm.
+        """
+        self._on_cancel = callback
+
+    def _fire_cancel_callback(self) -> None:
+        """Invoke the on_cancel callback if registered."""
+        cb = self._on_cancel
+        if cb is not None:
+            try:
+                cb()
+            except Exception as e:
+                logger.warning(f"on_cancel callback failed: {e}")
+
+    # =========================================================================
+    # Flush Ticker (for live metric updates)
+    # =========================================================================
+
+    def _start_flush_ticker(self):
+        """Start background flush ticker thread for live metric updates."""
+        self._stop_flush_ticker.clear()
+
+        def flush_ticker_loop():
+            while not self._stop_flush_ticker.wait(timeout=self._flush_interval):
+                try:
+                    self._maybe_flush_metrics()
+                except Exception as e:
+                    logger.debug(f"Flush ticker failed: {e}")
+
+        self._flush_ticker_thread = threading.Thread(
+            target=flush_ticker_loop,
+            name=f"gateway-flush-ticker-{self._run_id}",
+            daemon=True,
+        )
+        self._flush_ticker_thread.start()
+
+    # =========================================================================
+    # Signal Handling
+    # =========================================================================
+
+    def _register_signal_handlers(self):
+        """Register handlers for SIGINT/SIGTERM to cleanup on Ctrl+C."""
+
+        def signal_handler(signum, frame):
+            logger.info(f"Received signal {signum}, finishing run as cancelled")
+            self.finish(state="cancelled")
+
+            # Call original handler
+            if signum == signal.SIGINT and callable(self._original_sigint):
+                self._original_sigint(signum, frame)
+            elif signum == signal.SIGTERM and callable(self._original_sigterm):
+                self._original_sigterm(signum, frame)
+
+            sys.exit(128 + signum)
+
+        try:
+            self._original_sigint = signal.signal(signal.SIGINT, signal_handler)
+            self._original_sigterm = signal.signal(signal.SIGTERM, signal_handler)
+        except Exception as e:
+            # Signal handling may not work in all environments (e.g., threads)
+            logger.debug(f"Could not register signal handlers: {e}")
+
+    def _restore_signal_handlers(self):
+        """Restore original signal handlers."""
+        try:
+            if self._original_sigint is not None:
+                signal.signal(signal.SIGINT, self._original_sigint)
+            if self._original_sigterm is not None:
+                signal.signal(signal.SIGTERM, self._original_sigterm)
+        except Exception as e:
+            logger.debug(f"Could not restore signal handlers: {e}")
+
+    def _atexit_handler(self):
+        """Called on Python exit - mark as failed if not finished."""
+        if not self._finished:
+            logger.warning("Run not finished, marking as failed")
+            try:
+                self.finish(state="failed")
+            except Exception:
+                pass  # Best effort
+
+    # =========================================================================
+    # Logging Methods
+    # =========================================================================
+
+    def log(
+        self,
+        metrics: Union[Dict[str, Any], str],
+        value: Optional[float] = None,
+        *,
+        step: Optional[int] = None,
+        labels: Optional[Dict[str, str]] = None,
+        unit: str = "",
+        metric_type: str = "gauge",
+    ):
+        """Log metrics for this run.
+
+        Can be called with a dict of metrics::
+
+            run.log({"loss": 0.5, "accuracy": 0.92}, step=100)
+
+        Or with a single metric name and value::
+
+            run.log("loss", 0.5, step=100, unit="", metric_type="gauge")
+
+        Args:
+            metrics: Dict of {metric_name: value} or single metric name.
+            value: Metric value (only when metrics is a string).
+            step: Training step/iteration number.
+            labels: Per-metric labels for filtering.
+            unit: Unit of measurement (e.g., "ms", "tokens", "%").
+            metric_type: Type hint for visualization ("gauge", "counter", "histogram").
+        """
+        with self._lock:
+            if self._finished:
+                logger.warning("Cannot log to finished run")
+                return
+
+        if isinstance(metrics, str):
+            # Single metric
+            if value is None:
+                raise ValueError("value is required when logging a single metric")
+            metric_list = [
+                {
+                    "name": metrics,
+                    "value": float(value),
+                    "step": step,
+                    "labels": labels or {},
+                    "unit": unit,
+                    "metricType": metric_type,
+                }
+            ]
+        else:
+            # Dict of metrics
+            metric_list = [
+                {
+                    "name": name,
+                    "value": float(val),
+                    "step": step,
+                    "labels": labels or {},
+                    "unit": unit,
+                    "metricType": metric_type,
+                }
+                for name, val in metrics.items()
+            ]
+
+        # Admission and buffering are one lifecycle transaction. A concurrent
+        # finish either observes this complete batch or rejects it above.
+        with self._lock:
+            if self._finished:
+                logger.warning("Cannot log to finished run")
+                return
+            with self._metric_buffer_lock:
+                self._metric_buffer.extend(metric_list)
+                should_flush = len(self._metric_buffer) >= 100
+
+        # Flush if buffer is large enough
+        if should_flush:
+            self._flush_metrics()
+
+    def log_metrics(self, metrics: List[Metric]):
+        """Log a batch of Metric objects.
+
+        Args:
+            metrics: List of Metric objects with full metadata.
+        """
+        for m in metrics:
+            self.log(
+                m.name,
+                m.value,
+                step=m.step,
+                labels=m.labels,
+                unit=m.unit,
+                metric_type=m.metric_type,
+            )
+
+    def _maybe_flush_metrics(self):
+        """Flush metrics if enough time has passed."""
+        with self._lock:
+            if self._finished:
+                return
+        if time.time() - self._last_flush_time > self._flush_interval:
+            self._flush_metrics()
+
+    def _flush_metrics(self) -> bool:
+        """Send one buffered metric batch and report whether it was delivered."""
+        # Serialize the complete take/send/requeue operation. Buffer locking
+        # alone cannot keep a periodic flush from racing the final drain.
+        with self._metric_flush_lock:
+            with self._metric_buffer_lock:
+                if not self._metric_buffer:
+                    return True
+                metrics_to_send = self._metric_buffer.copy()
+                self._metric_buffer.clear()
+                self._last_flush_time = time.time()
+
+            if not self._run_id or self._run_id.startswith("local-"):
+                logger.debug(f"Skipping metric flush (local run): {len(metrics_to_send)} metrics")
+                return True
+
+            try:
+                response = requests.post(
+                    f"{self._api_url}/api/public/runs/{self._run_id}/metrics",
+                    json={
+                        "metrics": metrics_to_send,
+                    },
+                    headers=self._get_headers(),
+                    timeout=30,
+                )
+                response.raise_for_status()
+                logger.debug(f"Flushed {len(metrics_to_send)} metrics")
+                return True
+            except Exception as error:
+                if not getattr(self, "_metric_flush_warned", False):
+                    logger.warning("Failed to flush metrics (suppressing further warnings): " f"{error}")
+                    self._metric_flush_warned = True
+
+                # Requeueing can recover a transient error. Only permanently
+                # dropped accepted metrics become a sticky delivery failure.
+                with self._metric_buffer_lock:
+                    combined = metrics_to_send + self._metric_buffer
+                    if len(combined) > MAX_METRIC_BUFFER_SIZE:
+                        dropped = len(combined) - MAX_METRIC_BUFFER_SIZE
+                        self._dropped_metric_count += dropped
+                        logger.warning(
+                            f"Dropping {dropped} oldest metrics to prevent buffer overflow "
+                            f"(max: {MAX_METRIC_BUFFER_SIZE})"
+                        )
+                        combined = combined[-MAX_METRIC_BUFFER_SIZE:]
+                    self._metric_buffer = combined
+                return False
+
+    def flush(self) -> None:
+        """Manually flush buffered metrics to the backend.
+
+        Call this when you want immediate visibility of logged metrics
+        without waiting for the automatic flush interval.
+
+        Example:
+            run.log({"accuracy": 0.95})
+            run.flush()  # Metrics visible immediately in UI
+        """
+        self._flush_metrics()
+
+    # =========================================================================
+    # Live State (Tables-like feature for rich structured data)
+    # =========================================================================
+
+    def log_state(self, state: Dict[str, Any]) -> None:
+        """Log rich structured state for live visualization.
+
+        This is similar to W&B Tables - allows streaming structured data
+        (candidates, trees, pareto fronts) that can be visualized in real-time.
+
+        Unlike log() which is for numeric metrics, log_state() is for
+        complex objects like:
+        - Candidate prompts with full text
+        - Evolution tree structure (parent-child relationships)
+        - Pareto front membership
+        - Per-instance scores
+
+        Args:
+            state: Dictionary containing structured state data.
+                   Common keys for GEPA:
+                   - candidates: List of candidate dicts with idx, prompt, score, parent
+                   - parents: Parent indices for tree structure
+                   - scores: Validation scores per candidate
+                   - pareto_front: Dict mapping instance to best candidate indices
+                   - best_idx: Index of best overall candidate
+                   - generation: Current generation number
+
+        Example:
+            run.log_state({
+                "candidates": [
+                    {"idx": 0, "prompt": "You are...", "score": 0.72, "parent": None},
+                    {"idx": 1, "prompt": "As an expert...", "score": 0.78, "parent": 0},
+                ],
+                "pareto_front": {"instance_1": [1], "instance_2": [0, 1]},
+                "best_idx": 1,
+                "generation": 3,
+            })
+        """
+        if not self._run_id or self._run_id.startswith("local-"):
+            logger.info(f"[log_state] Skipping (local run): {self._run_id}")
+            return
+
+        logger.info(f"[log_state] Sending to {self._api_url}/api/public/runs/{self._run_id}/state")
+        try:
+            response = requests.patch(
+                f"{self._api_url}/api/public/runs/{self._run_id}/state",
+                json={"state": state},
+                headers=self._get_headers(),
+                timeout=30,
+            )
+            response.raise_for_status()
+            logger.info(f"[log_state] Success: {response.status_code}")
+        except Exception as e:
+            logger.warning(f"[log_state] Failed: {e}")
+
+    # =========================================================================
+    # Scores (Gateway Scores linked to traces)
+    # =========================================================================
+
+    def score(
+        self,
+        name: str,
+        value: Union[float, str],
+        *,
+        trace_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        observation_id: Optional[str] = None,
+        data_type: str = "NUMERIC",
+        comment: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Post a score to Gateway, linked to a trace.
+
+        Scores are evaluation results attached to traces (visible on the Scores
+        page and in trace detail views). Unlike run.log() which logs run-level
+        metrics, score() creates Gateway Score entities linked to specific traces.
+
+        Auto-detection:
+            - If trace_id is not provided, attempts to read it from the current
+              OTel span context.
+            - If session_id is not provided, reads it from the active Rollout.
+
+        Args:
+            name: Score name (e.g., "rubric", "accuracy", "task_completion").
+            value: Score value. Numeric for NUMERIC type, string for CATEGORICAL.
+            trace_id: Trace ID to link the score to. Auto-detected from OTel
+                context if not provided.
+            session_id: Session ID to link the score to. Auto-detected from
+                the active Rollout if not provided.
+            observation_id: Optional observation/span ID for span-level scores.
+            data_type: Score data type: "NUMERIC", "CATEGORICAL", or "BOOLEAN".
+            comment: Optional comment explaining the score.
+            metadata: Optional metadata dict.
+
+        Example::
+
+            with msk.run("experiment") as r:
+                with r.rollout("task-1") as rollout:
+                    result = await run_agent(task)
+                    grade = evaluate(result)
+                    # Score linked to trace + session automatically
+                    r.score("rubric", grade)
+
+                    # Or with explicit IDs
+                    r.score("accuracy", 0.95, trace_id="abc123")
+        """
+        with self._lock:
+            if self._finished:
+                logger.warning("Cannot score on a finished run")
+                return
+
+        # Auto-detect session_id from active Rollout
+        if session_id is None:
+            current_rollout = get_current_rollout()
+            if current_rollout is not None:
+                session_id = current_rollout.session_id
+
+        # Session scores are valid first-class scores and intentionally carry
+        # no trace/observation target. Only consult ambient trace context when
+        # no explicit or active Rollout session is available.
+        if not session_id and trace_id is None:
+            trace_id = self._get_current_trace_id()
+
+        if not session_id and not trace_id:
+            logger.warning(
+                "Cannot post score: no session_id or trace_id is available. "
+                "Call score() inside a Rollout/traced block or pass a target explicitly."
+            )
+            return
+
+        # Build score payload (matches Gateway POST /api/public/scores schema).
+        # The API requires EXACTLY ONE of: traceId, sessionId, or datasetRunId.
+        # Prefer sessionId (scores typically grade a whole session/rollout);
+        # fall back to traceId when no session context exists.
+        score_payload: Dict[str, Any] = {
+            "name": name,
+            "value": value,
+            "dataType": data_type,
+        }
+
+        if session_id:
+            score_payload["sessionId"] = session_id
+        elif trace_id:
+            score_payload["traceId"] = trace_id
+            if observation_id:
+                score_payload["observationId"] = observation_id
+        else:
+            logger.warning(f"Cannot post score '{name}': no sessionId or traceId available.")
+            return
+        if comment:
+            score_payload["comment"] = comment
+
+        # Merge metadata with run context
+        score_metadata: Dict[str, Any] = {
+            "source": "gatewaysdk",
+            "run_id": self._run_id,
+            "run_name": self._name,
+        }
+        if metadata:
+            score_metadata.update(metadata)
+        score_payload["metadata"] = score_metadata
+
+        # Send in a background thread (non-blocking)
+        self._send_score_async(score_payload)
+
+    def success(
+        self,
+        name: str,
+        value: Union[bool, float] = True,
+        *,
+        agent_id: Optional[str] = None,
+        metric_name: Optional[str] = None,
+        trace_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        observation_id: Optional[str] = None,
+        data_type: Optional[str] = None,
+        comment: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Emit a success signal that feeds a Gateway Success Metric.
+
+        A thin convenience wrapper over :meth:`score`. A boolean ``value`` is
+        recorded as a 1/0 BOOLEAN score (a pass/fail signal); a numeric value is
+        recorded as a NUMERIC score. The signal is tagged with the agent
+        identity (``gateway.agent.id``) and the metric name so the platform can
+        roll it up into that agent's Success Metric.
+
+        Args:
+            name: Signal/score name. Match it to the Success Metric (or a linked
+                eval's score name) you want it to count toward.
+            value: ``True``/``False`` for a pass/fail signal, or a number.
+            agent_id: The agent this signal belongs to (``gateway.agent.id``).
+                Defaults to the trace's agent when omitted.
+            metric_name: The Success Metric name; defaults to ``name``.
+            trace_id, session_id, observation_id: Linkage; auto-detected like
+                :meth:`score` when omitted.
+            data_type: Override the inferred "BOOLEAN"/"NUMERIC".
+            comment: Optional comment.
+            metadata: Extra metadata merged into the signal.
+
+        Example::
+
+            with gatewaysdk.run("prod") as r:
+                r.success("resolved_ticket", True, agent_id="support-triage")
+                r.success("handle_time_s", 42.0, agent_id="support-triage")
+        """
+        if isinstance(value, bool):
+            numeric_value: float = 1.0 if value else 0.0
+            inferred_type = "BOOLEAN"
+        else:
+            numeric_value = float(value)
+            inferred_type = "NUMERIC"
+
+        signal_metadata: Dict[str, Any] = {
+            "gateway.success.signal": True,
+            "gateway.success.metricName": metric_name or name,
+        }
+        if agent_id:
+            signal_metadata["gateway.agent.id"] = agent_id
+        if metadata:
+            signal_metadata.update(metadata)
+
+        self.score(
+            name,
+            numeric_value,
+            trace_id=trace_id,
+            session_id=session_id,
+            observation_id=observation_id,
+            data_type=data_type or inferred_type,
+            comment=comment,
+            metadata=signal_metadata,
+        )
+
+    def _get_current_trace_id(self) -> Optional[str]:
+        """Get the trace ID from the current OTel span context.
+
+        Returns the hex trace ID if there is an active span, None otherwise.
+        """
+        try:
+            from opentelemetry import trace as otel_trace
+
+            current_span = otel_trace.get_current_span()
+            span_context = current_span.get_span_context()
+            if span_context and span_context.is_valid:
+                return format(span_context.trace_id, "032x")
+        except Exception as e:
+            logger.debug(f"Could not get trace_id from OTel context: {e}")
+        return None
+
+    def _send_score_async(self, score_payload: Dict[str, Any]) -> None:
+        """Send a score to Gateway in a background thread.
+
+        Threads are tracked and joined in finish() to ensure delivery.
+        """
+        if not self._run_id or self._run_id.startswith("local-"):
+            logger.debug(f"Skipping score (local run): {score_payload.get('name')}")
+            return
+
+        headers = self._get_headers()
+        url = f"{self._api_url}/api/public/scores"
+
+        def _send():
+            try:
+                response = requests.post(
+                    url,
+                    json=score_payload,
+                    headers=headers,
+                    timeout=30,
+                )
+                if not response.ok:
+                    try:
+                        error_body = response.json()
+                    except Exception:
+                        error_body = response.text
+                    logger.warning(
+                        f"Failed to post score '{score_payload['name']}': "
+                        f"HTTP {response.status_code} - {error_body}"
+                    )
+                    self._record_score_delivery_failure(
+                        f"score '{score_payload['name']}' was rejected with " f"HTTP {response.status_code}"
+                    )
+                    return
+                logger.info(
+                    f"Score posted: {score_payload['name']}={score_payload['value']} "
+                    f"(trace={score_payload.get('traceId', 'N/A')[:12]}...)"
+                )
+            except Exception as error:
+                logger.warning(f"Failed to post score: {error}")
+                self._record_score_delivery_failure(f"score '{score_payload['name']}' delivery raised {error}")
+
+        thread = threading.Thread(
+            target=_send,
+            name=f"gateway-score-{score_payload.get('name', 'unknown')}",
+            daemon=True,
+        )
+
+        # Register and start under the lifecycle gate so finish cannot take a
+        # worker snapshot between acceptance and thread startup.
+        with self._lock:
+            if self._finished:
+                logger.warning("Cannot score on a finished run")
+                return
+            self._score_threads.append(thread)
+            try:
+                thread.start()
+            except Exception as error:
+                self._score_threads.remove(thread)
+                failure = f"score '{score_payload['name']}' worker could not start: {error}"
+                self._score_delivery_failures.append(failure)
+                logger.warning(failure)
+
+    def _record_score_delivery_failure(self, failure: str) -> None:
+        """Retain a terminal score outcome until finish decides final state."""
+        with self._lock:
+            self._score_delivery_failures.append(failure)
+
+    # =========================================================================
+    # Finish
+    # =========================================================================
+
+    def finish(self, state: RunState = "completed") -> None:
+        """Mark the run as finished.
+
+        Args:
+            state: Final state ("completed", "failed", "cancelled").
+        """
+        with self._lock:
+            if self._finished:
+                return
+            # This is the admission gate for both metrics and scores. Work
+            # registered before this point is included in the snapshots below;
+            # later calls retain the existing warning/no-op behavior.
+            self._finished = True
+            score_threads = list(self._score_threads)
+
+        delivery_failures: List[str] = []
+
+        # Stop periodic producers before the final drain. A flush already in
+        # progress is serialized with the final attempt by _metric_flush_lock.
+        self._stop_heartbeat.set()
+        self._stop_flush_ticker.set()
+
+        if self._heartbeat_thread:
+            self._heartbeat_thread.join(timeout=5)
+        if self._flush_ticker_thread:
+            self._flush_ticker_thread.join(timeout=2)
+
+        # Flush remaining metrics
+        try:
+            metrics_delivered = self._flush_metrics()
+        except Exception as error:
+            metrics_delivered = False
+            logger.warning(f"Final metric delivery raised unexpectedly: {error}")
+
+        if not metrics_delivered:
+            delivery_failures.append("metric delivery failed during the final flush")
+        if self._dropped_metric_count:
+            delivery_failures.append(
+                f"metric delivery permanently dropped {self._dropped_metric_count} " "accepted metrics"
+            )
+
+        # Wait for pending score submissions to complete
+        for thread in score_threads:
+            try:
+                thread.join(timeout=10)
+                if thread.is_alive():
+                    delivery_failures.append(
+                        f"score delivery worker {getattr(thread, 'name', '<unknown>')} "
+                        "did not finish before the join deadline"
+                    )
+            except Exception as error:
+                delivery_failures.append(f"score delivery worker could not be joined: {error}")
+
+        with self._lock:
+            self._score_threads.clear()
+            delivery_failures.extend(self._score_delivery_failures)
+
+        requested_success = state == "completed"
+        effective_state: RunState = "failed" if requested_success and delivery_failures else state
+        self._state = effective_state
+
+        self._last_finish_error = None
+        self._last_finish_exception = None
+        try:
+            finish_confirmed = self._send_finish(effective_state)
+        except Exception as error:
+            self._last_finish_error = str(error)
+            self._last_finish_exception = error
+            finish_confirmed = False
+            logger.warning(f"Failed to send finish: {error}")
+
+        completion_error: Optional[GatewayRunError] = None
+        if requested_success and delivery_failures:
+            reason = "; ".join(delivery_failures)
+            if not finish_confirmed:
+                reason = (
+                    f"{reason}; failed state was not confirmed: " f"{self._last_finish_error or 'unknown finish error'}"
+                )
+            completion_error = GatewayRunError(f"Run finish delivery failed: {reason}")
+        elif requested_success and not finish_confirmed:
+            self._state = "failed"
+            completion_error = GatewayRunError(
+                "Run finish could not be confirmed: " f"{self._last_finish_error or 'unknown finish error'}"
+            )
+
+        # Restore signal handlers
+        self._restore_signal_handlers()
+
+        # Unregister atexit
+        try:
+            atexit.unregister(self._atexit_handler)
+        except Exception:
+            pass
+
+        # Do not erase a different outer/sentinel Run context.
+        if get_current_run() is self:
+            _set_current_run(None)
+
+        if completion_error is not None:
+            logger.warning(f"Gateway Run finish failed: {self._name} ({self._state})")
+            if self._last_finish_exception is not None:
+                raise completion_error from self._last_finish_exception
+            raise completion_error
+
+        if delivery_failures or not finish_confirmed:
+            logger.warning("Gateway Run cleanup completed with delivery warnings: " f"{self._name} ({self._state})")
+        else:
+            logger.info(f"Gateway Run finished: {self._name} ({self._state})")
+
+    def _send_finish(self, state: RunState) -> bool:
+        """Send final state to the backend and report confirmation."""
+        self._last_finish_error = None
+        self._last_finish_exception = None
+        if not self._run_id or self._run_id.startswith("local-"):
+            return True
+
+        try:
+            response = requests.post(
+                f"{self._api_url}/api/public/runs/{self._run_id}/finish",
+                json={
+                    "state": state,
+                },
+                headers=self._get_headers(),
+                timeout=10,
+            )
+            response.raise_for_status()
+            return True
+        except Exception as error:
+            self._last_finish_error = str(error)
+            self._last_finish_exception = error
+            logger.warning(f"Failed to send finish: {error}")
+            return False
+
+    # =========================================================================
+    # Context Manager
+    # =========================================================================
+
+    def __enter__(self) -> "Run":
+        """Context manager entry - start the run."""
+        return self._start()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - finish with appropriate state."""
+        if exc_type is not None:
+            # Exception occurred - mark as failed
+            logger.error(f"Run failed with exception: {exc_val}")
+            self.finish(state="failed")
+        elif not self._finished:
+            # No exception, not manually finished - complete
+            self.finish(state="completed")
+
+        return False  # Don't suppress exceptions
+
+
+# =============================================================================
+# Convenience Functions
+# =============================================================================
+
+
+def init_run(
+    name: Optional[str] = None,
+    *,
+    project_id: Optional[str] = None,
+    heartbeat_interval: int = 30,
+    source: str = "sdk",
+    labels: Optional[Dict[str, Any]] = None,
+    config: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    **kwargs,
+) -> Run:
+    """Initialize and start a new Gateway Run.
+
+    This is the main entry point for creating runs. For the safest usage,
+    prefer the context manager form::
+
+        with gatewaysdk.run("my-experiment") as run:
+            run.log({"metric": value})
+
+    Args:
+        name: Human-readable name for the run.
+        project_id: Project ID (from env GATEWAY_PROJECT_ID if not provided).
+        heartbeat_interval: Seconds between heartbeats (default: 30).
+        source: Source identifier (e.g., "gepa", "trainer").
+        labels: Arbitrary key-value labels for filtering.
+        config: Hyperparameters/config dict.
+        tags: Simple string tags for categorization.
+
+    Returns:
+        An initialized and started Run instance.
+
+    Note:
+        Use run.rollout() to create individual episodes/rollouts within a run.
+        Each rollout gets its own session_id for trace grouping.
+    """
+    run = Run(
+        name=name,
+        project_id=project_id,
+        heartbeat_interval=heartbeat_interval,
+        source=source,
+        labels=labels,
+        config=config,
+        tags=tags,
+        **kwargs,
+    )
+    return run._start()
+
+
+@contextmanager
+def run(
+    name: Optional[str] = None,
+    **kwargs,
+):
+    """Context manager for a Gateway Run.
+
+    This is the recommended way to use runs as it automatically handles
+    crashes, signals, and exceptions::
+
+        with gatewaysdk.run("my-experiment", config={"lr": 0.01}) as r:
+            for epoch in range(100):
+                loss = train()
+                r.log({"loss": loss}, step=epoch)
+
+    Args:
+        name: Human-readable name for the run.
+        **kwargs: Additional arguments passed to Run constructor.
+
+    Yields:
+        A Run instance.
+    """
+    r = Run(name=name, **kwargs)
+    with r:
+        yield r
+
+
+def log(metrics: Dict[str, Any], step: Optional[int] = None, **kwargs):
+    """Log metrics to the current run.
+
+    Convenience function for logging to the currently active run.
+
+    Args:
+        metrics: Dict of {metric_name: value}.
+        step: Training step/iteration number.
+        **kwargs: Additional arguments passed to Run.log().
+
+    Raises:
+        RuntimeError: If no run is currently active.
+    """
+    current = get_current_run()
+    if current is None:
+        raise RuntimeError("No active run. Call gatewaysdk.init_run() or use 'with gatewaysdk.run(...)'")
+    current.log(metrics, step=step, **kwargs)
+
+
+def score(
+    name: str,
+    value: Union[float, str],
+    **kwargs,
+) -> None:
+    """Post a score to Gateway from the current run.
+
+    Convenience function for posting scores to the currently active run.
+    Auto-detects trace_id from OTel context and session_id from the active
+    Rollout.
+
+    Args:
+        name: Score name (e.g., "rubric", "accuracy").
+        value: Score value.
+        **kwargs: Additional arguments passed to Run.score().
+
+    Raises:
+        RuntimeError: If no run is currently active.
+    """
+    current = get_current_run()
+    if current is None:
+        raise RuntimeError("No active run. Call gatewaysdk.init_run() or use 'with gatewaysdk.run(...)'")
+    current.score(name, value, **kwargs)
+
+
+def success(
+    name: str,
+    value: Union[bool, float] = True,
+    **kwargs,
+) -> None:
+    """Emit a success signal from the current run.
+
+    Convenience function for the currently active run; see :meth:`Run.success`.
+
+    Raises:
+        RuntimeError: If no run is currently active.
+    """
+    current = get_current_run()
+    if current is None:
+        raise RuntimeError("No active run. Call gatewaysdk.init_run() or use 'with gatewaysdk.run(...)'")
+    current.success(name, value, **kwargs)
+
+
+def finish(state: RunState = "completed"):
+    """Finish the current run.
+
+    Convenience function for finishing the currently active run.
+
+    Args:
+        state: Final state ("completed", "failed", "cancelled").
+    """
+    current = get_current_run()
+    if current is None:
+        return
+    current.finish(state=state)
+
+
+# =============================================================================
+# Decorator-based API
+# =============================================================================
+
+
+def experiment(
+    name: Optional[str] = None,
+    *,
+    inject_run: bool = True,
+    **run_kwargs,
+) -> Callable[[F], F]:
+    """Decorator to wrap a function in a Gateway Run.
+
+    This provides an alternative to context managers for experiment tracking.
+    The run is automatically started before the function executes and finished
+    after it completes (or fails).
+
+    Example::
+
+        @gatewaysdk.experiment("gepa-v1", config={"lr": 0.01})
+        def train(run):
+            run.log({"started": True})
+            for candidate in population:
+                evaluate(candidate)
+
+        @gatewaysdk.experiment(inject_run=False)
+        def simple_experiment():
+            # Access run via get_current_run()
+            gatewaysdk.log({"metric": value})
+
+    Args:
+        name: Human-readable name for the run. Defaults to function name.
+        inject_run: If True (default), passes the Run instance as the first
+            argument to the decorated function. If False, access via
+            get_current_run().
+        **run_kwargs: Additional arguments passed to Run constructor
+            (config, labels, tags, source, etc).
+
+    Returns:
+        A decorator function.
+    """
+
+    def decorator(fn: F) -> F:
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            run_name = name or fn.__name__
+            r = Run(name=run_name, **run_kwargs)
+            with r:
+                if inject_run:
+                    return fn(r, *args, **kwargs)
+                return fn(*args, **kwargs)
+
+        @functools.wraps(fn)
+        async def async_wrapper(*args, **kwargs):
+            run_name = name or fn.__name__
+            r = Run(name=run_name, **run_kwargs)
+            with r:
+                if inject_run:
+                    return await fn(r, *args, **kwargs)
+                return await fn(*args, **kwargs)
+
+        if asyncio.iscoroutinefunction(fn):
+            return async_wrapper  # type: ignore
+        return wrapper  # type: ignore
+
+    return decorator
+
+
+def episode(
+    name: Optional[str] = None,
+    *,
+    name_fn: Optional[Callable[..., str]] = None,
+    inject_rollout: bool = False,
+) -> Callable[[F], F]:
+    """Decorator to wrap a function in a Rollout (episode).
+
+    This provides an alternative to context managers for episode tracking.
+    The rollout is automatically started before the function executes and
+    ended after it completes.
+
+    Requires an active run (created via @experiment decorator or context manager).
+
+    Example::
+
+        @gatewaysdk.experiment("gepa", config={"lr": 0.01})
+        def train(run):
+            for candidate in population:
+                evaluate(candidate)
+
+        @gatewaysdk.episode()
+        def evaluate(candidate):
+            result = agent.execute(task)
+            # Traces automatically get run_id + session_id
+
+        # With dynamic name based on arguments
+        @gatewaysdk.episode(name_fn=lambda c: f"candidate-{c.id}")
+        def evaluate_named(candidate):
+            result = agent.execute(task)
+
+    Args:
+        name: Static name for the rollout. Defaults to function name.
+        name_fn: Optional function that takes the same arguments as the
+            decorated function and returns a rollout name.
+        inject_rollout: If True, passes the Rollout instance as the first
+            argument to the decorated function. Default False.
+
+    Returns:
+        A decorator function.
+
+    Raises:
+        RuntimeError: If no active run exists when the function is called.
+    """
+
+    def decorator(fn: F) -> F:
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            current_run = get_current_run()
+            if current_run is None:
+                raise RuntimeError(
+                    "@episode requires an active run. " "Use @experiment decorator or 'with gatewaysdk.run(...)' first."
+                )
+
+            # Determine rollout name
+            if name_fn is not None:
+                rollout_name = name_fn(*args, **kwargs)
+            else:
+                rollout_name = name or fn.__name__
+
+            rollout = current_run.rollout(rollout_name)
+            with rollout:
+                if inject_rollout:
+                    return fn(rollout, *args, **kwargs)
+                return fn(*args, **kwargs)
+
+        @functools.wraps(fn)
+        async def async_wrapper(*args, **kwargs):
+            current_run = get_current_run()
+            if current_run is None:
+                raise RuntimeError(
+                    "@episode requires an active run. " "Use @experiment decorator or 'with gatewaysdk.run(...)' first."
+                )
+
+            # Determine rollout name
+            if name_fn is not None:
+                rollout_name = name_fn(*args, **kwargs)
+            else:
+                rollout_name = name or fn.__name__
+
+            rollout = current_run.rollout(rollout_name)
+            with rollout:
+                if inject_rollout:
+                    return await fn(rollout, *args, **kwargs)
+                return await fn(*args, **kwargs)
+
+        if asyncio.iscoroutinefunction(fn):
+            return async_wrapper  # type: ignore
+        return wrapper  # type: ignore
+
+    return decorator