server api

Author: Saba9

tests/unit/test_run.py

@@ -59,8 +59,12 @@ def test_run_log_calls_client_for_spaces(temp_dir):
     run.log(metrics)
 
     time.sleep(0.6)
-    _, kwargs = client.predict.call_args
-    assert kwargs["api_name"] == "/bulk_log"
+    bulk_log_calls = [
+        call for call in client.predict.call_args_list
+        if call.kwargs.get("api_name") == "/bulk_log"
+    ]
+    assert len(bulk_log_calls) >= 1
+    kwargs = bulk_log_calls[-1].kwargs
     assert len(kwargs["logs"]) == 1
     assert kwargs["logs"][0]["project"] == "proj"
     assert kwargs["logs"][0]["run"] == "run1"

trackio/__init__.py

@@ -105,7 +105,7 @@ def _cleanup_current_run():
     if run is not None:
         try:
             if not run._finished:
-                run.finish(status="failed")
+                run.finish(status="failed", _atexit=True)
         except Exception:
             pass
 

trackio/run.py

@@ -26,7 +26,13 @@
 from trackio.sqlite_storage import SQLiteStorage
 from trackio.table import Table
 from trackio.trace import Trace
-from trackio.typehints import AlertEntry, LogEntry, SystemLogEntry, UploadEntry
+from trackio.typehints import (
+    AlertEntry,
+    LogEntry,
+    RunStatusEntry,
+    SystemLogEntry,
+    UploadEntry,
+)
 from trackio.utils import MEDIA_DIR, _emit_nonfatal_warning, _get_default_namespace
 
 BATCH_SEND_INTERVAL = 0.5
@@ -139,6 +145,7 @@ def __init__(
         self._queued_system_logs: list[SystemLogEntry] = []
         self._queued_uploads: list[UploadEntry] = []
         self._queued_alerts: list[AlertEntry] = []
+        self._queued_run_status: list[RunStatusEntry] = []
         self._stop_flag = threading.Event()
         self._config_logged = False
         max_step = self._safe_get_max_step_for_run()
@@ -166,7 +173,20 @@ def __init__(
                 description="remote Trackio logging thread",
             )
 
-        SQLiteStorage.set_run_status(self.project, self.name, "running", run_id=self.id)
+        if self._is_local:
+            SQLiteStorage.set_run_status(
+                self.project, self.name, "running", run_id=self.id
+            )
+        else:
+            with self._client_lock:
+                self._queued_run_status.append(
+                    {
+                        "project": self.project,
+                        "run": self.name,
+                        "run_id": self.id,
+                        "status": "running",
+                    }
+                )
         self._finished = False
 
         self._gpu_monitor: "GpuMonitor | AppleGpuMonitor | None" = None
@@ -435,6 +455,7 @@ def _batch_sender(self):
             or len(self._queued_system_logs) > 0
             or len(self._queued_uploads) > 0
             or len(self._queued_alerts) > 0
+            or len(self._queued_run_status) > 0
             or self._has_local_buffer
         ):
             if not self._stop_flag.is_set():
@@ -522,6 +543,23 @@ def _batch_sender(self):
                             self._write_alerts_to_sqlite(alerts_to_send)
                             failed = True
 
+                    if self._queued_run_status:
+                        status_to_send = self._queued_run_status.copy()
+                        self._queued_run_status.clear()
+                        try:
+                            for entry in status_to_send:
+                                self._client.predict(
+                                    api_name="/set_run_status",
+                                    project=entry["project"],
+                                    run=entry["run"],
+                                    status=entry["status"],
+                                    run_id=entry.get("run_id"),
+                                    hf_token=self._hf_token_for_remote(),
+                                )
+                        except Exception:
+                            self._queued_run_status[0:0] = status_to_send
+                            failed = True
+
                     if failed:
                         consecutive_failures += 1
                     else:
@@ -965,11 +1003,24 @@ def log_system(self, metrics: dict):
         except Exception as e:
             _emit_nonfatal_warning(f"trackio.log_system() failed: {e}")
 
-    def finish(self, status: str = "finished"):
+    def finish(self, status: str = "finished", _atexit: bool = False):
         if self._finished:
             return
         self._finished = True
 
+        join_timeout = 2 if _atexit else 30
+
+        if not self._is_local:
+            with self._client_lock:
+                self._queued_run_status.append(
+                    {
+                        "project": self.project,
+                        "run": self.name,
+                        "run_id": self.id,
+                        "status": status,
+                    }
+                )
+
         try:
             if self._gpu_monitor is not None:
                 try:
@@ -984,24 +1035,28 @@ def finish(self, status: str = "finished"):
 
             if self._is_local:
                 if self._local_sender_thread is not None:
-                    print("* Run finished. Uploading logs to Trackio (please wait...)")
-                    self._local_sender_thread.join(timeout=30)
+                    if not _atexit:
+                        print(
+                            "* Run finished. Uploading logs to Trackio (please wait...)"
+                        )
+                    self._local_sender_thread.join(timeout=join_timeout)
                     if self._local_sender_thread.is_alive():
                         _emit_nonfatal_warning(
-                            "Could not flush all logs within 30s. Some data may be buffered locally."
+                            f"Could not flush all logs within {join_timeout}s. Some data may be buffered locally."
                         )
                 else:
                     with self._client_lock:
                         self._flush_queues_inline()
             else:
                 if self._client_thread is not None:
-                    print(
-                        "* Run finished. Uploading logs to the remote Trackio server (please wait...)"
-                    )
-                    self._client_thread.join(timeout=30)
+                    if not _atexit:
+                        print(
+                            "* Run finished. Uploading logs to the remote Trackio server (please wait...)"
+                        )
+                    self._client_thread.join(timeout=join_timeout)
                     if self._client_thread.is_alive():
                         _emit_nonfatal_warning(
-                            "Could not flush all logs within 30s. Some data may be buffered locally."
+                            f"Could not flush all logs within {join_timeout}s. Some data may be buffered locally."
                         )
                 else:
                     with self._client_lock:
@@ -1029,4 +1084,13 @@ def finish(self, status: str = "finished"):
         except Exception as e:
             _emit_nonfatal_warning(f"trackio.finish() failed: {e}")
 
-        SQLiteStorage.set_run_status(self.project, self.name, status, run_id=self.id)
+        if self._is_local:
+            SQLiteStorage.set_run_status(
+                self.project, self.name, status, run_id=self.id
+            )
+        elif self._queued_run_status:
+            self._warn_once(
+                "finish-remote-status",
+                f"trackio.finish() could not record run status '{status}' on the remote server. "
+                "The dashboard may not reflect the final state of this run.",
+            )

trackio/server.py

@@ -619,16 +619,19 @@ def bulk_alert(
                 "steps": [],
                 "timestamps": [],
                 "alert_ids": [],
+                "data_list": [],
             }
         alerts_by_run[key]["titles"].append(entry["title"])
         alerts_by_run[key]["texts"].append(entry.get("text"))
         alerts_by_run[key]["levels"].append(entry["level"])
         alerts_by_run[key]["steps"].append(entry.get("step"))
         alerts_by_run[key]["timestamps"].append(entry.get("timestamp"))
         alerts_by_run[key]["alert_ids"].append(entry.get("alert_id"))
+        alerts_by_run[key]["data_list"].append(entry.get("data"))
 
     for (project, run, run_id), data in alerts_by_run.items():
         has_alert_ids = any(aid is not None for aid in data["alert_ids"])
+        has_data = any(d is not None for d in data["data_list"])
         payload = dict(
             project=project,
             run=run,
@@ -639,13 +642,26 @@ def bulk_alert(
             steps=data["steps"],
             timestamps=data["timestamps"],
             alert_ids=data["alert_ids"] if has_alert_ids else None,
+            data_list=data["data_list"] if has_data else None,
         )
         try:
             SQLiteStorage.bulk_alert(**payload)
         except sqlite3.OperationalError:
             _enqueue_write("bulk_alert", payload)
 
 
+def set_run_status(
+    request: Request,
+    project: str,
+    run: str,
+    status: str,
+    run_id: str | None,
+    hf_token: str | None,
+) -> None:
+    assert_can_write_metrics(request, hf_token)
+    SQLiteStorage.set_run_status(project, run, status, run_id=run_id)
+
+
 def get_alerts(
     project: str,
     run: str | None = None,
@@ -949,6 +965,7 @@ def _api_registry() -> dict[str, Any]:
         "bulk_log": bulk_log,
         "bulk_log_system": bulk_log_system,
         "bulk_alert": bulk_alert,
+        "set_run_status": set_run_status,
         "get_alerts": get_alerts,
         "get_metric_values": get_metric_values,
         "get_runs_for_project": get_runs_for_project,

trackio/typehints.py

@@ -35,6 +35,13 @@ class AlertEntry(TypedDict, total=False):
     data: dict | None
 
 
+class RunStatusEntry(TypedDict, total=False):
+    project: str
+    run: str
+    run_id: str | None
+    status: str
+
+
 class UploadEntry(TypedDict):
     project: str
     run: str | None