Report measured usage costs and link completed runs to the annotation viewer

README.md

@@ -123,3 +123,13 @@ out. A refreshable dashboard lists each user's jobs, progress and cancellation.
 
 Hugging Face login is available on the deployed Space. Local mode is preview
 only: it cannot launch jobs using a developer's cached credentials.
+
+The Jobs dashboard refreshes every 15 seconds and shows usage-derived cost:
+provider-reported input/output/cache tokens at the pinned rates, plus HF-reported
+CPU runtime rounded to minutes. Unknown request charges remain separately
+reserved; this is not an invoice. Worker usage checkpoints are saved every ten
+seconds. The billing link opens the user's final HF charges.
+
+Finished annotated copies offer **Show annotations**, opening the first
+materialized episode with `?tab=annotations` in the visualizer. Private results
+require signing in there. JSONL-only exports retain their Results download link.

deploy.py

@@ -57,7 +57,7 @@ def main():
             "frontend/dist/**",
         ],
         delete_patterns=["app.py"],
-        commit_message="Discover dataset camera views before estimating annotation jobs",
+        commit_message="Report measured usage costs and link completed runs to the annotation viewer",
     )
     print(f"https://huggingface.co/spaces/{repo}\nRevision: {result.oid}")
 

frontend/src/main.tsx

@@ -117,7 +117,24 @@ type Progress = {
   expected: number;
   completed: number;
   failed: number;
-  usage: { measured_inference_usd?: number };
+  usage: {
+    measured_inference_usd?: number;
+    successful_calls?: number;
+    input_tokens?: number;
+    output_tokens?: number;
+    rate_limited_requests?: number;
+  };
+  cost?: {
+    inference_usd: number | null;
+    cpu_usd: number | null;
+    total_usd: number | null;
+    uncertain_inference_usd: number;
+    running_seconds: number | null;
+    basis: string;
+    billing_url: string;
+  };
+  annotation_url?: string | null;
+  output_mode?: string;
 };
 const EXPORT = "Annotation files only",
   COPY = "Annotated dataset copy";
@@ -348,7 +365,12 @@ function App() {
     } catch {}
   }, [dark]);
   useEffect(() => {
-    if (page === "jobs" && session?.user) void loadJobs();
+    if (page !== "jobs" || !session?.user) return;
+    void loadJobs();
+    const timer = window.setInterval(() => {
+      if (document.visibilityState === "visible") void loadJobs();
+    }, 15000);
+    return () => window.clearInterval(timer);
   }, [page, session?.user?.name]);
   useEffect(() => {
     if (review) modalRef.current?.showModal();
@@ -471,20 +493,22 @@ function App() {
     setJobsLoading(true);
     setError("");
     try {
-      setJobs(await api<Job[]>("/api/jobs"));
+      const rows = await api<Job[]>("/api/jobs");
+      setJobs(rows);
+      await Promise.allSettled(rows.map((job) => getProgress(job, true)));
     } catch (e) {
       setError((e as Error).message);
     } finally {
       setJobsLoading(false);
     }
   }
-  async function getProgress(job: Job) {
+  async function getProgress(job: Job, quiet = false) {
     try {
       setProgress((old) => ({ ...old }));
       const p = await api<Progress>(`/api/jobs/${job.id}/progress`);
       setProgress((old) => ({ ...old, [job.id]: p }));
     } catch (e) {
-      setNotice((e as Error).message);
+      if (!quiet) setNotice((e as Error).message);
     }
   }
   async function submit() {
@@ -1440,7 +1464,92 @@ function App() {
                           </small>
                         </div>
                       )}
+                      {progress[j.id]?.cost && (
+                        <div className="job-cost">
+                          <div className="job-cost-total">
+                            <div>
+                              <span>Usage cost</span>
+                              <strong>
+                                {progress[j.id].cost!.total_usd == null
+                                  ? "Pending"
+                                  : new Intl.NumberFormat("en-US", {
+                                      style: "currency",
+                                      currency: "USD",
+                                      minimumFractionDigits: 4,
+                                    }).format(progress[j.id].cost!.total_usd!)}
+                              </strong>
+                            </div>
+                            <span className="quiet-badge">
+                              {terminal.has(j.stage)
+                                ? "Run ended"
+                                : "Updates every 15 s"}
+                            </span>
+                          </div>
+                          <div className="job-cost-breakdown">
+                            <span>
+                              Model inference{" "}
+                              <b>
+                                {progress[j.id].cost!.inference_usd == null
+                                  ? "Pending"
+                                  : "$" +
+                                    progress[j.id].cost!.inference_usd!.toFixed(
+                                      4,
+                                    )}
+                              </b>
+                            </span>
+                            <span>
+                              CPU runtime{" "}
+                              <b>
+                                {progress[j.id].cost!.cpu_usd == null
+                                  ? "Pending"
+                                  : "$" +
+                                    progress[j.id].cost!.cpu_usd!.toFixed(4)}
+                              </b>
+                            </span>
+                            <span>
+                              Model calls{" "}
+                              <b>
+                                {num(
+                                  progress[j.id].usage.successful_calls || 0,
+                                )}
+                              </b>
+                            </span>
+                          </div>
+                          {progress[j.id].cost!.uncertain_inference_usd >
+                            0.000001 && (
+                            <p className="cost-uncertain">
+                              Additional uncertain inference charges: up to $
+                              {progress[
+                                j.id
+                              ].cost!.uncertain_inference_usd.toFixed(4)}
+                              . These are reserved separately.
+                            </p>
+                          )}
+                          <p>
+                            {progress[j.id].cost!.basis}{" "}
+                            <a
+                              href={progress[j.id].cost!.billing_url}
+                              target="_blank"
+                              rel="noreferrer"
+                            >
+                              View HF billing <ArrowUpRight size={10} />
+                            </a>
+                          </p>
+                        </div>
+                      )}
                       <div className="job-actions">
+                        {terminal.has(j.stage) &&
+                          progress[j.id]?.annotation_url && (
+                            <a
+                              className="button primary"
+                              href={progress[j.id].annotation_url!}
+                              target="_blank"
+                              rel="noreferrer"
+                            >
+                              <Layers3 size={14} /> Show annotations{" "}
+                              <ArrowUpRight size={13} />
+                            </a>
+                          )}
                         <button
                           className="button secondary"
                           onClick={() => getProgress(j)}
@@ -1480,9 +1589,11 @@ function App() {
                 </div>
               )}
               <p className="section-footnote">
-                Only jobs launched from Annotation Studio appear here. A
-                completed job can contain partial results; check the saved run
-                report.
+                Show annotations opens the visualizer’s Annotations tab for an
+                annotated dataset copy. Sign in there to access private results.
+                Annotation-file exports can be downloaded from Results. Only
+                jobs launched from Annotation Studio appear here. A completed
+                job can contain partial results; check the saved run report.
               </p>
             </>
           )}

frontend/src/style.css

@@ -2491,3 +2491,71 @@ footer > span {
   line-height: 1.6;
   margin-top: 6px;
 }
+
+.job-cost {
+  margin: 16px 0;
+  padding: 18px;
+  background: var(--surface-1);
+  border: 1px solid var(--line);
+  border-radius: 8px;
+}
+.job-cost-total {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 12px;
+}
+.job-cost-total > div {
+  display: flex;
+  align-items: baseline;
+  gap: 12px;
+}
+.job-cost-total strong {
+  font-size: 23px;
+  font-variant-numeric: tabular-nums;
+  letter-spacing: -0.5px;
+}
+.job-cost-total span {
+  font-size: 12px;
+  color: var(--fg-muted);
+}
+.job-cost-breakdown {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 12px 28px;
+  margin-top: 16px;
+  font-size: 12px;
+  color: var(--fg-muted);
+}
+.job-cost-breakdown b {
+  font-weight: 550;
+  color: var(--fg-soft);
+  margin-left: 8px;
+  font-variant-numeric: tabular-nums;
+}
+.job-cost p {
+  font-size: 11px;
+  color: var(--fg-faint);
+  line-height: 1.6;
+  margin-top: 14px;
+}
+.job-cost p a {
+  text-decoration: underline;
+  white-space: nowrap;
+}
+.job-cost .cost-uncertain {
+  color: var(--fg-soft);
+}
+@media (max-width: 520px) {
+  .job-cost-total {
+    align-items: flex-start;
+  }
+  .job-cost-total > div {
+    flex-direction: column;
+    gap: 5px;
+  }
+  .job-cost-breakdown {
+    flex-direction: column;
+    gap: 8px;
+  }
+}

provider.py

@@ -17,10 +17,16 @@ class BudgetExceededError(RuntimeError):
     pass
 
 
+class DeadlineExceededError(BudgetExceededError):
+    pass
+
+
 class Budget:
-    def __init__(self, limit, input_rate, output_rate):
+    def __init__(self, limit, input_rate, output_rate, cached_input_rate=None):
         self.limit, self.input_rate, self.output_rate = limit, input_rate, output_rate
         self.reserved, self.measured, self.calls = 0.0, 0.0, 0
+        self.cached_input_rate = input_rate if cached_input_rate is None else cached_input_rate
+        self.input_tokens = self.output_tokens = self.cached_tokens = self.attempts = self.rate_limits = 0
         self.lock = threading.Lock()
 
     def reserve(self, max_tokens):
@@ -30,6 +36,7 @@ class Budget:
             if self.reserved + amount > self.limit:
                 raise BudgetExceededError("Spending target reached; completed results are saved.")
             self.reserved += amount
+            self.attempts += 1
         return amount
 
     def reconcile(self, reservation, usage):
@@ -38,11 +45,22 @@ class Budget:
         values = [usage.get("prompt_tokens"), usage.get("completion_tokens")]
         if any(type(v) is not int or v < 0 for v in values):
             return
-        cost = (values[0] * self.input_rate + values[1] * self.output_rate) / 1e6
+        cached = usage.get(
+            "cached_tokens", (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0)
+        )
+        cached = min(values[0], max(0, cached)) if type(cached) is int else 0
+        cost = (
+            (values[0] - cached) * self.input_rate
+            + cached * self.cached_input_rate
+            + values[1] * self.output_rate
+        ) / 1e6
         with self.lock:
             self.reserved += cost - reservation
             self.measured += cost
             self.calls += 1
+            self.input_tokens += values[0]
+            self.output_tokens += values[1]
+            self.cached_tokens += cached
 
     def summary(self):
         with self.lock:
@@ -51,6 +69,12 @@ class Budget:
                 "reserved_including_uncertain_usd": self.reserved,
                 "successful_calls": self.calls,
                 "inference_target_usd": self.limit,
+                "input_tokens": self.input_tokens,
+                "output_tokens": self.output_tokens,
+                "cached_input_tokens": self.cached_tokens,
+                "attempted_requests": self.attempts,
+                "requests_without_usage": self.attempts - self.calls,
+                "rate_limited_requests": self.rate_limits,
             }
 
 
@@ -82,17 +106,25 @@ def compact_images(messages, max_bytes=2_800_000):
 
 
 class Provider:
-    def __init__(self, token, model, budget):
+    def __init__(self, token, model, budget, deadline=None):
         self.token, self.model, self.budget = token, model, budget
         self.semaphore = threading.BoundedSemaphore(2)
+        self.deadline = deadline
         self.client = _GenericTextClient(self.generate, VlmConfig(max_new_tokens=4096, temperature=0.2))
 
     def generate_json(self, *args, **kwargs):
         return self.client.generate_json(*args, **kwargs)
 
+    def remaining_time(self):
+        left = self.deadline - time.time() if self.deadline else 300
+        if left <= 1:
+            raise DeadlineExceededError("Runtime limit reached; saving completed episodes.")
+        return left
+
     def generate(self, batch, max_tokens, temperature):
         output = []
         for messages in batch:
+            self.remaining_time()
             content, _ = _to_openai_messages(messages)
             body = {
                 "model": self.model,
@@ -104,22 +136,26 @@ class Provider:
             }
             with self.semaphore:
                 for attempt in range(8):
+                    left = self.remaining_time()
                     reservation = self.budget.reserve(max_tokens)
                     try:
                         response = requests.post(
                             "https://router.huggingface.co/v1/chat/completions",
                             headers={"Authorization": "Bearer " + self.token},
                             json=body,
-                            timeout=(20, 300),
+                            timeout=(min(20, left), min(300, left)),
                         )
                     except requests.RequestException:
+                        self.remaining_time()
                         # The provider may have processed the request. Never release its reservation.
                         raise RuntimeError(
                             "Inference connection failed; charges are reserved conservatively."
                         ) from None
                     if response.status_code == 429:
                         # Rejections may be ambiguous; retain the reservation before retrying.
-                        time.sleep(min(2**attempt, 30))
+                        with self.budget.lock:
+                            self.budget.rate_limits += 1
+                        time.sleep(min(2**attempt, 30, self.remaining_time()))
                         continue
                     if response.status_code >= 400:
                         raise RuntimeError(f"Inference provider returned HTTP {response.status_code}.")

reporting.py

@@ -0,0 +1,48 @@
+"""Usage-derived cost reporting and links to materialized annotations."""
+
+import json
+import math
+from datetime import UTC, datetime
+
+from studio import COPY
+
+
+def usage_costs(job, report):
+    usage = report.get("usage") or {}
+    manifest = report.get("manifest") or {}
+    if not manifest:
+        try:
+            manifest = json.loads((getattr(job, "environment", None) or {}).get("ANNOTATION_MANIFEST", "{}"))
+        except (TypeError, ValueError):
+            manifest = {}
+    rate = manifest.get("costs", {}).get("cpu_rate")
+    seconds = getattr(getattr(job, "durations", None), "running_secs", None)
+    if seconds is None and getattr(job, "started_at", None):
+        end = getattr(job, "finished_at", None) or datetime.now(UTC)
+        seconds = max(0, (end - job.started_at).total_seconds())
+    cpu = None
+    if isinstance(seconds, (int, float)) and math.isfinite(seconds) and rate is not None:
+        cpu = math.ceil(max(0, seconds) / 60) * rate
+    inference = usage.get("measured_inference_usd")
+    return {
+        "inference_usd": inference,
+        "cpu_usd": cpu,
+        "total_usd": inference + cpu if inference is not None and cpu is not None else None,
+        "uncertain_inference_usd": max(
+            0, usage.get("reserved_including_uncertain_usd", 0) - (inference or 0)
+        ),
+        "running_seconds": seconds,
+        "basis": "Reported tokens × published rates, plus HF job runtime rounded to minutes. Final HF billing may differ.",
+        "billing_url": "https://huggingface.co/settings/billing",
+    }
+
+
+def annotation_url(report):
+    manifest = report.get("manifest", {})
+    written = report.get("written_episode_ids", [])
+    # Old completed copy reports predate the explicit materialization checkpoint.
+    if not written and report.get("status") == "completed":
+        written = [e["episode_index"] for e in report.get("episodes", []) if e["status"] == "completed"]
+    if manifest.get("mode") != COPY or not written:
+        return None
+    return f"https://lerobot-visualize-dataset.hf.space/{manifest['output']}/{min(written)}?tab=annotations"

server.py

@@ -18,6 +18,7 @@ from fastapi.staticfiles import StaticFiles
 from huggingface_hub import HfApi, hf_hub_download
 from prompt_config import public_catalog, validate_prompts
 from pydantic import BaseModel, Field
+from reporting import annotation_url, usage_costs
 from studio import EXPORT, FEATURES, launch, prepare, read_cameras
 
 app = FastAPI(title="LeRobot Annotation Studio", docs_url=None, redoc_url=None, openapi_url=None)
@@ -390,11 +391,17 @@ def progress(job_id: str, request: Request):
             report = json.loads(Path(path).read_text())
         episodes = report.get("episodes", [])
         return {
-            "status": report["status"],
+            "status": "interrupted"
+            if job.status.stage in {"COMPLETED", "ERROR", "CANCELED", "DELETED"}
+            and report["status"] == "running"
+            else report["status"],
             "expected": report["expected_episodes"],
             "completed": sum(e["status"] == "completed" for e in episodes),
             "failed": sum(e["status"] != "completed" for e in episodes),
             "usage": report.get("usage", {}),
+            "cost": usage_costs(job, report),
+            "annotation_url": annotation_url(report),
+            "output_mode": report.get("manifest", {}).get("mode"),
         }
     except Exception as exc:
         raise HTTPException(400, "No saved progress yet. The job may still be starting.") from exc

studio.py

@@ -19,6 +19,7 @@ from prompt_config import TEMPLATES, validate_prompts
 MODEL = "Qwen/Qwen3.8-Flash-Next:featherless-ai"
 INPUT_RATE = 0.15
 OUTPUT_RATE = 0.50
+CACHED_INPUT_RATE = 0.03
 PRICE_DATE = "2026-09-25"
 FEATURES = {
     "Subtasks": "subtask",
@@ -217,6 +218,7 @@ def estimate(selection, features, vqa_hz, hours, cpu_rate, prompts=None):
         "hours": hours,
         "input_rate": INPUT_RATE,
         "output_rate": OUTPUT_RATE,
+        "cached_input_rate": CACHED_INPUT_RATE,
         "price_date": PRICE_DATE,
     }
 
@@ -255,7 +257,7 @@ def prepare(dataset, start, count, features, camera, vqa_hz, hours, mode, token=
     }
 
 
-def job_command(space, revision):
+def job_command(space, revision, timeout_minutes=None):
     # Only trusted deployment values enter shell syntax. User options travel in JSON env.
     validate_repo_id(space)
     if not re.fullmatch(r"[a-f0-9]{40}", revision):
@@ -271,6 +273,9 @@ def job_command(space, revision):
         "\n".join(
             [
                 "set -euo pipefail",
+                f"export ANNOTATION_DEADLINE_EPOCH=$(( $(date +%s) + {timeout_minutes * 60 - 30} ))"
+                if timeout_minutes
+                else ":",
                 "mkdir -p /workspace/app /workspace/source",
                 "cd /workspace",
                 "pip install --quiet huggingface_hub==1.26.0",
@@ -285,7 +290,11 @@ def job_command(space, revision):
     ]
 
 
-def launch(plan, output, budget, acknowledged, token, space=None, api_factory=HfApi):
+def launch(plan, output, budget, acknowledged, token, space=None, api_factory=HfApi, timeout_minutes=None):
+    if timeout_minutes is not None and (
+        type(timeout_minutes) is not int or not 1 <= timeout_minutes <= plan["costs"]["hours"] * 60
+    ):
+        raise ValueError("Runtime limit must be within the quoted timeout.")
     if not token:
         raise ValueError("Sign in with Hugging Face before starting a job.")
     if not acknowledged:
@@ -343,13 +352,14 @@ def launch(plan, output, budget, acknowledged, token, space=None, api_factory=Hf
             "space": space,
             "space_revision": revision,
             "owner": username,
+            "runtime_limit_minutes": timeout_minutes,
         }
     )
     job = api.run_job(
         image="python:3.12-slim",
-        command=job_command(space, revision),
+        command=job_command(space, revision, timeout_minutes),
         flavor="cpu-upgrade",
-        timeout=f"{plan['costs']['hours']}h",
+        timeout=f"{timeout_minutes}m" if timeout_minutes else f"{plan['costs']['hours']}h",
         namespace=username,
         name=f"annotate-{plan['nonce'][:12]}",
         labels={

test_studio.py

@@ -465,3 +465,71 @@ def test_camera_discovery_rejects_unsupported_sources():
         studio.video_cameras({"codebase_version": "v3.0", "fps": 30, "features": {}})
     with pytest.raises(ValueError, match="v3"):
         studio.video_cameras({"codebase_version": "v2.1"})
+
+
+def test_usage_costs_use_reported_runtime_and_keep_unknown_charges_separate():
+    from reporting import usage_costs
+
+    job = SimpleNamespace(durations=SimpleNamespace(running_secs=61))
+    result = usage_costs(
+        job,
+        {
+            "manifest": {"costs": {"cpu_rate": 0.0005}},
+            "usage": {"measured_inference_usd": 0.2, "reserved_including_uncertain_usd": 0.24},
+        },
+    )
+    assert result["cpu_usd"] == 0.001
+    assert result["total_usd"] == pytest.approx(0.201)
+    assert result["uncertain_inference_usd"] == pytest.approx(0.04)
+    assert usage_costs(SimpleNamespace(), {})["total_usd"] is None
+
+
+def test_visualizer_links_only_target_materialized_copy_annotations():
+    from reporting import annotation_url
+
+    report = {
+        "manifest": {"mode": studio.COPY, "output": "alice/out"},
+        "status": "partial",
+        "written_episode_ids": [12, 5],
+    }
+    assert annotation_url(report) == "https://lerobot-visualize-dataset.hf.space/alice/out/5?tab=annotations"
+    report["manifest"]["mode"] = studio.EXPORT
+    assert annotation_url(report) is None
+    report["manifest"]["mode"] = studio.COPY
+    report["written_episode_ids"] = []
+    assert annotation_url(report) is None
+
+
+def test_cached_usage_and_deadline_stop_before_a_paid_request(monkeypatch):
+    import provider
+
+    budget = Budget(1, 1, 2, cached_input_rate=0.2)
+    reserved = budget.reserve(100)
+    budget.reconcile(reserved, {"prompt_tokens": 10, "completion_tokens": 20, "cached_tokens": 6})
+    assert budget.summary()["measured_inference_usd"] == pytest.approx(0.0000452)
+    assert budget.summary()["input_tokens"] == 10
+    assert budget.summary()["cached_input_tokens"] == 6
+    post = Mock()
+    monkeypatch.setattr(provider.requests, "post", post)
+    client = provider.Provider("token", studio.MODEL, budget, deadline=time.time() - 1)
+    with pytest.raises(provider.DeadlineExceededError):
+        client.generate([[{"role": "user", "content": "hello"}]], 100, 0.2)
+    post.assert_not_called()
+
+
+def test_five_minute_run_has_job_and_worker_deadlines():
+    api = api_mock()
+    studio.launch(
+        plan(),
+        "alice/output",
+        30,
+        True,
+        "token",
+        space="lerobot/annotation-studio",
+        api_factory=lambda **_: api,
+        timeout_minutes=5,
+    )
+    kwargs = api.run_job.call_args.kwargs
+    assert kwargs["timeout"] == "5m"
+    assert "ANNOTATION_DEADLINE_EPOCH" in str(kwargs["command"])
+    assert "270" in str(kwargs["command"])

worker.py

@@ -5,6 +5,8 @@ import json
 import math
 import os
 import shutil
+import threading
+import time
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from contextvars import ContextVar
@@ -13,7 +15,7 @@ from pathlib import Path
 import av
 from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, snapshot_download
 from prompt_config import validate_prompts
-from provider import Budget, BudgetExceededError, Provider
+from provider import Budget, BudgetExceededError, DeadlineExceededError, Provider
 from studio import COPY, FEATURES, read_selection, relative_path
 
 from lerobot.annotations.steerable_pipeline.config import InterjectionsConfig, PlanConfig, VqaConfig
@@ -302,18 +304,30 @@ def run(manifest, work=Path("/workspace/run")):
         allow_patterns=["meta/**"],
     )
     budget = Budget(
-        manifest["inference_budget"], manifest["costs"]["input_rate"], manifest["costs"]["output_rate"]
+        manifest["inference_budget"],
+        manifest["costs"]["input_rate"],
+        manifest["costs"]["output_rate"],
+        manifest["costs"].get("cached_input_rate"),
     )
-    provider = Provider(token, manifest["model"], budget)
+    deadline = float(os.environ.get("ANNOTATION_DEADLINE_EPOCH", "0")) or None
+    provider = Provider(token, manifest["model"], budget, deadline=deadline)
     progress = {
         "manifest": manifest,
         "pipeline": json.loads(Path(__file__).with_name("pipeline-version.json").read_text()),
         "status": "running",
         "episodes": [],
         "expected_episodes": len(ids),
+        "written_episode_ids": [],
+        "worker_started_at": time.time(),
     }
 
+    commit_lock = threading.RLock()
+
     def checkpoint(extra=()):
+        with commit_lock:
+            return _checkpoint(extra)
+
+    def _checkpoint(extra=()):
         progress["usage"] = budget.summary()
         operations = [
             CommitOperationAdd(
@@ -330,6 +344,17 @@ def run(manifest, work=Path("/workspace/run")):
         )
 
     checkpoint()
+    heartbeat_stop = threading.Event()
+
+    def heartbeat():
+        while not heartbeat_stop.wait(10):
+            try:
+                checkpoint()
+            except Exception:
+                print("Progress update delayed; will retry.", flush=True)
+
+    heartbeat_thread = threading.Thread(target=heartbeat, daemon=True)
+    heartbeat_thread.start()
     shards = defaultdict(list)
     info = selection["info"]
     for row in selection["episodes"]:
@@ -383,8 +408,13 @@ def run(manifest, work=Path("/workspace/run")):
                     try:
                         result = future.result()
                         completed.append(record)
-                    except BudgetExceededError:
-                        result = {"episode_index": record.episode_index, "status": "budget_stopped"}
+                    except BudgetExceededError as exc:
+                        result = {
+                            "episode_index": record.episode_index,
+                            "status": "time_stopped"
+                            if isinstance(exc, DeadlineExceededError)
+                            else "budget_stopped",
+                        }
                         stop = True
                         for pending in futures:
                             pending.cancel()
@@ -418,11 +448,16 @@ def run(manifest, work=Path("/workspace/run")):
                         for p in [*written, root / "meta/info.json"]
                     ]
                 )
+                progress["written_episode_ids"].extend(r.episode_index for r in completed)
+                checkpoint()
             for file in files:
                 (root / file).unlink(missing_ok=True)
             shutil.rmtree(staging, ignore_errors=True)
             if stop:
                 break
+        heartbeat_stop.set()
+        heartbeat_thread.join(timeout=15)
+        progress["worker_finished_at"] = time.time()
         completed_count = sum(e["status"] == "completed" for e in progress["episodes"])
         progress["status"] = "completed" if completed_count == len(ids) else "partial"
         final_commit = checkpoint()
@@ -434,7 +469,9 @@ def run(manifest, work=Path("/workspace/run")):
                 api.delete_tag(output, tag=version, repo_type="dataset")
             api.create_tag(output, tag=version, repo_type="dataset", revision=final_commit.oid)
     except Exception as exc:
-        progress.update(status="partial", error_type=type(exc).__name__)
+        heartbeat_stop.set()
+        heartbeat_thread.join(timeout=15)
+        progress.update(status="partial", error_type=type(exc).__name__, worker_finished_at=time.time())
         checkpoint()
         raise RuntimeError("Annotation stopped; check the saved run report for completed episodes.") from None
     print(f"Result: {progress['status']} — https://huggingface.co/datasets/{output}", flush=True)