Spaces:
Sleeping
Sleeping
zofiasmolenasana
fix(vertex): synchronous CustomJob.submit; train UI GCP arch sync
93e4108 unverified | """Load per-unit ``progress.json`` files from GCS for the local /train dashboard.""" | |
| from __future__ import annotations | |
| import json | |
| import re | |
| from typing import Any | |
| DIR_RE = re.compile(r"^(.+)_seed(\d+)_fold(\d+)$") | |
| def parse_gs_uri(uri: str) -> tuple[str, str]: | |
| u = uri.strip().rstrip("/") | |
| if not u.startswith("gs://"): | |
| raise ValueError("EXCEL_CHUNKER_GCS_EXPERIMENTS_URI must start with gs://") | |
| rest = u[5:] | |
| slash = rest.find("/") | |
| if slash == -1: | |
| raise ValueError("URI must include a prefix path after the bucket name") | |
| bucket = rest[:slash] | |
| prefix = rest[slash + 1 :].lstrip("/") | |
| if not prefix: | |
| raise ValueError("URI prefix must not be empty (point at .../experiments)") | |
| if not prefix.endswith("/"): | |
| prefix = prefix + "/" | |
| return bucket, prefix | |
| def fetch_suite_from_gcs(uri: str) -> dict[str, Any]: | |
| """Return a dict shaped like ``SuiteProgress.to_api_dict()`` from blob ``*/progress.json``.""" | |
| try: | |
| from google.cloud import storage | |
| except ImportError as e: | |
| raise RuntimeError( | |
| "Install google-cloud-storage (e.g. pip install google-cloud-storage)" | |
| ) from e | |
| bucket_name, prefix = parse_gs_uri(uri) | |
| client = storage.Client() | |
| bucket = client.bucket(bucket_name) | |
| units: list[dict[str, Any]] = [] | |
| for blob in bucket.list_blobs(prefix=prefix): | |
| if not blob.name.endswith("/progress.json"): | |
| continue | |
| try: | |
| raw = blob.download_as_text(timeout=60) | |
| row = json.loads(raw) | |
| except Exception: | |
| continue | |
| if not isinstance(row, dict): | |
| continue | |
| parent = blob.name.rsplit("/", 2)[-2] if "/" in blob.name else "" | |
| m = DIR_RE.match(parent) | |
| try: | |
| seed_v = int(row["seed"]) if row.get("seed") is not None else (int(m.group(2)) if m else 0) | |
| fold_v = int(row["fold"]) if row.get("fold") is not None else (int(m.group(3)) if m else 0) | |
| except (TypeError, ValueError): | |
| seed_v, fold_v = 0, 0 | |
| units.append( | |
| { | |
| "arch": row.get("arch") or (m.group(1) if m else ""), | |
| "seed": seed_v, | |
| "fold": fold_v, | |
| "run_tag": row.get("run_tag") or parent, | |
| "status": row.get("status") or "unknown", | |
| "epoch": int(row.get("epoch") or 0), | |
| "epochs": int(row.get("epochs") or 0), | |
| "loss": row.get("loss"), | |
| "val_acc": row.get("val_acc"), | |
| "best_acc": row.get("best_acc"), | |
| "message": row.get("message") or "", | |
| "error": row.get("error"), | |
| "updated_at": row.get("updated_at"), | |
| } | |
| ) | |
| def _sort_key(u: dict) -> tuple: | |
| return (u.get("seed", 0), u.get("fold", 0), u.get("arch", ""), u.get("run_tag", "")) | |
| # Stable fallback order (seed / fold / arch). | |
| units.sort(key=_sort_key) | |
| # Newest worker activity first so a fresh Vertex run is visible at the top. | |
| units.sort(key=lambda u: u.get("updated_at") or "", reverse=True) | |
| done_like = {"done", "skipped", "error"} | |
| completed = sum(1 for u in units if u.get("status") in done_like) | |
| running = any(u.get("status") == "running" for u in units) | |
| return { | |
| "running": running, | |
| "started_at": None, | |
| "finished_at": None, | |
| "total_units": len(units), | |
| "completed_units": completed, | |
| "units": units, | |
| "log_tail": [], | |
| "source": "gcs", | |
| "gcs_uri": uri, | |
| } | |