source/study/tracking.py

#!/usr/bin/env python3
"""Publish sanitized controller events to Trackio from a durable SQLite outbox."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import re
import signal
import sqlite3
import sys
import time


SPACE_ID = "burtenshaw/beam-pi-programbench"
PROJECT = "beam-pi-programbench-20261009"
LABEL = re.compile(r"^[A-Za-z0-9_.:/-]{1,180}$")
CONFIG_TEXT = {"model", "provider", "reasoning", "task", "condition", "kind", "phase",
               "task_id", "benchmark_commit", "test_dataset_commit", "stop_reason"}
KINDS = {"run_started", "progress", "score", "run_finished", "error", "budget_stop"}
STATUSES = {"pending", "running", "completed", "failed", "cancelled", "budget_exhausted",
            "wall_time_exhausted", "passed", "blocked", "validated", "initializing",
            "calibrating", "calibration_failed", "validation_failed", "interrupted"}


def numeric(value):
    return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value)


def safe_config(config):
    """Allow scalar numeric settings and explicitly named, short categorical strings."""
    result = {}
    if not isinstance(config, dict):
        return result
    for key, value in config.items():
        if not isinstance(key, str) or not LABEL.fullmatch(key):
            continue
        if any(word in key.lower() for word in ("secret", "token_key", "api_key", "password", "credential")):
            continue
        if isinstance(value, bool) or numeric(value):
            result[key] = value
        elif key in CONFIG_TEXT and isinstance(value, str) and LABEL.fullmatch(value):
            result[key] = value
    return result


def sanitize(event):
    run_id = event.get("run_id")
    if not isinstance(run_id, str) or not LABEL.fullmatch(run_id):
        raise ValueError("invalid run_id")
    kind = event.get("kind")
    if kind not in KINDS:
        raise ValueError("invalid event kind")
    config = safe_config(event.get("config", {}))
    for key in ("task", "condition"):
        value = event.get(key)
        if isinstance(value, str) and LABEL.fullmatch(value):
            config[key] = value
    config["phase"] = ("infrastructure_validation" if run_id.startswith("validation-")
                       else "study_summary" if run_id == "study-summary" else "benchmark")
    metrics = {key: value for key, value in event.get("metrics", {}).items()
               if isinstance(key, str) and LABEL.fullmatch(key) and numeric(value)}
    metrics[f"event/{kind}"] = 1
    status = event.get("status")
    if status in STATUSES:
        for label in sorted(STATUSES):
            metrics[f"status/{label}"] = int(label == status)
    # Deliberately omit message, traces, prompts, outputs, paths, and arbitrary text.
    return {"run_id": run_id, "config": config, "metrics": metrics}


class Tracker:
    def __init__(self, state_dir, space_id=SPACE_ID, project=PROJECT, client=None):
        self.state_dir = Path(state_dir)
        self.state_dir.mkdir(parents=True, exist_ok=True)
        os.chmod(self.state_dir, 0o700)
        self.db = sqlite3.connect(self.state_dir / "outbox.sqlite")
        self.db.execute("PRAGMA journal_mode=DELETE")
        self.db.execute("PRAGMA synchronous=FULL")
        self.db.executescript("""
            CREATE TABLE IF NOT EXISTS events (
                id INTEGER PRIMARY KEY, log_id TEXT UNIQUE NOT NULL, payload TEXT NOT NULL,
                remote_visible INTEGER NOT NULL DEFAULT 0);
            CREATE TABLE IF NOT EXISTS cursor (
                path TEXT PRIMARY KEY, identity TEXT NOT NULL, offset INTEGER NOT NULL);
        """)
        self.db.commit()
        self.space_id, self.project, self.client = space_id, project, client
        self.last_error = None

    def enqueue(self, event):
        payload = sanitize(event)
        # Hash the source event for stable IDs across replay/restart, without storing it.
        log_id = hashlib.sha256(json.dumps(event, sort_keys=True).encode()).hexdigest()
        self.db.execute("INSERT OR IGNORE INTO events(log_id,payload) VALUES (?,?)",
                        (log_id, json.dumps(payload, sort_keys=True)))
        self.db.commit()

    def ingest(self, events_path):
        path = Path(events_path)
        if not path.exists():
            return
        stat = path.stat()
        identity = f"{stat.st_dev}:{stat.st_ino}"
        saved = self.db.execute("SELECT identity,offset FROM cursor WHERE path=?",
                                (str(path.resolve()),)).fetchone()
        offset = saved[1] if saved and saved[0] == identity and saved[1] <= stat.st_size else 0
        with path.open("rb") as source:
            source.seek(offset)
            while True:
                line = source.readline()
                if not line or not line.endswith(b"\n"):
                    break
                self.enqueue(json.loads(line))
                self.db.execute("INSERT OR REPLACE INTO cursor VALUES (?,?,?)",
                                (str(path.resolve()), identity, source.tell()))
                self.db.commit()

    def pending(self):
        return self.db.execute("SELECT count(*) FROM events WHERE remote_visible=0").fetchone()[0]

    def _connect(self):
        if self.client is None:
            os.environ.setdefault("TRACKIO_DIR", str(self.state_dir / "trackio"))
            from huggingface_hub import get_token
            from trackio.remote_client import RemoteClient
            self.client = RemoteClient(self.space_id, hf_token=get_token(),
                                       httpx_kwargs={"timeout": 30}, verbose=False)
        return self.client

    def flush(self):
        rows = self.db.execute("SELECT id,log_id,payload FROM events WHERE remote_visible=0 ORDER BY id LIMIT 100").fetchall()
        if not rows:
            self.write_status()
            return True
        try:
            client = self._connect()
            from huggingface_hub import get_token
            logs = []
            by_run = {}
            for seq, log_id, payload_text in rows:
                payload = json.loads(payload_text)
                run_id = payload["run_id"]
                metrics = payload["metrics"] | {"tracking/event_sequence": seq}
                logs.append({"project": self.project, "run": run_id, "run_id": run_id,
                             "metrics": metrics, "step": seq, "log_id": log_id,
                             "config": payload["config"]})
                by_run.setdefault(run_id, set()).add(seq)
            client.predict(api_name="/bulk_log", logs=logs, hf_token=get_token())
            # An HTTP success can mean queued writes. Only read-back events are acknowledged.
            acknowledged = set()
            for run_id, sequences in by_run.items():
                values = client.predict(api_name="/get_metric_values", project=self.project,
                                        run=run_id, run_id=run_id,
                                        metric_name="tracking/event_sequence", max_points=None)
                acknowledged.update(item["value"] for item in values
                                    if item.get("value") in sequences)
            self.db.executemany("UPDATE events SET remote_visible=1 WHERE id=?",
                                [(seq,) for seq in acknowledged])
            self.db.commit()
            self.last_error = None if len(acknowledged) == len(rows) else "remote_readback_incomplete"
            self.write_status()
            return self.last_error is None
        except Exception as error:
            code = getattr(getattr(error, "response", None), "status_code", None)
            self.last_error = type(error).__name__ + (f"_HTTP_{code}" if code else "")
            print(f"TRACKING_UPLOAD_PENDING: {self.last_error}; {self.pending()} events retained locally",
                  file=sys.stderr, flush=True)
            self.client = None
            self.write_status()
            return False

    def write_status(self):
        total, visible = self.db.execute("SELECT count(*),coalesce(sum(remote_visible),0) FROM events").fetchone()
        status = {"timestamp": time.time(), "space_id": self.space_id, "project": self.project,
                  "local_durable_events": total, "remote_readback_events": visible,
                  "pending_events": total - visible, "last_error": self.last_error,
                  "remote_acknowledgement": "metric_readback", "local_storage": "sqlite_full_sync"}
        tmp = self.state_dir / "status.json.tmp"
        tmp.write_text(json.dumps(status, indent=2) + "\n")
        tmp.replace(self.state_dir / "status.json")

    def close(self):
        self.db.close()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--events", required=True, type=Path)
    parser.add_argument("--state-dir", required=True, type=Path)
    parser.add_argument("--space-id", default=SPACE_ID)
    parser.add_argument("--project", default=PROJECT)
    parser.add_argument("--once", action="store_true")
    parser.add_argument("--poll-seconds", type=float, default=15)
    args = parser.parse_args()
    tracker = Tracker(args.state_dir, args.space_id, args.project)
    stopped = False

    def stop(*_):
        nonlocal stopped
        stopped = True

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    try:
        while True:
            try:
                tracker.ingest(args.events)
            except Exception as error:
                tracker.last_error = "ingest_" + type(error).__name__
                tracker.write_status()
                print(f"TRACKING_INGEST_ERROR: {tracker.last_error}; inspect controller event schema",
                      file=sys.stderr, flush=True)
                if args.once or stopped:
                    return 2
            tracker.flush()
            if args.once or stopped:
                # Drain already-ingested batches once; stop if remote cannot confirm.
                while tracker.pending() and tracker.flush():
                    pass
                return 1 if tracker.pending() else 0
            deadline = time.monotonic() + args.poll_seconds
            while not stopped and time.monotonic() < deadline:
                time.sleep(min(1, max(0, deadline - time.monotonic())))
    finally:
        tracker.close()


if __name__ == "__main__":
    raise SystemExit(main())