davanstrien/embedding-fleet-dashboard

1"""Embedding-fleet control plane — run-level view of a fan-out embedding run.2 3Renders one run of launch-embedding-fleet.py (uv-scripts/embeddings): documents-processed4progress, tokens, live ~$ cost vs the hard ceiling, ETA, GPU utilization, replica health,5and a per-worker table. Polls the run bucket + Jobs API every few seconds.6"""7 8import os9 10import gradio as gr11 12from control_plane import RunView, list_runs, load_run13 14BUCKET = os.environ.get("FLEET_BUCKET", "davanstrien/embedding-runs")15NAMESPACE = os.environ.get("FLEET_NAMESPACE") or BUCKET.split("/")[0]16POLL_SECS = float(os.environ.get("FLEET_POLL_SECS", "5"))17 18CSS = """19:root {20  --ink: #1a1a1a; --ink-2: #555; --ink-3: #999;21  --surface: #fffef9; --card: #ffffff; --line: #e4e2da;22  --accent: #3d6ea5; --ok: #2e7d43; --bad: #b3382c;23}24@media (prefers-color-scheme: dark) {25  :root { --ink: #ececec; --ink-2: #b0b0b0; --ink-3: #7a7a7a;26          --surface: #131313; --card: #1c1c1c; --line: #333;27          --accent: #7ba7d4; --ok: #6fbf85; --bad: #e07a6e; }28}29.gradio-container { max-width: 1080px !important; }30#fleet-html h2 { font-weight: 600; margin: 0 0 2px; }31.fleet-head { color: var(--ink-2); font-size: 0.92rem; margin-bottom: 14px; }32.fleet-head a { color: var(--accent); text-decoration: none; }33.tiles { display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr));34         gap: 10px; margin: 12px 0 6px; }35.tile { background: var(--card); border: 1px solid var(--line); border-radius: 6px;36        padding: 10px 14px; }37.tile .k { font-size: 0.72rem; text-transform: uppercase; letter-spacing: 0.05em;38           color: var(--ink-3); }39.tile .v { font-size: 1.45rem; font-variant-numeric: tabular-nums; color: var(--ink); }40.tile .s { font-size: 0.78rem; color: var(--ink-2); }41.tile .v.ok { color: var(--ok); } .tile .v.bad { color: var(--bad); }42.bar-wrap { margin: 10px 0 2px; }43.bar-label { display: flex; justify-content: space-between; font-size: 0.85rem;44             color: var(--ink-2); margin-bottom: 4px; font-variant-numeric: tabular-nums; }45.bar { height: 10px; background: var(--line); border-radius: 5px; overflow: hidden; }46.bar > div { height: 100%; background: var(--accent); border-radius: 5px 0 0 5px;47             transition: width 0.6s ease; }48"""49 50 51def fmt_int(n):52    return f"{n:,}" if n is not None else "—"53 54 55def fmt_secs(s):56    if s is None:57        return "—"58    if s < 90:59        return f"{s:.0f}s"60    if s < 5400:61        return f"{s / 60:.0f} min"62    return f"{s / 3600:.1f} h"63 64 65def render_run(view: RunView) -> str:66    m = view.manifest67    pct = 100.0 * view.rows_done / view.rows_total if view.rows_total else 0.068    in_url = f"https://huggingface.co/datasets/{m['input_dataset']}"69    out_url = f"https://huggingface.co/datasets/{m['output_dataset']}"70    ceiling = f"ceiling ≤ ${view.cost_ceiling_usd:,.2f}" if view.cost_ceiling_usd else "est."71    gpu = f"{view.gpu_util:.0f}%" if view.gpu_util is not None else "—"72    health_cls = "bad" if view.errored else "ok"73    eta = "done" if view.eta_secs == 0 else fmt_secs(view.eta_secs)74    return f"""75<h2>{m["output_dataset"].split("/")[-1]}</h2>76<div class="fleet-head">77  <a href="{in_url}">{m["input_dataset"]}</a> → <a href="{out_url}">{m["output_dataset"]}</a>78  · <code>{m["model"].split("/")[-1]}</code>79  · {view.num_shards} × {m["flavor"]} · run <code>{view.run_id}</code>80</div>81<div class="bar-wrap">82  <div class="bar-label">83    <span>{fmt_int(view.rows_done)} of {fmt_int(view.rows_total)} documents</span>84    <span>{pct:.1f}%</span>85  </div>86  <div class="bar"><div style="width:{min(pct, 100):.2f}%"></div></div>87</div>88<div class="tiles">89  <div class="tile"><div class="k">Tokens (est.)</div>90    <div class="v">{fmt_int(view.tokens_done_est)}</div></div>91  <div class="tile"><div class="k">Cost</div>92    <div class="v">~${view.cost_usd:,.2f}</div><div class="s">{ceiling}</div></div>93  <div class="tile"><div class="k">ETA</div>94    <div class="v">{eta}</div></div>95  <div class="tile"><div class="k">GPU util</div>96    <div class="v">{gpu}</div><div class="s">mean of running</div></div>97  <div class="tile"><div class="k">Replicas healthy</div>98    <div class="v {health_cls}">{view.healthy}/{view.num_shards}</div>99    <div class="s">{view.done} done · {view.errored} error</div></div>100</div>101"""102 103 104def worker_table(view: RunView):105    rows = []106    for w in view.workers:107        rows.append([108            w.rank,109            w.stage + (f" ({w.state})" if w.state and w.state != "running" else ""),110            f"{w.rows_done:,}" + (f" / {w.rows_total:,}" if w.rows_total else ""),111            f"{w.rows_per_sec:,.0f}" if w.rows_per_sec else "—",112            f"{w.gpu_util:.0f}%" if w.gpu_util is not None else "—",113            f"~${w.cost_usd:.3f}" if w.cost_usd is not None else "—",114            w.job_id or "—",115        ])116    return rows117 118 119def refresh(run_id):120    if not run_id:121        return "<p>No runs found in the bucket yet.</p>", []122    view = load_run(BUCKET, run_id, namespace=NAMESPACE)123    if view is None:124        return f"<p>Run <code>{run_id}</code> has no manifest.</p>", []125    return render_run(view), worker_table(view)126 127 128def init(request: gr.Request):129    runs = list_runs(BUCKET)130    wanted = dict(request.query_params).get("run")131    selected = wanted if wanted in runs else (runs[0] if runs else None)132    html, table = refresh(selected)133    return gr.Dropdown(choices=runs, value=selected), html, table134 135 136with gr.Blocks(css=CSS, title="Embedding Fleet") as demo:137    with gr.Row():138        # allow_custom_value: ?run= deep links and API calls may reference runs139        # that appeared after the choices list was built.140        run_dd = gr.Dropdown(label="Run", choices=[], scale=3, allow_custom_value=True)141        reload_btn = gr.Button("Reload runs", scale=1)142    html = gr.HTML(elem_id="fleet-html")143    table = gr.Dataframe(144        headers=["rank", "stage", "rows", "rows/s", "gpu", "~$", "job"],145        interactive=False, label="Workers",146    )147    timer = gr.Timer(POLL_SECS)148 149    demo.load(init, inputs=None, outputs=[run_dd, html, table])150    timer.tick(refresh, inputs=run_dd, outputs=[html, table])151    run_dd.change(refresh, inputs=run_dd, outputs=[html, table])152    reload_btn.click(lambda: gr.Dropdown(choices=list_runs(BUCKET)), outputs=run_dd)153 154if __name__ == "__main__":155    demo.launch()156