Auto-deploy to Spaces

trackio/README.md

@@ -0,0 +1,5 @@
+---
+sdk: gradio
+sdk_version: 5.30.0
+app_file: ui.py
+---

trackio/__init__.py

@@ -1,9 +1,13 @@
 import contextvars
+import time
 import webbrowser
 from pathlib import Path
 
 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
@@ -24,24 +28,60 @@
 config = {}
 
 
-def init(project: str, name: str | None = None, config: dict | None = None) -> Run:
+def init(
+    project: str, name: str | None = None, space_id=None, config: dict | None = None
+) -> Run:
     if not current_server.get():
-        _, url, _ = demo.launch(
-            show_api=False, inline=False, quiet=True, prevent_thread_lock=True
-        )
+        if space_id is None:
+            _, url, _ = demo.launch(
+                show_api=False, inline=False, quiet=True, prevent_thread_lock=True
+            )
+        else:
+            url = space_id
         current_server.set(url)
     else:
         url = current_server.get()
+
+    client = None
+    try:
+        client = Client(url, verbose=False)
+    except RepositoryNotFoundError as e:
+        # if we're passed a space_id, ignore the error here, otherwise re-raise it
+        if space_id is None:
+            raise e
+    if client is None and space_id is not None:
+        # try to create the space on demand
+        deploy_as_space(space_id)
+        # try to wait until the Client can initialize
+        max_attempts = 30
+        attempts = 0
+        while client is None:
+            try:
+                client = Client(space_id, verbose=False)
+            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)
+            attempts += 1
+            if attempts >= max_attempts:
+                break
+
     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}")')
-
+        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:
+            print(
+                f"* Trackio metrics logged to: https://huggingface.co/spaces/{space_id}"
+            )
     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

trackio/deploy.py

@@ -0,0 +1,41 @@
+import os
+from importlib.resources import files
+
+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
+
+    hf_api.upload_folder(
+        repo_id=space_id,
+        repo_type="space",
+        folder_path=trackio_path,
+    )

trackio/sqlite_storage.py

@@ -2,7 +2,10 @@
 import os
 import sqlite3
 
-from trackio.utils import RESERVED_KEYS, TRACKIO_DIR
+try:
+    from trackio.utils import RESERVED_KEYS, TRACKIO_DIR
+except:  # noqa: E722
+    from utils import RESERVED_KEYS, TRACKIO_DIR
 
 
 class SQLiteStorage:

trackio/ui.py

@@ -3,8 +3,12 @@
 import gradio as gr
 import pandas as pd
 
-from trackio.sqlite_storage import SQLiteStorage
-from trackio.utils import RESERVED_KEYS, TRACKIO_LOGO_PATH
+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):