Merge pull request #11 from gradio-app/reorg

Add a separate `trackio show` command to launch the UI

Author: abidlabs

README.md

@@ -28,6 +28,8 @@ uv pip install trackio
 
 ## Usage
 
+The usage of `trackio` is designed to be a drop-in replacement for `wandb` in most cases:
+
 ```python
 import trackio as wandb
 import random
@@ -38,7 +40,7 @@ epochs = 8
 
 def simulate_multiple_runs():
     for run in range(runs):
-        wandb.init(project="minimal-train-loop2", config={
+        wandb.init(project="fake-training", config={
             "epochs": epochs,
             "learning_rate": 0.001,
             "batch_size": 64
@@ -66,6 +68,35 @@ def simulate_multiple_runs():
 simulate_multiple_runs()
 ```
 
+Running the above will print to the terminal instructions on launching the dashboard.
+
+# Dashboard
+
+You can launch the dashboard by running in your terminal:
+
+```bash
+$ trackio show
+```
+
+or, in Python:
+
+```py
+trackio.show()
+```
+
+You can also provide an optional `project` name as the argument to load a specific project directly:
+
+```bash
+$ trackio show --project "my project"
+```
+
+or, in Python:
+
+```py
+trackio.show(project="my project")
+```
+
+
 ![Screen Recording 2025-05-12 at 2 43 38 PM](https://github.com/user-attachments/assets/d627c9c3-7365-4250-839c-db67dde34a02)
 
 

pyproject.toml

@@ -32,6 +32,9 @@ dev = [
     "ruff==0.9.3"
 ]
 
+[project.scripts]
+trackio = "trackio.cli:main"
+
 [tool.hatch.build.targets.wheel]
 packages = ["trackio"]
 

trackio/__init__.py

@@ -1,40 +1,54 @@
 import contextvars
+import webbrowser
 from pathlib import Path
 
 from gradio_client import Client
 
 from trackio.run import Run
-from trackio.ui import launch_gradio
+from trackio.ui import demo
+from trackio.utils import TRACKIO_DIR, 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 = {}
 
 
-def init(project, name=None, config=None):
+def init(project: str, name: str | None = None, config: dict | None = None) -> Run:
     if not current_server.get():
-        url = launch_gradio(show_api=False, inline=False, quiet=True)
-        print(f"* Trackio server launched at: {url}")
-        print(f"** View project dashboard at: {url}?project={project}")
+        _, 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}")
+        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}")')
+
+    current_project.set(project)
     client = Client(url, verbose=False)
     run = Run(project=project, client=client, name=name, config=config)
     current_run.set(run)
     globals()["config"] = run.config
+    return run
 
 
-def log(metrics):
+def log(metrics: dict) -> None:
     if current_run.get() is None:
         raise RuntimeError("Call trackio.init() before log().")
     current_run.get().log(metrics)
@@ -44,3 +58,14 @@ def finish():
     if current_run.get() is None:
         raise RuntimeError("Call trackio.init() before finish().")
     current_run.get().finish()
+
+
+def show(project: str | None = None):
+    _, url, share_url = demo.launch(
+        show_api=False, quiet=True, inline=False, prevent_thread_lock=True
+    )
+    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()

trackio/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()

trackio/ui.py

@@ -11,14 +11,16 @@ def get_projects(request: gr.Request):
     storage = SQLiteStorage("", "", {})
     projects = storage.get_projects()
     if project := request.query_params.get("project"):
-        pass
+        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,
     )
 
 
@@ -63,9 +65,9 @@ def log(project: str, run: str, metrics: dict[str, Any]) -> None:
 
 
 with gr.Blocks(theme="citrus") as demo:
-    with gr.Sidebar():
+    with gr.Sidebar() as sidebar:
         gr.Markdown("# 🎯 Trackio Dashboard")
-        project_dd = gr.Dropdown(label="Project")
+        project_dd = gr.Dropdown(label="Project", allow_custom_value=True)
         gr.Markdown("### ⚙️ Settings")
         realtime_cb = gr.Checkbox(label="Refresh realtime", value=True)
     with gr.Row():
@@ -125,10 +127,5 @@ def update_dashboard(project, runs):
             )
 
 
-def launch_gradio(**kwargs) -> str:
-    _, url, _ = demo.launch(**kwargs)
-    return url
-
-
 if __name__ == "__main__":
-    launch_gradio()
+    demo.launch()

trackio/utils.py

@@ -1,5 +1,7 @@
 import os
 import random
+import sys
+import time
 
 from huggingface_hub.constants import HF_HOME
 
@@ -199,3 +201,14 @@ def generate_readable_name():
     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.")

trackio/version.txt

@@ -1 +1 @@
-0.0.5
+0.0.6