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,155 @@
+import contextvars
+import time
+import webbrowser
+from pathlib import Path
+
+import huggingface_hub
+from gradio_client import Client
+from httpx import ReadTimeout
+from huggingface_hub.errors import RepositoryNotFoundError
+
+from trackio.deploy import deploy_as_space
+from trackio.run import Run
+from trackio.ui import demo
+from trackio.utils import TRACKIO_DIR, TRACKIO_LOGO_PATH, block_except_in_notebook
+
+__version__ = Path(__file__).parent.joinpath("version.txt").read_text().strip()
+
+
+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
+)
+
+config = {}
+SPACE_URL = "https://huggingface.co/spaces/{space_id}"
+
+
+def init(
+    project: str,
+    name: str | None = None,
+    space_id: str | None = None,
+    config: dict | None = None,
+) -> 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". If the Space does not exist, it will be created. If the Space already exists, the project will be logged to it.
+        config: A dictionary of configuration options. Provided for compatibility with wandb.init()
+    """
+    if not current_server.get() and space_id is None:
+        _, url, _ = demo.launch(
+            show_api=False, inline=False, quiet=True, prevent_thread_lock=True
+        )
+        current_server.set(url)
+    else:
+        url = current_server.get()
+
+    if current_project.get() is None or current_project.get() != project:
+        print(f"* Trackio project initialized: {project}")
+
+        if space_id is None:
+            print(f"* Trackio metrics logged to: {TRACKIO_DIR}")
+            print(
+                f'\n* View dashboard by running in your terminal: trackio show --project "{project}"'
+            )
+            print(f'* or by running in Python: trackio.show(project="{project}")')
+        else:
+            create_space_if_not_exists(space_id)
+            print(
+                f"* View dashboard by going to: {SPACE_URL.format(space_id=space_id)}"
+            )
+    current_project.set(project)
+
+    space_or_url = space_id if space_id else url
+    client = Client(space_or_url, verbose=False)
+    run = Run(project=project, client=client, name=name, config=config)
+    current_run.set(run)
+    globals()["config"] = run.config
+    return run
+
+
+def create_space_if_not_exists(space_id: str) -> None:
+    """
+    Creates a new Hugging Face Space if it does not exist.
+
+    Args:
+        space_id: The ID of the Space to create.
+    """
+    if "/" not in space_id:
+        raise ValueError(
+            f"Invalid space ID: {space_id}. Must be in the format: username/reponame."
+        )
+    try:
+        huggingface_hub.repo_info(space_id, repo_type="space")
+        print(f"* Found existing space: {SPACE_URL.format(space_id=space_id)}")
+        return
+    except RepositoryNotFoundError:
+        pass
+
+    print(f"* Creating new space: {SPACE_URL.format(space_id=space_id)}")
+    deploy_as_space(space_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 log(metrics: dict) -> None:
+    """
+    Logs metrics to the current run.
+
+    Args:
+        metrics: A dictionary of metrics to log.
+    """
+    if current_run.get() is None:
+        raise RuntimeError("Call trackio.init() before log().")
+    current_run.get().log(metrics)
+
+
+def finish():
+    """
+    Finishes the current run.
+    """
+    if current_run.get() is None:
+        raise RuntimeError("Call trackio.init() before finish().")
+    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)
+    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()

deploy.py

@@ -0,0 +1,57 @@
+import io
+import os
+from importlib.resources import files
+from pathlib import Path
+
+import gradio
+import huggingface_hub
+
+
+def deploy_as_space(title: str):
+    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()
+
+    space_id = huggingface_hub.create_repo(
+        title,
+        space_sdk="gradio",
+        repo_type="space",
+        exist_ok=True,
+    ).repo_id
+    assert space_id == title  # not sure why these would differ
+
+    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"],
+    )

run.py

@@ -0,0 +1,25 @@
+from gradio_client import Client
+
+from trackio.utils import 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):
+        self.client.predict(
+            api_name="/log", project=self.project, run=self.name, metrics=metrics
+        )
+
+    def finish(self):
+        pass

sqlite_storage.py

@@ -0,0 +1,123 @@
+import json
+import os
+import sqlite3
+
+try:
+    from trackio.utils import RESERVED_KEYS, TRACKIO_DIR
+except:  # noqa: E722
+    from utils import RESERVED_KEYS, TRACKIO_DIR
+
+
+class SQLiteStorage:
+    def __init__(self, project: str, name: str, config: dict):
+        self.project = project
+        self.name = name
+        self.config = config
+        self.db_path = os.path.join(TRACKIO_DIR, "trackio.db")
+
+        os.makedirs(TRACKIO_DIR, exist_ok=True)
+
+        self._init_db()
+        self._save_config()
+
+    def _init_db(self):
+        """Initialize the SQLite database with required tables."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS metrics (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
+                    project_name TEXT NOT NULL,
+                    run_name TEXT NOT NULL,
+                    metrics TEXT NOT NULL
+                )
+            """)
+
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS configs (
+                    project_name TEXT NOT NULL,
+                    run_name TEXT NOT NULL,
+                    config TEXT NOT NULL,
+                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+                    PRIMARY KEY (project_name, run_name)
+                )
+            """)
+
+            conn.commit()
+
+    def _save_config(self):
+        """Save the run configuration to the database."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "INSERT OR REPLACE INTO configs (project_name, run_name, config) VALUES (?, ?, ?)",
+                (self.project, self.name, json.dumps(self.config)),
+            )
+            conn.commit()
+
+    def log(self, metrics: dict):
+        """Log metrics to the database."""
+        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}"
+                )
+
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                INSERT INTO metrics 
+                (project_name, run_name, metrics)
+                VALUES (?, ?, ?)
+                """,
+                (self.project, self.name, json.dumps(metrics)),
+            )
+            conn.commit()
+
+    def get_metrics(self, project: str, run: str) -> list[dict]:
+        """Retrieve metrics for a specific run."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                SELECT timestamp, metrics
+                FROM metrics
+                WHERE project_name = ? AND run_name = ?
+                ORDER BY timestamp
+                """,
+                (project, run),
+            )
+            rows = cursor.fetchall()
+
+            results = []
+            for row in rows:
+                timestamp, metrics_json = row
+                metrics = json.loads(metrics_json)
+                metrics["timestamp"] = timestamp
+                results.append(metrics)
+
+            return results
+
+    def get_projects(self) -> list[str]:
+        """Get list of all projects."""
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT DISTINCT project_name FROM metrics")
+            return [row[0] for row in cursor.fetchall()]
+
+    def get_runs(self, project: str) -> list[str]:
+        """Get list of all runs for a project."""
+        with sqlite3.connect(self.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,214 @@
+from typing import Any
+
+import gradio as gr
+import pandas as pd
+
+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
+
+
+def get_projects(request: gr.Request):
+    storage = SQLiteStorage("", "", {})
+    projects = storage.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,
+    )
+
+
+def get_runs(project):
+    if not project:
+        return []
+    storage = SQLiteStorage("", "", {})
+    return storage.get_runs(project)
+
+
+def load_run_data(project: str | None, run: str | None, smoothing: bool):
+    if not project or not run:
+        return None
+    storage = SQLiteStorage("", "", {})
+    metrics = storage.get_metrics(project, run)
+    if not metrics:
+        return None
+    df = pd.DataFrame(metrics)
+    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[numeric_cols] = df[numeric_cols].ewm(alpha=0.1).mean()
+    if "step" not in df.columns:
+        df["step"] = range(len(df))
+    return df
+
+
+def update_runs(project, filter_text, user_interacted_with_runs=False):
+    if project is None:
+        runs = []
+    else:
+        runs = get_runs(project)
+        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)
+    else:
+        return gr.CheckboxGroup(choices=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 toggle_timer(cb_value):
+    if cb_value:
+        return gr.Timer(active=True)
+    else:
+        return gr.Timer(active=False)
+
+
+def log(project: str, run: str, metrics: dict[str, Any]) -> None:
+    storage = SQLiteStorage(project, run, {})
+    storage.log(metrics)
+
+
+def configure(request: gr.Request):
+    if metrics := request.query_params.get("metrics"):
+        return metrics.split(",")
+    else:
+        return []
+
+
+with gr.Blocks(theme="citrus", title="Trackio Dashboard") as demo:
+    with gr.Sidebar() 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")
+        run_tb = gr.Textbox(label="Runs (99)", placeholder="Type to filter...")
+        run_cb = gr.CheckboxGroup(label="Runs", choices=[], interactive=True)
+    with gr.Sidebar(position="right", open=False) as settings_sidebar:
+        gr.Markdown("### ⚙️ Settings")
+        realtime_cb = gr.Checkbox(label="Refresh realtime", value=True)
+        smoothing_cb = gr.Checkbox(label="Smoothing", value=True)
+
+    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,
+    )
+    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,
+        show_progress="hidden",
+    )
+    gr.on(
+        [demo.load, project_dd.change],
+        fn=update_runs,
+        inputs=[project_dd, run_tb],
+        outputs=run_cb,
+        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=log,
+        api_name="log",
+    )
+
+    x_lim = gr.State(None)
+
+    def update_x_lim(select_data: gr.SelectData):
+        return select_data.index
+
+    @gr.render(
+        triggers=[
+            demo.load,
+            run_cb.change,
+            timer.tick,
+            smoothing_cb.change,
+            x_lim.change,
+        ],
+        inputs=[project_dd, run_cb, smoothing_cb, metrics_subset, x_lim],
+    )
+    def update_dashboard(project, runs, smoothing, metrics_subset, x_lim_value):
+        dfs = []
+        for run in runs:
+            df = load_run_data(project, run, smoothing)
+            if df is not None:
+                df["run"] = run
+                dfs.append(df)
+        if dfs:
+            master_df = pd.concat(dfs, ignore_index=True)
+        else:
+            master_df = pd.DataFrame()
+        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]
+        plots: list[gr.LinePlot] = []
+        for col in range(len(numeric_cols) // 2):
+            with gr.Row(key=f"row-{col}"):
+                for i in range(2):
+                    plot = gr.LinePlot(
+                        master_df,
+                        x="step",
+                        y=numeric_cols[2 * col + i],
+                        color="run" if "run" in master_df.columns else None,
+                        title=numeric_cols[2 * col + i],
+                        key=f"plot-{col}-{i}",
+                        preserved_by_key=None,
+                        x_lim=x_lim_value,
+                        y_lim=[
+                            min(master_df[numeric_cols[2 * col + i]]),
+                            max(master_df[numeric_cols[2 * col + i]]),
+                        ],
+                        show_fullscreen_button=True,
+                    )
+                    plots.append(plot)
+
+        for plot in plots:
+            plot.select(update_x_lim, outputs=x_lim)
+            plot.double_click(lambda: None, outputs=x_lim)
+
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=[TRACKIO_LOGO_PATH], show_api=False)

utils.py

@@ -0,0 +1,217 @@
+import os
+import random
+import sys
+import time
+from pathlib import Path
+
+from huggingface_hub.constants import HF_HOME
+
+RESERVED_KEYS = ["project", "run", "timestamp"]
+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.")

version.txt

@@ -0,0 +1 @@
+0.0.8