Fix/remove flaky E2E space tests (#485)

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

Author: abidlabs

.changeset/khaki-rockets-cough.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:Fix/remove flaky E2E space tests

tests/e2e-spaces/conftest.py

@@ -1,8 +1,10 @@
 import os
 import time
 
+import huggingface_hub
 import pytest
 from gradio_client import Client
+from huggingface_hub.errors import HfHubHTTPError, RepositoryNotFoundError
 
 from trackio import deploy, utils
 
@@ -12,6 +14,7 @@ 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")
+    space_id, _, _ = utils.preprocess_space_and_dataset_ids(space_id, None)
     return space_id
 
 
@@ -20,16 +23,64 @@ def _ensure_space_ready(test_space_id):
     space_id, dataset_id, bucket_id = utils.preprocess_space_and_dataset_ids(
         test_space_id, None
     )
+
+    _reset_test_space(space_id)
     deploy.create_space_if_not_exists(space_id, None, dataset_id, bucket_id, None)
+    _wait_for_space_ready(space_id)
+
+
+def _reset_test_space(space_id):
+    try:
+        huggingface_hub.delete_repo(space_id, repo_type="space")
+    except RepositoryNotFoundError:
+        return
+    except HfHubHTTPError as e:
+        if e.response is not None and e.response.status_code in (401, 403):
+            raise RuntimeError(
+                f"Cannot delete test Space {space_id}: check Hugging Face auth and write access."
+            ) from e
+        raise
+
+
+def _wait_for_space_ready(space_id, timeout=300):
+    hf_api = huggingface_hub.HfApi()
+    deadline = time.time() + timeout
+    delay = 2
+    last_error = None
+    last_stage = None
+
+    terminal_stages = {"NO_APP_FILE", "CONFIG_ERROR", "BUILD_ERROR", "RUNTIME_ERROR"}
 
-    deadline = time.time() + 300
     while time.time() < deadline:
         try:
-            Client(test_space_id, verbose=False)
-            return
-        except Exception:
-            time.sleep(10)
-    pytest.fail(f"Space {test_space_id} not ready after 5 minutes")
+            info = hf_api.space_info(space_id, timeout=30)
+            runtime = getattr(info, "runtime", None)
+            if runtime is not None:
+                last_stage = str(runtime.stage)
+                if last_stage in terminal_stages:
+                    raise RuntimeError(
+                        f"Space {space_id} entered terminal stage {last_stage}. "
+                        "Check the Hugging Face build logs."
+                    )
+                if last_stage == "RUNNING":
+                    Client(space_id, verbose=False, httpx_kwargs={"timeout": 30})
+                    return
+        except HfHubHTTPError as e:
+            last_error = e
+            if e.response is not None and e.response.status_code in (401, 403):
+                raise RuntimeError(
+                    f"Cannot access test Space {space_id}: check Hugging Face auth and permissions."
+                ) from e
+        except Exception as e:
+            last_error = e
+
+        time.sleep(delay)
+        delay = min(delay * 1.5, 15)
+
+    pytest.fail(
+        f"Space {space_id} not ready after {timeout}s. "
+        f"Last observed stage={last_stage!r}, last error={last_error!r}"
+    )
 
 
 @pytest.fixture

tests/e2e-spaces/test_data_robustness.py

@@ -164,7 +164,10 @@ def wrapped_predict(*args, **kwargs):
         f"Expected all 3 logs on Space after recovery, got {summary['num_logs']}"
     )
 
-    local_logs = SQLiteStorage.get_logs(project=project_name, run=run_name)
-    assert len(local_logs) == 0, (
-        f"Expected local buffer to be empty after flush, but found {len(local_logs)} rows"
+    deadline = time.time() + 60
+    while time.time() < deadline and SQLiteStorage.has_pending_data(project_name):
+        time.sleep(2)
+
+    assert not SQLiteStorage.has_pending_data(project_name), (
+        "Expected pending local buffer rows to be cleared after flush"
     )

tests/e2e-spaces/test_metrics_on_spaces.py

@@ -9,6 +9,46 @@
 from trackio import utils
 
 
+def _predict_run_summary(
+    test_space_id: str,
+    project_name: str,
+    run_name: str,
+    *,
+    min_num_logs: int = 0,
+    timeout: float = 240,
+):
+    deadline = time.time() + timeout
+    last_err: Exception | None = None
+    flush_attempted = False
+    while time.time() < deadline:
+        try:
+            client = Client(test_space_id, verbose=False)
+            summary = client.predict(
+                project=project_name, run=run_name, api_name="/get_run_summary"
+            )
+            if summary["num_logs"] >= min_num_logs:
+                return summary
+            last_err = None
+        except Exception as e:
+            last_err = e
+        if not flush_attempted and time.time() > deadline - max(timeout - 60, 0):
+            flush_run = trackio.init(
+                project=project_name,
+                name=f"flush_{secrets.token_urlsafe(4)}",
+                space_id=test_space_id,
+                auto_log_gpu=False,
+            )
+            flush_deadline = time.time() + 30
+            while flush_run._client is None and time.time() < flush_deadline:
+                time.sleep(0.1)
+            flush_run.finish()
+            flush_attempted = True
+        time.sleep(5)
+    if last_err is not None:
+        raise last_err
+    raise TimeoutError("get_run_summary timed out before logs appeared")
+
+
 def test_basic_logging(test_space_id):
     project_name = f"test_project_{secrets.token_urlsafe(8)}"
     run_name = "test_run"
@@ -18,15 +58,14 @@ def test_basic_logging(test_space_id):
     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"
+    summary = _predict_run_summary(
+        test_space_id, project_name, run_name, min_num_logs=2
     )
     assert summary["num_logs"] == 2
     assert "loss" in summary["metrics"]
     assert "acc" in summary["metrics"]
 
+    client = Client(test_space_id)
     loss_values = client.predict(
         project=project_name,
         run=run_name,
@@ -159,8 +198,8 @@ def test_bucket_space_preserves_logged_metrics_after_restart(test_space_id):
             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"
+    summary = _predict_run_summary(
+        test_space_id, project_name, run_name, min_num_logs=1, timeout=360
     )
     assert summary["num_logs"] == 1
     assert "loss" in summary["metrics"] and "acc" in summary["metrics"]

tests/e2e-spaces/test_sync_and_freeze.py

@@ -26,7 +26,9 @@ def _download_parquet_from_bucket(bucket_id, remote_name="metrics.parquet"):
     with tempfile.TemporaryDirectory() as tmp:
         local_path = Path(tmp) / remote_name
         huggingface_hub.download_bucket_files(
-            bucket_id, files=[(remote_name, str(local_path))]
+            bucket_id,
+            files=[(remote_name, str(local_path))],
+            token=huggingface_hub.utils.get_token(),
         )
         return pd.read_parquet(local_path)
 
@@ -40,11 +42,20 @@ def _cleanup_space(space_id):
 
 def _cleanup_bucket(bucket_id):
     try:
-        huggingface_hub.delete_bucket(bucket_id)
+        huggingface_hub.delete_bucket(
+            bucket_id, token=huggingface_hub.utils.get_token()
+        )
     except Exception:
         pass
 
 
+def _namespace_scoped_repo_id(test_space_id: str, repo_name: str) -> str:
+    if "/" in test_space_id:
+        namespace = test_space_id.split("/", 1)[0]
+        return f"{namespace}/{repo_name}"
+    return repo_name
+
+
 def test_sync_to_gradio_space(test_space_id, temp_dir):
     project_name = f"test_sync_gradio_{secrets.token_urlsafe(8)}"
     run_name = "run1"
@@ -76,11 +87,11 @@ def test_sync_to_gradio_space(test_space_id, temp_dir):
     assert loss_values[2]["value"] == 0.1
 
 
-def test_sync_to_static_space_incremental(temp_dir):
+def test_sync_to_static_space_incremental(test_space_id, temp_dir):
     project_name = f"test_sync_static_{secrets.token_urlsafe(8)}"
     run_name = "run1"
     suffix = secrets.token_urlsafe(6)
-    space_id = f"trackio-test-static-{suffix}"
+    space_id = _namespace_scoped_repo_id(test_space_id, f"trackio-test-static-{suffix}")
     space_id, _, bucket_id = utils.preprocess_space_and_dataset_ids(space_id, None)
 
     try:
@@ -127,7 +138,9 @@ def test_sync_gradio_then_freeze_to_static(test_space_id, temp_dir):
     time.sleep(5)
 
     suffix = secrets.token_urlsafe(6)
-    frozen_space_id = f"trackio-test-frozen-{suffix}"
+    frozen_space_id = _namespace_scoped_repo_id(
+        test_space_id, f"trackio-test-frozen-{suffix}"
+    )
     frozen_space_id, _, frozen_bucket_id = utils.preprocess_space_and_dataset_ids(
         frozen_space_id, None
     )

trackio/bucket_storage.py

@@ -64,6 +64,7 @@ def _download_db_from_bucket(
         huggingface_hub.download_bucket_files(
             bucket_id,
             files=[(remote_path, str(local_path))],
+            token=huggingface_hub.utils.get_token(),
         )
         return local_path.exists()
     except Exception:

trackio/deploy.py

@@ -7,6 +7,7 @@
 import tempfile
 import threading
 import time
+from collections import Counter
 from importlib.resources import files
 from pathlib import Path
 
@@ -498,11 +499,13 @@ def sync_incremental(
 
     client = Client(space_id, verbose=False, httpx_kwargs={"timeout": 90})
     hf_token = huggingface_hub.utils.get_token()
+    expected_run_counts: Counter[str] = Counter()
 
     if pending_only:
         pending_logs = SQLiteStorage.get_pending_logs(project)
         if pending_logs:
             logs = pending_logs["logs"]
+            expected_run_counts.update(log["run"] for log in logs)
             for i in range(0, len(logs), SYNC_BATCH_SIZE):
                 batch = logs[i : i + SYNC_BATCH_SIZE]
                 print(
@@ -550,6 +553,7 @@ def sync_incremental(
     else:
         all_logs = SQLiteStorage.get_all_logs_for_sync(project)
         if all_logs:
+            expected_run_counts.update(log["run"] for log in all_logs)
             for i in range(0, len(all_logs), SYNC_BATCH_SIZE):
                 batch = all_logs[i : i + SYNC_BATCH_SIZE]
                 print(
@@ -568,12 +572,52 @@ def sync_incremental(
                     api_name="/bulk_log_system", logs=batch, hf_token=hf_token
                 )
 
+    _wait_for_remote_sync(client, project, expected_run_counts)
     SQLiteStorage.set_project_metadata(project, "space_id", space_id)
     print(
         f"* Synced successfully to space: {_BOLD_ORANGE}{SPACE_URL.format(space_id=space_id)}{_RESET}"
     )
 
 
+def _wait_for_remote_sync(
+    client: Client,
+    project: str,
+    expected_run_counts: Counter[str],
+    timeout: int = 180,
+) -> None:
+    if not expected_run_counts:
+        return
+
+    deadline = time.time() + timeout
+    delay = 2
+    last_error: Exception | None = None
+    pending = dict(expected_run_counts)
+
+    while time.time() < deadline and pending:
+        completed = []
+        for run_name, expected_num_logs in pending.items():
+            try:
+                summary = client.predict(
+                    project=project, run=run_name, api_name="/get_run_summary"
+                )
+                if summary.get("num_logs") == expected_num_logs:
+                    completed.append(run_name)
+            except Exception as e:
+                last_error = e
+        for run_name in completed:
+            pending.pop(run_name, None)
+        if pending:
+            time.sleep(delay)
+            delay = min(delay * 1.5, 15)
+
+    if pending:
+        raise TimeoutError(
+            f"Remote sync for project '{project}' did not become visible for runs "
+            f"{sorted(pending.items())} within {timeout}s. "
+            f"Last error: {last_error!r}"
+        )
+
+
 def upload_dataset_for_static(
     project: str,
     dataset_id: str,
@@ -831,6 +875,14 @@ def _do_sync():
                 create_space_if_not_exists(
                     space_id, bucket_id=bucket_id, private=private
                 )
+                _wait_for_remote_sync(
+                    Client(space_id, verbose=False, httpx_kwargs={"timeout": 90}),
+                    project,
+                    Counter(
+                        log["run"]
+                        for log in SQLiteStorage.get_all_logs_for_sync(project)
+                    ),
+                )
             else:
                 sync_incremental(project, space_id, private=private, pending_only=False)
         SQLiteStorage.set_project_metadata(project, "space_id", space_id)