fix: replace hard timeout with inactivity timeout for model loading (#4707)

The 180s wall-clock timeout would kill model loads on slow connections
even when the download was actively progressing. Now the worker sends
heartbeat status messages every 30s during loading, and the orchestrator
resets its 300s deadline on each one — so it only times out when the
subprocess goes truly silent.

studio/backend/core/inference/orchestrator.py

@@ -263,12 +263,17 @@ class InferenceOrchestrator:
         except (EOFError, OSError, ValueError):
             return None
 
-    def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict:
+    def _wait_response(self, expected_type: str, timeout: float = 300.0) -> dict:
         """Block until a response of the expected type arrives.
 
         Also handles 'status' and 'error' events during the wait.
         Returns the matching response dict.
         Raises RuntimeError on timeout or subprocess crash.
+
+        The *timeout* is an **inactivity** timeout: it resets whenever the
+        subprocess sends a status message, so long-running operations (large
+        downloads, slow model loads) won't be killed as long as the subprocess
+        keeps reporting progress.
         """
         deadline = time.monotonic() + timeout
 
@@ -293,6 +298,8 @@ class InferenceOrchestrator:
 
             if rtype == "status":
                 logger.info("Subprocess status: %s", resp.get("message", ""))
+                # Reset deadline — subprocess is still alive and working
+                deadline = time.monotonic() + timeout
                 continue
 
             # Other response types during wait — skip
@@ -303,7 +310,8 @@ class InferenceOrchestrator:
             )
 
         raise RuntimeError(
-            f"Timeout waiting for '{expected_type}' response after {timeout}s"
+            f"Timeout waiting for '{expected_type}' response "
+            f"(no activity for {timeout}s)"
         )
 
     def _drain_queue(self) -> list:
@@ -625,7 +633,7 @@ class InferenceOrchestrator:
                 needed_major,
             )
             self._spawn_subprocess(sub_config)
-            resp = self._wait_response("loaded", timeout = 180)
+            resp = self._wait_response("loaded")
 
             # Update local state from response
             if resp.get("success"):
@@ -672,7 +680,7 @@ class InferenceOrchestrator:
                     "model_name": model_name,
                 }
             )
-            resp = self._wait_response("unloaded", timeout = 30)
+            resp = self._wait_response("unloaded")
 
             # Update local state
             self.models.pop(model_name, None)

studio/backend/core/inference/worker.py

@@ -22,6 +22,7 @@ from loggers import get_logger
 import os
 import queue as _queue
 import sys
+import threading
 import time
 import traceback
 from io import BytesIO
@@ -114,6 +115,29 @@ def _build_model_config(config: dict):
     return mc
 
 
+def _start_heartbeat(resp_queue: Any, interval: float = 30.0) -> threading.Event:
+    """Start a daemon thread that sends periodic status heartbeats.
+
+    Returns a stop event — set it to terminate the heartbeat thread.
+    """
+    stop = threading.Event()
+
+    def _beat():
+        while not stop.wait(interval):
+            _send_response(
+                resp_queue,
+                {
+                    "type": "status",
+                    "message": "Still loading model...",
+                    "ts": time.time(),
+                },
+            )
+
+    t = threading.Thread(target = _beat, daemon = True)
+    t.start()
+    return stop
+
+
 def _handle_load(backend, config: dict, resp_queue: Any) -> None:
     """Handle a load command: load a model into the backend."""
     try:
@@ -173,14 +197,20 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                     model_name,
                 )
 
-        success = backend.load_model(
-            config = mc,
-            max_seq_length = config.get("max_seq_length", 2048),
-            load_in_4bit = load_in_4bit,
-            hf_token = hf_token,
-            trust_remote_code = trust_remote_code,
-            gpu_ids = config.get("resolved_gpu_ids"),
-        )
+        # Send heartbeats every 30s so the orchestrator knows we're still alive
+        # (download / weight loading can take a long time on slow connections)
+        heartbeat_stop = _start_heartbeat(resp_queue, interval = 30.0)
+        try:
+            success = backend.load_model(
+                config = mc,
+                max_seq_length = config.get("max_seq_length", 2048),
+                load_in_4bit = load_in_4bit,
+                hf_token = hf_token,
+                trust_remote_code = trust_remote_code,
+                gpu_ids = config.get("resolved_gpu_ids"),
+            )
+        finally:
+            heartbeat_stop.set()
 
         if success:
             # Build model_info for the parent to mirror