changes

Author: abidlabs

tests/e2e-spaces/test_buckets_backend.py

@@ -2,9 +2,11 @@
 import time
 
 import huggingface_hub
+import pytest
 from gradio_client import Client
 
 import trackio
+from trackio import utils
 
 
 def test_basic_logging(test_space_id):
@@ -92,3 +94,60 @@ def test_runs_data_persisted_after_restart(test_space_id):
     lr = cfg.get("learning_rate")
     assert lr is not None and abs(float(lr) - 0.001) < 1e-6
     assert cfg.get("epochs") == 10
+
+
+def test_bucket_space_preserves_logged_metrics_after_restart(test_space_id):
+    _, dataset_id, bucket_id = utils.preprocess_space_and_dataset_ids(
+        test_space_id, None, None
+    )
+    if dataset_id is not None or bucket_id is None:
+        pytest.skip("Requires a Space deployed with bucket backend (no dataset_id).")
+
+    project_name = f"test_bucket_persist_{secrets.token_urlsafe(8)}"
+    run_name = "metrics_run"
+
+    trackio.init(project=project_name, name=run_name, space_id=test_space_id)
+    trackio.log(metrics={"loss": 0.42, "acc": 0.88})
+    trackio.finish()
+
+    client = Client(test_space_id)
+    client.predict(api_name="/force_sync")
+
+    huggingface_hub.add_space_variable(
+        test_space_id, "TRACKIO_TEST_RESTART", secrets.token_urlsafe(8)
+    )
+
+    time.sleep(10)
+    deadline = time.time() + 300
+    client = None
+    while time.time() < deadline:
+        try:
+            client = Client(test_space_id, verbose=False)
+            break
+        except Exception:
+            time.sleep(10)
+    assert client is not None, "Space did not come back up after restart"
+
+    summary = client.predict(
+        project=project_name, run=run_name, api_name="/get_run_summary"
+    )
+    assert summary["num_logs"] == 1
+    assert "loss" in summary["metrics"] and "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) == 1
+    assert abs(float(loss_values[0]["value"]) - 0.42) < 1e-6
+
+    acc_values = client.predict(
+        project=project_name,
+        run=run_name,
+        metric_name="acc",
+        api_name="/get_metric_values",
+    )
+    assert len(acc_values) == 1
+    assert abs(float(acc_values[0]["value"]) - 0.88) < 1e-6

tests/e2e-spaces/test_metrics_on_spaces.py

@@ -7,4 +7,4 @@ tags:
 {LINKED_HUB_METADATA}hf_oauth: true
 hf_oauth_scopes:
  - write-repos
-{BUCKET_MOUNT}---
+---

trackio/README.md

@@ -16,6 +16,7 @@
     import tomli as tomllib
 
 import gradio
+import httpx
 import huggingface_hub
 from gradio_client import Client, handle_file
 from httpx import ReadTimeout
@@ -34,13 +35,17 @@
 SPACE_URL = "https://huggingface.co/spaces/{space_id}"
 
 
-def _readme_linked_hub_yaml(dataset_id: str | None, bucket_id: str | None) -> str:
-    parts = []
+def _readme_linked_hub_yaml(dataset_id: str | None) -> str:
     if dataset_id is not None:
-        parts.append(f"datasets:\n - {dataset_id}\n")
-    if bucket_id is not None:
-        parts.append(f"buckets:\n - {bucket_id}\n")
-    return "".join(parts)
+        return f"datasets:\n - {dataset_id}\n"
+    return ""
+
+
+def _space_app_py_content(bucket_id: str | None) -> str:
+    if bucket_id is None:
+        return "import trackio\ntrackio.show()\n"
+    path = Path(files("trackio")) / "space_bucket_app.py"
+    return path.read_text(encoding="utf-8")
 
 
 def _retry_hf_write(op_name: str, fn, retries: int = 4, initial_delay: float = 1.5):
@@ -157,15 +162,8 @@ def deploy_as_space(
         readme_content = readme_content.replace("{GRADIO_VERSION}", gradio.__version__)
         readme_content = readme_content.replace("{APP_FILE}", "app.py")
         readme_content = readme_content.replace(
-            "{LINKED_HUB_METADATA}", _readme_linked_hub_yaml(dataset_id, bucket_id)
+            "{LINKED_HUB_METADATA}", _readme_linked_hub_yaml(dataset_id)
         )
-        if bucket_id is not None:
-            bucket_mount = (
-                f"hf_mount:\n - src: hf://buckets/{bucket_id}\n   dst: /data/trackio\n"
-            )
-        else:
-            bucket_mount = ""
-        readme_content = readme_content.replace("{BUCKET_MOUNT}", bucket_mount)
         readme_buffer = io.BytesIO(readme_content.encode("utf-8"))
         hf_api.upload_file(
             path_or_fileobj=readme_buffer,
@@ -217,8 +215,7 @@ def deploy_as_space(
             ],
         )
 
-    app_file_content = """import trackio
-trackio.show()"""
+    app_file_content = _space_app_py_content(bucket_id)
     app_file_buffer = io.BytesIO(app_file_content.encode("utf-8"))
     hf_api.upload_file(
         path_or_fileobj=app_file_buffer,
@@ -231,6 +228,7 @@ def deploy_as_space(
         huggingface_hub.add_space_secret(space_id, "HF_TOKEN", hf_token)
     if bucket_id is not None:
         huggingface_hub.add_space_variable(space_id, "TRACKIO_BUCKET_ID", bucket_id)
+        huggingface_hub.add_space_variable(space_id, "TRACKIO_DIR", "/data/trackio")
     elif 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"):
@@ -307,15 +305,32 @@ def _wait_until_space_running(space_id: str, timeout: int = 300) -> None:
     hf_api = huggingface_hub.HfApi()
     start = time.time()
     delay = 2
+    request_timeout = 45.0
+    failure_stages = frozenset(
+        ("NO_APP_FILE", "CONFIG_ERROR", "BUILD_ERROR", "RUNTIME_ERROR")
+    )
     while time.time() - start < timeout:
         try:
-            info = hf_api.space_info(space_id)
-            if info.runtime and info.runtime.stage == "RUNNING":
-                return
-        except (huggingface_hub.utils.HfHubHTTPError, ReadTimeout):
+            info = hf_api.space_info(space_id, timeout=request_timeout)
+            if info.runtime:
+                stage = str(info.runtime.stage)
+                if stage in failure_stages:
+                    raise RuntimeError(
+                        f"Space {space_id} entered terminal stage {stage}. "
+                        "Fix README.md or app files; see build logs on the Hub."
+                    )
+                if stage == "RUNNING":
+                    return
+        except RuntimeError:
+            raise
+        except (huggingface_hub.utils.HfHubHTTPError, httpx.RequestError):
             pass
         time.sleep(delay)
         delay = min(delay * 1.5, 15)
+    raise TimeoutError(
+        f"Space {space_id} did not reach RUNNING within {timeout}s. "
+        "Check status and build logs on the Hub."
+    )
 
 
 def wait_until_space_exists(
@@ -337,7 +352,7 @@ def wait_until_space_exists(
         try:
             hf_api.space_info(space_id)
             return
-        except (huggingface_hub.utils.HfHubHTTPError, ReadTimeout):
+        except (huggingface_hub.utils.HfHubHTTPError, httpx.RequestError):
             time.sleep(delay)
             delay = min(delay * 2, 60)
     raise TimeoutError("Waiting for space to exist took longer than expected")
@@ -569,7 +584,7 @@ def deploy_as_static_space(
         else:
             raise ValueError(f"Failed to create Space: {e}")
 
-    linked = _readme_linked_hub_yaml(dataset_id, bucket_id)
+    linked = _readme_linked_hub_yaml(dataset_id)
     readme_content = (
         "---\nsdk: static\npinned: false\ntags:\n - trackio\n"
         f"{linked}---\n"

trackio/deploy.py

@@ -0,0 +1,62 @@
+import importlib
+import os
+import platform
+import subprocess
+import sys
+import urllib.request
+
+_DEFAULT_MOUNT = "/data/trackio"
+_BIN_DIR = "/tmp/trackio-hf-mount-bin"
+
+
+def _platform_release_asset():
+    if sys.platform != "linux":
+        raise OSError(f"hf-mount on Spaces requires Linux, got {sys.platform}")
+    m = platform.machine().lower()
+    if m in ("x86_64", "amd64"):
+        return "x86_64", "linux"
+    if m in ("aarch64", "arm64"):
+        return "aarch64", "linux"
+    raise OSError(f"Unsupported machine for hf-mount: {m}")
+
+
+def _download_hf_mount_binaries():
+    arch, plat = _platform_release_asset()
+    base = "https://github.com/huggingface/hf-mount/releases/latest/download"
+    os.makedirs(_BIN_DIR, exist_ok=True)
+    for name in ("hf-mount", "hf-mount-nfs", "hf-mount-fuse"):
+        binary = f"{name}-{arch}-{plat}"
+        url = f"{base}/{binary}"
+        dest = os.path.join(_BIN_DIR, name)
+        with urllib.request.urlopen(url) as response:
+            with open(dest, "wb") as out:
+                out.write(response.read())
+        os.chmod(dest, 0o755)
+
+
+def start_hf_mount_for_trackio_bucket():
+    if os.environ.get("SYSTEM") != "spaces":
+        return
+    bucket_id = os.environ.get("TRACKIO_BUCKET_ID")
+    if not bucket_id:
+        return
+    mount_path = os.environ.get("TRACKIO_DIR", _DEFAULT_MOUNT)
+    parent = os.path.dirname(mount_path.rstrip("/")) or "/"
+    os.makedirs(parent, exist_ok=True)
+    os.makedirs(mount_path, exist_ok=True)
+    hf_mount = os.path.join(_BIN_DIR, "hf-mount")
+    if not os.path.isfile(hf_mount):
+        _download_hf_mount_binaries()
+    env = {**os.environ, "PATH": _BIN_DIR + os.pathsep + os.environ.get("PATH", "")}
+    subprocess.run(
+        [hf_mount, "start", "bucket", bucket_id, mount_path],
+        check=True,
+        env=env,
+        timeout=600,
+    )
+
+
+start_hf_mount_for_trackio_bucket()
+
+trackio = importlib.import_module("trackio")
+trackio.show()

trackio/space_bucket_app.py

@@ -134,9 +134,8 @@ def _get_trackio_dir() -> Path:
         return Path("/data/trackio")
     if os.environ.get("TRACKIO_DIR"):
         return Path(os.environ.get("TRACKIO_DIR"))
-    bucket_mount = Path("/data/trackio")
-    if os.environ.get("TRACKIO_BUCKET_ID") and bucket_mount.exists():
-        return bucket_mount
+    if os.environ.get("TRACKIO_BUCKET_ID"):
+        return Path("/data/trackio")
     return Path(HF_HOME) / "trackio"