Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+trackio_logo.png filter=lfs diff=lfs merge=lfs -text

__init__.py

@@ -0,0 +1,138 @@
+import os
+import webbrowser
+from pathlib import Path
+
+from gradio_client import Client
+
+from trackio import context_vars, deploy, utils
+from trackio.imports import import_csv
+from trackio.run import Run
+from trackio.sqlite_storage import SQLiteStorage
+from trackio.ui import demo
+from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_PATH
+
+__version__ = Path(__file__).parent.joinpath("version.txt").read_text().strip()
+
+__all__ = ["init", "log", "finish", "show", "import_csv"]
+
+
+config = {}
+
+
+def init(
+    project: str,
+    name: str | None = None,
+    space_id: str | None = None,
+    dataset_id: str | None = None,
+    config: dict | None = None,
+    resume: str = "never",
+) -> Run:
+    """
+    Creates a new Trackio project and returns a Run object.
+
+    Args:
+        project: The name of the project (can be an existing project to continue tracking or a new project to start tracking from scratch).
+        name: The name of the run (if not provided, a default name will be generated).
+        space_id: If provided, the project will be logged to a Hugging Face Space instead of a local directory. Should be a complete Space name like "username/reponame" or "orgname/reponame", or just "reponame" in which case the Space will be created in the currently-logged-in Hugging Face user's namespace. If the Space does not exist, it will be created. If the Space already exists, the project will be logged to it.
+        dataset_id: If provided, a persistent Hugging Face Dataset will be created and the metrics will be synced to it every 5 minutes. Should be a complete Dataset name like "username/datasetname" or "orgname/datasetname", or just "datasetname" in which case the Dataset will be created in the currently-logged-in Hugging Face user's namespace. If the Dataset does not exist, it will be created. If the Dataset already exists, the project will be appended to it. If not provided, the metrics will be logged to a local SQLite database, unless a `space_id` is provided, in which case a Dataset will be automatically created with the same name as the Space but with the "_dataset" suffix.
+        config: A dictionary of configuration options. Provided for compatibility with wandb.init()
+        resume: Controls how to handle resuming a run. Can be one of:
+            - "must": Must resume the run with the given name, raises error if run doesn't exist
+            - "allow": Resume the run if it exists, otherwise create a new run
+            - "never": Never resume a run, always create a new one
+    """
+    if not context_vars.current_server.get() and space_id is None:
+        _, url, _ = demo.launch(
+            show_api=False, inline=False, quiet=True, prevent_thread_lock=True
+        )
+        context_vars.current_server.set(url)
+    else:
+        url = context_vars.current_server.get()
+
+    space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
+
+    if (
+        context_vars.current_project.get() is None
+        or context_vars.current_project.get() != project
+    ):
+        print(f"* Trackio project initialized: {project}")
+
+        if dataset_id is not None:
+            os.environ["TRACKIO_DATASET_ID"] = dataset_id
+            print(
+                f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}"
+            )
+        if space_id is None:
+            print(f"* Trackio metrics logged to: {TRACKIO_DIR}")
+            utils.print_dashboard_instructions(project)
+        else:
+            deploy.create_space_if_not_exists(space_id, dataset_id)
+            print(
+                f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"
+            )
+    context_vars.current_project.set(project)
+
+    space_or_url = space_id if space_id else url
+    client = Client(space_or_url, verbose=False)
+
+    if resume == "must":
+        if name is None:
+            raise ValueError("Must provide a run name when resume='must'")
+        if name not in SQLiteStorage.get_runs(project):
+            raise ValueError(f"Run '{name}' does not exist in project '{project}'")
+    elif resume == "allow":
+        if name is not None and name in SQLiteStorage.get_runs(project):
+            print(f"* Resuming existing run: {name}")
+    elif resume == "never":
+        if name is not None and name in SQLiteStorage.get_runs(project):
+            name = None
+    else:
+        raise ValueError("resume must be one of: 'must', 'allow', or 'never'")
+
+    run = Run(project=project, client=client, name=name, config=config)
+    context_vars.current_run.set(run)
+    globals()["config"] = run.config
+    return run
+
+
+def log(metrics: dict) -> None:
+    """
+    Logs metrics to the current run.
+
+    Args:
+        metrics: A dictionary of metrics to log.
+    """
+    if context_vars.current_run.get() is None:
+        raise RuntimeError("Call trackio.init() before log().")
+    context_vars.current_run.get().log(metrics)
+
+
+def finish():
+    """
+    Finishes the current run.
+    """
+    if context_vars.current_run.get() is None:
+        raise RuntimeError("Call trackio.init() before finish().")
+    context_vars.current_run.get().finish()
+
+
+def show(project: str | None = None):
+    """
+    Launches the Trackio dashboard.
+
+    Args:
+        project: The name of the project whose runs to show. If not provided, all projects will be shown and the user can select one.
+    """
+    _, url, share_url = demo.launch(
+        show_api=False,
+        quiet=True,
+        inline=False,
+        prevent_thread_lock=True,
+        favicon_path=TRACKIO_LOGO_PATH,
+        allowed_paths=[TRACKIO_LOGO_PATH],
+    )
+    base_url = share_url + "/" if share_url else url
+    dashboard_url = base_url + f"?project={project}" if project else base_url
+    print(f"* Trackio UI launched at: {dashboard_url}")
+    webbrowser.open(dashboard_url)
+    utils.block_except_in_notebook()

cli.py

@@ -0,0 +1,26 @@
+import argparse
+
+from trackio import show
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Trackio CLI")
+    subparsers = parser.add_subparsers(dest="command")
+
+    ui_parser = subparsers.add_parser(
+        "show", help="Show the Trackio dashboard UI for a project"
+    )
+    ui_parser.add_argument(
+        "--project", required=False, help="Project name to show in the dashboard"
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "show":
+        show(args.project)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

context_vars.py

@@ -0,0 +1,19 @@
+import contextvars
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from trackio.run import Run
+    from trackio.sqlite_storage import CommitScheduler, DummyCommitScheduler
+
+current_run: contextvars.ContextVar["Run | None"] = contextvars.ContextVar(
+    "current_run", default=None
+)
+current_project: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "current_project", default=None
+)
+current_server: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "current_server", default=None
+)
+current_scheduler: contextvars.ContextVar[
+    "CommitScheduler | DummyCommitScheduler | None"
+] = contextvars.ContextVar("current_scheduler", default=None)

deploy.py

@@ -0,0 +1,138 @@
+import io
+import os
+import time
+from importlib.resources import files
+from pathlib import Path
+
+import gradio
+import huggingface_hub
+from gradio_client import Client, handle_file
+from httpx import ReadTimeout
+from huggingface_hub.errors import RepositoryNotFoundError
+
+from trackio.sqlite_storage import SQLiteStorage
+
+SPACE_URL = "https://huggingface.co/spaces/{space_id}"
+
+
+def deploy_as_space(
+    space_id: str,
+    dataset_id: str | None = None,
+):
+    if (
+        os.getenv("SYSTEM") == "spaces"
+    ):  # in case a repo with this function is uploaded to spaces
+        return
+
+    trackio_path = files("trackio")
+
+    hf_api = huggingface_hub.HfApi()
+    whoami = None
+    login = False
+    try:
+        whoami = hf_api.whoami()
+        if whoami["auth"]["accessToken"]["role"] != "write":
+            login = True
+    except OSError:
+        login = True
+    if login:
+        print("Need 'write' access token to create a Spaces repo.")
+        huggingface_hub.login(add_to_git_credential=False)
+        whoami = hf_api.whoami()
+
+    huggingface_hub.create_repo(
+        space_id,
+        space_sdk="gradio",
+        repo_type="space",
+        exist_ok=True,
+    )
+
+    with open(Path(trackio_path, "README.md"), "r") as f:
+        readme_content = f.read()
+        readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)
+        readme_buffer = io.BytesIO(readme_content.encode("utf-8"))
+        hf_api.upload_file(
+            path_or_fileobj=readme_buffer,
+            path_in_repo="README.md",
+            repo_id=space_id,
+            repo_type="space",
+        )
+
+    huggingface_hub.utils.disable_progress_bars()
+    hf_api.upload_folder(
+        repo_id=space_id,
+        repo_type="space",
+        folder_path=trackio_path,
+        ignore_patterns=["README.md"],
+    )
+
+    hf_token = huggingface_hub.utils.get_token()
+    if hf_token is not None:
+        huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)
+    if dataset_id is not None:
+        huggingface_hub.add_space_variable(space_id, "TRACKIO_DATASET_ID", dataset_id)
+
+
+def create_space_if_not_exists(
+    space_id: str,
+    dataset_id: str | None = None,
+) -> None:
+    """
+    Creates a new Hugging Face Space if it does not exist. If a dataset_id is provided, it will be added as a space variable.
+
+    Args:
+        space_id: The ID of the Space to create.
+        dataset_id: The ID of the Dataset to add to the Space.
+    """
+    if "/" not in space_id:
+        raise ValueError(
+            f"Invalid space ID: {space_id}. Must be in the format: username/reponame or orgname/reponame."
+        )
+    if dataset_id is not None and "/" not in dataset_id:
+        raise ValueError(
+            f"Invalid dataset ID: {dataset_id}. Must be in the format: username/datasetname or orgname/datasetname."
+        )
+    try:
+        huggingface_hub.repo_info(space_id, repo_type="space")
+        print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")
+        if dataset_id is not None:
+            huggingface_hub.add_space_variable(
+                space_id, "TRACKIO_DATASET_ID", dataset_id
+            )
+        return
+    except RepositoryNotFoundError:
+        pass
+
+    print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")
+    deploy_as_space(space_id, dataset_id)
+
+    client = None
+    for _ in range(30):
+        try:
+            client = Client(space_id, verbose=False)
+            if client:
+                break
+        except ReadTimeout:
+            print("* Space is not yet ready. Waiting 5 seconds...")
+            time.sleep(5)
+        except ValueError as e:
+            print(f"* Space gave error {e}. Trying again in 5 seconds...")
+            time.sleep(5)
+
+
+def upload_db_to_space(project: str, space_id: str) -> None:
+    """
+    Uploads the database of a local Trackio project to a Hugging Face Space.
+
+    Args:
+        project: The name of the project to upload.
+        space_id: The ID of the Space to upload to.
+    """
+    db_path = SQLiteStorage.get_project_db_path(project)
+    client = Client(space_id, verbose=False)
+    client.predict(
+        api_name="/upload_db_to_space",
+        project=project,
+        uploaded_db=handle_file(db_path),
+        hf_token=huggingface_hub.utils.get_token(),
+    )

dummy_commit_scheduler.py

@@ -0,0 +1,12 @@
+# A dummy object to fit the interface of huggingface_hub's CommitScheduler
+class DummyCommitSchedulerLock:
+    def __enter__(self):
+        return None
+
+    def __exit__(self, exception_type, exception_value, exception_traceback):
+        pass
+
+
+class DummyCommitScheduler:
+    def __init__(self):
+        self.lock = DummyCommitSchedulerLock()

imports.py

@@ -0,0 +1,118 @@
+import os
+from pathlib import Path
+
+import pandas as pd
+
+from trackio import deploy, utils
+from trackio.sqlite_storage import SQLiteStorage
+
+
+def import_csv(
+    csv_path: str,
+    project: str,
+    name: str | None = None,
+    space_id: str | None = None,
+    dataset_id: str | None = None,
+) -> None:
+    """
+    Imports a CSV file into a Trackio project. The CSV file must contain a "step" column, may optionally
+    contain a "timestamp" column, and any other columns will be treated as metrics. It should also include
+    a header row with the column names.
+
+    TODO: call init() and return a Run object so that the user can continue to log metrics to it.
+
+    Args:
+        csv_path: The str or Path to the CSV file to import.
+        project: The name of the project to import the CSV file into. Must not be an existing project.
+        name: The name of the Run to import the CSV file into. If not provided, a default name will be generated.
+        name: The name of the run (if not provided, a default name will be generated).
+        space_id: If provided, the project will be logged to a Hugging Face Space instead of a local directory. Should be a complete Space name like "username/reponame" or "orgname/reponame", or just "reponame" in which case the Space will be created in the currently-logged-in Hugging Face user's namespace. If the Space does not exist, it will be created. If the Space already exists, the project will be logged to it.
+        dataset_id: If provided, a persistent Hugging Face Dataset will be created and the metrics will be synced to it every 5 minutes. Should be a complete Dataset name like "username/datasetname" or "orgname/datasetname", or just "datasetname" in which case the Dataset will be created in the currently-logged-in Hugging Face user's namespace. If the Dataset does not exist, it will be created. If the Dataset already exists, the project will be appended to it. If not provided, the metrics will be logged to a local SQLite database, unless a `space_id` is provided, in which case a Dataset will be automatically created with the same name as the Space but with the "_dataset" suffix.
+    """
+    if SQLiteStorage.get_runs(project):
+        raise ValueError(
+            f"Project '{project}' already exists. Cannot import CSV into existing project."
+        )
+
+    csv_path = Path(csv_path)
+    if not csv_path.exists():
+        raise FileNotFoundError(f"CSV file not found: {csv_path}")
+
+    df = pd.read_csv(csv_path)
+    if df.empty:
+        raise ValueError("CSV file is empty")
+
+    column_mapping = utils.simplify_column_names(df.columns.tolist())
+    df = df.rename(columns=column_mapping)
+
+    step_column = None
+    for col in df.columns:
+        if col.lower() == "step":
+            step_column = col
+            break
+
+    if step_column is None:
+        raise ValueError("CSV file must contain a 'step' or 'Step' column")
+
+    if name is None:
+        name = csv_path.stem
+
+    metrics_list = []
+    steps = []
+    timestamps = []
+
+    numeric_columns = []
+    for column in df.columns:
+        if column == step_column:
+            continue
+        if column == "timestamp":
+            continue
+
+        try:
+            pd.to_numeric(df[column], errors="raise")
+            numeric_columns.append(column)
+        except (ValueError, TypeError):
+            continue
+
+    for _, row in df.iterrows():
+        metrics = {}
+        for column in numeric_columns:
+            if pd.notna(row[column]):
+                metrics[column] = float(row[column])
+
+        if metrics:
+            metrics_list.append(metrics)
+            steps.append(int(row[step_column]))
+
+            if "timestamp" in df.columns and pd.notna(row["timestamp"]):
+                timestamps.append(str(row["timestamp"]))
+            else:
+                timestamps.append("")
+
+    if metrics_list:
+        SQLiteStorage.bulk_log(
+            project=project,
+            run=name,
+            metrics_list=metrics_list,
+            steps=steps,
+            timestamps=timestamps,
+        )
+
+    print(
+        f"* Imported {len(metrics_list)} rows from {csv_path} into project '{project}' as run '{name}'"
+    )
+    print(f"* Metrics found: {', '.join(metrics_list[0].keys())}")
+
+    space_id, dataset_id = utils.preprocess_space_and_dataset_ids(space_id, dataset_id)
+    if dataset_id is not None:
+        os.environ["TRACKIO_DATASET_ID"] = dataset_id
+        print(f"* Trackio metrics will be synced to Hugging Face Dataset: {dataset_id}")
+
+    if space_id is None:
+        utils.print_dashboard_instructions(project)
+    else:
+        deploy.create_space_if_not_exists(space_id, dataset_id)
+        deploy.upload_db_to_space(project, space_id)
+        print(
+            f"* View dashboard by going to: {deploy.SPACE_URL.format(space_id=space_id)}"
+        )

run.py

@@ -0,0 +1,36 @@
+import huggingface_hub
+from gradio_client import Client
+
+from trackio.utils import RESERVED_KEYS, generate_readable_name
+
+
+class Run:
+    def __init__(
+        self,
+        project: str,
+        client: Client,
+        name: str | None = None,
+        config: dict | None = None,
+    ):
+        self.project = project
+        self.client = client
+        self.name = name or generate_readable_name()
+        self.config = config or {}
+
+    def log(self, metrics: dict):
+        for k in metrics.keys():
+            if k in RESERVED_KEYS or k.startswith("__"):
+                raise ValueError(
+                    f"Please do not use this reserved key as a metric: {k}"
+                )
+        self.client.predict(
+            api_name="/log",
+            project=self.project,
+            run=self.name,
+            metrics=metrics,
+            hf_token=huggingface_hub.utils.get_token(),
+        )
+
+    def finish(self):
+        """Cleanup when run is finished."""
+        pass

sqlite_storage.py

@@ -0,0 +1,241 @@
+import glob
+import json
+import os
+import sqlite3
+from datetime import datetime
+
+from huggingface_hub import CommitScheduler
+
+try:
+    from trackio.context_vars import current_scheduler
+    from trackio.dummy_commit_scheduler import DummyCommitScheduler
+    from trackio.utils import TRACKIO_DIR
+except:  # noqa: E722
+    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:
+        """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")
+
+    @staticmethod
+    def init_db(project: str) -> str:
+        """
+        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)
+        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
+                    )
+                """)
+                conn.commit()
+        return db_path
+
+    @staticmethod
+    def get_scheduler():
+        """
+        Get the scheduler for the database based on the environment variables.
+        This applies to both local and Spaces.
+        """
+        if current_scheduler.get() is not None:
+            return current_scheduler.get()
+        hf_token = os.environ.get("HF_TOKEN")
+        dataset_id = os.environ.get("TRACKIO_DATASET_ID")
+        if dataset_id is None:
+            scheduler = DummyCommitScheduler()
+        else:
+            scheduler = CommitScheduler(
+                repo_id=dataset_id,
+                repo_type="dataset",
+                folder_path=TRACKIO_DIR,
+                private=True,
+                squash_history=True,
+                token=hf_token,
+            )
+        current_scheduler.set(scheduler)
+        return scheduler
+
+    @staticmethod
+    def log(project: str, run: str, metrics: dict):
+        """
+        Safely log metrics to the database. Before logging, this method will ensure the database exists
+        and is set up with the correct tables. It also uses the scheduler to lock the database so
+        that there is no race condition when logging / syncing to the Hugging Face Dataset.
+        """
+        db_path = SQLiteStorage.init_db(project)
+
+        with SQLiteStorage.get_scheduler().lock:
+            with sqlite3.connect(db_path) as conn:
+                cursor = conn.cursor()
+
+                cursor.execute(
+                    """
+                    SELECT MAX(step) 
+                    FROM metrics 
+                    WHERE project_name = ? AND run_name = ?
+                    """,
+                    (project, run),
+                )
+                last_step = cursor.fetchone()[0]
+                current_step = 0 if last_step is None else last_step + 1
+
+                current_timestamp = datetime.now().isoformat()
+
+                cursor.execute(
+                    """
+                    INSERT INTO metrics 
+                    (timestamp, project_name, run_name, step, metrics)
+                    VALUES (?, ?, ?, ?, ?)
+                    """,
+                    (
+                        current_timestamp,
+                        project,
+                        run,
+                        current_step,
+                        json.dumps(metrics),
+                    ),
+                )
+                conn.commit()
+
+    @staticmethod
+    def bulk_log(
+        project: str,
+        run: str,
+        metrics_list: list[dict],
+        steps: list[int] | None = None,
+        timestamps: list[str] | None = None,
+    ):
+        """Bulk log metrics to the database with specified steps and timestamps."""
+        if not metrics_list:
+            return
+
+        if steps is None:
+            steps = list(range(len(metrics_list)))
+
+        if timestamps is None:
+            timestamps = [datetime.now().isoformat()] * len(metrics_list)
+
+        if len(metrics_list) != len(steps) or len(metrics_list) != len(timestamps):
+            raise ValueError(
+                "metrics_list, steps, and timestamps must have the same length"
+            )
+
+        db_path = SQLiteStorage.init_db(project)
+        with SQLiteStorage.get_scheduler().lock:
+            with sqlite3.connect(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),
+                        )
+                    )
+
+                cursor.executemany(
+                    """
+                    INSERT INTO metrics 
+                    (timestamp, project_name, run_name, step, metrics)
+                    VALUES (?, ?, ?, ?, ?)
+                    """,
+                    data,
+                )
+                conn.commit()
+
+    @staticmethod
+    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):
+            return []
+
+        with sqlite3.connect(db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                SELECT timestamp, step, metrics
+                FROM metrics
+                WHERE project_name = ? AND run_name = ?
+                ORDER BY timestamp
+                """,
+                (project, run),
+            )
+            rows = cursor.fetchall()
+
+            results = []
+            for row in rows:
+                timestamp, step, metrics_json = row
+                metrics = json.loads(metrics_json)
+                metrics["timestamp"] = timestamp
+                metrics["step"] = 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
+
+        return list(set(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):
+            return []
+
+        with sqlite3.connect(db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT DISTINCT run_name FROM metrics WHERE project_name = ?",
+                (project,),
+            )
+            return [row[0] for row in cursor.fetchall()]
+
+    def finish(self):
+        """Cleanup when run is finished."""
+        pass

ui.py

@@ -0,0 +1,503 @@
+import os
+import shutil
+from typing import Any
+
+import gradio as gr
+import huggingface_hub as hf
+import pandas as pd
+
+HfApi = hf.HfApi()
+
+try:
+    from trackio.sqlite_storage import SQLiteStorage
+    from trackio.utils import RESERVED_KEYS, TRACKIO_LOGO_PATH
+except:  # noqa: E722
+    from sqlite_storage import SQLiteStorage
+    from utils import RESERVED_KEYS, TRACKIO_LOGO_PATH
+
+css = """
+#run-cb .wrap {
+    gap: 2px;
+}
+#run-cb .wrap label {
+    line-height: 1;
+    padding: 6px;
+}
+"""
+
+COLOR_PALETTE = [
+    "#3B82F6",
+    "#EF4444",
+    "#10B981",
+    "#F59E0B",
+    "#8B5CF6",
+    "#EC4899",
+    "#06B6D4",
+    "#84CC16",
+    "#F97316",
+    "#6366F1",
+]
+
+
+def get_color_mapping(runs: list[str], smoothing: bool) -> dict[str, str]:
+    """Generate color mapping for runs, with transparency for original data when smoothing is enabled."""
+    color_map = {}
+
+    for i, run in enumerate(runs):
+        base_color = COLOR_PALETTE[i % len(COLOR_PALETTE)]
+
+        if smoothing:
+            color_map[f"{run}_smoothed"] = base_color
+            color_map[f"{run}_original"] = base_color + "4D"
+        else:
+            color_map[run] = base_color
+
+    return color_map
+
+
+def get_projects(request: gr.Request):
+    dataset_id = os.environ.get("TRACKIO_DATASET_ID")
+    projects = SQLiteStorage.get_projects()
+    if project := request.query_params.get("project"):
+        interactive = False
+    else:
+        interactive = True
+        project = projects[0] if projects else None
+    return gr.Dropdown(
+        label="Project",
+        choices=projects,
+        value=project,
+        allow_custom_value=True,
+        interactive=interactive,
+        info=f"&#x21bb; Synced to <a href='https://huggingface.co/datasets/{dataset_id}' target='_blank'>{dataset_id}</a> every 5 min"
+        if dataset_id
+        else None,
+    )
+
+
+def get_runs(project) -> list[str]:
+    if not project:
+        return []
+    return SQLiteStorage.get_runs(project)
+
+
+def get_available_metrics(project: str, runs: list[str]) -> list[str]:
+    """Get all available metrics across all runs for x-axis selection."""
+    if not project or not runs:
+        return ["step", "time"]
+
+    all_metrics = set()
+    for run in runs:
+        metrics = SQLiteStorage.get_metrics(project, run)
+        if metrics:
+            df = pd.DataFrame(metrics)
+            numeric_cols = df.select_dtypes(include="number").columns
+            numeric_cols = [c for c in numeric_cols if c not in RESERVED_KEYS]
+            all_metrics.update(numeric_cols)
+
+    # Always include step and time as options
+    all_metrics.add("step")
+    all_metrics.add("time")
+
+    # Sort metrics by prefix
+    sorted_metrics = sort_metrics_by_prefix(list(all_metrics))
+
+    # Put step and time at the beginning
+    result = ["step", "time"]
+    for metric in sorted_metrics:
+        if metric not in result:
+            result.append(metric)
+
+    return result
+
+
+def load_run_data(project: str | None, run: str | None, smoothing: bool, x_axis: str):
+    if not project or not run:
+        return None
+    metrics = SQLiteStorage.get_metrics(project, run)
+    if not metrics:
+        return None
+    df = pd.DataFrame(metrics)
+
+    if "step" not in df.columns:
+        df["step"] = range(len(df))
+
+    if x_axis == "time" and "timestamp" in df.columns:
+        df["timestamp"] = pd.to_datetime(df["timestamp"])
+        first_timestamp = df["timestamp"].min()
+        df["time"] = (df["timestamp"] - first_timestamp).dt.total_seconds()
+        x_column = "time"
+    elif x_axis == "step":
+        x_column = "step"
+    else:
+        x_column = x_axis
+
+    if smoothing:
+        numeric_cols = df.select_dtypes(include="number").columns
+        numeric_cols = [c for c in numeric_cols if c not in RESERVED_KEYS]
+
+        df_original = df.copy()
+        df_original["run"] = f"{run}_original"
+        df_original["data_type"] = "original"
+
+        df_smoothed = df.copy()
+        window_size = max(3, min(10, len(df) // 10))  # Adaptive window size
+        df_smoothed[numeric_cols] = (
+            df_smoothed[numeric_cols]
+            .rolling(window=window_size, center=True, min_periods=1)
+            .mean()
+        )
+        df_smoothed["run"] = f"{run}_smoothed"
+        df_smoothed["data_type"] = "smoothed"
+
+        combined_df = pd.concat([df_original, df_smoothed], ignore_index=True)
+        combined_df["x_axis"] = x_column
+        return combined_df
+    else:
+        df["run"] = run
+        df["data_type"] = "original"
+        df["x_axis"] = x_column
+        return df
+
+
+def update_runs(project, filter_text, user_interacted_with_runs=False):
+    if project is None:
+        runs = []
+        num_runs = 0
+    else:
+        runs = get_runs(project)
+        num_runs = len(runs)
+        if filter_text:
+            runs = [r for r in runs if filter_text in r]
+    if not user_interacted_with_runs:
+        return gr.CheckboxGroup(choices=runs, value=runs), gr.Textbox(
+            label=f"Runs ({num_runs})"
+        )
+    else:
+        return gr.CheckboxGroup(choices=runs), gr.Textbox(label=f"Runs ({num_runs})")
+
+
+def filter_runs(project, filter_text):
+    runs = get_runs(project)
+    runs = [r for r in runs if filter_text in r]
+    return gr.CheckboxGroup(choices=runs, value=runs)
+
+
+def update_x_axis_choices(project, runs):
+    """Update x-axis dropdown choices based on available metrics."""
+    available_metrics = get_available_metrics(project, runs)
+    return gr.Dropdown(
+        label="X-axis",
+        choices=available_metrics,
+        value="step",
+    )
+
+
+def toggle_timer(cb_value):
+    if cb_value:
+        return gr.Timer(active=True)
+    else:
+        return gr.Timer(active=False)
+
+
+def check_auth(hf_token: str | None) -> None:
+    if os.getenv("SYSTEM") == "spaces":  # if we are running in Spaces
+        # check auth token passed in
+        if hf_token is None:
+            raise PermissionError(
+                "Expected a HF_TOKEN to be provided when logging to a Space"
+            )
+        who = HfApi.whoami(hf_token)
+        access_token = who["auth"]["accessToken"]
+        owner_name = os.getenv("SPACE_AUTHOR_NAME")
+        repo_name = os.getenv("SPACE_REPO_NAME")
+        # make sure the token user is either the author of the space,
+        # or is a member of an org that is the author.
+        orgs = [o["name"] for o in who["orgs"]]
+        if owner_name != who["name"] and owner_name not in orgs:
+            raise PermissionError(
+                "Expected the provided hf_token to be the user owner of the space, or be a member of the org owner of the space"
+            )
+        # reject fine-grained tokens without specific repo access
+        if access_token["role"] == "fineGrained":
+            matched = False
+            for item in access_token["fineGrained"]["scoped"]:
+                if (
+                    item["entity"]["type"] == "space"
+                    and item["entity"]["name"] == f"{owner_name}/{repo_name}"
+                    and "repo.write" in item["permissions"]
+                ):
+                    matched = True
+                    break
+                if (
+                    item["entity"]["type"] == "user"
+                    and item["entity"]["name"] == owner_name
+                    and "repo.write" in item["permissions"]
+                ):
+                    matched = True
+                    break
+            if not matched:
+                raise PermissionError(
+                    "Expected the provided hf_token with fine grained permissions to provide write access to the space"
+                )
+        # reject read-only tokens
+        elif access_token["role"] != "write":
+            raise PermissionError(
+                "Expected the provided hf_token to provide write permissions"
+            )
+
+
+def upload_db_to_space(
+    project: str, uploaded_db: gr.FileData, hf_token: str | None
+) -> None:
+    check_auth(hf_token)
+    db_project_path = SQLiteStorage.get_project_db_path(project)
+    if os.path.exists(db_project_path):
+        raise gr.Error(
+            f"Trackio database file already exists for project {project}, cannot overwrite."
+        )
+    os.makedirs(os.path.dirname(db_project_path), exist_ok=True)
+    shutil.copy(uploaded_db["path"], db_project_path)
+
+
+def log(
+    project: str,
+    run: str,
+    metrics: dict[str, Any],
+    hf_token: str | None,
+) -> None:
+    check_auth(hf_token)
+    SQLiteStorage.log(project=project, run=run, metrics=metrics)
+
+
+def sort_metrics_by_prefix(metrics: list[str]) -> list[str]:
+    """
+    Sort metrics by grouping prefixes together.
+    Metrics without prefixes come first, then grouped by prefix.
+
+    Example:
+    Input: ["train/loss", "loss", "train/acc", "val/loss"]
+    Output: ["loss", "train/acc", "train/loss", "val/loss"]
+    """
+    no_prefix = []
+    with_prefix = []
+
+    for metric in metrics:
+        if "/" in metric:
+            with_prefix.append(metric)
+        else:
+            no_prefix.append(metric)
+
+    no_prefix.sort()
+
+    prefix_groups = {}
+    for metric in with_prefix:
+        prefix = metric.split("/")[0]
+        if prefix not in prefix_groups:
+            prefix_groups[prefix] = []
+        prefix_groups[prefix].append(metric)
+
+    sorted_with_prefix = []
+    for prefix in sorted(prefix_groups.keys()):
+        sorted_with_prefix.extend(sorted(prefix_groups[prefix]))
+
+    return no_prefix + sorted_with_prefix
+
+
+def configure(request: gr.Request):
+    sidebar_param = request.query_params.get("sidebar")
+    match sidebar_param:
+        case "collapsed":
+            sidebar = gr.Sidebar(open=False, visible=True)
+        case "hidden":
+            sidebar = gr.Sidebar(visible=False)
+        case _:
+            sidebar = gr.Sidebar(visible=True)
+
+    if metrics := request.query_params.get("metrics"):
+        return metrics.split(","), sidebar
+    else:
+        return [], sidebar
+
+
+with gr.Blocks(theme="citrus", title="Trackio Dashboard", css=css) as demo:
+    with gr.Sidebar(visible=False) as sidebar:
+        gr.Markdown(
+            f"<div style='display: flex; align-items: center; gap: 8px;'><img src='/gradio_api/file={TRACKIO_LOGO_PATH}' width='32' height='32'><span style='font-size: 2em; font-weight: bold;'>Trackio</span></div>"
+        )
+        project_dd = gr.Dropdown(label="Project", allow_custom_value=True)
+        run_tb = gr.Textbox(label="Runs", placeholder="Type to filter...")
+        run_cb = gr.CheckboxGroup(
+            label="Runs", choices=[], interactive=True, elem_id="run-cb"
+        )
+        gr.HTML("<hr>")
+        realtime_cb = gr.Checkbox(label="Refresh metrics realtime", value=True)
+        smoothing_cb = gr.Checkbox(label="Smooth metrics", value=True)
+        x_axis_dd = gr.Dropdown(
+            label="X-axis",
+            choices=["step", "time"],
+            value="step",
+        )
+
+    timer = gr.Timer(value=1)
+    metrics_subset = gr.State([])
+    user_interacted_with_run_cb = gr.State(False)
+
+    gr.on([demo.load], fn=configure, outputs=[metrics_subset, sidebar])
+    gr.on(
+        [demo.load],
+        fn=get_projects,
+        outputs=project_dd,
+        show_progress="hidden",
+    )
+    gr.on(
+        [timer.tick],
+        fn=update_runs,
+        inputs=[project_dd, run_tb, user_interacted_with_run_cb],
+        outputs=[run_cb, run_tb],
+        show_progress="hidden",
+    )
+    gr.on(
+        [demo.load, project_dd.change],
+        fn=update_runs,
+        inputs=[project_dd, run_tb],
+        outputs=[run_cb, run_tb],
+        show_progress="hidden",
+    )
+    gr.on(
+        [demo.load, project_dd.change, run_cb.change],
+        fn=update_x_axis_choices,
+        inputs=[project_dd, run_cb],
+        outputs=x_axis_dd,
+        show_progress="hidden",
+    )
+
+    realtime_cb.change(
+        fn=toggle_timer,
+        inputs=realtime_cb,
+        outputs=timer,
+        api_name="toggle_timer",
+    )
+    run_cb.input(
+        fn=lambda: True,
+        outputs=user_interacted_with_run_cb,
+    )
+    run_tb.input(
+        fn=filter_runs,
+        inputs=[project_dd, run_tb],
+        outputs=run_cb,
+    )
+
+    gr.api(
+        fn=upload_db_to_space,
+        api_name="upload_db_to_space",
+    )
+    gr.api(
+        fn=log,
+        api_name="log",
+    )
+
+    x_lim = gr.State(None)
+    last_steps = gr.State({})
+
+    def update_x_lim(select_data: gr.SelectData):
+        return select_data.index
+
+    def update_last_steps(project, runs):
+        """Update the last step from all runs to detect when new data is available."""
+        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
+
+    timer.tick(
+        fn=update_last_steps,
+        inputs=[project_dd, run_cb],
+        outputs=last_steps,
+        show_progress="hidden",
+    )
+
+    @gr.render(
+        triggers=[
+            demo.load,
+            run_cb.change,
+            last_steps.change,
+            smoothing_cb.change,
+            x_lim.change,
+            x_axis_dd.change,
+        ],
+        inputs=[project_dd, run_cb, smoothing_cb, metrics_subset, x_lim, x_axis_dd],
+        show_progress="hidden",
+    )
+    def update_dashboard(project, runs, smoothing, metrics_subset, x_lim_value, x_axis):
+        dfs = []
+        original_runs = runs.copy()
+
+        for run in runs:
+            df = load_run_data(project, run, smoothing, x_axis)
+            if df is not None:
+                dfs.append(df)
+
+        if dfs:
+            master_df = pd.concat(dfs, ignore_index=True)
+        else:
+            master_df = pd.DataFrame()
+
+        if master_df.empty:
+            return
+
+        x_column = "step"
+        if dfs and not dfs[0].empty and "x_axis" in dfs[0].columns:
+            x_column = dfs[0]["x_axis"].iloc[0]
+
+        numeric_cols = master_df.select_dtypes(include="number").columns
+        numeric_cols = [c for c in numeric_cols if c not in RESERVED_KEYS]
+        if metrics_subset:
+            numeric_cols = [c for c in numeric_cols if c in metrics_subset]
+
+        numeric_cols = sort_metrics_by_prefix(list(numeric_cols))
+        color_map = get_color_mapping(original_runs, smoothing)
+
+        with gr.Row(key="row"):
+            for metric_idx, metric_name in enumerate(numeric_cols):
+                metric_df = master_df.dropna(subset=[metric_name])
+                if not metric_df.empty:
+                    plot = gr.LinePlot(
+                        metric_df,
+                        x=x_column,
+                        y=metric_name,
+                        color="run" if "run" in metric_df.columns else None,
+                        color_map=color_map,
+                        title=metric_name,
+                        key=f"plot-{metric_idx}",
+                        preserved_by_key=None,
+                        x_lim=x_lim_value,
+                        y_lim=[
+                            metric_df[metric_name].min(),
+                            metric_df[metric_name].max(),
+                        ],
+                        show_fullscreen_button=True,
+                        min_width=400,
+                    )
+                plot.select(update_x_lim, outputs=x_lim, key=f"select-{metric_idx}")
+                plot.double_click(
+                    lambda: None, outputs=x_lim, key=f"double-{metric_idx}"
+                )
+
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=[TRACKIO_LOGO_PATH], show_api=False, show_error=True)

utils.py

@@ -0,0 +1,278 @@
+import os
+import random
+import re
+import sys
+import time
+from pathlib import Path
+
+import huggingface_hub
+from huggingface_hub.constants import HF_HOME
+
+RESERVED_KEYS = ["project", "run", "timestamp", "step", "time"]
+TRACKIO_DIR = os.path.join(HF_HOME, "trackio")
+
+TRACKIO_LOGO_PATH = str(Path(__file__).parent.joinpath("trackio_logo.png"))
+
+
+def generate_readable_name():
+    """
+    Generates a random, readable name like "dainty-sunset-1"
+    """
+    adjectives = [
+        "dainty",
+        "brave",
+        "calm",
+        "eager",
+        "fancy",
+        "gentle",
+        "happy",
+        "jolly",
+        "kind",
+        "lively",
+        "merry",
+        "nice",
+        "proud",
+        "quick",
+        "silly",
+        "tidy",
+        "witty",
+        "zealous",
+        "bright",
+        "shy",
+        "bold",
+        "clever",
+        "daring",
+        "elegant",
+        "faithful",
+        "graceful",
+        "honest",
+        "inventive",
+        "jovial",
+        "keen",
+        "lucky",
+        "modest",
+        "noble",
+        "optimistic",
+        "patient",
+        "quirky",
+        "resourceful",
+        "sincere",
+        "thoughtful",
+        "upbeat",
+        "valiant",
+        "warm",
+        "youthful",
+        "zesty",
+        "adventurous",
+        "breezy",
+        "cheerful",
+        "delightful",
+        "energetic",
+        "fearless",
+        "glad",
+        "hopeful",
+        "imaginative",
+        "joyful",
+        "kindly",
+        "luminous",
+        "mysterious",
+        "neat",
+        "outgoing",
+        "playful",
+        "radiant",
+        "spirited",
+        "tranquil",
+        "unique",
+        "vivid",
+        "wise",
+        "zany",
+        "artful",
+        "bubbly",
+        "charming",
+        "dazzling",
+        "earnest",
+        "festive",
+        "gentlemanly",
+        "hearty",
+        "intrepid",
+        "jubilant",
+        "knightly",
+        "lively",
+        "magnetic",
+        "nimble",
+        "orderly",
+        "peaceful",
+        "quick-witted",
+        "robust",
+        "sturdy",
+        "trusty",
+        "upstanding",
+        "vibrant",
+        "whimsical",
+    ]
+    nouns = [
+        "sunset",
+        "forest",
+        "river",
+        "mountain",
+        "breeze",
+        "meadow",
+        "ocean",
+        "valley",
+        "sky",
+        "field",
+        "cloud",
+        "star",
+        "rain",
+        "leaf",
+        "stone",
+        "flower",
+        "bird",
+        "tree",
+        "wave",
+        "trail",
+        "island",
+        "desert",
+        "hill",
+        "lake",
+        "pond",
+        "grove",
+        "canyon",
+        "reef",
+        "bay",
+        "peak",
+        "glade",
+        "marsh",
+        "cliff",
+        "dune",
+        "spring",
+        "brook",
+        "cave",
+        "plain",
+        "ridge",
+        "wood",
+        "blossom",
+        "petal",
+        "root",
+        "branch",
+        "seed",
+        "acorn",
+        "pine",
+        "willow",
+        "cedar",
+        "elm",
+        "falcon",
+        "eagle",
+        "sparrow",
+        "robin",
+        "owl",
+        "finch",
+        "heron",
+        "crane",
+        "duck",
+        "swan",
+        "fox",
+        "wolf",
+        "bear",
+        "deer",
+        "moose",
+        "otter",
+        "beaver",
+        "lynx",
+        "hare",
+        "badger",
+        "butterfly",
+        "bee",
+        "ant",
+        "beetle",
+        "dragonfly",
+        "firefly",
+        "ladybug",
+        "moth",
+        "spider",
+        "worm",
+        "coral",
+        "kelp",
+        "shell",
+        "pebble",
+        "boulder",
+        "cobble",
+        "sand",
+        "wavelet",
+        "tide",
+        "current",
+    ]
+    adjective = random.choice(adjectives)
+    noun = random.choice(nouns)
+    number = random.randint(1, 99)
+    return f"{adjective}-{noun}-{number}"
+
+
+def block_except_in_notebook():
+    in_notebook = bool(getattr(sys, "ps1", sys.flags.interactive))
+    if in_notebook:
+        return
+    try:
+        while True:
+            time.sleep(0.1)
+    except (KeyboardInterrupt, OSError):
+        print("Keyboard interruption in main thread... closing dashboard.")
+
+
+def simplify_column_names(columns: list[str]) -> dict[str, str]:
+    """
+    Simplifies column names to first 10 alphanumeric or "/" characters with unique suffixes.
+
+    Args:
+        columns: List of original column names
+
+    Returns:
+        Dictionary mapping original column names to simplified names
+    """
+    simplified_names = {}
+    used_names = set()
+
+    for col in columns:
+        alphanumeric = re.sub(r"[^a-zA-Z0-9/]", "", col)
+        base_name = alphanumeric[:10] if alphanumeric else f"col_{len(used_names)}"
+
+        final_name = base_name
+        suffix = 1
+        while final_name in used_names:
+            final_name = f"{base_name}_{suffix}"
+            suffix += 1
+
+        simplified_names[col] = final_name
+        used_names.add(final_name)
+
+    return simplified_names
+
+
+def print_dashboard_instructions(project: str) -> None:
+    """
+    Prints instructions for viewing the Trackio dashboard.
+
+    Args:
+        project: The name of the project to show dashboard for.
+    """
+    YELLOW = "\033[93m"
+    BOLD = "\033[1m"
+    RESET = "\033[0m"
+
+    print("* View dashboard by running in your terminal:")
+    print(f'{BOLD}{YELLOW}trackio show --project "{project}"{RESET}')
+    print(f'* or by running in Python: trackio.show(project="{project}")')
+
+
+def preprocess_space_and_dataset_ids(
+    space_id: str | None, dataset_id: str | None
+) -> tuple[str | None, str | None]:
+    if space_id is not None and "/" not in space_id:
+        username = huggingface_hub.whoami()["name"]
+        space_id = f"{username}/{space_id}"
+    if dataset_id is not None and "/" not in dataset_id:
+        username = huggingface_hub.whoami()["name"]
+        dataset_id = f"{username}/{dataset_id}"
+    if space_id is not None and dataset_id is None:
+        dataset_id = f"{space_id}_dataset"
+    return space_id, dataset_id

version.txt

@@ -0,0 +1 @@
+0.0.17