fix: size remote deadlines for worker device

backend/worker/deadlines.py

@@ -126,7 +126,9 @@ class Deadlines:
         }
 
 
-def _base_execution_seconds(text: Optional[str]) -> float:
+def _base_execution_seconds(
+    text: Optional[str], *, execution_device: Optional[str] = None
+) -> float:
     """Delegate to model_manager's budget; fall back to its formula.
 
     The lazy import keeps this module usable in a process that has no torch —
@@ -135,14 +137,18 @@ def _base_execution_seconds(text: Optional[str]) -> float:
     try:
         from services import model_manager  # noqa: PLC0415 — intentionally lazy
 
-        return float(model_manager.generate_timeout_s(text))
+        return float(
+            model_manager.generate_timeout_s(
+                text, execution_device=execution_device
+            )
+        )
     except Exception:
         base = _GENERATE_TIMEOUT_S
         try:
             from core.device_caps import detect_host_caps  # noqa: PLC0415
 
             if (
-                detect_host_caps().family == "cpu"
+                (execution_device or detect_host_caps().family) == "cpu"
                 and "OMNIVOICE_GENERATE_TIMEOUT_S" not in os.environ
             ):
                 base = _CPU_GENERATE_TIMEOUT_S
@@ -163,6 +169,7 @@ def for_task(
     model_resident: bool = False,
     model_downloaded: bool = True,
     input_seconds: float = 0.0,
+    execution_device: Optional[str] = None,
 ) -> Deadlines:
     """Compute the deadlines for one attempt.
 
@@ -174,7 +181,9 @@ def for_task(
     op = Operation.coerce(operation)
     multiplier, grace = _PROFILE[op]
 
-    execution = _base_execution_seconds(text) * multiplier
+    execution = _base_execution_seconds(
+        text, execution_device=execution_device
+    ) * multiplier
     # Media-length operations scale on duration, not characters.
     if input_seconds > 0:
         execution = max(execution, input_seconds * multiplier)

backend/worker/scheduler.py

@@ -506,6 +506,7 @@ class Scheduler:
             model_resident=worker.is_warm(task.engine, task.model_id),
             model_downloaded=True,
             input_seconds=float(task.params.get("input_seconds") or 0.0),
+            execution_device=worker.capacity.backend,
         )
         attempt.renew_lease(budget.accept_seconds, now=now)
         self._save(task, now=now)
@@ -997,6 +998,7 @@ class Scheduler:
             text=task.params.get("text"),
             model_resident=bool(worker and worker.is_warm(task.engine, task.model_id)),
             input_seconds=float(task.params.get("input_seconds") or 0.0),
+            execution_device=worker.capacity.backend if worker else None,
         )
 
 

tests/test_worker_deadlines.py

@@ -44,6 +44,29 @@ def test_fallback_formula_matches_when_model_manager_is_unavailable(monkeypatch)
     assert deadlines._base_execution_seconds("x" * 4000) == pytest.approx(real)
 
 
+def test_execution_budget_uses_target_worker_device(monkeypatch):
+    """A CPU controller must not enlarge a remote CUDA worker's budget."""
+    from services import model_manager
+
+    monkeypatch.setattr(model_manager, "GPU_JOB_TIMEOUT_S", 300.0)
+    monkeypatch.setattr(model_manager, "CPU_JOB_TIMEOUT_S", 777.0)
+    monkeypatch.setattr(model_manager, "_GENERATE_TIMEOUT_EXPLICIT", False)
+    monkeypatch.setattr(model_manager, "_CONFIGURED_GPU_JOB_TIMEOUT_S", 300.0)
+
+    assert for_task("tts", text="short", execution_device="cuda").execution_seconds == 300
+    assert for_task("tts", text="short", execution_device="cpu").execution_seconds == 777
+
+
+def test_universal_timeout_override_wins_on_cpu_worker(monkeypatch):
+    from services import model_manager
+
+    monkeypatch.setattr(model_manager, "GPU_JOB_TIMEOUT_S", 444.0)
+    monkeypatch.setattr(model_manager, "CPU_JOB_TIMEOUT_S", 777.0)
+    monkeypatch.setattr(model_manager, "_GENERATE_TIMEOUT_EXPLICIT", True)
+
+    assert for_task("tts", text="short", execution_device="cpu").execution_seconds == 444
+
+
 def test_accept_is_generous_enough_for_a_busy_worker():
     """The old 2s accept would time out against a worker mid-inference holding
     the GIL, then penalise it for being busy."""

tests/test_worker_scheduler.py

@@ -11,6 +11,7 @@ import asyncio
 
 import pytest
 
+from worker import deadlines as deadline_policy
 from worker.capacity import ModelSlot
 from worker.errors import ErrorClass, WorkerError
 from worker.identity import issue_session
@@ -381,6 +382,25 @@ def test_assignment_reserves_capacity_and_sets_deadlines():
     assert assignment.deadlines.accept_seconds > 0
 
 
+def test_assignment_deadline_uses_selected_workers_device(monkeypatch):
+    pool = _pool(_record("w1"))
+    worker = pool.get("w1")
+    worker.capacity.backend = "cuda"
+    seen = []
+    real = deadline_policy.for_task
+
+    def recording_for_task(*args, **kwargs):
+        seen.append(kwargs.get("execution_device"))
+        return real(*args, **kwargs)
+
+    monkeypatch.setattr(deadline_policy, "for_task", recording_for_task)
+    sched = _scheduler(pool)
+    _submit(sched)
+    sched.next_assignment(now=1000.0)
+
+    assert seen == ["cuda"]
+
+
 def test_no_capable_worker_fails_the_task_rather_than_ageing_it_out():
     sched = _scheduler(_pool(_record("w1")))
     task = _submit(sched, engine="nope")