changes

Author: abidlabs

trackio/__init__.py

@@ -260,7 +260,9 @@ def init(
             space_url = deploy.SPACE_HOST_URL.format(
                 user_name=user_name, space_name=space_name
             )
-            print(f"* View dashboard by going to: {space_url}")
+            print(
+                f"* View dashboard by going to: {deploy._BOLD_ORANGE}{space_url}{deploy._RESET}"
+            )
             if utils.is_in_notebook() and embed:
                 utils.embed_url_in_notebook(space_url)
     context_vars.current_project.set(project)

trackio/bucket_storage.py

@@ -6,8 +6,6 @@
 from trackio.sqlite_storage import SQLiteStorage
 from trackio.utils import MEDIA_DIR, TRACKIO_DIR
 
-DB_EXT = ".db"
-
 
 def create_bucket_if_not_exists(bucket_id: str, private: bool | None = None) -> None:
     huggingface_hub.create_bucket(bucket_id, private=private or False, exist_ok=True)
@@ -40,14 +38,3 @@ def upload_project_to_bucket(project: str, bucket_id: str) -> None:
                 files_to_add.append((str(media_file), str(rel)))
 
     huggingface_hub.batch_bucket_files(bucket_id, add=files_to_add)
-
-
-def upload_all_projects_to_bucket(bucket_id: str) -> None:
-    if not TRACKIO_DIR.exists():
-        return
-    for db_file in TRACKIO_DIR.glob(f"*{DB_EXT}"):
-        project = db_file.stem
-        try:
-            upload_project_to_bucket(project, bucket_id)
-        except FileNotFoundError:
-            continue

trackio/deploy.py

@@ -37,6 +37,8 @@
 
 SPACE_HOST_URL = "https://{user_name}-{space_name}.hf.space/"
 SPACE_URL = "https://huggingface.co/spaces/{space_id}"
+_BOLD_ORANGE = "\033[1m\033[38;5;208m"
+_RESET = "\033[0m"
 
 
 def _readme_linked_hub_yaml(dataset_id: str | None) -> str:
@@ -317,7 +319,9 @@ def create_space_if_not_exists(
         )
     try:
         huggingface_hub.repo_info(space_id, repo_type="space")
-        print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")
+        print(
+            f"* Found existing space: {_BOLD_ORANGE}{SPACE_URL.format(space_id=space_id)}{_RESET}"
+        )
         return
     except RepositoryNotFoundError:
         pass
@@ -328,7 +332,9 @@ def create_space_if_not_exists(
         else:
             raise ValueError(f"Failed to create Space: {e}")
 
-    print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")
+    print(
+        f"* Creating new space: {_BOLD_ORANGE}{SPACE_URL.format(space_id=space_id)}{_RESET}"
+    )
     deploy_as_space(
         space_id,
         space_storage,
@@ -540,7 +546,9 @@ def sync_incremental(
                 )
 
     SQLiteStorage.set_project_metadata(project, "space_id", space_id)
-    print(f"* Synced successfully to space: {SPACE_URL.format(space_id=space_id)}")
+    print(
+        f"* Synced successfully to space: {_BOLD_ORANGE}{SPACE_URL.format(space_id=space_id)}{_RESET}"
+    )
 
 
 def upload_dataset_for_static(
@@ -693,7 +701,9 @@ def deploy_as_static_space(
             ),
         )
 
-    print(f"* Static Space deployed: {SPACE_URL.format(space_id=space_id)}")
+    print(
+        f"* Static Space deployed: {_BOLD_ORANGE}{SPACE_URL.format(space_id=space_id)}{_RESET}"
+    )
 
 
 def sync(

trackio/server.py

@@ -1,11 +1,15 @@
 """The main API layer for the Trackio UI."""
 
 import base64
+import logging
 import os
 import re
 import secrets
 import shutil
+import sqlite3
+import threading
 import time
+from collections import deque
 from functools import lru_cache
 from typing import Any
 from urllib.parse import urlencode
@@ -23,6 +27,66 @@
 
 HfApi = hf.HfApi()
 
+logger = logging.getLogger("trackio")
+
+_write_queue: deque[tuple[str, Any]] = deque()
+_flush_thread: threading.Thread | None = None
+_flush_lock = threading.Lock()
+_FLUSH_INTERVAL = 2.0
+_MAX_RETRIES = 30
+
+
+def _enqueue_write(kind: str, payload: Any) -> None:
+    _write_queue.append((kind, payload))
+    _ensure_flush_thread()
+
+
+def _ensure_flush_thread() -> None:
+    global _flush_thread
+    with _flush_lock:
+        if _flush_thread is not None and _flush_thread.is_alive():
+            return
+        _flush_thread = threading.Thread(target=_flush_loop, daemon=True)
+        _flush_thread.start()
+
+
+def _flush_loop() -> None:
+    retries = 0
+    while _write_queue and retries < _MAX_RETRIES:
+        kind, payload = _write_queue[0]
+        try:
+            if kind == "bulk_log":
+                SQLiteStorage.bulk_log(**payload)
+            elif kind == "bulk_log_system":
+                SQLiteStorage.bulk_log_system(**payload)
+            elif kind == "bulk_alert":
+                SQLiteStorage.bulk_alert(**payload)
+            _write_queue.popleft()
+            retries = 0
+        except sqlite3.OperationalError as e:
+            msg = str(e).lower()
+            if "disk i/o error" in msg or "readonly" in msg:
+                retries += 1
+                logger.warning(
+                    "write queue: flush failed (%s), retry %d/%d",
+                    e,
+                    retries,
+                    _MAX_RETRIES,
+                )
+                time.sleep(min(_FLUSH_INTERVAL * retries, 15.0))
+            else:
+                logger.error("write queue: non-retryable error (%s), dropping entry", e)
+                _write_queue.popleft()
+                retries = 0
+    if _write_queue:
+        logger.error(
+            "write queue: giving up after %d retries, %d entries dropped",
+            _MAX_RETRIES,
+            len(_write_queue),
+        )
+        _write_queue.clear()
+
+
 write_token = secrets.token_urlsafe(32)
 
 OAUTH_CALLBACK_PATH = "/login/callback"
@@ -345,14 +409,18 @@ def bulk_log(
 
     for (project, run), data in logs_by_run.items():
         has_log_ids = any(lid is not None for lid in data["log_ids"])
-        SQLiteStorage.bulk_log(
+        payload = dict(
             project=project,
             run=run,
             metrics_list=data["metrics"],
             steps=data["steps"],
             config=data["config"],
             log_ids=data["log_ids"] if has_log_ids else None,
         )
+        try:
+            SQLiteStorage.bulk_log(**payload)
+        except sqlite3.OperationalError:
+            _enqueue_write("bulk_log", payload)
 
 
 def bulk_log_system(
@@ -372,13 +440,17 @@ def bulk_log_system(
 
     for (project, run), data in logs_by_run.items():
         has_log_ids = any(lid is not None for lid in data["log_ids"])
-        SQLiteStorage.bulk_log_system(
+        payload = dict(
             project=project,
             run=run,
             metrics_list=data["metrics"],
             timestamps=data["timestamps"],
             log_ids=data["log_ids"] if has_log_ids else None,
         )
+        try:
+            SQLiteStorage.bulk_log_system(**payload)
+        except sqlite3.OperationalError:
+            _enqueue_write("bulk_log_system", payload)
 
 
 def bulk_alert(
@@ -408,7 +480,7 @@ def bulk_alert(
 
     for (project, run), data in alerts_by_run.items():
         has_alert_ids = any(aid is not None for aid in data["alert_ids"])
-        SQLiteStorage.bulk_alert(
+        payload = dict(
             project=project,
             run=run,
             titles=data["titles"],
@@ -418,6 +490,10 @@ def bulk_alert(
             timestamps=data["timestamps"],
             alert_ids=data["alert_ids"] if has_alert_ids else None,
         )
+        try:
+            SQLiteStorage.bulk_alert(**payload)
+        except sqlite3.OperationalError:
+            _enqueue_write("bulk_alert", payload)
 
 
 def get_alerts(

trackio/space_volumes.py

@@ -41,28 +41,6 @@ def get_space_volumes(
     return list(vols) if vols else []
 
 
-def create_bucket(
-    namespace: str,
-    short_name: str,
-    *,
-    private: bool = True,
-    token: str | None = None,
-) -> None:
-    t = _token(token)
-    base = ENDPOINT.rstrip("/")
-    url = f"{base}/api/buckets/{namespace}/{short_name}"
-    headers = {"Authorization": f"Bearer {t}"}
-    r = get_session().post(
-        url,
-        headers=headers,
-        json={"private": private},
-        timeout=60.0,
-    )
-    if r.status_code == 409:
-        return
-    hf_raise_for_status(r)
-
-
 def set_space_volumes(
     space_id: str,
     volumes: list[dict[str, Any]],