Fixing some issues related to deployed Trackio Spaces (#386)

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

Author: abidlabs

.changeset/shy-cobras-sneeze.md

@@ -0,0 +1,5 @@
+---
+"trackio": patch
+---
+
+feat:Fixing some issues related to deployed Trackio Spaces

docs/source/environment_variables.md

@@ -23,6 +23,8 @@ export TRACKIO_LOGO_LIGHT_URL="https://example.com/logo-light.png"
 export TRACKIO_LOGO_DARK_URL="https://example.com/logo-dark.png"
 ```
 
+> **Note:** For remote Trackio Spaces, these environment variables are only applied when the Space is first created via `trackio.init(space_id=...)`. To change logos on an existing Space, update the Space variables directly in the Hugging Face Space settings.
+
 ### `TRACKIO_PLOT_ORDER`
 
 Controls the ordering of plots and metric groups in the Trackio dashboard. The value should be a comma-separated list of metric patterns that specify the desired order. Groups are preserved - if `train/loss` is specified first, all other `train/*` metrics will appear together in the train group, with `train/loss` appearing first within that group.
@@ -44,6 +46,8 @@ export TRACKIO_PLOT_ORDER="train/loss,val/loss"
 - Groups appear in the order of their first matching pattern
 - Unspecified metrics appear in alphabetical order after specified ones
 
+> **Note:** For remote Trackio Spaces, this environment variable is only applied when the Space is first created via `trackio.init(space_id=...)`. To change the plot order on an existing Space, update the `TRACKIO_PLOT_ORDER` Space variable directly in the Hugging Face Space settings.
+
 ### `TRACKIO_THEME`
 
 Sets the theme for the Trackio dashboard. Can be a built-in Gradio theme name or a theme from the Hugging Face Hub.
@@ -59,6 +63,8 @@ export TRACKIO_THEME="gstaff/xkcd"
 export TRACKIO_THEME="ParityError/Anime"
 ```
 
+> **Note:** For remote Trackio Spaces, this environment variable is only applied when the Space is first created via `trackio.init(space_id=...)`. To change the theme on an existing Space, update the `TRACKIO_THEME` Space variable directly in the Hugging Face Space settings.
+
 ### `TRACKIO_COLOR_PALETTE`
 
 Customizes the color palette used for plot lines in the Trackio dashboard. The value should be a comma-separated list of hex color codes. These colors will be cycled through when plotting multiple runs.

tests/conftest.py

@@ -1,3 +1,4 @@
+import os
 import tempfile
 from pathlib import Path
 
@@ -8,6 +9,14 @@
 from trackio.media import write_audio, write_video
 
 
+@pytest.fixture
+def test_space_id():
+    space_id = os.environ.get("TEST_SPACE_ID")
+    if not space_id:
+        pytest.skip("TEST_SPACE_ID environment variable not set")
+    return space_id
+
+
 @pytest.fixture
 def temp_dir(monkeypatch):
     """Fixture that creates a temporary TRACKIO_DIR."""

tests/e2e-spaces/test_metrics_available_on_spaces.py

@@ -0,0 +1,81 @@
+import secrets
+
+import huggingface_hub
+from gradio_client import Client
+
+import trackio
+
+
+def test_basic_logging(test_space_id):
+    project_name = f"test_project_{secrets.token_urlsafe(8)}"
+    run_name = "test_run"
+
+    trackio.init(project=project_name, name=run_name, space_id=test_space_id)
+    trackio.log(metrics={"loss": 0.1})
+    trackio.log(metrics={"loss": 0.2, "acc": 0.9})
+    trackio.finish()
+
+    client = Client(test_space_id)
+
+    summary = client.predict(
+        project=project_name, run=run_name, api_name="/get_run_summary"
+    )
+    assert summary["num_logs"] == 2
+    assert "loss" in summary["metrics"]
+    assert "acc" in summary["metrics"]
+
+    loss_values = client.predict(
+        project=project_name,
+        run=run_name,
+        metric_name="loss",
+        api_name="/get_metric_values",
+    )
+    assert len(loss_values) == 2
+    assert loss_values[0]["value"] == 0.1
+    assert loss_values[0]["step"] == 0
+    assert loss_values[1]["value"] == 0.2
+    assert loss_values[1]["step"] == 1
+
+    acc_values = client.predict(
+        project=project_name,
+        run=run_name,
+        metric_name="acc",
+        api_name="/get_metric_values",
+    )
+    assert len(acc_values) == 1
+    assert acc_values[0]["value"] == 0.9
+    assert acc_values[0]["step"] == 1
+
+
+def test_runs_data_persisted_after_restart(test_space_id):
+    """Test that runs with configs are correctly restored after Space restart."""
+    project_name = f"test_project_{secrets.token_urlsafe(8)}"
+    run_name = "test_run_with_config"
+
+    trackio.init(
+        project=project_name,
+        name=run_name,
+        space_id=test_space_id,
+        config={"learning_rate": 0.001, "epochs": 10},
+    )
+    trackio.log(metrics={"loss": 0.5})
+    trackio.finish()
+
+    client = Client(test_space_id)
+
+    client.predict(api_name="/force_sync")
+
+    # This will force a restart of the Space
+    huggingface_hub.add_space_variable(
+        test_space_id, "TRACKIO_TEST_RESTART", secrets.token_urlsafe(8)
+    )
+
+    client = Client(test_space_id)
+
+    headers, rows, run_names = client.predict(
+        project=project_name, api_name="/get_runs_data"
+    )
+
+    assert run_name in run_names
+    assert any("0.001" in str(row) for row in rows)
+    assert any("10" in str(row) for row in rows)

tests/e2e-spaces/test_metrics_on_spaces.py

@@ -160,7 +160,6 @@ def deploy_as_space(
         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)
-
     if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):
         huggingface_hub.add_space_variable(
             space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url
@@ -169,13 +168,10 @@ def deploy_as_space(
         huggingface_hub.add_space_variable(
             space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url
         )
-
     if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):
         huggingface_hub.add_space_variable(space_id, "TRACKIO_PLOT_ORDER", plot_order)
-
     if theme := os.environ.get("TRACKIO_THEME"):
         huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)
-
     huggingface_hub.add_space_variable(space_id, "GRADIO_MCP_SERVER", "True")
 
 
@@ -211,36 +207,13 @@ def create_space_if_not_exists(
     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
-            )
-        if logo_light_url := os.environ.get("TRACKIO_LOGO_LIGHT_URL"):
-            huggingface_hub.add_space_variable(
-                space_id, "TRACKIO_LOGO_LIGHT_URL", logo_light_url
-            )
-        if logo_dark_url := os.environ.get("TRACKIO_LOGO_DARK_URL"):
-            huggingface_hub.add_space_variable(
-                space_id, "TRACKIO_LOGO_DARK_URL", logo_dark_url
-            )
-
-        if plot_order := os.environ.get("TRACKIO_PLOT_ORDER"):
-            huggingface_hub.add_space_variable(
-                space_id, "TRACKIO_PLOT_ORDER", plot_order
-            )
-
-        if theme := os.environ.get("TRACKIO_THEME"):
-            huggingface_hub.add_space_variable(space_id, "TRACKIO_THEME", theme)
         return
     except RepositoryNotFoundError:
         pass
     except HfHubHTTPError as e:
         if e.response.status_code in [401, 403]:  # unauthorized or forbidden
             print("Need 'write' access token to create a Spaces repo.")
             huggingface_hub.login(add_to_git_credential=False)
-            huggingface_hub.add_space_variable(
-                space_id, "TRACKIO_DATASET_ID", dataset_id
-            )
         else:
             raise ValueError(f"Failed to create Space: {e}")
 

trackio/deploy.py

@@ -183,6 +183,7 @@ def export_to_parquet():
         """
         Exports all projects' DB files as Parquet under the same path but with extension ".parquet".
         Also exports system_metrics to separate parquet files with "_system.parquet" suffix.
+        Also exports configs to separate parquet files with "_configs.parquet" suffix.
         """
         if not SQLiteStorage._dataset_import_attempted:
             return
@@ -196,6 +197,7 @@ def export_to_parquet():
             parquet_path = db_path.with_suffix(".parquet")
             system_parquet_path = db_path.with_suffix("") / ""
             system_parquet_path = TRACKIO_DIR / (db_path.stem + "_system.parquet")
+            configs_parquet_path = TRACKIO_DIR / (db_path.stem + "_configs.parquet")
             if (not parquet_path.exists()) or (
                 db_path.stat().st_mtime > parquet_path.stat().st_mtime
             ):
@@ -243,6 +245,31 @@ def export_to_parquet():
                         use_content_defined_chunking=True,
                     )
 
+            if (not configs_parquet_path.exists()) or (
+                db_path.stat().st_mtime > configs_parquet_path.stat().st_mtime
+            ):
+                with sqlite3.connect(str(db_path)) as conn:
+                    try:
+                        configs_df = pd.read_sql("SELECT * FROM configs", conn)
+                    except Exception:
+                        configs_df = pd.DataFrame()
+                if not configs_df.empty:
+                    config_data = configs_df["config"].copy()
+                    config_data = pd.DataFrame(
+                        config_data.apply(
+                            lambda x: deserialize_values(orjson.loads(x))
+                        ).values.tolist(),
+                        index=configs_df.index,
+                    )
+                    del configs_df["config"]
+                    for col in config_data.columns:
+                        configs_df[col] = config_data[col]
+                    configs_df.to_parquet(
+                        configs_parquet_path,
+                        write_page_index=True,
+                        use_content_defined_chunking=True,
+                    )
+
     @staticmethod
     def _cleanup_wal_sidecars(db_path: Path) -> None:
         """Remove leftover -wal/-shm files for a DB basename (prevents disk I/O errors)."""
@@ -259,6 +286,7 @@ def import_from_parquet():
         """
         Imports to all DB files that have matching files under the same path but with extension ".parquet".
         Also imports system_metrics from "_system.parquet" files.
+        Also imports configs from "_configs.parquet" files.
         """
         if not TRACKIO_DIR.exists():
             return
@@ -267,7 +295,9 @@ def import_from_parquet():
         parquet_names = [
             f
             for f in all_paths
-            if f.endswith(".parquet") and not f.endswith("_system.parquet")
+            if f.endswith(".parquet")
+            and not f.endswith("_system.parquet")
+            and not f.endswith("_configs.parquet")
         ]
         for pq_name in parquet_names:
             parquet_path = TRACKIO_DIR / pq_name
@@ -310,6 +340,29 @@ def import_from_parquet():
                 df.to_sql("system_metrics", conn, if_exists="replace", index=False)
                 conn.commit()
 
+        configs_parquet_names = [f for f in all_paths if f.endswith("_configs.parquet")]
+        for pq_name in configs_parquet_names:
+            parquet_path = TRACKIO_DIR / pq_name
+            db_name = pq_name.replace("_configs.parquet", DB_EXT)
+            db_path = TRACKIO_DIR / db_name
+
+            df = pd.read_parquet(parquet_path)
+            if "config" not in df.columns:
+                config_data = df.copy()
+                other_cols = ["id", "run_name", "created_at"]
+                df = df[[c for c in other_cols if c in df.columns]]
+                for col in other_cols:
+                    if col in config_data.columns:
+                        del config_data[col]
+                config_data = orjson.loads(config_data.to_json(orient="records"))
+                df["config"] = [
+                    orjson.dumps(serialize_values(row)) for row in config_data
+                ]
+
+            with sqlite3.connect(str(db_path), timeout=30.0) as conn:
+                df.to_sql("configs", conn, if_exists="replace", index=False)
+                conn.commit()
+
     @staticmethod
     def get_scheduler():
         """
@@ -330,7 +383,12 @@ def get_scheduler():
                     repo_type="dataset",
                     folder_path=TRACKIO_DIR,
                     private=True,
-                    allow_patterns=["*.parquet", "*_system.parquet", "media/**/*"],
+                    allow_patterns=[
+                        "*.parquet",
+                        "*_system.parquet",
+                        "*_configs.parquet",
+                        "media/**/*",
+                    ],
                     squash_history=True,
                     token=hf_token,
                     on_before_commit=SQLiteStorage.export_to_parquet,