From d6d3f599848b2001bc2a3711496c2b7bd17586ec Mon Sep 17 00:00:00 2001
From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Date: Tue, 31 Mar 2026 07:35:04 +0400
Subject: [PATCH] fix: replace hard timeout with inactivity timeout for model
 loading (#4707)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 | 16 +++++--
 studio/backend/core/inference/worker.py       | 46 +++++++++++++++----
 2 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index 78cc60e1c6..0366524078 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/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)
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index b3ce43795c..9948142afc 100644
--- a/studio/backend/core/inference/worker.py
+++ b/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