fix(workers): retain granted deadlines across disconnects and restart

backend/core/db.py

@@ -274,7 +274,8 @@ _BASE_SCHEMA = """
         started_at REAL,
         finished_at REAL,
         lease_expires_at REAL,
-        grace_expires_at REAL
+        grace_expires_at REAL,
+        deadlines_json TEXT
     );
     CREATE INDEX IF NOT EXISTS idx_remote_attempts_task ON remote_task_attempts(task_id);
     CREATE INDEX IF NOT EXISTS idx_remote_attempts_worker ON remote_task_attempts(worker_id, state);

backend/migrations/versions/0011_remote_attempt_deadlines.py

@@ -0,0 +1,18 @@
+"""Retain the dispatch-time deadline policy across worker/control-plane loss."""
+from alembic import op
+import sqlalchemy as sa
+
+revision = "0011_remote_attempt_deadlines"
+down_revision = "0010_remote_worker_schema"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    columns = op.get_bind().execute(sa.text("PRAGMA table_info(remote_task_attempts)"))
+    if not any(row[1] == "deadlines_json" for row in columns):
+        op.add_column("remote_task_attempts", sa.Column("deadlines_json", sa.Text(), nullable=True))
+
+
+def downgrade() -> None:
+    op.drop_column("remote_task_attempts", "deadlines_json")

backend/worker/lifecycle.py

@@ -32,6 +32,7 @@ import uuid
 from dataclasses import dataclass, field
 from typing import Iterable, Optional
 
+from worker.deadlines import Deadlines
 from worker.clock import resolve
 from worker.errors import ErrorClass, WorkerError
 
@@ -224,6 +225,9 @@ class Attempt:
     stage: str = ""
     error: Optional[WorkerError] = None
 
+    # Snapshot the lease policy granted at dispatch, including after restart.
+    deadlines: Optional[Deadlines] = None
+
     def matches(self, *, session_epoch: Optional[int] = None) -> bool:
         """Fence check: reject messages from a superseded session."""
         if session_epoch is None:

backend/worker/scheduler.py

@@ -656,6 +656,7 @@ class Scheduler:
                 task.engine, task.model_id, task.operation
             ),
         )
+        attempt.deadlines = budget
         attempt.renew_lease(budget.accept_seconds, now=now)
         self._save(task, now=now)
         self._emit("assigned", task)
@@ -1298,6 +1299,9 @@ class Scheduler:
 
     def _budget_for(self, task: Task) -> deadline_policy.Deadlines:
         attempt = task.active_attempt
+        if attempt is not None and attempt.deadlines is not None:
+            return attempt.deadlines
+        # Legacy stored attempts have no snapshot; retain their prior policy.
         worker = self.pool.get(attempt.worker_id) if attempt else None
         return deadline_policy.for_task(
             task.operation,

backend/worker/task_store.py

@@ -33,6 +33,7 @@ from core.db import db_conn
 from core.path_security import UnsafePath, resolve_within, safe_filename
 from worker.clock import resolve
 from worker.errors import ErrorClass, WorkerError
+from worker.deadlines import Deadlines
 from worker.lifecycle import Attempt, AttemptState, PriorityClass, Task, TaskState
 
 logger = logging.getLogger("omnivoice.worker")
@@ -67,6 +68,8 @@ def _row_to_attempt(row) -> Attempt:
         state=AttemptState(row["state"]),
         created_at=float(row["created_at"]),
     )
+    if row["deadlines_json"]:
+        attempt.deadlines = Deadlines(**json.loads(row["deadlines_json"]))
     attempt.accepted_at = row["accepted_at"]
     attempt.started_at = row["started_at"]
     attempt.finished_at = row["finished_at"]
@@ -635,12 +638,13 @@ def _upsert_attempts(conn, task: Task) -> None:
             "INSERT INTO remote_task_attempts "
             "(id, task_id, worker_id, session_epoch, attempt_number, state, progress, stage, "
             " error_json, created_at, accepted_at, started_at, finished_at, lease_expires_at, "
-            " grace_expires_at) "
-            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "
+            " grace_expires_at, deadlines_json) "
+            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "
             "ON CONFLICT(id) DO UPDATE SET state=excluded.state, progress=excluded.progress, "
             " stage=excluded.stage, error_json=excluded.error_json, accepted_at=excluded.accepted_at, "
             " started_at=excluded.started_at, finished_at=excluded.finished_at, "
-            " lease_expires_at=excluded.lease_expires_at, grace_expires_at=excluded.grace_expires_at",
+            " lease_expires_at=excluded.lease_expires_at, grace_expires_at=excluded.grace_expires_at, "
+            " deadlines_json=excluded.deadlines_json",
             (
                 attempt.attempt_id,
                 attempt.task_id,
@@ -657,6 +661,7 @@ def _upsert_attempts(conn, task: Task) -> None:
                 attempt.finished_at,
                 attempt.lease_expires_at,
                 attempt.grace_expires_at,
+                json.dumps(attempt.deadlines.to_dict()) if attempt.deadlines else None,
             ),
         )
 

docs/remote-workers.md

@@ -245,6 +245,10 @@ grace window to come back, and if it returns carrying a finished result, that
 result is used — the task is never run twice just because a network blip
 happened. Only when the window expires is the task retried elsewhere.
 
+Each attempt retains the deadline budget granted at dispatch, including after a
+worker disconnect or control-plane restart; changed worker availability cannot
+shorten an in-flight attempt’s execution allowance.
+
 **A worker fails repeatedly.** After three consecutive failures that are
 actually its fault, it is paused for a minute, then automatically given one
 task to prove itself. Repeated trips back off further, up to thirty minutes.

tests/test_under_provisioned_gpu_budget_1804.py

@@ -320,21 +320,19 @@ def test_the_task_deadline_still_covers_the_raised_execution_budget(device):
     assert ceiling >= corrected.total_seconds
 
 
-def test_losing_the_worker_never_shortens_an_under_provisioned_budget():
-    """`Scheduler._budget_for` recomputes with no worker once one disconnects,
-    so `under_provisioned` goes False there. That must not shorten anything: no
-    worker means no `execution_device`, which `_base_execution_seconds` already
-    coerces to "cpu" — the very budget the floor raises an under-provisioned
-    card to.
-
-    Driven through a real scheduler rather than the policy alone (CodeRabbit on
-    the PR): the coercion lives in the disconnect path, so a test that only
-    called `for_task` would pass even if that path stopped doing it."""
+@pytest.mark.parametrize("gpu_seconds", [300, 900])
+def test_losing_the_worker_never_shortens_an_under_provisioned_budget(
+    gpu_seconds, monkeypatch, model_manager,
+):
+    """Worker loss must retain the granted budget, including GPU > CPU overrides."""
     from worker import deadlines
     from worker.identity import issue_session
     from worker.pool import WorkerPool
     from worker.scheduler import Scheduler
 
+    monkeypatch.setattr(model_manager, "GPU_JOB_TIMEOUT_S", gpu_seconds)
+    monkeypatch.setattr(model_manager, "CPU_JOB_TIMEOUT_S", 600.0)
+    monkeypatch.setattr(model_manager, "_CPU_GENERATE_TIMEOUT_EXPLICIT", True)
     now = 1000.0
     worker = _worker()  # 4 GB card, 6 GB engine
     pool = WorkerPool()
@@ -357,7 +355,7 @@ def test_losing_the_worker_never_shortens_an_under_provisioned_budget():
     # Bound: the worker is present, so its own verdict raises the budget.
     bound = sched._budget_for(task)
     on_cpu = deadlines.for_task("tts", text="short", execution_device="cpu")
-    assert bound.execution_seconds == on_cpu.execution_seconds
+    assert bound.execution_seconds == max(gpu_seconds, on_cpu.execution_seconds)
 
     # …and it survives the worker vanishing.
     pool.disconnect(worker.record.id)

tests/test_worker_remote_migration.py

@@ -36,3 +36,19 @@ def test_remote_schema_upgrades_from_previous_head(tmp_path, monkeypatch):
         } <= tables
         columns = {row[1] for row in conn.execute("PRAGMA table_info(remote_tasks)")}
         assert "pinned_worker_id" in columns
+
+
+def test_attempt_deadlines_upgrade_preserves_existing_attempt(tmp_path, monkeypatch):
+    db_path = tmp_path / "existing-remote.db"
+    with sqlite3.connect(db_path) as conn:
+        conn.execute("CREATE TABLE alembic_version (version_num VARCHAR(64) NOT NULL)")
+        conn.execute("INSERT INTO alembic_version VALUES ('0010_remote_worker_schema')")
+        conn.execute("CREATE TABLE remote_task_attempts (id TEXT PRIMARY KEY, state TEXT)")
+        conn.execute("INSERT INTO remote_task_attempts VALUES ('existing', 'running')")
+    monkeypatch.setenv("OMNIVOICE_DB_PATH", str(db_path))
+    _upgrade(str(db_path))
+    _upgrade(str(db_path))
+    with sqlite3.connect(db_path) as conn:
+        assert conn.execute(
+            "SELECT id, state, deadlines_json FROM remote_task_attempts"
+        ).fetchone() == ("existing", "running", None)

tests/test_worker_scheduler.py

@@ -418,8 +418,8 @@ def test_assignment_deadline_uses_selected_workers_device(monkeypatch):
     _submit(sched)
     assignment = sched.next_assignment(now=1000.0)
 
-    sched._budget_for(assignment.task)
-    assert seen == ["cuda", "cuda"]
+    assert sched._budget_for(assignment.task) is assignment.deadlines
+    assert seen == ["cuda"]  # Keep the granted policy instead of recomputing.
 
 
 def test_cpu_fallback_capability_overrides_machine_cuda(monkeypatch):

tests/test_worker_task_store.py

@@ -397,3 +397,19 @@ def test_result_directory_delete_is_durable_before_its_row_is_forgotten(
         == 1
     )
     assert task_store.get("t1") is None
+
+
+def test_dispatch_budget_survives_reload_without_worker(db):
+    from worker.deadlines import Deadlines
+    from worker.pool import WorkerPool
+    from worker.scheduler import Scheduler
+
+    task = _task()
+    task_store.create(task, now=1000.0)
+    attempt = task.assign(worker_id="w1", session_epoch=1, now=1001.0)
+    budget = Deadlines(20, 1800, 900, 120, 900, 75)
+    attempt.deadlines = budget
+    task_store.save(task, now=1002.0)
+    loaded = task_store.get(task.task_id)
+    assert loaded.active_attempt.deadlines == budget
+    assert Scheduler(WorkerPool(), persist=False)._budget_for(loaded) == budget