feat(trackio): capture experiments + record the hosted URL on the DAG ([70])

roar/analyzers/experiment_trackers.py

@@ -2,6 +2,7 @@
 
 import json
 import re
+import sqlite3
 from pathlib import Path
 from typing import Any, ClassVar
 
@@ -18,6 +19,7 @@ class ExperimentTrackerAnalyzer(Analyzer):
     TRACKER_PATTERNS: ClassVar[dict[str, list[str]]] = {
         "wandb": ["wandb/", ".wandb"],
         "mlflow": ["mlruns/", "mlartifacts/"],
+        "trackio": ["huggingface/trackio", "/trackio/"],
         "neptune": [".neptune/"],
         "tensorboard": ["/runs/", "events.out.tfevents"],
     }
@@ -72,7 +74,12 @@ def analyze(self, context: dict) -> dict | None:
         for tracker in trackers_found:
             run_info = self._extract_run_info(tracker, written, env)
             if run_info:
-                results["runs"].append(run_info)
+                # trackio reports one run per project DB (a run may touch
+                # several); the other trackers return a single run dict.
+                if isinstance(run_info, list):
+                    results["runs"].extend(run_info)
+                else:
+                    results["runs"].append(run_info)
 
             # Add ignore patterns for this tracker
             for pattern in self.IGNORE_PATTERNS:
@@ -88,18 +95,26 @@ def analyze(self, context: dict) -> dict | None:
             results["ignore_patterns"].extend(["mlruns/*", "mlartifacts/*"])
         if "neptune" in trackers_found:
             results["ignore_patterns"].append(".neptune/*")
+        if "trackio" in trackers_found:
+            results["ignore_patterns"].append("*/trackio/*")
 
         # Dedupe
         results["ignore_patterns"] = sorted(set(results["ignore_patterns"]))
 
         return results if results["trackers_detected"] else None
 
-    def _extract_run_info(self, tracker: str, written_files: list, env: dict) -> dict | None:
-        """Extract run URL and metadata for a specific tracker."""
+    def _extract_run_info(self, tracker: str, written_files: list, env: dict) -> dict | list | None:
+        """Extract run URL and metadata for a specific tracker.
+
+        Most trackers return a single run dict; ``trackio`` returns a list of
+        per-project run dicts (or ``None`` when nothing usable was found).
+        """
         if tracker == "wandb":
             return self._extract_wandb_info(written_files, env)
         elif tracker == "mlflow":
             return self._extract_mlflow_info(written_files, env)
+        elif tracker == "trackio":
+            return self._extract_trackio_info(written_files, env)
         elif tracker == "neptune":
             return self._extract_neptune_info(written_files, env)
         return None
@@ -198,6 +213,91 @@ def _extract_wandb_info(self, written_files: list, env: dict) -> dict | None:
 
         return info if len(info) > 1 else None
 
+    def _extract_trackio_info(self, written_files: list, env: dict) -> list | None:
+        """Extract trackio run info + the hosted-dashboard URL — scrape only.
+
+        trackio stores one SQLite DB per project at ``.../huggingface/trackio/
+        <project>.db`` and, when a run syncs to a HF Space, records the Space in
+        ``project_metadata`` (``key='space_id'``). We read that link + the run
+        identity, exactly as the W&B extractor reads wandb's on-disk run URL — no
+        roar-side config, no env, and **never the metrics**: TReqs stores no
+        customer data, so the DAG carries a link to the externally-hosted
+        experiment, and the metrics live on the dashboard.
+
+        Returns ONE record per project DB (never splicing one project's identity
+        onto another's Space), or ``None`` if no DB yielded anything usable. Each
+        record's URL is ``https://huggingface.co/spaces/<space_id>?project=<project>``.
+        (``env`` is unused — the space comes from what trackio itself persisted.)
+        """
+        del env
+        db_paths = sorted({p for p in written_files if "trackio/" in p and p.endswith(".db")})
+
+        runs: list[dict[str, Any]] = []
+        for db_path in db_paths:
+            path = Path(db_path)
+            if not path.exists():
+                continue
+            info: dict[str, Any] = {"tracker": "trackio", "project": path.stem}
+            try:
+                con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
+            except sqlite3.Error:
+                continue
+            try:
+                # Two INDEPENDENT lookups: an older/imported trackio schema may
+                # lack a ``run_id`` column or the ``project_metadata`` table, so
+                # a failure of one must not discard the other.
+                run = self._trackio_run_identity(con)
+                space = self._trackio_space_id(con)
+            finally:
+                con.close()
+
+            if run:
+                info["run_id"], info["run_name"] = run
+            if space:
+                info["space_id"] = space
+                info["url"] = f"https://huggingface.co/spaces/{space}?project={info['project']}"
+
+            # Keep only DBs that yielded a run identity and/or a hosted Space
+            # (more than the always-present tracker+project keys).
+            if len(info) > 2:
+                runs.append(info)
+
+        return runs or None
+
+    @staticmethod
+    def _trackio_run_identity(con: "sqlite3.Connection") -> tuple | None:
+        """``(run_id, run_name)`` of the most recent metric row, tolerating
+        older schemas whose ``metrics`` table lacks a ``run_id`` column (trackio
+        produces these when importing TensorBoard/older runs)."""
+        try:
+            cols = {row[1] for row in con.execute("PRAGMA table_info(metrics)").fetchall()}
+        except sqlite3.Error:
+            return None
+        if not cols or not ({"run_id", "run_name"} & cols):
+            return None
+        run_id_expr = "run_id" if "run_id" in cols else "NULL"
+        run_name_expr = "run_name" if "run_name" in cols else "NULL"
+        try:
+            row = con.execute(
+                f"SELECT {run_id_expr}, {run_name_expr} FROM metrics ORDER BY rowid DESC LIMIT 1"
+            ).fetchone()
+        except sqlite3.Error:
+            return None
+        if not row or (row[0] is None and row[1] is None):
+            return None
+        return row[0], row[1]
+
+    @staticmethod
+    def _trackio_space_id(con: "sqlite3.Connection") -> str | None:
+        """The hosted HF Space id trackio persisted once a run synced, or None."""
+        try:
+            row = con.execute(
+                "SELECT value FROM project_metadata WHERE key = 'space_id' LIMIT 1"
+            ).fetchone()
+        except sqlite3.Error:
+            return None
+        return row[0] if row and row[0] else None
+
     def _extract_mlflow_info(self, written_files: list, env: dict) -> dict | None:
         """Extract MLflow run info from local files."""
         info = {"tracker": "mlflow"}

tests/unit/test_experiment_trackers.py

@@ -0,0 +1,143 @@
+"""Unit tests for the trackio branch of ExperimentTrackerAnalyzer ([70]).
+
+roar *scrapes* trackio exactly like W&B: it reads the run identity + the hosted HF
+Space URL that trackio itself persisted (``project_metadata.space_id``), and never
+the metrics (TReqs stores no customer data). No roar-side config, env, or flags.
+
+The extractor returns ONE record per project DB (a run may touch several projects),
+and tolerates older/imported trackio schemas that lack a ``run_id`` column or the
+``project_metadata`` table.
+"""
+
+import sqlite3
+from pathlib import Path
+
+from roar.analyzers.experiment_trackers import ExperimentTrackerAnalyzer
+
+
+def _make_trackio_db(
+    dirpath: Path,
+    project: str,
+    space_id: str | None = None,
+    *,
+    with_project_metadata: bool = True,
+    with_run_id: bool = True,
+    run_id: str = "run123",
+    run_name: str = "brave-run-0",
+) -> str:
+    """Create a minimal trackio-shaped SQLite DB, return its path.
+
+    Mirrors the real schema: a ``metrics`` table and a ``project_metadata`` table
+    into which trackio writes ``space_id`` once a run syncs to a HF Space. The
+    flags let us reproduce older/imported schemas: ``with_run_id=False`` drops the
+    ``run_id`` column (trackio's TensorBoard import), ``with_project_metadata=False``
+    drops the whole table (a never-synced/older DB).
+    """
+    db = dirpath / "huggingface" / "trackio" / f"{project}.db"
+    db.parent.mkdir(parents=True, exist_ok=True)
+    con = sqlite3.connect(str(db))
+    if with_run_id:
+        con.execute(
+            "CREATE TABLE metrics (id INTEGER PRIMARY KEY, run_id TEXT, timestamp TEXT, "
+            "run_name TEXT, step INTEGER, metrics TEXT, log_id TEXT, space_id TEXT)"
+        )
+        for step in range(3):
+            con.execute(
+                "INSERT INTO metrics (run_id, run_name, step, metrics) VALUES (?,?,?,?)",
+                (run_id, run_name, step, f'{{"loss": {round(1.0 / (step + 1), 4)}}}'),
+            )
+    else:
+        con.execute(
+            "CREATE TABLE metrics (id INTEGER PRIMARY KEY, timestamp TEXT, "
+            "run_name TEXT, step INTEGER, metrics TEXT)"
+        )
+        for step in range(3):
+            con.execute(
+                "INSERT INTO metrics (run_name, step, metrics) VALUES (?,?,?)",
+                (run_name, step, f'{{"loss": {round(1.0 / (step + 1), 4)}}}'),
+            )
+    if with_project_metadata:
+        con.execute("CREATE TABLE project_metadata (key TEXT, value TEXT)")
+        if space_id:
+            con.execute(
+                "INSERT INTO project_metadata (key, value) VALUES ('space_id', ?)", (space_id,)
+            )
+    con.commit()
+    con.close()
+    return str(db)
+
+
+class TestTrackioExtractor:
+    def test_url_from_persisted_space_metadata(self, tmp_path: Path):
+        db = _make_trackio_db(tmp_path, "minimind-o", space_id="reproducible-ai/experiments")
+        infos = ExperimentTrackerAnalyzer()._extract_trackio_info([db], {})
+        assert infos is not None and len(infos) == 1
+        info = infos[0]
+        assert info["tracker"] == "trackio"
+        assert info["project"] == "minimind-o"
+        assert info["run_id"] == "run123"
+        assert info["run_name"] == "brave-run-0"
+        assert (
+            info["url"]
+            == "https://huggingface.co/spaces/reproducible-ai/experiments?project=minimind-o"
+        )
+
+    def test_no_space_metadata_records_identity_but_no_url(self, tmp_path: Path):
+        db = _make_trackio_db(tmp_path, "cosyvoice")  # local-only run, never synced
+        infos = ExperimentTrackerAnalyzer()._extract_trackio_info([db], {})
+        assert infos is not None and len(infos) == 1
+        assert infos[0]["project"] == "cosyvoice"
+        assert infos[0]["run_id"] == "run123"
+        assert "url" not in infos[0]
+
+    def test_detection_and_analyze_end_to_end(self, tmp_path: Path):
+        db = _make_trackio_db(tmp_path, "sam2", space_id="reproducible-ai/experiments")
+        analyzer = ExperimentTrackerAnalyzer()
+        ctx = {"tracer_data": {"written_files": [db]}, "env": {}}
+        assert analyzer.relevant(ctx) is True
+        res = analyzer.analyze(ctx)
+        assert res is not None
+        assert "trackio" in res["trackers_detected"]
+        run = next(r for r in res["runs"] if r["tracker"] == "trackio")
+        assert run["url"].endswith("?project=sam2")
+
+    def test_url_only_no_metric_values_leak(self, tmp_path: Path):
+        """No-data-storage invariant: metric VALUES must never reach lineage."""
+        db = _make_trackio_db(tmp_path, "m", space_id="reproducible-ai/experiments")
+        infos = ExperimentTrackerAnalyzer()._extract_trackio_info([db], {})
+        assert "loss" not in str(infos)
+
+    def test_multi_project_no_cross_contamination(self, tmp_path: Path):
+        """F1: a run touching >1 project must never splice one project's Space
+        onto another project's name."""
+        alpha = _make_trackio_db(tmp_path, "alpha", space_id="team/alpha-space")
+        zebra = _make_trackio_db(tmp_path, "zebra")  # local-only, no Space
+        infos = ExperimentTrackerAnalyzer()._extract_trackio_info([alpha, zebra], {})
+        assert infos is not None and len(infos) == 2
+        by_project = {i["project"]: i for i in infos}
+        assert (
+            by_project["alpha"]["url"]
+            == "https://huggingface.co/spaces/team/alpha-space?project=alpha"
+        )
+        # zebra was never synced → no URL, and alpha's Space never leaks onto it.
+        assert "url" not in by_project["zebra"]
+        assert not any("project=zebra" in i.get("url", "") for i in infos)
+
+    def test_missing_project_metadata_keeps_run_identity(self, tmp_path: Path):
+        """F2: a DB without the project_metadata table must still record the run
+        identity (the two lookups are independent)."""
+        db = _make_trackio_db(tmp_path, "beta", with_project_metadata=False)
+        infos = ExperimentTrackerAnalyzer()._extract_trackio_info([db], {})
+        assert infos is not None and len(infos) == 1
+        assert infos[0]["run_name"] == "brave-run-0"
+        assert "url" not in infos[0]
+
+    def test_runid_less_schema_still_records_url(self, tmp_path: Path):
+        """F2: an older/imported schema without a run_id column must still yield
+        the hosted-Space URL from project_metadata."""
+        db = _make_trackio_db(tmp_path, "omega", space_id="team/omega-space", with_run_id=False)
+        infos = ExperimentTrackerAnalyzer()._extract_trackio_info([db], {})
+        assert infos is not None and len(infos) == 1
+        assert infos[0]["url"] == "https://huggingface.co/spaces/team/omega-space?project=omega"
+        assert infos[0]["run_name"] == "brave-run-0"
+        assert infos[0].get("run_id") is None