davanstrien/embedding-fleet-dashboard

1"""Data layer for the embedding-fleet control plane.2 3One poll tick = bucket reads (run manifest + worker heartbeats) + Jobs API reads4(stage, durations, one metrics sample per running job). All aggregation to the5run level happens here; app.py only renders.6 7Cost figures are client-side estimates (flavor unit price x running time), NOT8billing — always presented as "~$".9"""10 11from __future__ import annotations12 13import json14import logging15import tempfile16import time17from concurrent.futures import ThreadPoolExecutor18from dataclasses import dataclass, field19from pathlib import Path20 21from huggingface_hub import (22    download_bucket_files,23    fetch_job_metrics,24    list_bucket_tree,25    list_jobs,26    list_jobs_hardware,27)28 29TERMINAL_OK = {"COMPLETED"}30TERMINAL_BAD = {"ERROR", "CANCELED", "DELETED"}31 32 33def _stage_name(job) -> str:34    stage = job.status.stage if job.status else None35    return getattr(stage, "value", None) or str(stage or "UNKNOWN")36 37 38@dataclass39class WorkerRow:40    rank: int41    job_id: str | None = None42    stage: str = "UNKNOWN"43    rows_done: int = 044    rows_total: int | None = None45    rows_per_sec: float = 0.046    tokens_done_est: int | None = None47    gpu_util: float | None = None48    cost_usd: float | None = None49    state: str | None = None  # worker-reported: running/done/error50 51 52@dataclass53class RunView:54    run_id: str55    manifest: dict56    workers: list[WorkerRow] = field(default_factory=list)57 58    # run-level aggregates59    rows_done: int = 060    rows_total: int = 061    tokens_done_est: int = 062    cost_usd: float = 0.063    cost_ceiling_usd: float | None = None64    eta_secs: float | None = None65    gpu_util: float | None = None66    healthy: int = 067    errored: int = 068    done: int = 069    num_shards: int = 070 71 72_PRICING: dict | None = None73 74 75def pricing() -> dict:76    global _PRICING77    if _PRICING is None:78        _PRICING = {hw.name: hw for hw in list_jobs_hardware()}79    return _PRICING80 81 82def parse_timeout_secs(timeout) -> float | None:83    """'20m' / '1h' / '90s' / plain seconds -> seconds."""84    if timeout is None:85        return None86    s = str(timeout).strip().lower()87    try:88        mult = {"s": 1, "m": 60, "h": 3600, "d": 86400}.get(s[-1])89        return float(s[:-1]) * mult if mult else float(s)90    except (ValueError, IndexError):91        return None92 93 94def list_runs(bucket: str) -> list[str]:95    """Run ids under runs/, newest first (ids are timestamp-prefixed)."""96    try:97        ids = [Path(e.path.rstrip("/")).name98               for e in list_bucket_tree(bucket, prefix="runs/", recursive=False)99               if e.__class__.__name__ == "BucketFolder"]100        # Timestamp-prefixed ids first (newest first), ad-hoc ids after.101        return sorted(set(ids), key=lambda r: (r[:8].isdigit(), r), reverse=True) if ids else []102    except Exception:103        return []104 105 106def _read_bucket_json(bucket: str, paths: list[str]) -> dict[str, dict]:107    """Fetch small JSON files from the bucket; missing files are skipped."""108    out: dict[str, dict] = {}109    if not paths:110        return out111    with tempfile.TemporaryDirectory() as td:112        pairs = [(p, Path(td) / p.replace("/", "__")) for p in paths]113        try:114            download_bucket_files(bucket, pairs, raise_on_missing_files=False)115        except Exception:116            return out117        for src, dst in pairs:118            if dst.exists():119                try:120                    out[src] = json.loads(dst.read_text())121                except (json.JSONDecodeError, OSError):122                    pass123    return out124 125 126def _sample_gpu_util(job_id: str, timeout: float = 3.0) -> float | None:127    """One metrics sample -> mean GPU utilization. Never blocks past `timeout`."""128 129    def _one():130        gen = iter(fetch_job_metrics(job_id=job_id))131        try:132            raw = next(gen)133        finally:134            getattr(gen, "close", lambda: None)()135        gpus = raw.get("gpus") or {}136        utils = [float(g.get("utilization") or 0) for g in gpus.values()]137        return sum(utils) / len(utils) if utils else None138 139    with ThreadPoolExecutor(max_workers=1) as pool:140        fut = pool.submit(_one)141        try:142            return fut.result(timeout=timeout)143        except Exception:144            return None145 146 147def _accrued_cost(job, hw_pricing: dict) -> float | None:148    flavor = getattr(job.flavor, "value", None) or (str(job.flavor) if job.flavor else None)149    hw = hw_pricing.get(flavor)150    if not hw:151        return None152    secs = job.durations.running_secs if job.durations else None153    if not secs and job.started_at and _stage_name(job) == "RUNNING":154        secs = time.time() - job.started_at.timestamp()155    if not secs:156        return None157    return secs / 60.0 * hw.unit_cost_usd158 159 160def load_run(bucket: str, run_id: str, namespace: str | None = None) -> RunView | None:161    """One full poll tick: manifest + heartbeats + job stages + metrics samples -> RunView."""162    manifest = _read_bucket_json(bucket, [f"runs/{run_id}/run.json"]).get(f"runs/{run_id}/run.json")163    if not manifest:164        return None165    n = manifest["num_shards"]166    view = RunView(run_id=run_id, manifest=manifest, num_shards=n,167                   rows_total=manifest.get("rows_total") or 0)168 169    status_paths = [f"runs/{run_id}/status/{i:05d}.json" for i in range(n)]170    statuses = _read_bucket_json(bucket, status_paths)171 172    # Jobs by label (server-side filter); fall back to manifest job_ids via list comprehension.173    jobs_by_id = {}174    try:175        for j in list_jobs(labels={"embedding-fleet-run": run_id}, namespace=namespace):176            jobs_by_id[j.id] = j177    except Exception as e:178        logging.getLogger("control-plane").warning(f"list_jobs failed: {e!r}")179    manifest_job_ids = manifest.get("job_ids") or []180 181    hw_pricing = pricing()182    workers: list[WorkerRow] = []183    running_job_ids: list[str] = []184    for rank in range(n):185        row = WorkerRow(rank=rank)186        st = statuses.get(f"runs/{run_id}/status/{rank:05d}.json")187        if st:188            row.state = st.get("state")189            row.rows_done = st.get("rows_done") or 0190            row.rows_total = st.get("rows_total")191            row.rows_per_sec = st.get("rows_per_sec") or 0.0192            row.tokens_done_est = st.get("tokens_done_est")193            row.job_id = st.get("job_id")194        if row.job_id is None and rank < len(manifest_job_ids):195            row.job_id = manifest_job_ids[rank]196        job = jobs_by_id.get(row.job_id)197        if job is None and str(rank) in {j.labels.get("rank") for j in jobs_by_id.values() if j.labels}:198            job = next(j for j in jobs_by_id.values() if (j.labels or {}).get("rank") == str(rank))199        if job is not None:200            row.stage = _stage_name(job)201            row.cost_usd = _accrued_cost(job, hw_pricing)202            if row.stage == "RUNNING":203                running_job_ids.append(row.job_id)204        workers.append(row)205 206    # One GPU sample per running job, in parallel, bounded.207    if running_job_ids:208        with ThreadPoolExecutor(max_workers=min(8, len(running_job_ids))) as pool:209            samples = dict(zip(running_job_ids,210                               pool.map(_sample_gpu_util, running_job_ids)))211        for row in workers:212            if row.job_id in samples:213                row.gpu_util = samples[row.job_id]214 215    # Consolidator cost (labeled role=consolidate) counts toward the run.216    consolidator_cost = sum(217        _accrued_cost(j, hw_pricing) or 0.0218        for j in jobs_by_id.values()219        if (j.labels or {}).get("role") == "consolidate"220    )221 222    # ---- aggregate ----223    view.workers = workers224    view.rows_done = sum(w.rows_done for w in workers)225    view.tokens_done_est = sum(w.tokens_done_est or 0 for w in workers)226    view.cost_usd = sum(w.cost_usd or 0.0 for w in workers) + consolidator_cost227    view.done = sum(1 for w in workers if w.state == "done" or w.stage in TERMINAL_OK)228    view.errored = sum(1 for w in workers if w.state == "error" or w.stage in TERMINAL_BAD)229    view.healthy = n - view.errored230    gpu_vals = [w.gpu_util for w in workers if w.gpu_util is not None]231    view.gpu_util = sum(gpu_vals) / len(gpu_vals) if gpu_vals else None232 233    timeout_secs = parse_timeout_secs(manifest.get("timeout"))234    hw = hw_pricing.get(manifest.get("flavor"))235    if timeout_secs and hw:236        view.cost_ceiling_usd = n * timeout_secs / 60.0 * hw.unit_cost_usd237 238    active_rps = sum(w.rows_per_sec for w in workers239                     if w.state == "running" and w.stage not in TERMINAL_BAD)240    remaining = max((view.rows_total or 0) - view.rows_done, 0)241    if active_rps > 0 and remaining > 0:242        view.eta_secs = remaining / active_rps243    elif remaining == 0 and view.rows_total:244        view.eta_secs = 0.0245    return view246