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