Allow deleting runs on Spaces (via Oauth) and on Colab (#268)

abidlabs · web-flow · commit 83cf08f214bd · 2025-10-02T11:54:35.000-07:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ requires-python = ">=3.10"
 dependencies = [
     "pandas<3.0.0",
     "huggingface-hub<1.0.0",
-    "gradio>=5.48.0,<6.0.0",
+    "gradio[oauth]>=5.48.0,<6.0.0",
     "numpy<3.0.0",
     "pillow<12.0.0",
     "orjson>=3.0,<4.0.0"
diff --git a/trackio/README.md b/trackio/README.md
@@ -4,4 +4,7 @@ sdk_version: {GRADIO_VERSION}
 app_file: ui/main.py
 tags:
  - trackio
+hf_oauth: true
+hf_oauth_scopes:
+ - write-repos
 ---
diff --git a/trackio/__init__.py b/trackio/__init__.py
@@ -23,6 +23,13 @@
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
+warnings.filterwarnings(
+    "ignore",
+    message="Empty session being created. Install gradio\\[oauth\\]",
+    category=UserWarning,
+    module="gradio.helpers",
+)
+
 __version__ = Path(__file__).parent.joinpath("version.txt").read_text().strip()
 
 __all__ = [
diff --git a/trackio/ui/fns.py b/trackio/ui/fns.py
@@ -3,6 +3,7 @@
 import os
 
 import gradio as gr
+import huggingface_hub as hf
 
 try:
     import trackio.utils as utils
@@ -21,9 +22,12 @@
 CONFIG_COLUMN_MAPPINGS_REVERSE = {v: k for k, v in CONFIG_COLUMN_MAPPINGS.items()}
 
 
+HfApi = hf.HfApi()
+
+
 def get_project_info() -> str | None:
     dataset_id = os.environ.get("TRACKIO_DATASET_ID")
-    space_id = os.environ.get("SPACE_ID")
+    space_id = utils.get_space()
     if utils.persistent_storage_enabled():
         return "&#10024; Persistent Storage is enabled, logs are stored directly in this Space."
     if dataset_id:
@@ -58,15 +62,107 @@ def get_projects(request: gr.Request):
     )
 
 
-def update_navbar_value(project_dd):
+def update_navbar_value(project_dd, request: gr.Request):
+    write_token = None
+    if hasattr(request, "query_params") and request.query_params:
+        write_token = request.query_params.get("write_token")
+
+    metrics_url = f"?selected_project={project_dd}"
+    runs_url = f"runs?selected_project={project_dd}"
+
+    if write_token:
+        metrics_url += f"&write_token={write_token}"
+        runs_url += f"&write_token={write_token}"
+
     return gr.Navbar(
         value=[
-            ("Metrics", f"?selected_project={project_dd}"),
-            ("Runs", f"runs?selected_project={project_dd}"),
+            ("Metrics", metrics_url),
+            ("Runs", runs_url),
         ]
     )
 
 
+def check_hf_token_has_write_access(hf_token: str | None) -> None:
+    """
+    Checks to see if the provided hf_token is valid and has write access to the Space
+    that Trackio is running in. If the hf_token is valid or if Trackio is not running
+    on a Space, this function does nothing. Otherwise, it raises a PermissionError.
+    """
+    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)
+        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
+        access_token = who["auth"]["accessToken"]
+        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"
+                        or item["entity"]["type"] == "org"
+                    )
+                    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 check_oauth_token_has_write_access(oauth_token: str | None) -> None:
+    """
+    Checks to see if the oauth token provided via Gradio's OAuth is valid and has write access
+    to the Space that Trackio is running in. If the oauth token is valid or if Trackio is not running
+    on a Space, this function does nothing. Otherwise, it raises a PermissionError.
+    """
+    if not os.getenv("SYSTEM") == "spaces":
+        return
+    if oauth_token is None:
+        raise PermissionError(
+            "Expected an oauth to be provided when logging to a Space"
+        )
+    who = HfApi.whoami(oauth_token)
+    user_name = who["name"]
+    owner_name = os.getenv("SPACE_AUTHOR_NAME")
+    if user_name == owner_name:
+        return
+    # check if user is a member of an org that owns the space with write permissions
+    for org in who["orgs"]:
+        if org["name"] == owner_name and org["roleInOrg"] == "write":
+            return
+    raise PermissionError(
+        "Expected the oauth token to be the user owner of the space, or be a member of the org owner of the space"
+    )
+
+
 def get_group_by_fields(project: str):
     configs = SQLiteStorage.get_all_run_configs(project) if project else {}
     keys = set()
diff --git a/trackio/ui/main.py b/trackio/ui/main.py
@@ -8,12 +8,9 @@
 from typing import Any
 
 import gradio as gr
-import huggingface_hub as hf
 import numpy as np
 import pandas as pd
 
-HfApi = hf.HfApi()
-
 try:
     import trackio.utils as utils
     from trackio.file_storage import FileStorage
@@ -267,63 +264,13 @@ def toggle_timer(cb_value):
         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"
-                        or item["entity"]["type"] == "org"
-                    )
-                    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:
     """
     Uploads the database of a local Trackio project to a Hugging Face Space.
     """
-    check_auth(hf_token)
+    fns.check_hf_token_has_write_access(hf_token)
     db_project_path = SQLiteStorage.get_project_db_path(project)
     if os.path.exists(db_project_path):
         raise gr.Error(
@@ -337,7 +284,7 @@ def bulk_upload_media(uploads: list[UploadEntry], hf_token: str | None) -> None:
     """
     Uploads media files to a Trackio dashboard. Each entry in the list is a tuple of the project, run, and media file to be uploaded.
     """
-    check_auth(hf_token)
+    fns.check_hf_token_has_write_access(hf_token)
     for upload in uploads:
         media_path = FileStorage.init_project_media_path(
             upload["project"], upload["run"], upload["step"]
@@ -357,7 +304,7 @@ def log(
     is kept for backwards compatibility for users who are connecting to a newer version of
     a Trackio Spaces dashboard with an older version of Trackio installed locally.
     """
-    check_auth(hf_token)
+    fns.check_hf_token_has_write_access(hf_token)
     SQLiteStorage.log(project=project, run=run, metrics=metrics, step=step)
 
 
@@ -368,7 +315,7 @@ def bulk_log(
     """
     Logs a list of metrics to a Trackio dashboard. Each entry in the list is a dictionary of the project, run, a dictionary of metrics, and optionally, a step and config.
     """
-    check_auth(hf_token)
+    fns.check_hf_token_has_write_access(hf_token)
 
     logs_by_run = {}
     for log_entry in logs:
@@ -627,12 +574,17 @@ def create_media_section(media_by_run: dict[str, dict[str, list[MediaData]]]):
     
     if (writeToken) {
         setCookie('trackio_write_token', writeToken, 7);
-        
-        urlParams.delete('write_token');
-        const newUrl = window.location.pathname + 
-            (urlParams.toString() ? '?' + urlParams.toString() : '') + 
-            window.location.hash;
-        window.history.replaceState({}, document.title, newUrl);
+                
+        // Only remove write_token from URL if not in iframe
+        // In iframes, keep it in URL as cookies may be blocked
+        const inIframe = window.self !== window.top;
+        if (!inIframe) {
+            urlParams.delete('write_token');
+            const newUrl = window.location.pathname + 
+                (urlParams.toString() ? '?' + urlParams.toString() : '') + 
+                window.location.hash;
+            window.history.replaceState({}, document.title, newUrl);
+        }
     }
 })();
 </script>
@@ -968,7 +920,7 @@ def update_dashboard(
             master_df = pd.DataFrame()
 
         if master_df.empty:
-            if space_id := os.environ.get("SPACE_ID"):
+            if space_id := utils.get_space():
                 gr.Markdown(INSTRUCTIONS_SPACES.format(space_id))
             else:
                 gr.Markdown(INSTRUCTIONS_LOCAL)
diff --git a/trackio/ui/runs.py b/trackio/ui/runs.py
diff --git a/trackio/utils.py b/trackio/utils.py
