diff --git a/roar/analyzers/experiment_trackers.py b/roar/analyzers/experiment_trackers.py index 541181ea..c6cdbe8e 100644 --- a/roar/analyzers/experiment_trackers.py +++ b/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/ + .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/?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"} diff --git a/tests/unit/test_experiment_trackers.py b/tests/unit/test_experiment_trackers.py new file mode 100644 index 00000000..52da312d --- /dev/null +++ b/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