refactors

Author: Saba9

tests/e2e-local/test_cli_agent_commands.py

@@ -68,6 +68,54 @@ def seeded_dir():
         mm.MEDIA_DIR, mu.MEDIA_DIR, tu.MEDIA_DIR, ss.MEDIA_DIR = orig_media
 
 
+@pytest.fixture(scope="module")
+def seeded_dir_with_unfinished():
+    """Fixture with two finished runs and one still-running run."""
+    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdir:
+        import trackio.media.media as mm
+        import trackio.media.utils as mu
+        import trackio.sqlite_storage as ss
+        import trackio.utils as tu
+
+        orig_trackio = ss.TRACKIO_DIR
+        orig_media = [mm.MEDIA_DIR, mu.MEDIA_DIR, tu.MEDIA_DIR, ss.MEDIA_DIR]
+        ss.TRACKIO_DIR = Path(tmpdir)
+        mm.MEDIA_DIR = mu.MEDIA_DIR = tu.MEDIA_DIR = ss.MEDIA_DIR = (
+            Path(tmpdir) / "media"
+        )
+
+        context_vars.current_run.set(None)
+        context_vars.current_project.set(None)
+        context_vars.current_server.set(None)
+        context_vars.current_space_id.set(None)
+
+        proj = "filter_test"
+
+        trackio.init(project=proj, name="done-run", config={"lr": 0.01})
+        for step in range(5):
+            trackio.log({"val/loss": 1.0 - step * 0.1}, step=step)
+        trackio.finish()
+
+        trackio.init(project=proj, name="also-done", config={"lr": 0.1})
+        for step in range(5):
+            trackio.log({"val/loss": 2.0 - step * 0.1}, step=step)
+        trackio.finish()
+
+        trackio.init(project=proj, name="still-running", config={"lr": 1.0})
+        for step in range(5):
+            trackio.log({"val/loss": 5.0 + step * 0.5}, step=step)
+
+        context_vars.current_run.set(None)
+        context_vars.current_project.set(None)
+        context_vars.current_server.set(None)
+        context_vars.current_space_id.set(None)
+
+        yield (tmpdir, proj)
+
+        ss.TRACKIO_DIR = orig_trackio
+        mm.MEDIA_DIR, mu.MEDIA_DIR, tu.MEDIA_DIR, ss.MEDIA_DIR = orig_media
+
+
 def _cli(args, env_dir):
     env = os.environ.copy()
     env["TRACKIO_DIR"] = env_dir
@@ -108,6 +156,29 @@ def test_best(seeded_dir):
     assert json.loads(r2.stdout)["best_run"] == "run-lr0.01"
 
 
+def test_best_excludes_unfinished_by_default(seeded_dir_with_unfinished):
+    tmpdir, proj = seeded_dir_with_unfinished
+    r = _cli(["best", "--project", proj, "--metric", "val/loss", "--json"], tmpdir)
+    assert r.returncode == 0
+    data = json.loads(r.stdout)
+    run_names = [e["run"] for e in data["ranking"]]
+    assert "still-running" not in run_names
+    assert len(run_names) == 2
+
+
+def test_best_include_all(seeded_dir_with_unfinished):
+    tmpdir, proj = seeded_dir_with_unfinished
+    r = _cli(
+        ["best", "--project", proj, "--metric", "val/loss", "--include-all", "--json"],
+        tmpdir,
+    )
+    assert r.returncode == 0
+    data = json.loads(r.stdout)
+    run_names = [e["run"] for e in data["ranking"]]
+    assert "still-running" in run_names
+    assert len(run_names) == 3
+
+
 def test_compare(seeded_dir):
     r = _cli(
         ["compare", "--project", PROJECT, "--metrics", "val/loss,accuracy", "--json"],
@@ -136,6 +207,37 @@ def test_compare(seeded_dir):
     assert len(json.loads(r2.stdout)["runs"]) == 2
 
 
+def test_compare_excludes_unfinished_by_default(seeded_dir_with_unfinished):
+    tmpdir, proj = seeded_dir_with_unfinished
+    r = _cli(["compare", "--project", proj, "--metrics", "val/loss", "--json"], tmpdir)
+    assert r.returncode == 0
+    data = json.loads(r.stdout)
+    run_names = [e["run"] for e in data["runs"]]
+    assert "still-running" not in run_names
+    assert len(run_names) == 2
+
+
+def test_compare_include_all(seeded_dir_with_unfinished):
+    tmpdir, proj = seeded_dir_with_unfinished
+    r = _cli(
+        [
+            "compare",
+            "--project",
+            proj,
+            "--metrics",
+            "val/loss",
+            "--include-all",
+            "--json",
+        ],
+        tmpdir,
+    )
+    assert r.returncode == 0
+    data = json.loads(r.stdout)
+    run_names = [e["run"] for e in data["runs"]]
+    assert "still-running" in run_names
+    assert len(run_names) == 3
+
+
 def test_summary(seeded_dir):
     r = _cli(
         ["summary", "--project", PROJECT, "--metric", "val/loss", "--json"], seeded_dir
@@ -155,6 +257,20 @@ def test_summary(seeded_dir):
         } <= run_entry.keys()
 
 
+def test_list_runs_json_includes_status(seeded_dir):
+    r = _cli(["list", "runs", "--project", PROJECT, "--json"], seeded_dir)
+    assert r.returncode == 0
+    data = json.loads(r.stdout)
+    assert "runs" in data
+    for entry in data["runs"]:
+        assert "name" in entry
+        assert "status" in entry
+    statuses = {e["name"]: e["status"] for e in data["runs"]}
+    assert statuses.get("run-lr0.01") == "finished"
+    assert statuses.get("run-lr0.1") == "finished"
+    assert statuses.get("run-lr1.0") == "finished"
+
+
 def test_best_error_cases(seeded_dir):
     assert (
         _cli(

tests/e2e-local/test_run_status.py

@@ -41,6 +41,33 @@ def test_api_run_status(temp_dir):
     assert run.status == "finished"
 
 
+def test_api_run_final_metrics(temp_dir):
+    trackio.init(project="final_metrics_test", name="run1")
+    trackio.log({"loss": 1.0, "acc": 0.5}, step=0)
+    trackio.log({"loss": 0.5, "acc": 0.8}, step=1)
+    trackio.finish()
+
+    run = trackio.Api().runs("final_metrics_test")[0]
+    fm = run.final_metrics
+    assert abs(fm["loss"] - 0.5) < 1e-6
+    assert abs(fm["acc"] - 0.8) < 1e-6
+
+
+def test_api_run_history_with_metric_filter(temp_dir):
+    trackio.init(project="history_test", name="run1")
+    for step in range(5):
+        trackio.log({"loss": 1.0 - step * 0.1, "acc": step * 0.1}, step=step)
+    trackio.finish()
+
+    run = trackio.Api().runs("history_test")[0]
+    full = run.history()
+    assert len(full) == 5
+
+    loss_history = run.history(metric="loss")
+    assert len(loss_history) == 5
+    assert all("value" in entry for entry in loss_history)
+
+
 def test_status_survives_multiple_runs(temp_dir):
     run1 = trackio.init(project="multi_status", name="run1")
     trackio.log({"loss": 0.5}, step=0)

tests/e2e-local/test_watchers.py

@@ -1,11 +1,11 @@
-from trackio.watchers import MetricWatcher, WatcherManager
+from trackio.watchers import AlertReason, MetricWatcher, WatcherManager
 
 
 def test_nan_inf_triggers_stop():
     w = MetricWatcher("loss", nan=True)
     alerts = w.check(float("nan"), step=10)
     assert len(alerts) == 1
-    assert alerts[0]["data"]["reason"] == "nan_inf"
+    assert alerts[0]["data"]["reason"] == AlertReason.NAN_INF
     assert w.should_stop
 
     w2 = MetricWatcher("loss", nan=False)
@@ -17,7 +17,7 @@ def test_max_value_with_dedup():
     assert len(w.check(5.0, step=0)) == 0
     alerts = w.check(15.0, step=1)
     assert len(alerts) == 1
-    assert alerts[0]["data"]["reason"] == "max_exceeded"
+    assert alerts[0]["data"]["reason"] == AlertReason.MAX_EXCEEDED
     assert w.should_stop
     assert len(w.check(15.0, step=2)) == 0
     w.check(5.0, step=3)
@@ -27,7 +27,9 @@ def test_max_value_with_dedup():
 def test_min_value_with_dedup():
     w = MetricWatcher("acc", min_value=0.5)
     assert len(w.check(0.8, step=0)) == 0
-    assert len(w.check(0.3, step=1)) == 1
+    alerts = w.check(0.3, step=1)
+    assert len(alerts) == 1
+    assert alerts[0]["data"]["reason"] == AlertReason.MIN_EXCEEDED
     assert len(w.check(0.3, step=2)) == 0
     w.check(0.8, step=3)
     assert len(w.check(0.3, step=4)) == 1
@@ -39,13 +41,23 @@ def test_spike_detection_with_dedup_and_reset():
         w.check(1.0, step=i)
     alerts = w.check(10.0, step=3)
     assert len(alerts) == 1
-    assert alerts[0]["data"]["reason"] == "spike"
+    assert alerts[0]["data"]["reason"] == AlertReason.SPIKE
     assert len(w.check(10.0, step=4)) == 0
     for i in range(3):
         w.check(1.0, step=5 + i)
     assert len(w.check(10.0, step=8)) == 1
 
 
+def test_spike_detection_works_for_negative_metrics():
+    w = MetricWatcher("reward", spike_factor=3.0, window=3)
+    for i in range(3):
+        w.check(-1.0, step=i)
+    assert len(w.check(-1.1, step=3)) == 0
+    alerts = w.check(2.0, step=4)
+    assert len(alerts) == 1
+    assert alerts[0]["data"]["reason"] == AlertReason.SPIKE
+
+
 def test_patience_min_mode():
     w = MetricWatcher("loss", patience=3, mode="min")
     w.check(1.0, step=0)
@@ -54,7 +66,7 @@ def test_patience_min_mode():
     w.check(0.95, step=3)
     alerts = w.check(0.95, step=4)
     assert len(alerts) == 1
-    assert alerts[0]["data"]["reason"] == "stagnation"
+    assert alerts[0]["data"]["reason"] == AlertReason.STAGNATION
     assert w.should_stop
     assert len(w.check(0.95, step=5)) == 0
 

trackio/__init__.py

@@ -39,7 +39,7 @@
 from trackio.table import Table
 from trackio.typehints import UploadEntry
 from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_DIR, _emit_nonfatal_warning
-from trackio.watchers import MetricWatcher, WatcherManager
+from trackio.watchers import AlertReason, MetricWatcher, WatcherManager
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
@@ -63,6 +63,7 @@ def __repr__(self) -> str:
     "watch",
     "should_stop",
     "AlertLevel",
+    "AlertReason",
     "show",
     "sync",
     "freeze",
@@ -97,8 +98,6 @@ def _cleanup_current_run():
         try:
             if not run._finished:
                 run.finish(status="failed")
-            else:
-                run.finish()
         except Exception:
             pass
 

trackio/api.py

@@ -27,40 +27,40 @@ def status(self) -> str | None:
         return SQLiteStorage.get_run_status(self.project, self.name, run_id=self.id)
 
     @property
-    def summary(self) -> dict:
-        logs = SQLiteStorage.get_logs(self.project, self.name)
-        final_values = {}
-        for log_entry in logs:
-            for key, value in log_entry.items():
-                if key not in ("timestamp", "step") and isinstance(value, (int, float)):
-                    final_values[key] = value
-        return final_values
+    def final_metrics(self) -> dict:
+        """Last recorded value for each numeric metric, keyed by metric name."""
+        metric_names = SQLiteStorage.get_all_metrics_for_run(self.project, self.name)
+        result = {}
+        for m in metric_names:
+            rows = SQLiteStorage.get_final_metric_for_runs(
+                self.project, m, mode="last", run_names=[self.name], status_filter=None
+            )
+            if rows:
+                result[m] = rows[0]["value"]
+        return result
 
     def metrics(self) -> list[str]:
         return SQLiteStorage.get_all_metrics_for_run(self.project, self.name)
 
-    def history(self, metric: str | None = None) -> list[dict]:
-        if metric is not None:
-            return SQLiteStorage.get_metric_values(self.project, self.name, metric)
-        return SQLiteStorage.get_logs(self.project, self.name)
-
-    def get_metric(
+    def history(
         self,
-        name: str,
+        metric: str | None = None,
         step: int | None = None,
         around_step: int | None = None,
         at_time: str | None = None,
-        window: int | float | None = None,
+        window: int | None = None,
     ) -> list[dict]:
-        return SQLiteStorage.get_metric_values(
-            self.project,
-            self.name,
-            name,
-            step=step,
-            around_step=around_step,
-            at_time=at_time,
-            window=window,
-        )
+        if metric is not None:
+            return SQLiteStorage.get_metric_values(
+                self.project,
+                self.name,
+                metric,
+                step=step,
+                around_step=around_step,
+                at_time=at_time,
+                window=window,
+            )
+        return SQLiteStorage.get_logs(self.project, self.name)
 
     def alerts(self, level: str | None = None, since: str | None = None) -> list[dict]:
         return SQLiteStorage.get_alerts(