extty 0.1.0__py3-none-any.whl

extty/__init__.py

@@ -0,0 +1,897 @@
+"""extty: Terminal-native ML experiment tracker."""
+
+import atexit
+import copy
+import dataclasses
+import functools
+import os
+import warnings
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+from typing import Any, Callable, ParamSpec, TypeVar
+
+from extty._logger import log as logger
+from extty.arg_parser import run_experiments_parser
+from extty.artifact import (
+    ArtifactMeta,
+    artifact_exists,
+    delete_artifact,
+    get_artifact,
+    list_artifacts,
+    load_artifact,
+)
+from extty.artifact import (
+    save_artifact as _save_artifact_raw,
+)
+from extty.chart import Chart
+from extty.checkpoints import (
+    DELETABLE_LOCALLY,
+    CheckpointStatus,
+    LocalCopy,
+    checkpoint_status,
+    local_checkpoint_dir,
+    read_local_checkpoint,
+    read_local_copy,
+)
+from extty.compare import compare, config_diff, plot_metric, reduce_metric
+from extty.confusion import ConfusionMatrix
+from extty.example import BatchExample, Example
+from extty.image import Image
+from extty.query import RunData, get_run, get_runs
+from extty.run import NoOpRun, Run
+from extty.s3 import S3Config
+from extty.storage import (
+    ChartRecord,
+    Checkpoint,
+    CheckpointFile,
+    ConfusionMatrixRecord,
+    ExampleRecord,
+    ImageRecord,
+    MetricPoint,
+    SystemMetricPoint,
+    get_runs_dir,
+)
+
+_logger = logger
+
+__all__ = [
+    "init",
+    "log",
+    "save_checkpoint",
+    "load_checkpoint",
+    "load_checkpoint_from",
+    "delete_local_checkpoint",
+    "finish",
+    "Run",
+    "RunData",
+    "get_runs",
+    "get_run",
+    "push",
+    "list_local_runs",
+    "compare",
+    "config_diff",
+    "plot_metric",
+    "reduce_metric",
+    "Example",
+    "BatchExample",
+    "ConfusionMatrix",
+    "ConfusionMatrixRecord",
+    "Chart",
+    "ChartRecord",
+    "Image",
+    "ImageRecord",
+    "MetricPoint",
+    "SystemMetricPoint",
+    "ExampleRecord",
+    "Checkpoint",
+    "CheckpointFile",
+    "has_active_run",
+    "NoOpRun",
+    "ArtifactMeta",
+    "save_artifact",
+    "list_artifacts",
+    "load_artifact",
+    "get_artifact",
+    "artifact_exists",
+    "delete_artifact",
+    "logger",
+    "run_experiments_parser",
+]
+
+# Version is managed by setuptools_scm
+try:
+    __version__ = version("extty")
+except PackageNotFoundError:
+    # Package is not installed, use a default version
+    __version__ = "0.0.0+unknown"
+
+_active_run: Run | None = None
+_atexit_registered: bool = False
+
+
+def has_active_run() -> bool:
+    return _active_run is not None
+
+
+def _atexit_finish_active_run() -> None:
+    """Best-effort flush of any still-active run when the interpreter exits.
+
+    Daemon writer threads in :class:`extty.async_sink.AsyncSink` are killed
+    without draining at interpreter shutdown, so this handler protects users
+    who forget to call :func:`extty.finish`.
+    """
+    global _active_run
+    run = _active_run
+    if run is None:
+        return
+    try:
+        run.finish()
+    except Exception:
+        _logger.exception("extty atexit handler failed while finishing run")
+    finally:
+        _active_run = None
+
+
+def _is_main_process(rank: int | None) -> bool:
+    if rank is not None:
+        return rank == 0
+    env_rank = os.environ.get("RANK")
+    if env_rank is not None:
+        try:
+            return int(env_rank) == 0
+        except ValueError:
+            return True
+    return True
+
+
+def init(
+    project: str,
+    *,
+    name: str | None = None,
+    config: dict[str, Any] | None = None,
+    system_metrics: bool = True,
+    rank: int | None = None,
+) -> Run:
+    """
+    Initialize a new experiment run.
+
+    On non-main processes in distributed training (detected via the RANK
+    environment variable or the explicit ``rank`` parameter), returns a
+    :class:`NoOpRun` that silently discards all operations.
+
+    Parameters
+    ----------
+    project : str
+        Name of the project/experiment group.
+    name : str, optional
+        Name for this specific run. Auto-generated if not provided.
+    config : dict, optional
+        Hyperparameters and configuration to log.
+    system_metrics : bool, default True
+        Whether to automatically collect system metrics (RAM, GPU).
+    rank : int, optional
+        Explicit process rank. When provided, overrides the RANK environment
+        variable. Only rank 0 creates a real run.
+
+    Returns
+    -------
+    Run
+        The initialized run object (or :class:`NoOpRun` on non-main processes).
+    """
+    global _active_run, _atexit_registered
+    if _active_run is not None:
+        _active_run.finish()
+
+    if _is_main_process(rank):
+        if not _atexit_registered:
+            atexit.register(_atexit_finish_active_run)
+            _atexit_registered = True
+        _active_run = Run(
+            project,
+            name=name,
+            config=config,
+            system_metrics=system_metrics,
+        )
+        _logger.info(
+            "initialized run %s/%s, writing to %s",
+            project,
+            _active_run.name,
+            _active_run.run_dir,
+        )
+    else:
+        _active_run = NoOpRun(project, name=name, config=config)
+
+    return _active_run
+
+
+def log(metrics: dict[str, Any], *, step: int) -> None:
+    """
+    Log metrics or structured examples for the current step.
+
+    Parameters
+    ----------
+    metrics : dict[str, Any]
+        Dictionary of metric names to values. Values can be:
+        - float/int: logged as metric
+        - Example: single prompt with grouped responses
+        - BatchExample: batch of prompts with grouped responses
+        - ConfusionMatrix: N×N matrix with class labels
+        - Chart: a 2D chart of (x, y) points with axis names
+        - Image: a PIL image wrapped in extty.Image, stored as PNG
+    step : int
+        The current training step.
+    """
+    if _active_run is None:
+        raise RuntimeError("No active run. Call extty.init() first.")
+    _active_run.log(metrics, step=step)
+
+
+def save_checkpoint(
+    step: int,
+    *,
+    path: str | None = None,
+    state_dict: Any = None,
+    optimizer_state_dict: Any = None,
+    keep_local: bool = False,
+) -> None:
+    """
+    Save a checkpoint for the active run.
+
+    Saved to the run directory, or uploaded to S3 when configured; see
+    :meth:`extty.Run.save_checkpoint`.
+
+    Parameters
+    ----------
+    step : int
+        The training step for this checkpoint.
+    path : str or None
+        Path to an existing checkpoint file to save.
+    state_dict : Any or None
+        Model state dict to serialize with torch.save.
+    optimizer_state_dict : Any or None
+        Optimizer state dict to include when using state_dict.
+    keep_local : bool, default False
+        Keep the local copy after a successful S3 upload.
+    """
+    if _active_run is None:
+        raise RuntimeError("No active run. Call extty.init() first.")
+    _active_run.save_checkpoint(
+        step,
+        path=path,
+        state_dict=state_dict,
+        optimizer_state_dict=optimizer_state_dict,
+        keep_local=keep_local,
+    )
+
+
+def load_checkpoint(
+    step: int,
+    *,
+    load_optimizer: bool = True,
+) -> dict[str, Any]:
+    """
+    Load a checkpoint for the active run, downloading from S3 if needed.
+
+    Parameters
+    ----------
+    step : int
+        The training step to load.
+    load_optimizer : bool, default True
+        Whether to include the optimizer state in the result.
+
+    Returns
+    -------
+    dict[str, Any]
+        Contains ``"model_state_dict"`` and optionally
+        ``"optimizer_state_dict"``.
+    """
+    if _active_run is None:
+        raise RuntimeError("No active run. Call extty.init() first.")
+    return _active_run.load_checkpoint(step, load_optimizer=load_optimizer)
+
+
+def load_checkpoint_from(
+    project: str,
+    run_name: str,
+    step: int,
+    *,
+    load_optimizer: bool = True,
+    s3_config: S3Config | None = None,
+) -> dict[str, Any]:
+    """
+    Load a checkpoint from any run, downloading from S3 if needed.
+
+    A complete copy in the local runs directory is used without contacting
+    S3.
+
+    Parameters
+    ----------
+    project : str
+        The project name.
+    run_name : str
+        The run name.
+    step : int
+        The training step to load.
+    load_optimizer : bool, default True
+        Whether to include the optimizer state in the result.
+    s3_config : S3Config, optional
+        S3 configuration. Loaded from environment if not provided.
+
+    Returns
+    -------
+    dict[str, Any]
+        Contains ``"model_state_dict"`` and optionally
+        ``"optimizer_state_dict"``.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the checkpoint is not available locally and S3 is not configured
+        or does not have it.
+    """
+    from extty.s3 import S3Storage
+
+    local_dir = local_checkpoint_dir(project, run_name, step)
+    local = read_local_checkpoint(local_dir, load_optimizer=load_optimizer)
+    if local is not None:
+        return local
+    if s3_config is None:
+        s3_config = S3Config.load()
+    if s3_config is None:
+        raise FileNotFoundError(
+            f"Checkpoint {project}/{run_name} step {step} not found in {local_dir} "
+            "and S3 is not configured."
+        )
+    storage = S3Storage(s3_config, project, run_name)
+    return storage.load_checkpoint(step, load_optimizer=load_optimizer)
+
+
+def delete_local_checkpoint(
+    project: str,
+    run_name: str,
+    step: int,
+    *,
+    force: bool = False,
+    s3_config: S3Config | None = None,
+) -> bool:
+    """
+    Delete a checkpoint's local files, leaving any S3 copy intact.
+
+    The local copy may be the only one (no S3 configured, its upload failed,
+    or S3 has since been given a different save of the step), so unless
+    ``force`` is set this deletes only when S3 holds the same save as the
+    local copy. This is the inverse of :func:`load_checkpoint_from`, using
+    the same addressing.
+
+    Parameters
+    ----------
+    project : str
+        The project name.
+    run_name : str
+        The run name.
+    step : int
+        The training step whose local files should be removed.
+    force : bool, default False
+        Delete even if S3 does not hold this save of the checkpoint.
+    s3_config : S3Config, optional
+        S3 configuration. Loaded from environment if not provided.
+
+    Returns
+    -------
+    bool
+        True if a local directory was found and removed, False if there was
+        nothing locally.
+
+    Raises
+    ------
+    RuntimeError
+        If ``force`` is not set and S3 cannot confirm it has this save of the
+        checkpoint (not configured, unreachable, missing it, or holding a
+        different save of the step).
+    """
+    import shutil
+
+    local_dir = local_checkpoint_dir(project, run_name, step)
+    local = read_local_copy(local_dir)
+    if local is None:
+        return False
+    if not force:
+        _require_remote_copy(project, run_name, step, local, local_dir, s3_config)
+    shutil.rmtree(local_dir)
+    return True
+
+
+_KEEP_REASONS = {
+    CheckpointStatus.LOCAL_ONLY: "It is not in S3.",
+    CheckpointStatus.UNTRACKED: "It is not in S3.",
+    CheckpointStatus.DIVERGED: "S3 has a different save of this step.",
+}
+
+
+def _require_remote_copy(
+    project: str,
+    run_name: str,
+    step: int,
+    local: LocalCopy,
+    local_dir: Path,
+    s3_config: S3Config | None,
+) -> None:
+    from extty.s3 import _S3_ERRORS, S3Storage
+
+    refusal = (
+        f"Refusing to delete {local_dir}: it may be the only copy of checkpoint "
+        f"{project}/{run_name} step {step}. Pass force=True to delete it anyway."
+    )
+    if s3_config is None:
+        s3_config = S3Config.load()
+    if s3_config is None:
+        raise RuntimeError(f"{refusal} (S3 is not configured.)")
+    try:
+        remote = S3Storage(s3_config, project, run_name).find_checkpoint(step)
+    except (ImportError, *_S3_ERRORS) as e:
+        raise RuntimeError(f"{refusal} (Could not check S3: {e})") from e
+    status = checkpoint_status(local, remote)
+    if status not in DELETABLE_LOCALLY:
+        raise RuntimeError(f"{refusal} ({_KEEP_REASONS[status]})")
+
+
+def finish() -> None:
+    """Finish the current run and flush all data."""
+    global _active_run
+    if _active_run is None:
+        raise RuntimeError("No active run to finish.")
+    _active_run.finish()
+    _active_run = None
+
+
+def save_artifact(
+    name: str,
+    path: str | Path,
+    *,
+    description: str = "",
+    metadata: dict[str, Any] | None = None,
+    s3_config: S3Config | None = None,
+) -> ArtifactMeta:
+    """
+    Save an artifact to S3.
+
+    If an active run exists, the artifact is automatically associated
+    with it.
+
+    Parameters
+    ----------
+    name : str
+        Unique name for this artifact.
+    path : str or Path
+        Local file or directory to upload.
+    description : str
+        Human-readable description.
+    metadata : dict[str, Any], optional
+        User-defined metadata (arbitrary JSON-serializable dict).
+    s3_config : S3Config, optional
+        S3 configuration. Loaded from environment if not provided.
+
+    Returns
+    -------
+    ArtifactMeta
+        Metadata for the saved artifact.
+    """
+    if isinstance(_active_run, NoOpRun):
+        return _active_run.save_artifact(
+            name, path, description=description, metadata=metadata, s3_config=s3_config
+        )
+    run_project = _active_run.project if _active_run is not None else None
+    run_name = _active_run.name if _active_run is not None else None
+    return _save_artifact_raw(
+        name,
+        path,
+        description=description,
+        metadata=metadata,
+        s3_config=s3_config,
+        run_project=run_project,
+        run_name=run_name,
+    )
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _sanitize_config(config: dict[str, Any]) -> dict[str, Any]:
+    def _convert(val: Any) -> Any:
+        if dataclasses.is_dataclass(val) and not isinstance(val, type):
+            return dataclasses.asdict(val)
+        if isinstance(val, dict):
+            return {k: _convert(v) for k, v in val.items()}
+        if isinstance(val, list):
+            return [_convert(v) for v in val]
+        return val
+
+    return {k: _convert(v) for k, v in config.items()}
+
+
+def experiment(
+    project: str,
+    name: str | None = None,
+    name_kwarg: str | None = None,
+    conf_kwargs: list[str] | None = None,
+    non_conf_kwargs: list[str] | None = None,
+    system_metrics: bool = True,
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    if conf_kwargs is not None and non_conf_kwargs is not None:
+        raise ValueError(
+            "cannot pass values for both `conf_kwargs` and `non_conf_kwargs`"
+        )
+
+    def dec(fn: Callable[P, T]) -> Callable[P, T]:
+        @functools.wraps(fn)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            if len(args) > 0:
+                warnings.warn(
+                    f"non-keyword args passed to {fn} will not be logged to `extty`."
+                )
+            try:
+                kwargs_copy = copy.deepcopy(kwargs)
+                if conf_kwargs is not None:
+                    kwargs_copy = {k: kwargs_copy[k] for k in conf_kwargs}
+                if non_conf_kwargs is not None:
+                    kwargs_copy = {
+                        k: v for k, v in kwargs_copy.items() if k not in non_conf_kwargs
+                    }
+
+                kwargs_copy = _sanitize_config(kwargs_copy)
+
+                run_name = name
+
+                if run_name is None and name_kwarg is not None:
+                    run_name = kwargs_copy.pop(name_kwarg)
+
+                init(
+                    project=project,
+                    name=run_name,
+                    config=kwargs_copy,
+                    system_metrics=system_metrics,
+                )
+                return fn(*args, **kwargs)
+            finally:
+                global _active_run
+                if _active_run is not None:
+                    finish()
+
+        return wrapper
+
+    return dec
+
+
+def list_local_runs(project: str | None = None) -> list[tuple[str, str]]:
+    """
+    List all local runs.
+
+    Parameters
+    ----------
+    project : str, optional
+        Filter to runs in this project only.
+
+    Returns
+    -------
+    list[tuple[str, str]]
+        List of (project, run_name) tuples, sorted by start time
+        (most recent first).
+    """
+    return [(r.project, r.name) for r in get_runs(project=project)]
+
+
+def push(
+    target: str | None = None,
+    *,
+    s3_config: S3Config | None = None,
+    force: bool = False,
+    dry_run: bool = False,
+) -> list[str]:
+    """
+    Push local runs to S3.
+
+    Parameters
+    ----------
+    target : str, optional
+        Target to push. Format: "project/run-name" for single run,
+        "project/" for all runs in project, or None for all runs.
+    s3_config : S3Config, optional
+        S3 configuration. If not provided, reads from environment variables.
+    force : bool, default False
+        Overwrite remote data without merging.
+    dry_run : bool, default False
+        Show what would be pushed without actually pushing.
+
+    Returns
+    -------
+    list[str]
+        List of pushed run paths (project/run_name format).
+    """
+    if s3_config is None:
+        s3_config = S3Config.load()
+    if s3_config is None:
+        raise ValueError(
+            "S3 configuration required. Set EXTTY_S3_BUCKET environment variable "
+            "or provide s3_config parameter."
+        )
+
+    try:
+        import boto3
+    except ImportError:
+        raise ImportError(
+            "boto3 is required for S3 operations. Install with: pip install extty[s3]"
+        )
+
+    kwargs: dict = {}
+    if s3_config.region:
+        kwargs["region_name"] = s3_config.region
+    if s3_config.access_key_id and s3_config.secret_access_key:
+        kwargs["aws_access_key_id"] = s3_config.access_key_id
+        kwargs["aws_secret_access_key"] = s3_config.secret_access_key
+    if s3_config.endpoint_url:
+        kwargs["endpoint_url"] = s3_config.endpoint_url
+
+    client = boto3.client("s3", **kwargs)
+
+    project_filter = None
+    run_filter = None
+    if target:
+        if target.endswith("/"):
+            project_filter = target.rstrip("/")
+        elif "/" in target:
+            project_filter, run_filter = target.split("/", 1)
+        else:
+            project_filter = target
+
+    runs_to_push = []
+    for proj, run_name in list_local_runs(project_filter):
+        if run_filter is not None and run_name != run_filter:
+            continue
+        runs_to_push.append((proj, run_name))
+
+    if dry_run:
+        for proj, run_name in runs_to_push:
+            _logger.info("would push: %s/%s", proj, run_name)
+        return [f"{p}/{r}" for p, r in runs_to_push]
+
+    pushed = []
+    runs_dir = get_runs_dir()
+    for proj, run_name in runs_to_push:
+        local_run_dir = runs_dir / proj / run_name
+        s3_prefix = f"{s3_config.prefix}/runs/{proj}/{run_name}"
+
+        _push_run_to_s3(client, s3_config.bucket, s3_prefix, local_run_dir, force=force)
+        pushed.append(f"{proj}/{run_name}")
+        _logger.info("pushed: %s/%s", proj, run_name)
+
+    return pushed
+
+
+def _push_run_to_s3(
+    client,
+    bucket: str,
+    s3_prefix: str,
+    local_run_dir,
+    force: bool = False,
+) -> None:
+    """Push a single run directory to S3."""
+    import json
+
+    meta_path = local_run_dir / "meta.json"
+    if meta_path.exists():
+        with open(meta_path) as f:
+            meta_content = f.read()
+        if force:
+            client.put_object(
+                Bucket=bucket,
+                Key=f"{s3_prefix}/meta.json",
+                Body=meta_content.encode("utf-8"),
+                ContentType="application/json",
+            )
+        else:
+            try:
+                response = client.get_object(
+                    Bucket=bucket, Key=f"{s3_prefix}/meta.json"
+                )
+                remote_meta = json.loads(response["Body"].read().decode("utf-8"))
+                local_meta = json.loads(meta_content)
+                merged_meta = _merge_meta(local_meta, remote_meta)
+                client.put_object(
+                    Bucket=bucket,
+                    Key=f"{s3_prefix}/meta.json",
+                    Body=json.dumps(merged_meta, indent=2).encode("utf-8"),
+                    ContentType="application/json",
+                )
+            except client.exceptions.NoSuchKey:
+                client.put_object(
+                    Bucket=bucket,
+                    Key=f"{s3_prefix}/meta.json",
+                    Body=meta_content.encode("utf-8"),
+                    ContentType="application/json",
+                )
+
+    metrics_dir = local_run_dir / "metrics"
+    if metrics_dir.exists():
+        for csv_file in metrics_dir.rglob("*.csv"):
+            relative_path = csv_file.relative_to(metrics_dir)
+            s3_key = f"{s3_prefix}/metrics/{relative_path}"
+            _push_csv_file(client, bucket, s3_key, csv_file, force=force)
+
+    examples_dir = local_run_dir / "examples"
+    if examples_dir.exists():
+        for jsonl_file in examples_dir.rglob("*.jsonl"):
+            relative_path = jsonl_file.relative_to(examples_dir)
+            s3_key = f"{s3_prefix}/examples/{relative_path}"
+            _push_jsonl_file(client, bucket, s3_key, jsonl_file, force=force)
+
+    confusion_dir = local_run_dir / "confusion_matrices"
+    if confusion_dir.exists():
+        for jsonl_file in confusion_dir.rglob("*.jsonl"):
+            relative_path = jsonl_file.relative_to(confusion_dir)
+            s3_key = f"{s3_prefix}/confusion_matrices/{relative_path}"
+            _push_jsonl_file(client, bucket, s3_key, jsonl_file, force=force)
+
+    charts_dir = local_run_dir / "charts"
+    if charts_dir.exists():
+        for jsonl_file in charts_dir.rglob("*.jsonl"):
+            relative_path = jsonl_file.relative_to(charts_dir)
+            s3_key = f"{s3_prefix}/charts/{relative_path}"
+            _push_jsonl_file(client, bucket, s3_key, jsonl_file, force=force)
+
+    images_dir = local_run_dir / "images"
+    if images_dir.exists():
+        for jsonl_file in images_dir.rglob("*.jsonl"):
+            relative_path = jsonl_file.relative_to(images_dir)
+            s3_key = f"{s3_prefix}/images/{relative_path}"
+            _push_jsonl_file(client, bucket, s3_key, jsonl_file, force=force)
+        for png_file in images_dir.rglob("*.png"):
+            relative_path = png_file.relative_to(images_dir)
+            s3_key = f"{s3_prefix}/images/{relative_path}"
+            _push_binary_file(
+                client, bucket, s3_key, png_file, "image/png", force=force
+            )
+
+    system_csv = local_run_dir / "system.csv"
+    if system_csv.exists():
+        s3_key = f"{s3_prefix}/system.csv"
+        _push_csv_file(client, bucket, s3_key, system_csv, force=force)
+
+
+def _merge_meta(local: dict, remote: dict) -> dict:
+    """Merge meta.json from local and remote."""
+    if local.get("status") == "completed":
+        return local
+    if remote.get("status") == "completed":
+        return remote
+    local_finished = local.get("finished_at")
+    remote_finished = remote.get("finished_at")
+    if local_finished and remote_finished:
+        return local if local_finished >= remote_finished else remote
+    if local_finished:
+        return local
+    if remote_finished:
+        return remote
+    return local
+
+
+def _push_csv_file(client, bucket: str, s3_key: str, local_path, force: bool) -> None:
+    """Push a CSV file to S3, optionally merging with existing data."""
+    import csv
+    import io
+
+    with open(local_path) as f:
+        reader = csv.DictReader(f)
+        local_rows = list(reader)
+        if not local_rows:
+            return
+        fieldnames = reader.fieldnames or []
+
+    if force:
+        with open(local_path, "rb") as f:
+            client.put_object(
+                Bucket=bucket, Key=s3_key, Body=f.read(), ContentType="text/csv"
+            )
+        return
+
+    existing_keys: set[tuple] = set()
+    existing_rows: list[dict] = []
+    try:
+        response = client.get_object(Bucket=bucket, Key=s3_key)
+        content = response["Body"].read().decode("utf-8")
+        reader = csv.DictReader(io.StringIO(content))
+        for row in reader:
+            key = (row.get("step", ""), row.get("timestamp", ""))
+            existing_keys.add(key)
+            existing_rows.append(row)
+    except client.exceptions.NoSuchKey:
+        pass
+
+    for row in local_rows:
+        key = (row.get("step", ""), row.get("timestamp", ""))
+        if key not in existing_keys:
+            existing_keys.add(key)
+            existing_rows.append(row)
+
+    existing_rows.sort(
+        key=lambda r: (float(r.get("step", 0)), float(r.get("timestamp", 0)))
+    )
+
+    output = io.StringIO()
+    writer = csv.DictWriter(output, fieldnames=fieldnames)
+    writer.writeheader()
+    writer.writerows(existing_rows)
+
+    client.put_object(
+        Bucket=bucket,
+        Key=s3_key,
+        Body=output.getvalue().encode("utf-8"),
+        ContentType="text/csv",
+    )
+
+
+def _push_binary_file(
+    client, bucket: str, s3_key: str, local_path, content_type: str, force: bool
+) -> None:
+    """Push an immutable binary file to S3, skipping if the key already exists."""
+    if not force:
+        try:
+            client.head_object(Bucket=bucket, Key=s3_key)
+            return
+        except Exception:
+            pass
+    with open(local_path, "rb") as f:
+        client.put_object(
+            Bucket=bucket, Key=s3_key, Body=f.read(), ContentType=content_type
+        )
+
+
+def _push_jsonl_file(client, bucket: str, s3_key: str, local_path, force: bool) -> None:
+    """Push a JSONL file to S3, optionally merging with existing data."""
+    import json
+
+    with open(local_path) as f:
+        local_records = [json.loads(line) for line in f if line.strip()]
+
+    if force:
+        with open(local_path, "rb") as f:
+            client.put_object(
+                Bucket=bucket,
+                Key=s3_key,
+                Body=f.read(),
+                ContentType="application/x-ndjson",
+            )
+        return
+
+    existing_keys: set[tuple] = set()
+    existing_records: list[dict] = []
+    try:
+        response = client.get_object(Bucket=bucket, Key=s3_key)
+        content = response["Body"].read().decode("utf-8")
+        for line in content.strip().split("\n"):
+            if line:
+                record = json.loads(line)
+                key = (record.get("step"), record.get("timestamp"))
+                existing_keys.add(key)
+                existing_records.append(record)
+    except client.exceptions.NoSuchKey:
+        pass
+
+    for record in local_records:
+        key = (record.get("step"), record.get("timestamp"))
+        if key not in existing_keys:
+            existing_keys.add(key)
+            existing_records.append(record)
+
+    existing_records.sort(key=lambda r: (r.get("step", 0), r.get("timestamp", 0)))
+
+    output = "\n".join(json.dumps(r) for r in existing_records)
+    if output:
+        output += "\n"
+
+    client.put_object(
+        Bucket=bucket,
+        Key=s3_key,
+        Body=output.encode("utf-8"),
+        ContentType="application/x-ndjson",
+    )