Optimizations for the exponential backoff and for indexing the sqlite db (#70)

* Refactor utils and improve imports

* Optimize database and UI performance

* Format run and sqlite_storage

* revert utils changes for now

* revert slots

* revert imports changes

* revert ui changes

* changes

* readme

* changes

* projects

* docstring

* changes

* changes

* rm project name from db

* use Path

* changes

* remove cache

* revert

* changes

* use index

* fix tests

* version bump

---------

Co-authored-by: Abubakar Abid <abubakar@huggingface.co>

Author: stabgan

README.md

@@ -145,6 +145,12 @@ To get started and see basic examples of usage, see these files:
 * [Persisting metrics in a Hugging Face Dataset](https://github.com/gradio-app/trackio/blob/main/examples/persist-dataset.py)
 * [Deploying the dashboard to Spaces](https://github.com/gradio-app/trackio/blob/main/examples/deploy-on-spaces.py)
 
+## Note: Trackio is in Beta (DB Schema May Change)
+
+Note that Trackio is in pre-release right now and we may release breaking changes. In particular, the schema of the Trackio sqlite database may change, which may require migrating or deleting existing database files (located by default at: `~/.cache/huggingface/trackio`).  
+
+Since Trackio is in beta, your feedback is welcome! Please create issues with bug reports or feature requests.
+
 ## License
 
 MIT License 

tests/__pycache__/conftest.cpython-312-pytest-8.3.4.pyc

@@ -1,4 +1,5 @@
 import tempfile
+from pathlib import Path
 
 import pytest
 
@@ -7,5 +8,5 @@
 def temp_db(monkeypatch):
     """Fixture that creates a temporary directory for database storage and patches the TRACKIO_DIR."""
     with tempfile.TemporaryDirectory() as tmpdir:
-        monkeypatch.setattr("trackio.sqlite_storage.TRACKIO_DIR", tmpdir)
+        monkeypatch.setattr("trackio.sqlite_storage.TRACKIO_DIR", Path(tmpdir))
         yield tmpdir

tests/__pycache__/test_run.cpython-312-pytest-8.3.4.pyc

@@ -128,16 +128,14 @@ def wait_until_space_exists(
     Args:
         space_id: The ID of the Space to wait for.
     """
-    client = None
-    for _ in range(30):
+    delay = 1
+    for _ in range(10):
         try:
-            client = Client(space_id, verbose=False)
-            if client:
-                break
-        except ReadTimeout:
-            time.sleep(5)
-        except ValueError:
-            time.sleep(5)
+            Client(space_id, verbose=False)
+            return
+        except (ReadTimeout, ValueError):
+            time.sleep(delay)
+            delay = min(delay * 2, 30)
     raise TimeoutError("Waiting for space to exist took longer than expected")
 
 

tests/__pycache__/test_sqlite_storage.cpython-312-pytest-8.3.4.pyc

@@ -1,53 +1,64 @@
-import glob
 import json
 import os
 import sqlite3
 from datetime import datetime
+from pathlib import Path
 
 from huggingface_hub import CommitScheduler
 
-try:
+try:  # absolute imports when installed
     from trackio.context_vars import current_scheduler
     from trackio.dummy_commit_scheduler import DummyCommitScheduler
     from trackio.utils import TRACKIO_DIR
-except:  # noqa: E722
+except Exception:  # relative imports for local execution on Spaces
     from context_vars import current_scheduler
     from dummy_commit_scheduler import DummyCommitScheduler
     from utils import TRACKIO_DIR
 
 
 class SQLiteStorage:
     @staticmethod
-    def get_project_db_path(project: str) -> str:
+    def _get_connection(db_path: Path) -> sqlite3.Connection:
+        conn = sqlite3.connect(str(db_path))
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    @staticmethod
+    def get_project_db_path(project: str) -> Path:
         """Get the database path for a specific project."""
         safe_project_name = "".join(
             c for c in project if c.isalnum() or c in ("-", "_")
         ).rstrip()
         if not safe_project_name:
             safe_project_name = "default"
-        return os.path.join(TRACKIO_DIR, f"{safe_project_name}.db")
+        return TRACKIO_DIR / f"{safe_project_name}.db"
 
     @staticmethod
-    def init_db(project: str) -> str:
+    def init_db(project: str) -> Path:
         """
         Initialize the SQLite database with required tables.
         Returns the database path.
         """
         db_path = SQLiteStorage.get_project_db_path(project)
-        os.makedirs(os.path.dirname(db_path), exist_ok=True)
+        db_path.parent.mkdir(parents=True, exist_ok=True)
         with SQLiteStorage.get_scheduler().lock:
             with sqlite3.connect(db_path) as conn:
                 cursor = conn.cursor()
                 cursor.execute("""
                     CREATE TABLE IF NOT EXISTS metrics (
                         id INTEGER PRIMARY KEY AUTOINCREMENT,
                         timestamp TEXT NOT NULL,
-                        project_name TEXT NOT NULL,
                         run_name TEXT NOT NULL,
                         step INTEGER NOT NULL,
                         metrics TEXT NOT NULL
                     )
                 """)
+                cursor.execute(
+                    """
+                    CREATE INDEX IF NOT EXISTS idx_metrics_run_step
+                    ON metrics(run_name, step)
+                    """
+                )
                 conn.commit()
         return db_path
 
@@ -85,16 +96,16 @@ def log(project: str, run: str, metrics: dict):
         db_path = SQLiteStorage.init_db(project)
 
         with SQLiteStorage.get_scheduler().lock:
-            with sqlite3.connect(db_path) as conn:
+            with SQLiteStorage._get_connection(db_path) as conn:
                 cursor = conn.cursor()
 
                 cursor.execute(
                     """
                     SELECT MAX(step) 
                     FROM metrics 
-                    WHERE project_name = ? AND run_name = ?
+                    WHERE run_name = ?
                     """,
-                    (project, run),
+                    (run,),
                 )
                 last_step = cursor.fetchone()[0]
                 current_step = 0 if last_step is None else last_step + 1
@@ -103,13 +114,12 @@ def log(project: str, run: str, metrics: dict):
 
                 cursor.execute(
                     """
-                    INSERT INTO metrics 
-                    (timestamp, project_name, run_name, step, metrics)
-                    VALUES (?, ?, ?, ?, ?)
+                    INSERT INTO metrics
+                    (timestamp, run_name, step, metrics)
+                    VALUES (?, ?, ?, ?)
                     """,
                     (
                         current_timestamp,
-                        project,
                         run,
                         current_step,
                         json.dumps(metrics),
@@ -142,15 +152,14 @@ def bulk_log(
 
         db_path = SQLiteStorage.init_db(project)
         with SQLiteStorage.get_scheduler().lock:
-            with sqlite3.connect(db_path) as conn:
+            with SQLiteStorage._get_connection(db_path) as conn:
                 cursor = conn.cursor()
 
                 data = []
                 for i, metrics in enumerate(metrics_list):
                     data.append(
                         (
                             timestamps[i],
-                            project,
                             run,
                             steps[i],
                             json.dumps(metrics),
@@ -159,9 +168,9 @@ def bulk_log(
 
                 cursor.executemany(
                     """
-                    INSERT INTO metrics 
-                    (timestamp, project_name, run_name, step, metrics)
-                    VALUES (?, ?, ?, ?, ?)
+                    INSERT INTO metrics
+                    (timestamp, run_name, step, metrics)
+                    VALUES (?, ?, ?, ?)
                     """,
                     data,
                 )
@@ -171,71 +180,85 @@ def bulk_log(
     def get_metrics(project: str, run: str) -> list[dict]:
         """Retrieve metrics for a specific run. The metrics also include the step count (int) and the timestamp (datetime object)."""
         db_path = SQLiteStorage.get_project_db_path(project)
-        if not os.path.exists(db_path):
+        if not db_path.exists():
             return []
 
-        with sqlite3.connect(db_path) as conn:
+        with SQLiteStorage._get_connection(db_path) as conn:
             cursor = conn.cursor()
             cursor.execute(
                 """
                 SELECT timestamp, step, metrics
                 FROM metrics
-                WHERE project_name = ? AND run_name = ?
+                WHERE run_name = ?
                 ORDER BY timestamp
                 """,
-                (project, run),
+                (run,),
             )
-            rows = cursor.fetchall()
 
+            rows = cursor.fetchall()
             results = []
             for row in rows:
-                timestamp, step, metrics_json = row
-                metrics = json.loads(metrics_json)
-                metrics["timestamp"] = timestamp
-                metrics["step"] = step
+                metrics = json.loads(row["metrics"])
+                metrics["timestamp"] = row["timestamp"]
+                metrics["step"] = row["step"]
                 results.append(metrics)
+
             return results
 
     @staticmethod
     def get_projects() -> list[str]:
-        """Get list of all projects by scanning database files."""
-        projects = []
-        if not os.path.exists(TRACKIO_DIR):
-            return projects
-
-        db_files = glob.glob(os.path.join(TRACKIO_DIR, "*.db"))
-
-        for db_file in db_files:
-            try:
-                with sqlite3.connect(db_file) as conn:
-                    cursor = conn.cursor()
-                    cursor.execute(
-                        "SELECT name FROM sqlite_master WHERE type='table' AND name='metrics'"
-                    )
-                    if cursor.fetchone():
-                        cursor.execute("SELECT DISTINCT project_name FROM metrics")
-                        project_names = [row[0] for row in cursor.fetchall()]
-                        projects.extend(project_names)
-            except sqlite3.Error:
-                continue
+        """
+        Get list of all projects by scanning the database files in the trackio directory.
+        """
+        projects: set[str] = set()
+        if not TRACKIO_DIR.exists():
+            return []
 
-        return list(set(projects))
+        for db_file in TRACKIO_DIR.glob("*.db"):
+            project_name = db_file.stem
+            projects.add(project_name)
+        return sorted(projects)
 
     @staticmethod
     def get_runs(project: str) -> list[str]:
         """Get list of all runs for a project."""
         db_path = SQLiteStorage.get_project_db_path(project)
-        if not os.path.exists(db_path):
+        if not db_path.exists():
             return []
 
-        with sqlite3.connect(db_path) as conn:
+        with SQLiteStorage._get_connection(db_path) as conn:
             cursor = conn.cursor()
             cursor.execute(
-                "SELECT DISTINCT run_name FROM metrics WHERE project_name = ?",
-                (project,),
+                "SELECT DISTINCT run_name FROM metrics",
             )
             return [row[0] for row in cursor.fetchall()]
 
+    @staticmethod
+    def get_max_steps_for_runs(project: str, runs: list[str]) -> dict[str, int]:
+        """Efficiently get the maximum step for multiple runs in a single query."""
+        db_path = SQLiteStorage.get_project_db_path(project)
+        if not db_path.exists():
+            return {run: 0 for run in runs}
+
+        with SQLiteStorage._get_connection(db_path) as conn:
+            cursor = conn.cursor()
+            placeholders = ",".join("?" * len(runs))
+            cursor.execute(
+                f"""
+                SELECT run_name, MAX(step) as max_step
+                FROM metrics
+                WHERE run_name IN ({placeholders})
+                GROUP BY run_name
+                """,
+                runs,
+            )
+
+            results = {run: 0 for run in runs}  # Default to 0 for runs with no data
+            for row in cursor.fetchall():
+                results[row["run_name"]] = row["max_step"]
+
+            return results
+
     def finish(self):
         """Cleanup when run is finished."""
         pass

tests/conftest.py

@@ -408,21 +408,7 @@ def update_last_steps(project, runs):
         if not project or not runs:
             return {}
 
-        last_steps = {}
-        for run in runs:
-            metrics = SQLiteStorage.get_metrics(project, run)
-            if metrics:
-                df = pd.DataFrame(metrics)
-                if "step" not in df.columns:
-                    df["step"] = range(len(df))
-                if not df.empty:
-                    last_steps[run] = df["step"].max().item()
-                else:
-                    last_steps[run] = 0
-            else:
-                last_steps[run] = 0
-
-        return last_steps
+        return SQLiteStorage.get_max_steps_for_runs(project, runs)
 
     timer.tick(
         fn=update_last_steps,

trackio/deploy.py

@@ -1,4 +1,3 @@
-import os
 import random
 import re
 import sys
@@ -9,7 +8,7 @@
 from huggingface_hub.constants import HF_HOME
 
 RESERVED_KEYS = ["project", "run", "timestamp", "step", "time"]
-TRACKIO_DIR = os.path.join(HF_HOME, "trackio")
+TRACKIO_DIR = Path(HF_HOME) / "trackio"
 
 TRACKIO_LOGO_PATH = str(Path(__file__).parent.joinpath("trackio_logo.png"))
 

trackio/sqlite_storage.py

@@ -1 +1 @@
-0.0.19
+0.0.20