diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 671ef363f5..31bbbdc748 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -132,6 +132,11 @@ class ExportOrchestrator: """True while an export / load / cleanup command is running.""" return self._export_active + def is_worker_alive(self) -> bool: + """True while the persistent export subprocess is running (op or idle).""" + proc = self._proc + return proc is not None and proc.is_alive() + def was_cancelled(self) -> bool: """True if the in-flight (or most recent) run was cancelled by the user.""" return self._cancel_requested @@ -204,6 +209,23 @@ class ExportOrchestrator: def _spawn_subprocess(self, config: dict) -> None: """Spawn a new export subprocess.""" + # Last-resort recheck for spawns outside an active op. Inside an op, _export_active is set and + # load_checkpoint already rechecked, so a reservation here is an install about to observe + # is_export_active() and abort; raising would kill this export for an install that never proceeds. + from utils.transformers_version import sidecar_swap_in_progress + + from utils.transformers_version import sidecar_swap_kind + + _swap_kind = sidecar_swap_kind() + # Inside an active op an INSTALL reservation is about to abort on the + # is_export_active check, but a lazy REPAIR has no such check and can be + # rebuilding the sidecar right now, so it must always refuse the spawn. + if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active): + from utils.transformers_version import SidecarSwapInProgress + raise SidecarSwapInProgress( + "A transformers installation is replacing the latest sidecar; " + "retry when it completes." + ) from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -231,11 +253,17 @@ class ExportOrchestrator: adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Export subprocess started (pid=%s)", self._proc.pid) - def _shutdown_subprocess(self, timeout: float = 10.0) -> None: - """Gracefully shut down the export subprocess.""" + def _shutdown_subprocess(self, timeout: float = 10.0) -> bool: + """Gracefully shut down the export subprocess. + + Returns True only once the worker is confirmed dead. If it survives + terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives + SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the + pre-swap liveness guard can still observe the survivor instead of a cleared + handle and refuse the destructive sidecar swap.""" if self._proc is None or not self._proc.is_alive(): self._proc = None - return + return True self._drain_queue() @@ -265,10 +293,20 @@ class ExportOrchestrator: except Exception: pass + if self._proc is not None and self._proc.is_alive(): + # Survived SIGKILL (uninterruptible syscall): keep the handle so callers + # and the pre-swap guard see a live worker rather than a nulled one. + logger.error( + "Export subprocess still alive after terminate/kill; " + "preserving its handle for the pre-swap liveness check" + ) + return False + self._proc = None self._cmd_queue = None self._resp_queue = None logger.info("Export subprocess shut down") + return True def _cleanup(self): """atexit handler.""" @@ -409,14 +447,44 @@ class ExportOrchestrator: self._export_active = True op_success, op_message = False, "" try: + # Handshake with the sidecar install route: _export_active is set above, so either this + # recheck refuses BEFORE tearing down the old worker (keeping the loaded checkpoint), or + # the install sees is_export_active() and 409s. The spawn-time recheck stays as a last resort. + from utils.transformers_version import sidecar_swap_in_progress + + if sidecar_swap_in_progress(): + from utils.transformers_version import SidecarSwapInProgress + op_message = ( + "A transformers installation is replacing the latest " + "sidecar; retry when it completes." + ) + raise SidecarSwapInProgress(op_message) # Always kill any existing subprocess and spawn fresh. if self._ensure_subprocess_alive(): - self._shutdown_subprocess() + if self._shutdown_subprocess() is False: + # Survivor still holds GPU memory (a wedged CUDA syscall outliving + # SIGKILL); its handle is kept so is_worker_alive() and the pre-swap + # guard still see it. Do not spawn a second worker over it -- fail so + # the load can retry once it exits. + op_message = ( + "The current export worker did not exit and still holds GPU " + "memory; not starting a new checkpoint load over it. Retry shortly." + ) + return False, op_message elif self._proc is not None: self._shutdown_subprocess(timeout = 2) logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path) - self._spawn_subprocess(sub_config) + try: + self._spawn_subprocess(sub_config) + except Exception: + # The old worker is already gone; a stale current_checkpoint + # would make the Export page claim a loaded checkpoint that + # the next op then fails on with "no subprocess running". + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + raise try: resp = self._wait_response("loaded") @@ -560,6 +628,18 @@ class ExportOrchestrator: self._export_active = True op_success, op_message, op_output_path = False, "", None try: + # Handshake with the sidecar install route (see load_checkpoint): _export_active is set + # above, so this recheck refuses before the command is sent, or the install sees the active + # op and 409s. Without it, an install would block in cleanup_memory behind a long export op. + from utils.transformers_version import sidecar_swap_in_progress + + if sidecar_swap_in_progress(): + from utils.transformers_version import SidecarSwapInProgress + op_message = ( + "A transformers installation is replacing the latest " + "sidecar; retry when it completes." + ) + raise SidecarSwapInProgress(op_message) cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 7828116236..08993a9a08 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -236,6 +236,17 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: checkpoint_path = cmd["checkpoint_path"] max_seq_length = cmd.get("max_seq_length", 2048) load_in_4bit = cmd.get("load_in_4bit", True) + # Latest-sidecar checkpoints load 16-bit here too: bnb 4-bit feeds quantized + # expert weights into unvalidated paths (same flip as the chat worker). + if load_in_4bit: + from utils.transformers_version import latest_tier_active_for + if latest_tier_active_for(checkpoint_path, cmd.get("hf_token")): + load_in_4bit = False + logger.info( + "Latest-transformers sidecar active for %s - forcing a 16-bit " + "export load (4-bit is disabled for brand-new architectures)", + checkpoint_path, + ) trust_remote_code = cmd.get("trust_remote_code", False) # Auto-enable trust_remote_code for NemotronH/Nano models. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 2b1ceca75a..6d0b13ced9 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -174,6 +174,21 @@ class InferenceOrchestrator: def _spawn_subprocess(self, config: dict) -> None: """Spawn a new inference subprocess.""" + # Same recheck as the training/export spawns, REPAIR reservations only: a + # repair swaps without holding the lifecycle gate this load's caller owns, + # while an install cannot swap until this gate is released (and then its + # queued-load snapshot aborts it), so tolerating installs here lets the + # load win instead of failing both sides. Also covers the OpenAI + # auto-switch path, which enters _load_model_impl without route guards. + from utils.transformers_version import ( + SidecarSwapInProgress, + sidecar_swap_kind, + ) + + if sidecar_swap_kind() == "repair": + raise SidecarSwapInProgress( + "A transformers repair is replacing the latest sidecar; retry when it completes." + ) from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -210,12 +225,24 @@ class InferenceOrchestrator: if self._cancel_event is not None: self._cancel_event.set() - def _shutdown_subprocess(self, timeout: float = 10.0) -> None: - """Gracefully shut down the inference subprocess.""" + def is_worker_alive(self) -> bool: + """True while the inference subprocess is running, even with no model + active (a failed load can leave a live worker holding sidecar modules).""" + proc = self._proc + return proc is not None and proc.is_alive() + + def _shutdown_subprocess(self, timeout: float = 10.0) -> bool: + """Gracefully shut down the inference subprocess. + + Returns True only once the worker is confirmed dead. If it survives + terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives + SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the + pre-swap liveness guard can still observe the survivor instead of a cleared + handle and refuse the destructive sidecar swap.""" self._stop_dispatcher() # before killing subprocess if self._proc is None or not self._proc.is_alive(): self._proc = None - return + return True # 1. Cancel any ongoing generation first (instant via mp.Event) self._cancel_generation() @@ -252,12 +279,22 @@ class InferenceOrchestrator: except Exception: pass + if self._proc is not None and self._proc.is_alive(): + # Survived SIGKILL (uninterruptible syscall): keep the handle so callers + # and the pre-swap guard see a live worker rather than a nulled one. + logger.error( + "Inference subprocess still alive after terminate/kill; " + "preserving its handle for the pre-swap liveness check" + ) + return False + self._proc = None self._cmd_queue = None self._resp_queue = None self._cancel_event = None self._drain_event = None logger.info("Inference subprocess shut down") + return True def _cleanup(self): """atexit handler.""" @@ -882,6 +919,13 @@ class InferenceOrchestrator: # Public API — same interface as InferenceBackend # ------------------------------------------------------------------ + # Monotonic count of PUBLISHED loads; lets the install route detect a load + # (including a same-model reload) that completed while it waited on the gate. + # Bumped when the load result is published, not at load start: a start-time + # bump is already visible when the installer snapshots mid-load, so the + # completed reload would look unchanged and get unloaded by the swap. + load_generation: int = 0 + def load_model( self, config, # ModelConfig @@ -935,13 +979,36 @@ class InferenceOrchestrator: sub_config["resolved_gpu_ids"] = resolved_gpu_ids sub_config["gpu_selection"] = gpu_selection + # Recheck the sidecar reservation BEFORE tearing the old worker down, + # for REPAIRS only: an install holds this same lifecycle gate, so it + # cannot swap while this load runs, and its queued-load snapshot + # aborts it after this load publishes -- the load wins cleanly. + # Raising here (repair) keeps the current model loaded. + from utils.transformers_version import ( + SidecarSwapInProgress, + sidecar_swap_kind, + ) + + if sidecar_swap_kind() == "repair": + raise SidecarSwapInProgress( + "A transformers repair is replacing the latest sidecar; " + "retry when it completes." + ) + # Always kill the existing subprocess and spawn fresh: reusing one # after unsloth patches torch internals breaks getsource on reload. if self._ensure_subprocess_alive(): self._cancel_generation() time.sleep(0.3) - self._shutdown_subprocess() - + if self._shutdown_subprocess() is False: + # The worker survived terminate/kill (e.g. a wedged CUDA syscall that + # outlives SIGKILL). Its handle is kept, so is_worker_alive() and the + # pre-swap guard still see it; do not spawn a second worker over one + # still holding GPU memory. Fail so the load can retry once it exits. + raise RuntimeError( + "The current inference worker did not exit and still holds GPU " + "memory; not starting a new model over it. Retry shortly." + ) elif self._proc is not None: self._shutdown_subprocess(timeout = 2) @@ -1030,6 +1097,7 @@ class InferenceOrchestrator: return False model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) + self.load_generation += 1 # A load always spawns a fresh subprocess holding only this model, so # mirror that. A lingering stale name would pass unload_model's "not in # self.models" guard, and the worker's absent-name fallback would unload @@ -1061,8 +1129,15 @@ class InferenceOrchestrator: self.models.clear() raise Exception(error) - except Exception: + except Exception as exc: self.loading_models.discard(model_name) + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive(): + # Raised before the old worker was torn down: the previous model + # is still live, so keep the mirrors (clearing them would let the + # installer treat the worker as inactive and kill it unreported). + raise self.active_model_name = None self.models.clear() raise diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index d56353ee56..e4628dcea8 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -291,6 +291,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: hf_token = _clean_token(config.get("hf_token")) load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True)) + # Latest-transformers sidecar models load 16-bit: bnb 4-bit feeds quantized + # expert weights into unvalidated paths (e.g. grouped-MoE torch._grouped_mm). + if load_in_4bit: + from utils.transformers_version import latest_tier_active_for + if latest_tier_active_for(config["model_name"], hf_token): + load_in_4bit = False + logger.info( + "Latest-transformers sidecar active for %s - forcing a 16-bit " + "load (4-bit is disabled for brand-new architectures)", + config["model_name"], + ) + trust_remote_code = config.get("trust_remote_code", False) if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token): trust_remote_code = True diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 6b32ec873d..406c780e81 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -881,93 +881,125 @@ class TrainingBackend: else: defer_auto_selection = True - # Synchronous validation passed -> free VRAM (export + chat) now, before - # auto-selection and the spawn, so placement sees the freed memory. - if before_spawn is not None: - try: - before_spawn() - except Exception: - logger.warning("before_spawn hook failed; continuing", exc_info = True) + # Handshake with the sidecar install route: mark the spawn in progress BEFORE rechecking + # the reservation, so either this recheck aborts, or the install's is_training_active() + # sees this flag (or the recorded proc) and refuses. + from utils.transformers_version import sidecar_swap_in_progress - if defer_auto_selection: - resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs) - config["resolved_gpu_ids"] = resolved_gpu_ids - config["gpu_selection"] = gpu_selection - - from .worker import run_training_process + self._spawn_in_progress = True + if sidecar_swap_in_progress(): + self._spawn_in_progress = False + from utils.transformers_version import SidecarSwapInProgress + raise SidecarSwapInProgress( + "A transformers installation is replacing the latest sidecar; " + "retry when it completes." + ) + # Any exception between the handshake above and the flag reset below would + # otherwise leave _spawn_in_progress latched, wedging is_training_active + # (and the install route) until restart. try: - with native_path_secret_removed_for_child_start(): - event_queue = _CTX.Queue() - stop_queue = _CTX.Queue() + # Synchronous validation passed -> free VRAM (export + chat) now, before + # auto-selection and the spawn, so placement sees the freed memory. Runs AFTER the handshake + # so a lost race to an install can't tear down chat/export for a training run that never spawns. + if before_spawn is not None: + try: + before_spawn() + except Exception: + logger.warning("before_spawn hook failed; continuing", exc_info = True) - proc = _CTX.Process( - target = run_without_native_path_secret, - args = (run_training_process,), - kwargs = { - "event_queue": event_queue, - "stop_queue": stop_queue, - "config": config, - }, - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid + if defer_auto_selection: + try: + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + None, **gpu_selection_kwargs + ) + except Exception: + # Flag is already set; a failed GPU selection must not leave is_training_active stuck True. + self._spawn_in_progress = False + raise + config["resolved_gpu_ids"] = resolved_gpu_ids + config["gpu_selection"] = gpu_selection + + from .worker import run_training_process + + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + + proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) + except Exception: + logger.error("Failed to start training subprocess", exc_info = True) + self._spawn_in_progress = False + return False + + logger.info("Training subprocess started (pid=%s)", proc.pid) + + # Reset state (old pump thread dead, proc.start() succeeded). + self.current_job_id = job_id + self._should_stop = False + self._cancel_requested = False + self._complete_seen.clear() + self._progress = TrainingProgress( + is_training = True, status_message = "Initializing training..." + ) + self.loss_history.clear() + self.lr_history.clear() + self.step_history.clear() + self.grad_norm_history.clear() + self.grad_norm_step_history.clear() + self.eval_loss_history.clear() + self.eval_step_history.clear() + self.eval_enabled = False + self._output_dir = None + self._metric_buffer.clear() + self._run_finalized = False + self._db_run_created = False + self._db_create_in_progress = False # a stale watchdog create can't block this run + self._db_total_steps_set = False + self._db_config = _sanitize_db_config(config) + self._db_started_at = datetime.now(timezone.utc).isoformat() + # Start each job Xet-first; keep config so a stall can respawn over HTTP. + self._last_full_config = config + self._in_model_load = False + self._xet_fallback_used = False + self._needs_xet_respawn = False + + # Create the DB run row before the pump can consume events, so it appears + # in history during model loading and a fast terminal worker can't race the + # pump into a duplicate create/finalize. From here the pump only finalizes. + self._ensure_db_run_created() + + # Assign handles and start the pump together under the lock so a concurrent + # poll can't see a live _proc with no pump and spawn a duplicate. + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._pump_running = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = proc + self._pump_thread = new_pump + new_pump.start() + self._spawn_in_progress = False + + return True - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) except Exception: - logger.error("Failed to start training subprocess", exc_info = True) - return False - - logger.info("Training subprocess started (pid=%s)", proc.pid) - - # Reset state (old pump thread dead, proc.start() succeeded). - self.current_job_id = job_id - self._should_stop = False - self._cancel_requested = False - self._complete_seen.clear() - self._progress = TrainingProgress( - is_training = True, status_message = "Initializing training..." - ) - self.loss_history.clear() - self.lr_history.clear() - self.step_history.clear() - self.grad_norm_history.clear() - self.grad_norm_step_history.clear() - self.eval_loss_history.clear() - self.eval_step_history.clear() - self.eval_enabled = False - self._output_dir = None - self._metric_buffer.clear() - self._run_finalized = False - self._db_run_created = False - self._db_create_in_progress = False # a stale watchdog create can't block this run - self._db_total_steps_set = False - self._db_config = _sanitize_db_config(config) - self._db_started_at = datetime.now(timezone.utc).isoformat() - # Start each job Xet-first; keep config so a stall can respawn over HTTP. - self._last_full_config = config - self._in_model_load = False - self._xet_fallback_used = False - self._needs_xet_respawn = False - - # Create the DB run row before the pump can consume events, so it appears - # in history during model loading and a fast terminal worker can't race the - # pump into a duplicate create/finalize. From here the pump only finalizes. - self._ensure_db_run_created() - - # Assign handles and start the pump together under the lock so a concurrent - # poll can't see a live _proc with no pump and spawn a duplicate. - new_pump = threading.Thread(target = self._pump_loop, daemon = True) - with self._lock: - self._pump_running = False - self._event_queue = event_queue - self._stop_queue = stop_queue - self._proc = proc - self._pump_thread = new_pump - new_pump.start() - - return True + self._spawn_in_progress = False + raise def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" @@ -1266,50 +1298,84 @@ class TrainingBackend: from .worker import run_training_process - try: - with native_path_secret_removed_for_child_start(): - event_queue = _CTX.Queue() - stop_queue = _CTX.Queue() - new_proc = _CTX.Process( - target = run_without_native_path_secret, - args = (run_training_process,), - kwargs = { - "event_queue": event_queue, - "stop_queue": stop_queue, - "config": config, - }, - daemon = True, - ) - new_proc.start() - from utils.process_lifetime import adopt_pid + # This run is active, so an install request 409s rather than proceeds: a reservation seen here + # is transient (an aborting install or short lazy repair). Wait it out instead of stranding the + # stalled run; only a wedged reservation fails the respawn. + from utils.transformers_version import sidecar_swap_in_progress - adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep) - except Exception: - logger.error("Failed to respawn training subprocess", exc_info = True) - with self._lock: - # No replacement pump will run; clear the flag so a later run can't - # inherit a stale _pump_running=True and spawn a duplicate. - self._pump_running = False - self._progress.is_training = False - self._progress.error = "Failed to recover stalled model download" - self._ensure_db_run_created() - self._finalize_run_in_db( - status = "error", - error_message = "Failed to recover stalled model download", + self._spawn_in_progress = True + _swap_wait_deadline = time.time() + 120 + while sidecar_swap_in_progress() and time.time() < _swap_wait_deadline: + time.sleep(1) + if sidecar_swap_in_progress(): + # Raising here would land in the pump's broad finalization catch and + # strand the run in a training state with no worker: finalize it as a + # failure explicitly instead. + self._spawn_in_progress = False + msg = ( + "A transformers installation is replacing the latest sidecar; " + "cannot respawn the training worker." ) + logger.error(msg) + with self._lock: + self._progress.is_training = False + self._progress.error = msg + self._ensure_db_run_created() + self._finalize_run_in_db(status = "error", error_message = msg) return - logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) - new_pump = threading.Thread(target = self._pump_loop, daemon = True) - with self._lock: - self._in_model_load = False - self._event_queue = event_queue - self._stop_queue = stop_queue - self._proc = new_proc - self._pump_thread = new_pump - # Start under the lock so _ensure_pump_alive can never observe the - # new pump as a not-yet-started (dead) thread and spawn a duplicate. - new_pump.start() + # Reset the handshake flag on any unexpected failure past this point, so a + # crashed respawn cannot wedge is_training_active until restart. + try: + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + new_proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + new_proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep) + except Exception: + logger.error("Failed to respawn training subprocess", exc_info = True) + self._spawn_in_progress = False + with self._lock: + # No replacement pump will run; clear the flag so a later run can't + # inherit a stale _pump_running=True and spawn a duplicate. + self._pump_running = False + self._progress.is_training = False + self._progress.error = "Failed to recover stalled model download" + self._ensure_db_run_created() + self._finalize_run_in_db( + status = "error", + error_message = "Failed to recover stalled model download", + ) + return + + logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._in_model_load = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = new_proc + self._spawn_in_progress = False + self._pump_thread = new_pump + # Start under the lock so _ensure_pump_alive can never observe the + # new pump as a not-yet-started (dead) thread and spawn a duplicate. + new_pump.start() + except Exception: + self._spawn_in_progress = False + raise def _ensure_pump_alive(self) -> bool: """Restart the event pump if it crashed, even after the worker exited. @@ -1342,6 +1408,10 @@ class TrainingBackend: def is_training_active(self) -> bool: """Check if training is currently active.""" + # A spawn past its sidecar-swap recheck counts as active even before _proc is recorded, + # so an install cannot slip in mid-spawn. + if getattr(self, "_spawn_in_progress", False): + return True # Self-heal a crashed pump first: a dead pump must never leave the worker # training invisibly behind a frozen UI. Cheap enough for per-second polls. self._ensure_pump_alive() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 130e6ece64..1c84b8268f 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -3019,11 +3019,24 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ), xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) + # Latest-sidecar models load 16-bit here too: bnb 4-bit feeds quantized + # expert weights into unvalidated paths (same flip as the chat worker). + _train_load_in_4bit = config["load_in_4bit"] + if _train_load_in_4bit: + from utils.transformers_version import latest_tier_active_for + if latest_tier_active_for(model_name, hf_token): + _train_load_in_4bit = False + logger.info( + "Latest-transformers sidecar active for %s - forcing a 16-bit " + "training load (4-bit is disabled for brand-new architectures)", + model_name, + ) + try: success = trainer.load_model( model_name = model_name, max_seq_length = config["max_seq_length"], - load_in_4bit = config["load_in_4bit"], + load_in_4bit = _train_load_in_4bit, full_finetuning = not use_lora, hf_token = hf_token, is_dataset_image = config.get("is_dataset_image", False), diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 53441bdb3a..2bdc00dce1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -140,6 +140,27 @@ class ValidateModelRequest(BaseModel): ) +class TransformersUpgradeInfo(BaseModel): + """A model architecture no installed transformers ships, but a newer release does.""" + + model_type: str = Field( + ..., description = "config.json model_type unknown to every installed transformers" + ) + pypi_version: Optional[str] = Field( + None, description = "Latest transformers release on PyPI at check time" + ) + supported_in_pypi: bool = Field( + False, + description = "True if the latest PyPI release ships this model_type; Studio can " + "install it into a persistent sidecar after user consent.", + ) + supported_in_main: bool = Field( + False, + description = "True if transformers GitHub main ships this model_type (dev-only; " + "not installable through Studio yet).", + ) + + class ValidateModelResponse(BaseModel): """Result of model validation. @@ -167,6 +188,48 @@ class ValidateModelResponse(BaseModel): description = "Native training context length, read from the GGUF header when the file " "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", ) + # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. + requires_transformers_upgrade: bool = Field( + False, + description = "True when the model's architecture is unknown to every installed " + "transformers but a newer transformers ships it; the UI should offer the " + "install-latest-transformers consent dialog (or the dev-only notice).", + ) + transformers_upgrade: Optional[TransformersUpgradeInfo] = Field( + None, + description = "Details for the transformers-upgrade dialog; set only when " + "requires_transformers_upgrade is true.", + ) + + +class InstallLatestTransformersRequest(BaseModel): + """Consented request to install the latest transformers release into a sidecar.""" + + version: str = Field( + ..., + min_length = 1, + max_length = 64, + description = "Exact transformers version to install; must match the current " + "latest PyPI release reported by /validate.", + ) + + +class InstallLatestTransformersResponse(BaseModel): + """Result of the consented latest-transformers sidecar install.""" + + success: bool = Field(..., description = "Whether the sidecar was provisioned") + version: str = Field(..., description = "The requested transformers version") + message: str = Field(..., description = "Human-readable result") + model_unloaded: bool = Field( + False, + description = "Whether the active chat model was unloaded before the swap " + "(reported even on failure, so the client can restore its state)", + ) + latest_version: Optional[str] = Field( + None, + description = "On a version-mismatch failure: the release that superseded " + "the requested one, so the client can retry with it", + ) class GenerateRequest(BaseModel): diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index a7fd7cbec7..d44e2ac021 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -51,7 +51,17 @@ def _ensure_export_supported() -> None: Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + Also refuses (409) while a latest-transformers install is swapping .venv_t5_latest: an + export worker spawned mid-swap could activate a half-replaced sidecar. """ + from utils.transformers_latest import is_install_in_progress + + if is_install_in_progress(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + from utils.hardware import export_capability cap = export_capability() @@ -97,6 +107,11 @@ async def load_checkpoint( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error loading checkpoint: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -308,6 +323,11 @@ async def export_merged_model( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting merged model: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -347,6 +367,11 @@ async def export_base_model( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting base model: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -388,6 +413,11 @@ async def export_gguf( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting GGUF model: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -428,6 +458,11 @@ async def export_lora_adapter( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True) raise HTTPException( status_code = 500, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8141ff073a..073cd64400 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1693,6 +1693,9 @@ from models.inference import ( CompletionUsage, ValidateModelRequest, ValidateModelResponse, + TransformersUpgradeInfo, + InstallLatestTransformersRequest, + InstallLatestTransformersResponse, TextContentPart, ImageContentPart, ImageUrl, @@ -3836,11 +3839,25 @@ async def load_model( GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ + # A sidecar install that has reserved the swap must not lose to a load that + # then gets unloaded by the pre-swap teardown. Rechecked under the gate: an + # install can reserve while this request queues on the gate, so the pre-gate + # check alone is only a fast path. + from core.inference.llama_keepwarm import inference_lifecycle_gate + from utils.transformers_version import sidecar_swap_in_progress + + _swap_409 = HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + if sidecar_swap_in_progress(): + raise _swap_409 # Hold the lifecycle gate across the load so idle auto-unload can't unload the # model mid-load. Auto-switch calls _load_model_impl directly since it already # holds this gate. - from core.inference.llama_keepwarm import inference_lifecycle_gate async with inference_lifecycle_gate(): + if sidecar_swap_in_progress(): + raise _swap_409 return await _load_model_impl(request, fastapi_request, current_subject) @@ -4037,6 +4054,17 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre f"Resolved load_in_4bit={effective_load_in_4bit} for '{model_log_label}' " f"from adapter_config.json / base model (requested {request.load_in_4bit})" ) + # Latest-sidecar models load 16-bit (worker refuses bnb 4-bit); size the guard + # to match. Off-loop: tier resolution reads configs. + if effective_load_in_4bit and not config.is_gguf: + from utils.transformers_version import latest_tier_active_for + if await asyncio.to_thread(latest_tier_active_for, config.identifier, request.hf_token): + effective_load_in_4bit = False + logger.info( + f"Latest-transformers sidecar active for '{model_log_label}' - " + "sizing and loading in 16-bit (4-bit is disabled for brand-new " + "architectures)" + ) # Refuse a load that would OOM active training, before the unload step below # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work. @@ -4470,6 +4498,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e) raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Lost the spawn-time race to a sidecar install/repair: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) # Friendlier message for models Unsloth cannot load. if native_grant_backed: redacted_msg = redact_native_paths(str(e)) @@ -4598,16 +4631,6 @@ async def validate_model( detail = "gpu_ids is not supported for GGUF models yet.", ) effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) - # Off-loop: guard does sync nvidia-smi / HF work. - await asyncio.to_thread( - _guard_chat_load_against_training, - config, - model_identifier = model_identifier, - hf_token = request.hf_token, - load_in_4bit = effective_load_in_4bit, - max_seq_length = request.max_seq_length, - requested_gpu_ids = effective_gpu_ids, - ) # Both checks cover the [adapter, base] set (matching the scan route and workers): # either repo can ship auto_map code or a poisoned pickle. @@ -4624,16 +4647,69 @@ async def validate_model( security_targets = list(dict.fromkeys(security_targets)) is_gguf = getattr(config, "is_gguf", False) - # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a - # mixed repo are inert for this load, so gating on them is a false positive. Only - # run the remote-code/security preflight for non-GGUF loads. + # Does a newer transformers ship this model_type? Static overlay first, cached + # PyPI/main snapshot only for unknown types. Never fails validation; run before + # the training guard so an installable upgrade sizes as 16-bit. + transformers_upgrade: Optional[TransformersUpgradeInfo] = None + if not is_gguf: + from utils.transformers_latest import check_upgrade_for_model + + # Cover [adapter, base]: the worker activates transformers for the base model. + for _target in security_targets: + _upgrade = await asyncio.to_thread( + check_upgrade_for_model, _target, request.hf_token + ) + if _upgrade is not None: + transformers_upgrade = TransformersUpgradeInfo(**_upgrade) + break + + # Whether the model can load on the CURRENT transformers through its own remote + # code (auto_map, or the YAML trust default). Computed before the 16-bit flip + # because a model with this fallback still loads 4-bit without the offered install, + # exactly as /load does. requires_trust_remote_code = False - requires_security_review = False if not is_gguf: requires_trust_remote_code = any( _requires_trust_remote_code_for_model(_t, request.hf_token) for _t in security_targets ) + + # Mirror /load's latest-sidecar 16-bit flip so the guard sizes it the same way. An + # ALREADY-ACTIVE latest sidecar always forces 16-bit (the worker will). A merely + # OFFERED (not yet installed) upgrade forces 16-bit only when the model has NO + # custom-code fallback: with auto_map it still loads 4-bit on the current + # transformers (as /load does without a successful install), and the install route + # refuses while training is active, so sizing 16-bit here would 409 the only viable + # 4-bit path. /load re-sizes 16-bit after a successful install and re-guards there. + if effective_load_in_4bit and not is_gguf: + from utils.transformers_version import latest_tier_active_for + _install_only_upgrade = ( + transformers_upgrade is not None + and transformers_upgrade.supported_in_pypi + and transformers_upgrade.pypi_version + and not requires_trust_remote_code + ) + if _install_only_upgrade or await asyncio.to_thread( + latest_tier_active_for, config.identifier, request.hf_token + ): + effective_load_in_4bit = False + # Off-loop: guard does sync nvidia-smi / HF work. + await asyncio.to_thread( + _guard_chat_load_against_training, + config, + model_identifier = model_identifier, + hf_token = request.hf_token, + load_in_4bit = effective_load_in_4bit, + max_seq_length = request.max_seq_length, + requested_gpu_ids = effective_gpu_ids, + ) + + # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a + # mixed repo are inert for this load, so gating on them is a false positive. Only + # run the security preflight for non-GGUF loads (requires_trust_remote_code was + # already resolved above for the sizing flip). + requires_security_review = False + if not is_gguf: requires_security_review = any( _requires_security_review_for_model(_t, request.hf_token) for _t in security_targets ) @@ -4676,6 +4752,8 @@ async def validate_model( requires_trust_remote_code = requires_trust_remote_code, requires_security_review = requires_security_review, context_length = context_length, + requires_transformers_upgrade = transformers_upgrade is not None, + transformers_upgrade = transformers_upgrade, ) except HTTPException: @@ -4720,6 +4798,217 @@ async def validate_model( ) +# studio_router only: admin action, kept off the OpenAI-compatible /v1 mount. +@studio_router.post( + "/install-latest-transformers", response_model = InstallLatestTransformersResponse +) +async def install_latest_transformers_route( + request: InstallLatestTransformersRequest, current_subject: str = Depends(get_current_subject) +): + """ + Consented install of the latest transformers release into the persistent + .venv_t5_latest sidecar. + + Called after the user confirms the transformers-upgrade dialog raised by /validate + (requires_transformers_upgrade). The requested version must match the current latest + PyPI release (re-verified server-side); the sidecar then participates in routing on + this and every future start. A pip install runs off-loop, so this can take a minute. + """ + from utils.transformers_latest import install_latest_transformers + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + # The install stage-and-swaps .venv_t5_latest in place; a live worker would + # lazy-import from the new version mid-run, mixing incompatible modules. Gate on + # worker LIVENESS not tier (no HF token here, so tier re-resolution is unreliable + # for gated repos): training and export are refused, the chat model unloaded. + # Reserve the swap FIRST, before any await: training/export starts check this + # reservation, so raising it after the gate wait would let a worker slip in. + if not try_begin_sidecar_swap(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is already in progress.", + ) + # Until the installer thread takes over, this coroutine owns the reservation + # and must release it on any early exit (the 409 refusals below). + owns_reservation = True + try: + from core.export import get_export_backend + from core.training import get_training_backend + + if get_training_backend().is_training_active(): + raise HTTPException( + status_code = 409, + detail = ( + "A training run is active. Wait for it to finish before " + "installing a new transformers version." + ), + ) + _export = get_export_backend() + if _export.is_export_active(): + raise HTTPException( + status_code = 409, + detail = ( + "An export is running. Wait for it to finish before " + "installing a new transformers version." + ), + ) + # A loaded (idle) export checkpoint would be torn down by the pre-swap + # cleanup; if the swap then failed, that state would be silently lost + # with no rollback signal. Make the user unload it deliberately first. + if getattr(_export, "current_checkpoint", None): + raise HTTPException( + status_code = 409, + detail = ( + "An export checkpoint is loaded. Unload it from the Export " + "page before installing a new transformers version." + ), + ) + # In-flight streams passed the middleware already, so the lifecycle gate can't + # protect them and the swap's unload would kill them mid-stream; mirror the + # auto-switch busy check. This route is not middleware-counted and pending + # requests stay blocked in the middleware, so neither is subtracted here. + from core.inference.llama_keepwarm import ( + inference_lifecycle_gate, + note_model_unloaded, + other_inference_request_count, + ) + + if other_inference_request_count(current_request_counted = False, include_pending = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Another inference request is in progress. Wait for it to " + "finish before installing a new transformers version." + ), + ) + + # Hold the lifecycle gate /load holds so no HF worker can start (or be mid-load + # with active_model_name unset) while the sidecar is swapped. Teardown runs via + # before_swap, only once the staged install succeeded: a failed pip/compat check + # must not leave the user with their model gone. GGUF stays loaded (llama-server + # never imports transformers). + backend = get_inference_backend() + export_backend = get_export_backend() + + unloaded_chat = {"v": False} + + def _unload_before_swap() -> None: + # Runs on the install thread, inside the gate held by _gated_install. Any + # failure raises so the previous sidecar stays untouched (a worker that did + # not tear down cleanly may still lazy-import from it). Export teardown runs + # FIRST so its failure aborts while the chat model is still loaded; + # cleanup_memory shuts the subprocess down even when its command fails, so + # judge by worker liveness, not its return value. + export_backend.cleanup_memory() + export_alive = getattr(export_backend, "is_worker_alive", None) + if callable(export_alive) and export_alive(): + raise RuntimeError("Export worker still alive before the transformers swap") + active = getattr(backend, "active_model_name", None) + if active: + if not backend.unload_model(active): + # A failed unload still clears the orchestrator's model state, + # so the model is gone from the parent's view even though the + # swap aborts: report it so the client rolls back instead of + # pointing at an unloaded model. + if getattr(backend, "active_model_name", None) != active: + unloaded_chat["v"] = True + note_model_unloaded() + raise RuntimeError(f"Could not unload '{active}' before the transformers swap") + note_model_unloaded() + unloaded_chat["v"] = True + logger.info( + "Unloaded '%s' before swapping in transformers %s", + active, + request.version, + ) + # A failed load can leave a live worker with no active model that + # still holds sidecar modules (and blocks the rename on Windows). + worker_alive = getattr(backend, "is_worker_alive", None) + if callable(worker_alive) and worker_alive(): + # _shutdown_subprocess keeps the handle when the worker outlives SIGKILL, + # so both its False result and the liveness recheck catch a survivor + # rather than the recheck being fooled by a nulled handle. + stopped = backend._shutdown_subprocess() + if not stopped or worker_alive(): + raise RuntimeError("Inference worker still alive before the transformers swap") + + def _run_install() -> dict: + # Owns the reservation from here: releasing in the thread, not the route, + # keeps it held if the request is cancelled while the install still stages. + try: + return install_latest_transformers(request.version, _unload_before_swap, True) + finally: + end_sidecar_swap() + + # Snapshot before waiting on the gate: a /load already holding it can + # complete meanwhile (including a same-model reload with new settings), + # and the installer must not unload a model whose successful LoadResponse + # the client is about to render. The generation counter catches reloads + # the name alone would miss. + active_before_gate = ( + getattr(backend, "active_model_name", None), + getattr(backend, "load_generation", 0), + ) + + async def _gated_install() -> dict: + # Held by THIS task, not the request coroutine: a cancelled POST unwinding an + # `async with` here would drop the only guard /load honors mid-install. + async with inference_lifecycle_gate(): + _active_now = ( + getattr(backend, "active_model_name", None), + getattr(backend, "load_generation", 0), + ) + if _active_now != active_before_gate: + end_sidecar_swap() + raise HTTPException( + status_code = 409, + detail = ( + "A model load completed while the install was waiting. " + "Retry the install." + ), + ) + # Recheck under the gate: new streams bump their in-flight count while + # holding it, so once held nothing slips past (the pre-gate check is only + # a fast path and can be outlasted by a wait on a long /load). + if ( + other_inference_request_count( + current_request_counted = False, include_pending = False + ) + > 0 + ): + end_sidecar_swap() + raise HTTPException( + status_code = 409, + detail = ( + "Another inference request is in progress. Wait for " + "it to finish before installing a new transformers " + "version." + ), + ) + return await asyncio.to_thread(_run_install) + + install_task = asyncio.ensure_future(_gated_install()) + owns_reservation = False + # shield: a cancelled request stops waiting, but the installer runs to + # completion (holding the gate) instead of being torn down mid-swap. + result = await asyncio.shield(install_task) + finally: + if owns_reservation: + end_sidecar_swap() + if not result["success"]: + if result.get("latest_version"): + # Structured failure so the dialog can update to the newer release + # and offer a retry that can actually succeed. + return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"]) + if unloaded_chat["v"]: + # The chat model is already gone even though the swap failed; return a + # structured failure (not a bare 400) so the client can restore its + # model state instead of pointing at an unloaded model. + return InstallLatestTransformersResponse(**result, model_unloaded = True) + raise HTTPException(status_code = 400, detail = result["message"]) + return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"]) + + @router.post("/unload", response_model = UnloadResponse) async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)): """ diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1da1c4f425..5e633f4896 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -146,6 +146,16 @@ async def start_training( # No in-process ensure_transformers_version(): the subprocess # (worker.py) activates the correct version before importing ML libs. + # A consented latest-transformers install stage-and-swaps .venv_t5_latest; + # a worker spawned mid-swap could activate a half-replaced sidecar. + from utils.transformers_latest import is_install_in_progress + + if is_install_in_progress(): + raise HTTPException( + status_code = 409, + detail = ("A transformers installation is in progress. Retry when it completes."), + ) + backend = get_training_backend() # S3 dataset loading needs the optional boto3 dependency. Reject early @@ -341,6 +351,24 @@ async def start_training( "s3_config": request.s3_config.model_dump() if request.s3_config else None, } + # Latest-sidecar models size and train 16-bit (same flip as chat load): + # 4-bit is disabled for brand-new architectures, so VRAM coexistence + # checks must not underestimate against a load the worker will refuse. + if training_kwargs["load_in_4bit"]: + from utils.transformers_version import latest_tier_active_for + if await asyncio.to_thread( + latest_tier_active_for, + training_kwargs["model_name"], + training_kwargs["hf_token"] or None, + ): + training_kwargs["load_in_4bit"] = False + logger.info( + "Latest-transformers sidecar active for %s - sizing and " + "training in 16-bit (4-bit is disabled for brand-new " + "architectures)", + training_kwargs["model_name"], + ) + # Training page has no trust_remote_code toggle, so honor the YAML default # -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a # local path or a name merely starting with "unsloth/". @@ -426,9 +454,16 @@ async def start_training( logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e) # The hook runs only once start guards pass -> VRAM freed iff training starts. - success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs - ) + from utils.transformers_version import SidecarSwapInProgress + + try: + success = backend.start_training( + job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + ) + except SidecarSwapInProgress as exc: + # Expected loss of the race against a sidecar install: a retryable + # 409 matching the route-entry guard, not an internal error. + raise HTTPException(status_code = 409, detail = str(exc)) if not success: progress_error = backend.trainer.training_progress.error diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 487c3c7ce0..63dba8579c 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -651,11 +651,19 @@ class TestLoadModelGuardIntegration(unittest.TestCase): inf._shutdown_subprocess = MagicMock() llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None) llama.unload_model = MagicMock() - cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None) + cfg = SimpleNamespace( + is_gguf = False, + is_lora = False, + path = None, + base_model = None, + identifier = "unsloth/Qwen3-1.7B", + ) request = LoadRequest(model_path = "unsloth/Qwen3-1.7B") info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"} with ( + # Pin the latest-sidecar tier check so the guard path stays offline. + patch("utils.transformers_version.latest_tier_active_for", return_value = False), patch.object(self.route, "validate_extra_args", return_value = None), patch.object( self.route, diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index fe3c6d5a0d..fb80b6d061 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -874,6 +874,35 @@ def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): assert o.models == {} +def test_load_model_aborts_when_old_worker_survives_shutdown(monkeypatch): + # A wedged worker that outlives terminate/kill makes _shutdown_subprocess return + # False. load_model must not spawn a second worker over it (double GPU allocation + + # the survivor's handle is lost); it aborts so the load can retry once it exits. + import types + + from utils import transformers_version as tv + + o = _bare_orchestrator() + o.active_model_name = "old" + o.models = {"old": {}} + o.loading_models = set() + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([0], "sel")) + monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_cancel_generation", lambda: None) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: False) # survivor + monkeypatch.setattr( + o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn over a live survivor") + ) + + with pytest.raises(RuntimeError, match = "did not exit"): + o.load_model(types.SimpleNamespace(identifier = "new", gguf_variant = None)) + # The except path cleared the loading marker and mirrors. + assert "new" not in o.loading_models + assert o.active_model_name is None + + def test_load_model_proceeds_when_not_cancelled(monkeypatch): # Guard against a false abort: an uncancelled load keeps its marker and spawns. o = _bare_orchestrator() diff --git a/studio/backend/tests/test_shutdown_preserves_live_worker.py b/studio/backend/tests/test_shutdown_preserves_live_worker.py new file mode 100644 index 0000000000..faf273411c --- /dev/null +++ b/studio/backend/tests/test_shutdown_preserves_live_worker.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""_shutdown_subprocess returns whether the worker actually died, and preserves the +live handle when it survives terminate/kill. + +A GPU worker wedged in an uninterruptible CUDA syscall can outlive SIGKILL. If shutdown +nulled its handle anyway, is_worker_alive() would report False and the pre-swap liveness +guard would let the destructive .venv_t5_latest rename proceed while a live worker still +holds sidecar transformers modules (breaking the rename on Windows). The methods must keep +the handle and return False so callers can refuse the swap. +""" + +import pytest + +from core.export.orchestrator import ExportOrchestrator +from core.inference.orchestrator import InferenceOrchestrator + + +class _FakeProc: + """A subprocess handle that dies only on the requested step (or never).""" + + def __init__(self, dies_on = None): + self._alive = True + self._dies_on = dies_on # None | "join" | "terminate" | "kill" + + def is_alive(self): + return self._alive + + def join(self, timeout = None): + if self._dies_on == "join": + self._alive = False + + def terminate(self): + if self._dies_on == "terminate": + self._alive = False + + def kill(self): + if self._dies_on == "kill": + self._alive = False + + +def _bare_inference(): + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._stop_dispatcher = lambda: None + o._cancel_generation = lambda: None + o._drain_queue = lambda: [] + + class _Q: + def put(self, *a, **k): + pass + + o._cmd_queue = _Q() + o._resp_queue = _Q() + o._cancel_event = None + o._drain_event = None + return o + + +def _bare_export(): + o = ExportOrchestrator.__new__(ExportOrchestrator) + o._drain_queue = lambda: [] + + class _Q: + def put(self, *a, **k): + pass + + o._cmd_queue = _Q() + o._resp_queue = _Q() + return o + + +@pytest.fixture(autouse = True) +def _no_sleep(monkeypatch): + # _shutdown_subprocess sleeps 0.5s after cancelling; keep the tests instant. + import core.inference.orchestrator as inf_mod + monkeypatch.setattr(inf_mod.time, "sleep", lambda *_a, **_k: None) + + +class TestInferenceShutdownReturn: + def test_worker_that_dies_returns_true_and_clears_handle(self): + o = _bare_inference() + o._proc = _FakeProc(dies_on = "terminate") + assert o._shutdown_subprocess(timeout = 0.01) is True + assert o._proc is None + assert o.is_worker_alive() is False + + def test_survivor_returns_false_and_keeps_handle(self): + o = _bare_inference() + o._proc = _FakeProc(dies_on = None) # outlives terminate AND kill + assert o._shutdown_subprocess(timeout = 0.01) is False + assert o._proc is not None + # is_worker_alive stays truthful, so the pre-swap guard can refuse the swap. + assert o.is_worker_alive() is True + + def test_already_dead_returns_true(self): + o = _bare_inference() + o._proc = _FakeProc(dies_on = "join") + o._proc._alive = False + assert o._shutdown_subprocess(timeout = 0.01) is True + assert o._proc is None + + +class TestExportShutdownReturn: + def test_worker_that_dies_returns_true_and_clears_handle(self): + o = _bare_export() + o._proc = _FakeProc(dies_on = "terminate") + assert o._shutdown_subprocess(timeout = 0.01) is True + assert o._proc is None + assert o.is_worker_alive() is False + + def test_survivor_returns_false_and_keeps_handle(self): + o = _bare_export() + o._proc = _FakeProc(dies_on = None) + assert o._shutdown_subprocess(timeout = 0.01) is False + assert o._proc is not None + assert o.is_worker_alive() is True + + +class TestSpawnPathsHonorFailedShutdown: + """A fresh-load path must not spawn a second worker over one that outlived + terminate/kill: the survivor still holds GPU memory and its handle would be lost.""" + + def test_export_load_checkpoint_aborts_when_worker_survives(self, monkeypatch): + import threading + + import utils.transformers_version as tv + + o = ExportOrchestrator.__new__(ExportOrchestrator) + o._lock = threading.RLock() + o._proc = _FakeProc(dies_on = None) # survivor + o.clear_logs = lambda: None + o._cancel_requested = False + o._active_op_kind = None + o._export_active = False + o._ensure_subprocess_alive = lambda: True + o._shutdown_subprocess = lambda *a, **k: False + o._spawn_subprocess = lambda cfg: pytest.fail("must not spawn over a live survivor") + o._record_op_finished = lambda *a, **k: None + monkeypatch.setattr(tv, "sidecar_swap_in_progress", lambda: False) + + ok, msg = o.load_checkpoint(checkpoint_path = "ckpt") + + assert ok is False + assert "did not exit" in msg + # The finally cleared the op flags even though we returned early. + assert o._export_active is False diff --git a/studio/backend/tests/test_transformers_latest.py b/studio/backend/tests/test_transformers_latest.py new file mode 100644 index 0000000000..20616dccba --- /dev/null +++ b/studio/backend/tests/test_transformers_latest.py @@ -0,0 +1,1099 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the latest-transformers support check and the consented sidecar install.""" + +import ast +import json +import os +import textwrap +import time +import pytest +from pathlib import Path + + +# The backend uses "from utils..." imports; ensure the backend dir is on sys.path. +import sys + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the custom logger before importing the modules under test. +import types as _types + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import utils.transformers_latest as tl +import utils.transformers_version as tv +from utils.transformers_latest import ( + check_upgrade_for_model, + install_latest_transformers, + latest_transformers_supports, + _fetch_remote_model_types, + _model_types_from_config, +) +from utils.transformers_version import ( + _config_mapping_cache, + _config_json_cache, + _higher_tier, + _is_valid_version_string, + _model_types_from_source, + _tier_from_config_mapping, + _venv_t5_latest_packages, + activate_transformers_for_subprocess, + ensure_latest_transformers_venv, + get_transformers_tier, + latest_venv_pinned_version, +) + + +# A CONFIG_MAPPING_NAMES source exercising every construct the AST extractor supports. +_MAPPING_SOURCE = """ +from collections import OrderedDict +CONFIG_MAPPING_NAMES = OrderedDict( + [ + ("llama", "LlamaConfig"), + ("gemma4", "Gemma4Config"), + ], + **{"qwen3_moe": "Qwen3MoeConfig"}, +) +CONFIG_MAPPING_NAMES.update({"brandnew_arch": "BrandNewConfig"}) +""" + +_MAIN_ONLY_SOURCE = """ +CONFIG_MAPPING_NAMES = { + "llama": "LlamaConfig", + "gemma4": "Gemma4Config", + "qwen3_moe": "Qwen3MoeConfig", + "brandnew_arch": "BrandNewConfig", + "dev_only_arch": "DevOnlyConfig", +} +""" + + +class _FakeResponse: + def __init__(self, body: bytes): + self._body = body + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +def _fake_urlopen_factory(counter: dict): + """urlopen stub serving the PyPI JSON and both refs' mapping sources.""" + + def _fake_urlopen(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + counter[url] = counter.get(url, 0) + 1 + counter["__total__"] = counter.get("__total__", 0) + 1 + if url == tl._PYPI_JSON_URL: + return _FakeResponse(json.dumps({"info": {"version": "5.13.0"}}).encode()) + if "/v5.13.0/" in url and url.endswith("auto_mappings.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + if "/v5.13.0/" in url and url.endswith("configuration_auto.py"): + return _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n") + if "/main/" in url and url.endswith("auto_mappings.py"): + return _FakeResponse(_MAIN_ONLY_SOURCE.encode()) + if "/main/" in url and url.endswith("configuration_auto.py"): + return _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n") + raise AssertionError(f"unexpected URL fetched: {url}") + + return _fake_urlopen + + +@pytest.fixture(autouse = True) +def _isolated_caches(tmp_path: Path, monkeypatch): + """Fresh in-memory + on-disk caches per test; no accidental real studio_root writes.""" + tl.clear_caches() + monkeypatch.setattr(tl, "_cache_file", lambda: tmp_path / "transformers_latest_check.json") + # The sidecar swap reservation writes a lock file next to the venv dir; + # point it at tmp so tests never touch the real studio root. + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", raising = False) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + tl.clear_caches() + + +def _no_network(monkeypatch, exc = None): + """Fail every urlopen and return a counter; tests assert n == 0 to prove no fetch + happened (check_upgrade_for_model swallows exceptions, so a raising stub alone + cannot prove the negative).""" + calls = {"n": 0} + + def _raise(*args, **kwargs): + calls["n"] += 1 + raise (exc or OSError("network fetch attempted")) + + monkeypatch.setattr("urllib.request.urlopen", _raise) + return calls + + +# --- AST extraction shared with the static router --- + + +class TestModelTypesFromSource: + def test_ordereddict_update_and_unpacking(self): + keys = _model_types_from_source(_MAPPING_SOURCE) + assert keys == {"llama", "gemma4", "qwen3_moe", "brandnew_arch"} + + def test_plain_dict_literal(self): + keys = _model_types_from_source(_MAIN_ONLY_SOURCE) + assert "dev_only_arch" in keys and "llama" in keys + + def test_syntax_error_raises_for_caller_to_handle(self): + with pytest.raises(SyntaxError): + _model_types_from_source("def broken(:\n") + + +class TestFetchRemoteModelTypes: + def test_merges_both_auto_files(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + keys = _fetch_remote_model_types("v5.13.0") + assert keys is not None and "brandnew_arch" in keys + + def test_all_fetches_failing_returns_none(self, monkeypatch): + _no_network(monkeypatch, exc = OSError("no route")) + assert _fetch_remote_model_types("main") is None + + def test_empty_mapping_treated_as_failure(self, monkeypatch): + monkeypatch.setattr( + "urllib.request.urlopen", + lambda req, timeout = None: _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n"), + ) + assert _fetch_remote_model_types("main") is None + + def test_transient_failure_of_one_file_fails_whole_lookup(self, monkeypatch): + # One file times out: the partial map must not be returned and cached. + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + raise OSError("timed out") + + monkeypatch.setattr("urllib.request.urlopen", _fake) + assert _fetch_remote_model_types("main") is None + + def test_missing_auto_mappings_404_still_succeeds(self, monkeypatch): + # Pre-5.10 tags have no auto_mappings.py; a 404 must not fail the lookup. + import urllib.error + + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + raise urllib.error.HTTPError(url, 404, "Not Found", None, None) + + monkeypatch.setattr("urllib.request.urlopen", _fake) + keys = _fetch_remote_model_types("v5.9.0") + assert keys is not None and "brandnew_arch" in keys + + def test_unparseable_file_fails_whole_lookup(self, monkeypatch): + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + return _FakeResponse(b"def broken(:\n") + + monkeypatch.setattr("urllib.request.urlopen", _fake) + assert _fetch_remote_model_types("main") is None + + +# --- latest_transformers_supports: snapshot, cache, offline, kill switch --- + + +class TestLatestTransformersSupports: + def test_supported_in_pypi(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("brandnew_arch") + assert result == { + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + } + + def test_dev_only_arch_reported_main_only(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("dev_only_arch") + assert result["supported_in_pypi"] is False + assert result["supported_in_main"] is True + + def test_unknown_everywhere(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("no_such_arch") + assert result["supported_in_pypi"] is False and result["supported_in_main"] is False + + def test_network_failure_returns_none(self, monkeypatch): + _no_network(monkeypatch, exc = OSError("down")) + assert latest_transformers_supports("brandnew_arch") is None + + def test_offline_returns_none_without_fetch(self, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + calls = _no_network(monkeypatch) + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == 0 + + def test_kill_switch_returns_none_without_fetch(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + calls = _no_network(monkeypatch) + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == 0 + + def test_memory_cache_hit_avoids_refetch(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + first_total = counter["__total__"] + latest_transformers_supports("some_other_arch") + assert counter["__total__"] == first_total + + def test_disk_cache_survives_restart(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + # Simulate a restart: memory gone, disk snapshot stays, network unavailable. + tl.clear_caches() + _no_network(monkeypatch) + result = latest_transformers_supports("brandnew_arch") + assert result is not None and result["supported_in_pypi"] is True + + def test_expired_snapshot_refetches(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + stale = dict(tl._memory_snapshot, fetched_at = time.time() - tl._CACHE_TTL_SECONDS - 1) + tl.clear_caches() + tl._save_snapshot_file(stale) + first_total = counter["__total__"] + latest_transformers_supports("brandnew_arch") + assert counter["__total__"] > first_total + + def test_corrupt_disk_cache_ignored(self, monkeypatch, tmp_path: Path): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + tl._cache_file().write_text("{not json", encoding = "utf-8") + result = latest_transformers_supports("brandnew_arch") + assert result is not None and counter["__total__"] > 0 + + def test_failure_backoff_skips_immediate_retry(self, monkeypatch): + calls = {"n": 0} + + def _fail(*args, **kwargs): + calls["n"] += 1 + raise OSError("down") + + monkeypatch.setattr("urllib.request.urlopen", _fail) + assert latest_transformers_supports("brandnew_arch") is None + first = calls["n"] + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == first # backed off, no second network attempt + + +# --- check_upgrade_for_model: the tier hook --- + + +def _local_model(tmp_path: Path, model_type: str) -> str: + d = tmp_path / f"model_{model_type}" + d.mkdir() + (d / "config.json").write_text(json.dumps({"model_type": model_type})) + return str(d) + + +_FAKE_OVERLAYS = { + "default": frozenset({"llama", "bert", "gpt2"}), + "530": frozenset({"qwen3_moe", "qwen3_next"}), + "550": frozenset({"gemma4"}), + "510": frozenset({"gemma4_unified"}), + "latest": frozenset(), +} + + +def _fake_overlays(monkeypatch, overlays = None): + overlays = overlays or _FAKE_OVERLAYS + fake = lambda tier: overlays.get(tier, frozenset()) + monkeypatch.setattr(tv, "_config_model_types", fake) + monkeypatch.setattr(tl, "_config_model_types", fake) + + +class TestCheckUpgradeForModel: + def test_unknown_type_supported_in_pypi_signals(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) + assert result == { + "model_type": "brandnew_arch", + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + } + + def test_dev_only_type_signals_main_only(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = check_upgrade_for_model(_local_model(tmp_path, "dev_only_arch")) + assert result["supported_in_pypi"] is False and result["supported_in_main"] is True + + def test_unknown_everywhere_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + assert check_upgrade_for_model(_local_model(tmp_path, "no_such_arch")) is None + + def test_offline_falls_through_without_fetch(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + assert calls["n"] == 0 + + def test_network_failure_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch, exc = OSError("down")) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + + def test_known_default_type_never_fetches(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "llama")) is None + assert calls["n"] == 0 + + def test_known_sidecar_type_never_fetches(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "gemma4_unified")) is None + assert calls["n"] == 0 + + def test_hardcoded_tier_type_never_fetches_even_without_overlays( + self, tmp_path: Path, monkeypatch + ): + # Sidecar overlays unreadable, but the hardcoded tables route it. + _fake_overlays( + monkeypatch, + {"default": frozenset({"llama"})}, + ) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "qwen3_5_moe")) is None + assert calls["n"] == 0 + + def test_unreadable_default_overlay_bails_out(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch, {"default": frozenset()}) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + assert calls["n"] == 0 + + def test_no_model_type_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch) + d = tmp_path / "no_type" + d.mkdir() + (d / "config.json").write_text(json.dumps({"architectures": ["Whatever"]})) + assert check_upgrade_for_model(str(d)) is None + + def test_nested_model_type_is_used(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + d = tmp_path / "nested" + d.mkdir() + (d / "config.json").write_text(json.dumps({"text_config": {"model_type": "brandnew_arch"}})) + result = check_upgrade_for_model(str(d)) + assert result is not None and result["model_type"] == "brandnew_arch" + + def test_never_raises_on_internal_error(self, monkeypatch): + monkeypatch.setattr( + tl, "_load_config_json", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) + ) + assert check_upgrade_for_model("some/model") is None + + +class TestNestedModelTypeExtraction: + def test_top_level_wins(self): + assert _model_types_from_config( + {"model_type": "a", "text_config": {"model_type": "b"}} + ) == ["a", "b"] + + def test_nested_fallback(self): + assert _model_types_from_config({"llm_config": {"model_type": "b"}}) == ["b"] + + def test_missing_returns_none(self): + assert _model_types_from_config({}) == [] + + +# --- Routing parity: overlay-shipped model_types route as before, never remote-check --- + + +class TestRoutingParity: + def test_all_overlay_types_route_identically_and_never_check(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + expected_tier = { + "llama": "default", + "bert": "default", + "gpt2": "default", + "qwen3_moe": "530", + "qwen3_next": "530", + "gemma4": "550", + "gemma4_unified": "510", + } + for model_type, tier in expected_tier.items(): + cfg = {"model_type": model_type} + assert _tier_from_config_mapping(cfg) == tier, model_type + assert check_upgrade_for_model(_local_model(tmp_path, model_type)) is None + assert calls["n"] == 0 + + def test_real_installed_mappings_route_without_checker(self, monkeypatch, tmp_path: Path): + """Parity over the REAL installed overlays (base + any provisioned sidecar): + every shipped model_type resolves statically, so the remote checker never + fires and routing is byte-identical with the feature enabled.""" + _no_network(monkeypatch) + seen = 0 + for tier in ("default", "530", "550", "510"): + types = tv._config_model_types(tier) + if not types: + continue # overlay not provisioned in this environment + for model_type in types: + assert _tier_from_config_mapping({"model_type": model_type}) is not None + seen += 1 + if seen == 0: + pytest.skip("no transformers overlay available in this environment") + + def test_get_tier_unchanged_by_kill_switch(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch) + path = _local_model(tmp_path, "no_such_arch") + _config_json_cache.clear() + tier_default = get_transformers_tier(path, probe = False) + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + _config_json_cache.clear() + assert get_transformers_tier(path, probe = False) == tier_default == "default" + + +# --- .venv_t5_latest provisioning and routing participation --- + + +class TestLatestVenvProvisioning: + def test_version_string_validation(self): + assert _is_valid_version_string("5.13.0") + assert _is_valid_version_string("5.14.0rc1") + assert not _is_valid_version_string("5.13.0; rm -rf /") + assert not _is_valid_version_string("git+https://evil") + assert not _is_valid_version_string("") + + def test_packages_pin_exact_version(self): + pkgs = _venv_t5_latest_packages("5.13.0") + assert pkgs[0] == "transformers==5.13.0" + assert any(p.startswith("huggingface_hub==") for p in pkgs) + + def test_ensure_latest_writes_pin_and_invalidates_cache(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + recorded = {} + + def _fake_ensure(dir_, packages, label): + recorded["dir"] = dir_ + recorded["packages"] = packages + Path(dir_).mkdir(parents = True, exist_ok = True) + return True + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + _config_mapping_cache["latest"] = frozenset({"stale"}) + assert ensure_latest_transformers_venv("5.13.0") is True + # Stage-and-swap: pip installs into staging, the live dir is the swap result. + assert recorded["dir"] == str(venv_dir) + ".staging" + assert "transformers==5.13.0" in recorded["packages"] + assert venv_dir.is_dir() + assert not Path(str(venv_dir) + ".staging").exists() + assert latest_venv_pinned_version() == "5.13.0" + assert "latest" not in _config_mapping_cache + + def test_ensure_latest_upgrade_failure_keeps_old_sidecar(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + venv_dir.mkdir(parents = True) + (venv_dir / tv._LATEST_PIN_MARKER).write_text( + json.dumps({"version": "5.12.0", "packages": ["transformers==5.12.0"]}) + ) + (venv_dir / "transformers").mkdir() + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a, **k: True) + # Install fails mid-flight: the previous sidecar and pin survive. + monkeypatch.setattr(tv, "_ensure_venv_dir", lambda *a, **k: False) + assert ensure_latest_transformers_venv("5.13.0") is False + assert latest_venv_pinned_version() == "5.12.0" + assert (venv_dir / "transformers").is_dir() + assert not Path(str(venv_dir) + ".staging").exists() + + def test_ensure_latest_rejects_bad_version(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert ensure_latest_transformers_venv("5.13.0 && curl evil") is False + + def test_ensure_latest_offline_refuses(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert ensure_latest_transformers_venv("5.13.0") is False + + def test_unpinned_sidecar_never_installs(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert tv._ensure_venv_t5_latest_exists() is False + + def test_pinned_sidecar_repairs_with_same_version(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a: False) + recorded = {} + + def _fake_ensure(dir_, packages, label): + recorded["dir"] = dir_ + recorded["packages"] = packages + Path(dir_).mkdir(parents = True, exist_ok = True) + return True + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + assert tv._ensure_venv_t5_latest_exists() is True + # Repair also stage-and-swaps, never installing into the live dir. + assert recorded["dir"] == str(venv_dir) + ".staging" + assert "transformers==5.13.0" in recorded["packages"] + assert latest_venv_pinned_version() == "5.13.0" + + +class TestLatestTierRouting: + def test_latest_outranks_510(self): + assert _higher_tier("latest", "510") == "latest" + assert _higher_tier("510", "latest") == "latest" + + def test_tier_from_mapping_prefers_lowest_but_reaches_latest(self, monkeypatch): + overlays = dict(_FAKE_OVERLAYS) + overlays["latest"] = frozenset({"brandnew_arch"}) + _fake_overlays(monkeypatch, overlays) + assert _tier_from_config_mapping({"model_type": "brandnew_arch"}) == "latest" + # Anything a lower tier ships stays on the lower tier. + assert _tier_from_config_mapping({"model_type": "qwen3_moe"}) == "530" + + def test_overlay_dir_for_latest(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + (venv_dir / "transformers").mkdir(parents = True) + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + # Unpinned dir is ignored: activation refuses an unpinned sidecar. + assert tv._overlay_transformers_dir("latest") is None + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") + + def test_probe_order_excludes_unprovisioned_latest(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + + def test_probe_order_includes_provisioned_latest(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + ("latest",) + + def test_activation_prepends_latest_dir(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", lambda: True) + old_sys_path = list(sys.path) + old_pp = os.environ.get("PYTHONPATH") + try: + activate_transformers_for_subprocess("some/brand-new-model") + assert sys.path[0] == str(venv_dir) + assert os.environ["PYTHONPATH"].split(os.pathsep)[0] == str(venv_dir) + finally: + sys.path[:] = old_sys_path + if old_pp is None: + os.environ.pop("PYTHONPATH", None) + else: + os.environ["PYTHONPATH"] = old_pp + + def test_activation_raises_when_latest_missing(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + with pytest.raises(RuntimeError, match = "venv_t5_latest"): + activate_transformers_for_subprocess("some/brand-new-model") + + +# --- install_latest_transformers: the consent endpoint helper --- + + +class TestInstallLatestTransformers: + def test_success_path(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + recorded = {} + + def _fake_ensure( + version, + extra_packages = (), + before_swap = None, + ): + recorded["args"] = (version, extra_packages) + return True + + monkeypatch.setattr(tl, "ensure_latest_transformers_venv", _fake_ensure) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + result = install_latest_transformers("5.13.0") + assert result["success"] is True and result["version"] == "5.13.0" + assert recorded["args"] == ("5.13.0", ()) + + def test_version_mismatch_rejected(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + ) + result = install_latest_transformers("4.99.0") + assert result["success"] is False and "not the latest" in result["message"] + + def test_offline_rejected(self, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _no_network(monkeypatch) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "offline" in result["message"].lower() + + def test_kill_switch_rejected(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + _no_network(monkeypatch) + result = install_latest_transformers("5.13.0") + assert result["success"] is False + + def test_install_failure_reported(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (), before_swap = None: False, + ) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "failed" in result["message"] + + def test_blocked_by_incompatible_deps(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), ["numpy>=99.0"])) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + ) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "numpy>=99.0" in result["message"] + + def test_compat_shadows_passed_to_installer(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: (("tokenizers==0.23.0",), [])) + recorded = {} + + def _fake_ensure( + version, + extra_packages = (), + before_swap = None, + ): + recorded["extras"] = extra_packages + return True + + monkeypatch.setattr(tl, "ensure_latest_transformers_venv", _fake_ensure) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + result = install_latest_transformers("5.13.0") + assert result["success"] is True + assert recorded["extras"] == ("tokenizers==0.23.0",) + + +class TestCompatPlan: + def _patch_env(self, monkeypatch, requires, installed): + monkeypatch.setattr(tl, "_fetch_requires_dist", lambda v: requires) + + def _ver(name): + from importlib.metadata import PackageNotFoundError + + key = name.lower().replace("_", "-") + if key not in installed: + raise PackageNotFoundError(name) + return installed[key] + + monkeypatch.setattr("importlib.metadata.version", _ver) + + def test_satisfied_env_needs_nothing(self, monkeypatch): + self._patch_env( + monkeypatch, + ["tokenizers<=0.23.0,>=0.22.0", "safetensors>=0.8.0", "numpy>=1.17"], + {"tokenizers": "0.22.2", "safetensors": "0.8.0", "numpy": "2.4.4"}, + ) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + def test_unsatisfied_shadowable_dep_pinned(self, monkeypatch): + self._patch_env( + monkeypatch, + ["tokenizers>=0.24.0"], + {"tokenizers": "0.22.2"}, + ) + monkeypatch.setattr(tl, "_resolve_exact_version", lambda name, spec: "0.24.1") + extras, blockers = tl.compat_plan("5.99.0") + assert extras == ("tokenizers==0.24.1",) and blockers == [] + + def test_unsatisfied_non_shadowable_dep_blocks(self, monkeypatch): + self._patch_env(monkeypatch, ["numpy>=99.0"], {"numpy": "2.4.4"}) + extras, blockers = tl.compat_plan("5.99.0") + assert extras == () and blockers == ["numpy>=99.0"] + + def test_cli_only_dep_ignored(self, monkeypatch): + self._patch_env(monkeypatch, ["typer"], {}) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + def test_sidecar_provided_hub_checked_against_recipe_pin(self, monkeypatch): + self._patch_env(monkeypatch, ["huggingface-hub<2.0,>=1.5.0"], {"huggingface-hub": "0.36.2"}) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] # 1.8.0 sidecar pin satisfies it + + def test_sidecar_provided_hub_out_of_range_blocks(self, monkeypatch): + self._patch_env(monkeypatch, ["huggingface-hub>=2.1"], {"huggingface-hub": "0.36.2"}) + extras, blockers = tl.compat_plan("5.99.0") + assert blockers == ["huggingface-hub>=2.1"] + + def test_unfetchable_requires_dist_blocks_install(self, monkeypatch): + # Proceeding unverified could pin a sidecar whose imports crash workers. + monkeypatch.setattr(tl, "_fetch_requires_dist", lambda v: None) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and len(blockers) == 1 and "retry" in blockers[0] + + def test_extra_marker_requirements_skipped(self, monkeypatch): + self._patch_env( + monkeypatch, + ['torch>=99.0; extra == "torch"', 'pytest; python_version < "3.0"'], + {}, + ) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + +def test_get_snapshot_dedupes_concurrent_fetch(monkeypatch): + """While one thread is fetching, other callers return None instead of stacking fetches.""" + with tl._lock: + tl._is_fetching = True + calls = {"n": 0} + + def boom(): + calls["n"] += 1 + raise AssertionError("must not fetch while another fetch is in flight") + + monkeypatch.setattr(tl, "_refresh_snapshot", boom) + assert tl._get_snapshot() is None + assert calls["n"] == 0 + tl.clear_caches() + + +def test_install_serialized(): + """A second install call while one is in progress gets a structured refusal.""" + from utils.transformers_version import try_begin_sidecar_swap + + assert try_begin_sidecar_swap() is True + out = tl.install_latest_transformers("5.13.0") + assert out["success"] is False + assert "already in progress" in out["message"] + tl.clear_caches() + + +def test_install_in_progress_reflects_reservation(): + """is_install_in_progress mirrors the shared sidecar swap reservation, so a + lazy repair (which takes the same reservation) also blocks worker starts.""" + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + assert tl.is_install_in_progress() is False + assert try_begin_sidecar_swap() is True + try: + assert tl.is_install_in_progress() is True + finally: + end_sidecar_swap() + assert tl.is_install_in_progress() is False + + +def test_upgrade_check_sees_nested_model_types(monkeypatch): + """A supported wrapper with a brand-new nested backbone must still signal.""" + cfg = { + "model_type": "llava", # in every installed overlay + "text_config": {"model_type": "zz_brand_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_brand_new_llm", + "supported_in_main": mt == "zz_brand_new_llm", + }, + ) + out = tl.check_upgrade_for_model("some-org/wrapped-new-backbone") + assert out is not None + assert out["model_type"] == "zz_brand_new_llm" + + +def test_upgrade_check_ignores_nested_known_types(monkeypatch): + """All nested types known to installed overlays -> no signal, no remote call.""" + cfg = { + "model_type": "llava", + "text_config": {"model_type": "llama"}, + "vision_config": {"model_type": "clip_vision_model"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + calls = [] + monkeypatch.setattr(tl, "latest_transformers_supports", lambda mt: calls.append(mt) or None) + assert tl.check_upgrade_for_model("some-org/normal-vlm") is None + assert calls == [] + + +def test_upgrade_check_requires_primary_supported(monkeypatch): + """Latest supporting only a nested type must not prompt: routing still + cannot load the primary, so the install would not fix the model.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_llm", + "supported_in_main": mt == "zz_new_llm", + }, + ) + assert tl.check_upgrade_for_model("some-org/half-supported") is None + + +def test_upgrade_check_requires_every_missing_type(monkeypatch): + """Primary supported but a nested backbone missing from latest -> no prompt + (CONFIG_MAPPING would still fail on the sub-config); all supported -> signal + carries the primary type.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_wrapper", + "supported_in_main": mt == "zz_new_wrapper", + }, + ) + assert tl.check_upgrade_for_model("some-org/half-supported") is None + + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + }, + ) + out = tl.check_upgrade_for_model("some-org/fully-supported") + assert out is not None and out["model_type"] == "zz_new_wrapper" + + +def test_install_success_invalidates_capability_caches(monkeypatch): + """A successful install must drop tier probes, the latest mapping, and the + vision-detection cache so the new sidecar takes effect without a restart.""" + from utils.models import model_config as mc + + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + monkeypatch.setattr( + tl, "ensure_latest_transformers_venv", lambda v, extra_packages = (), before_swap = None: True + ) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + + tv._probe_tier_cache["stale/model"] = "default" + tv._config_mapping_cache["latest"] = frozenset({"stale_type"}) + tv._config_mapping_cache["default"] = frozenset({"llama"}) + mc._vision_detection_cache[("stale/model", None, False)] = False + + result = install_latest_transformers("5.13.0") + assert result["success"] is True + assert tv._probe_tier_cache == {} + assert "latest" not in tv._config_mapping_cache + assert tv._config_mapping_cache.get("default") == frozenset({"llama"}) # untouched + assert mc._vision_detection_cache == {} + + tv._probe_tier_cache.clear() + tv._config_mapping_cache.clear() + tl.clear_caches() + + +def test_vision_subprocess_unions_sidecar_registry(): + """The embedded vision-check script must extend the inlined parent sets with + the ACTIVE sidecar's registry so sidecar-only architectures classify.""" + from utils.models import model_config as mc + + script = mc._VISION_CHECK_SCRIPT + ast.parse(script) + stub_registry = { + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES": { + "zz_sidecar_vlm": "ZzSidecarForConditionalGeneration" + }, + } + ns = {} + # Exec only the registry-union block against a stubbed sidecar registry. + body = script.split("from transformers import AutoConfig", 1)[1] + body = body.split("kwargs = {", 1)[0] + helpers = script.split("sys.path.insert(0, backend_dir)", 1)[1] + helpers = helpers.split("try:", 1)[0] + exec(helpers, ns) + + class _FakeMa: + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = stub_registry[ + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES" + ] + + import sys as _sys + import types as _types + + fake_pkg = _types.ModuleType("transformers.models.auto") + fake_pkg.modeling_auto = _FakeMa + saved = { + k: _sys.modules.get(k) + for k in ("transformers.models.auto", "transformers.models.auto.modeling_auto") + } + _sys.modules["transformers.models.auto"] = fake_pkg + _sys.modules["transformers.models.auto.modeling_auto"] = _FakeMa + try: + exec(textwrap.dedent(body), ns) + finally: + for k, v in saved.items(): + if v is None: + _sys.modules.pop(k, None) + else: + _sys.modules[k] = v + + assert "zz_sidecar_vlm" in ns["_VLM_MODEL_TYPES"] + assert "ZzSidecarForConditionalGeneration" in ns["_VLM_CLASS_NAMES"] + + class _Cfg: + architectures = ["ZzSidecarForConditionalGeneration"] + model_type = "zz_sidecar_vlm" + + assert ns["_is_vlm"](_Cfg()) is True + + +def test_upgrade_check_mixed_pypi_main_reports_dev_only(monkeypatch): + """Primary in the PyPI release but a nested type only on main: no install + may be offered (CONFIG_MAPPING would fail on the nested sub-config), so the + aggregate must read as main-only.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_wrapper", + "supported_in_main": True, + }, + ) + out = tl.check_upgrade_for_model("some-org/mixed-support") + assert out is not None + assert out["model_type"] == "zz_new_wrapper" + assert out["supported_in_pypi"] is False # no install offered + assert out["supported_in_main"] is True + + +def test_install_endpoint_not_mounted_on_v1(): + """The consented pip-install endpoint is a Studio admin action; it must live + on studio_router (kept off the OpenAI-compatible /v1 mount), not router.""" + from routes import inference as ri + + path = "/install-latest-transformers" + assert path in [r.path for r in ri.studio_router.routes] + assert path not in [r.path for r in ri.router.routes] + + +def test_kill_switch_removes_provisioned_latest_from_routing(tmp_path, monkeypatch): + """UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS must roll back a provisioned latest + sidecar: no overlay mapping, no probe participation, no file deletion needed.""" + venv_dir = tmp_path / ".venv_t5_latest" + (venv_dir / "transformers").mkdir(parents = True) + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + + assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + ("latest",) + + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + tv._config_mapping_cache.pop("latest", None) + assert tv._overlay_transformers_dir("latest") is None + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + tv._config_mapping_cache.pop("latest", None) + + +def test_repair_failure_preserves_pin_and_live_dir(tmp_path, monkeypatch): + """A failed lazy repair must not delete the incomplete-but-pinned live + sidecar: the pin survives so a later attempt can still repair it.""" + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + (venv_dir / "partial_file").write_text("x") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a: False) + monkeypatch.setattr(tv, "_ensure_venv_dir", lambda *a, **k: False) + + from utils.transformers_version import latest_venv_pinned_version + + assert tv._ensure_venv_t5_latest_exists() is False + assert venv_dir.is_dir() + assert (venv_dir / "partial_file").exists() + assert latest_venv_pinned_version() == "5.13.0" + assert not (tmp_path / ".venv_t5_latest.staging").exists() + + +def test_failed_staging_install_removes_staging_dir(tmp_path, monkeypatch): + """A pip failure inside _ensure_venv_dir returns False without raising, so + the except cleanup never runs; the partial staging dir must still go.""" + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + + def _fake_ensure(dir_, packages, label): + Path(dir_).mkdir(parents = True, exist_ok = True) + (Path(dir_) / "partial").write_text("x") + return False + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + assert ensure_latest_transformers_venv("5.13.0") is False + assert not Path(str(venv_dir) + ".staging").exists() diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index b9b5abb9e5..a6e6803a5c 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -2550,3 +2550,649 @@ class TestHfEndpointUnreachable: t0 = time.time() result = hf_endpoint_unreachable(timeout = 2) assert result is True and (time.time() - t0) < 6.0 + + +class TestLatestTierActiveFor: + """latest_tier_active_for: the 16-bit guard for the consented latest sidecar.""" + + @staticmethod + def _pin( + monkeypatch, + tv, + version = "5.13.1", + ): + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: version) + monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: None) + + def test_true_when_tier_latest(self, monkeypatch): + import utils.transformers_version as tv + + self._pin(monkeypatch, tv) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is True + + def test_false_for_fixed_tiers(self, monkeypatch): + import utils.transformers_version as tv + self._pin(monkeypatch, tv) + for tier in ("default", "530", "550", "510"): + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, _t = tier, **k: _t) + assert tv.latest_tier_active_for("some/model") is False + + def test_false_without_pin_and_no_resolution(self, monkeypatch): + """No sidecar pin returns False before any tier or network resolution.""" + import utils.transformers_version as tv + + def _boom(*a, **k): + raise AssertionError("must not resolve without a pin") + + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None) + monkeypatch.setattr(tv, "_remote_lora_base", _boom) + monkeypatch.setattr(tv, "get_transformers_tier", _boom) + assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is False + + def test_never_raises(self, monkeypatch): + import utils.transformers_version as tv + + def _boom(*a, **k): + raise RuntimeError("tier resolution exploded") + + self._pin(monkeypatch, tv) + monkeypatch.setattr(tv, "get_transformers_tier", _boom) + assert tv.latest_tier_active_for("some/model") is False + + def test_remote_lora_base_is_resolved(self, monkeypatch): + """A remote adapter is judged by its base model, like worker activation.""" + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.13.1") + monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: "Zyphra/ZAYA1-8B") + tiers = {"Zyphra/ZAYA1-8B": "latest"} + monkeypatch.setattr( + tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default") + ) + assert tv.latest_tier_active_for("someuser/zaya-lora") is True + + def test_local_checkpoint_config_upgrades(self, monkeypatch, tmp_path): + """An adapter dir with its own config.json merges tiers like activation does.""" + import utils.transformers_version as tv + + adapter = tmp_path / "ckpt" + adapter.mkdir() + (adapter / "adapter_config.json").write_text("{}") + (adapter / "adapter_model.safetensors").write_text("x") + (adapter / "config.json").write_text("{}") + self._pin(monkeypatch, tv) + monkeypatch.setattr(tv, "_resolve_base_model", lambda name: "base/model") + tiers = {"base/model": "default", str(adapter): "latest"} + monkeypatch.setattr( + tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default") + ) + assert tv.latest_tier_active_for(str(adapter)) is True + + +class TestLatestTierForces16Bit: + """The inference worker and load route refuse bnb 4-bit on the latest sidecar.""" + + def _read(self, rel): + backend_dir = Path(__file__).resolve().parent.parent + return (backend_dir / rel).read_text() + + def test_worker_guard_present(self): + src = self._read("core/inference/worker.py") + assert "latest_tier_active_for" in src, ( + "core/inference/worker.py must force load_in_4bit=False when " + "latest_tier_active_for(model) is true: transformers' grouped-MoE " + "kernels crash on bnb-quantized expert weights for brand-new " + "architectures." + ) + + def test_route_guard_present(self): + src = self._read("routes/inference.py") + assert "latest_tier_active_for" in src, ( + "routes/inference.py must size the VRAM guard with the same 16-bit " + "flip the worker applies for latest-sidecar models." + ) + + def test_validate_route_mirrors_16bit_flip(self): + # Without the same flip, /validate sizes 4-bit and /load then 409s. + src = self._read("routes/inference.py") + body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0] + assert "latest_tier_active_for" in body, ( + "validate_model must apply the latest-sidecar 16-bit flip before " + "_guard_chat_load_against_training so /validate and /load agree." + ) + # First-time loads have no pin yet, so an installable upgrade must also size 16-bit. + assert body.index("check_upgrade_for_model") < body.index( + "_guard_chat_load_against_training" + ), "the upgrade check must run before the training guard" + assert ( + "supported_in_pypi" in body.split("_guard_chat_load_against_training")[0] + ), "an installable upgrade must force 16-bit sizing for the guard" + + def test_validate_offered_upgrade_preserves_custom_code_4bit(self): + # A merely-offered (not installed) upgrade must NOT force 16-bit sizing when the + # model has a custom-code (auto_map) fallback: /load loads it 4-bit without the + # install, and the install route refuses during active training, so 16-bit sizing + # here would 409 the only viable 4-bit path. + src = self._read("routes/inference.py") + body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0] + flip = body.split("Mirror /load's latest-sidecar 16-bit flip", 1)[1].split( + "_guard_chat_load_against_training", 1 + )[0] + assert "not requires_trust_remote_code" in flip, ( + "the offered-upgrade 16-bit flip must be gated on the absence of a custom-code " + "fallback so /validate does not 409 a 4-bit load /load would allow" + ) + # requires_trust_remote_code must be resolved before the flip consumes it. + assert body.index("requires_trust_remote_code = any(") < body.index( + "not requires_trust_remote_code" + ) + + def test_install_route_guards_active_latest_workers(self): + # Stage-and-swap replaces .venv_t5_latest in place, so a live worker on the + # old sidecar would lazy-import files from the new version. + src = self._read("routes/inference.py") + body = src.split("async def install_latest_transformers_route", 1)[1].split( + "\nasync def ", 1 + )[0] + assert ( + "is_training_active" in body + and "is_export_active" in body + and "inference_lifecycle_gate" in body + ), ( + "install_latest_transformers_route must refuse while training or export " + "runs, and hold the lifecycle gate while unloading the chat model and " + "swapping the sidecar." + ) + # The unload (via before_swap so failed installs keep the model), the export-worker + # teardown, and the install must all sit INSIDE the gate so no /load interleaves. + assert "unload_model(active)" in body + assert "cleanup_memory()" in body + # Export teardown precedes the chat unload so its failure aborts with the model still loaded. + assert body.index("cleanup_memory()") < body.index("unload_model(active)") + assert "install_latest_transformers(" in body and "_unload_before_swap" in body + # The gate must be owned by the shielded task, not the request coroutine: a cancelled + # POST unwinding an async-with would release the only guard /load honors mid-install. + gated_task = body.split("async def _gated_install", 1)[1] + assert "inference_lifecycle_gate():" in gated_task + assert "asyncio.to_thread(_run_install)" in gated_task + # The reservation must be taken BEFORE the (awaitable) gate wait, or a + # training/export start could slip in while this request queues on the gate. + assert body.index("try_begin_sidecar_swap()") < body.index( + "inference_lifecycle_gate():" + ), "the swap reservation must be raised before waiting on the lifecycle gate" + # A failed teardown must abort the swap (raise), not fall through to it. + assert body.count("raise RuntimeError") >= 3, ( + "export, chat-unload, and idle-worker teardown failures must raise so " + "the staged install never swaps under a live worker" + ) + # The installer thread owns (and releases) the reservation, shielded from + # request cancellation, so a cancelled POST cannot unlock a live swap. + assert "asyncio.shield" in body and "end_sidecar_swap()" in body + # In-flight generation streams predate the gate; the route refuses rather than kill them + # via the before_swap unload. The count is rechecked UNDER the gate, since a wait on a + # long /load outlasts the pre-gate fast path and streams take this same gate. + assert "other_inference_request_count" in body + gated_task = body.split("async def _gated_install", 1)[1] + assert "other_inference_request_count" in gated_task + + def test_start_routes_refuse_during_install(self): + # A worker spawned mid-swap could activate a half-replaced sidecar. + training = self._read("routes/training.py") + start = training.split("async def start_training", 1)[1].split("\nasync def ", 1)[0] + assert ( + "is_install_in_progress" in start + ), "training /start must refuse while a transformers install is in progress" + export = self._read("routes/export.py") + helper = export.split("def _ensure_export_supported", 1)[1].split("\ndef ", 1)[0] + assert ( + "is_install_in_progress" in helper + ), "mutating export routes must refuse while a transformers install is in progress" + + def test_spawn_sites_recheck_reservation(self): + # The route-level guards are one-shot; validation between them and the + # actual spawn can outlast an install's start, so the spawn itself rechecks. + training = self._read("core/training/training.py") + assert ( + training.count("sidecar_swap_in_progress()") >= 2 + ), "both training spawn sites must recheck the sidecar swap reservation" + export = self._read("core/export/orchestrator.py") + spawn = export.split("def _spawn_subprocess", 1)[1].split("\n def ", 1)[0] + assert ( + "sidecar_swap_kind()" in spawn + ), "the export subprocess spawn must recheck the sidecar swap reservation" + # Training marks the spawn active BEFORE its recheck, so either side sees the other: + # is_training_active covers the window between proc.start() and the _proc assignment. + assert training.index("self._spawn_in_progress = True") < training.index( + "if sidecar_swap_in_progress():" + ) + active = training.split("def is_training_active", 1)[1].split("\n def ", 1)[0] + assert "_spawn_in_progress" in active + # Export load-checkpoint refuses BEFORE tearing down the old worker, so a + # lost race against an install keeps the loaded checkpoint (no bare 500). + loadck = export.split("def load_checkpoint", 1)[1].split("\n def ", 1)[0] + assert loadck.index("sidecar_swap_in_progress()") < loadck.index("_shutdown_subprocess()") + # The training handshake precedes the VRAM-freeing before_spawn hook, so + # losing the race never tears down chat/export for a run that won't spawn. + assert training.index("self._spawn_in_progress = True") < training.index("before_spawn()") + # The spawn-time export check is op-aware for installs (the install side + # aborts on is_export_active) but always refuses for repairs, which have + # no such abort and can be rebuilding the sidecar right now. + assert ( + '_swap_kind == "repair" or (_swap_kind is not None and not self._export_active)' + in spawn + ) + + +class TestSidecarSwapReservation: + """The lazy repair takes the same reservation the install route and worker starts use.""" + + def _repair_setup(self, monkeypatch, tmp_path): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.setattr( + tv, + "_latest_pin_data", + lambda: { + "version": "5.99.0", + "packages": ["transformers==5.99.0"], + }, + ) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: False) + monkeypatch.setattr(tv, "_env_offline", lambda: False) + return tv + + def test_repair_holds_reservation_during_swap(self, monkeypatch, tmp_path): + tv = self._repair_setup(monkeypatch, tmp_path) + seen = {} + + def _fake_swap( + version, + packages, + before_swap = None, + ): + seen["active_during_swap"] = tv.sidecar_swap_in_progress() + return True + + monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _fake_swap) + assert tv._ensure_venv_t5_latest_exists() is True + assert seen["active_during_swap"] is True + assert tv.sidecar_swap_in_progress() is False + + def test_foreign_process_lock_file_visible(self, monkeypatch, tmp_path): + """A repair in a LIVE worker subprocess is seen (via the lock file) by this + process, and its lock is never broken while the owner is alive.""" + import os + import time + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + lock = tv._swap_lock_path() + lock.parent.mkdir(parents = True, exist_ok = True) + # A live owner (this process): visible and never reclaimed, even once aged past + # the cutoff -- a slow but live pip install must keep its lock. + lock.write_text('{"pid": %d}' % os.getpid()) + assert tv.sidecar_swap_in_progress() is True + assert tv.try_begin_sidecar_swap() is False + old_ts = time.time() - 3 * 60 * 60 + os.utime(lock, (old_ts, old_ts)) + assert tv.sidecar_swap_in_progress() is True + assert tv.try_begin_sidecar_swap() is False + + def test_dead_owner_lock_reclaimed_promptly(self, monkeypatch, tmp_path): + """A fresh lock whose recorded owner is dead is reclaimed at once, not after the + long cutoff: a crash mid-install must not wedge loads/training/export for hours.""" + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + lock = tv._swap_lock_path() + lock.parent.mkdir(parents = True, exist_ok = True) + # 999999 is not a live PID: a fresh dead-owner lock is immediately stale. + lock.write_text('{"pid": 999999, "kind": "install"}') + assert tv._pid_alive(999999) is False + assert tv.sidecar_swap_in_progress() is False + assert tv.try_begin_sidecar_swap() is True + try: + assert lock.is_file() + finally: + tv.end_sidecar_swap() + assert not lock.exists() + + def test_unreadable_pid_lock_uses_age_cutoff(self, monkeypatch, tmp_path): + """A lock with no readable owner PID (mid create-before-write, or corrupt) is not + reclaimed while fresh -- only after the long cutoff -- so a lock a live owner just + created is not stolen before its PID lands.""" + import os + import time + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + lock = tv._swap_lock_path() + lock.parent.mkdir(parents = True, exist_ok = True) + lock.write_text("") # created but metadata not yet written + assert tv.sidecar_swap_in_progress() is True + old_ts = time.time() - (tv._SWAP_LOCK_STALE_SECS + 60) + os.utime(lock, (old_ts, old_ts)) + assert tv.sidecar_swap_in_progress() is False + + def test_repair_refused_while_install_holds_reservation(self, monkeypatch, tmp_path): + tv = self._repair_setup(monkeypatch, tmp_path) + + def _must_not_run(*a, **k): + raise AssertionError("repair must not swap while an install is in progress") + + monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _must_not_run) + assert tv.try_begin_sidecar_swap() is True + try: + assert tv._ensure_venv_t5_latest_exists() is False + finally: + tv.end_sidecar_swap() + + +class TestRecoverStrandedSidecar: + """A swap whose activation rename AND rollback both fail strands the previous sidecar + at .old with no live dir (its pin marker went with it). Reading the pin self-heals it, + but never while a swap legitimately holds the reservation.""" + + def _setup(self, monkeypatch, tmp_path): + import utils.transformers_version as tv + + live = str(tmp_path / "venv_t5_latest") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", live) + # Stranded state: live gone, previous sidecar (with its marker) sits at .old. + retired = Path(live + ".old") + retired.mkdir(parents = True) + (retired / tv._LATEST_PIN_MARKER).write_text( + '{"version": "5.99.0", "packages": ["transformers==5.99.0"]}' + ) + return tv, Path(live), retired + + def test_stranded_old_recovered_on_pin_read(self, monkeypatch, tmp_path): + tv, live, retired = self._setup(monkeypatch, tmp_path) + data = tv._latest_pin_data() + assert live.is_dir() + assert not retired.exists() + assert data is not None and data["version"] == "5.99.0" + + def test_stranded_recovery_skipped_during_swap(self, monkeypatch, tmp_path): + tv, live, retired = self._setup(monkeypatch, tmp_path) + assert tv.try_begin_sidecar_swap() is True + try: + # A swap holds the reservation and may be mid-rename; do not race it. + assert tv._latest_pin_data() is None + assert not live.exists() + assert retired.is_dir() + finally: + tv.end_sidecar_swap() + # Once the swap is done, the next pin read recovers the stranded sidecar. + assert tv._latest_pin_data() is not None + assert live.is_dir() + + +class TestCachedLatestMappingRevalidated: + """A cached 'latest' mapping is dropped and re-resolved when the sidecar since broke + in-process, so routing self-heals instead of trusting a mapping parsed from a sidecar + that no longer exists (which would keep routing latest-only models to a broken tier).""" + + def test_broken_sidecar_drops_cached_latest_mapping(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: False) + seen = {"n": 0} + + def _fake_overlay(tier): + seen["n"] += 1 + return None # broken/unavailable -> empty, uncached + + monkeypatch.setattr(tv, "_overlay_transformers_dir", _fake_overlay) + assert tv._config_model_types("latest") == frozenset() + assert seen["n"] == 1 # re-resolved, not served from the stale cache + assert "latest" not in tv._config_mapping_cache + + def test_intact_sidecar_serves_cached_latest_mapping(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True) + monkeypatch.setattr( + tv, + "_overlay_transformers_dir", + lambda tier: pytest.fail("intact sidecar must serve the cache without re-resolving"), + ) + assert tv._config_model_types("latest") == frozenset({"brandnew"}) + + def test_non_latest_cache_not_revalidated(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_config_mapping_cache", {"530": frozenset({"gemma3"})}) + monkeypatch.setattr( + tv, + "_latest_sidecar_intact", + lambda: pytest.fail("non-latest tiers must not pay the sidecar-intact check"), + ) + assert tv._config_model_types("530") == frozenset({"gemma3"}) + + def test_deleted_pin_drops_cached_latest_mapping(self, monkeypatch, tmp_path): + # A pin marker deleted after the mapping was cached makes _latest_pin_data None; + # the cache must be dropped (not trusted), so routing re-resolves to no latest tier + # rather than routing to a latest tier that then fails worker activation. + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False) + monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + # No pin marker on disk -> _latest_pin_data() is None -> not intact. + assert tv._latest_sidecar_intact() is False + assert tv._config_model_types("latest") == frozenset() + assert "latest" not in tv._config_mapping_cache + + +class TestOverlayRepairsIncompleteSidecar: + """Routing self-heals a pinned latest sidecar that is present but incomplete, + not only one whose transformers/ dir vanished: workers refuse parent-only + repairs, so a sidecar missing a pinned package would fail every load.""" + + def _setup(self, monkeypatch, tmp_path, valid): + import utils.transformers_version as tv + + live = tmp_path / "venv_t5_latest" + (live / "transformers").mkdir(parents = True) + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live)) + monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False) + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0") + monkeypatch.setattr( + tv, + "_latest_pin_data", + lambda: {"version": "5.99.0", "packages": ["transformers==5.99.0", "tiktoken"]}, + ) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: valid) + monkeypatch.setattr(tv, "_latest_repair_failed_at", 0.0) + return tv + + def test_incomplete_sidecar_triggers_repair(self, monkeypatch, tmp_path): + tv = self._setup(monkeypatch, tmp_path, valid = False) + called = {"n": 0} + + def _fake_repair(): + called["n"] += 1 + return True + + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair) + src = tv._overlay_transformers_dir("latest") + assert called["n"] == 1 + assert src == str(tmp_path / "venv_t5_latest" / "transformers") + + def test_intact_sidecar_skips_repair(self, monkeypatch, tmp_path): + tv = self._setup(monkeypatch, tmp_path, valid = True) + + def _must_not_run(): + raise AssertionError("intact sidecar must not trigger a repair") + + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _must_not_run) + assert tv._overlay_transformers_dir("latest") == str( + tmp_path / "venv_t5_latest" / "transformers" + ) + + def test_failed_repair_backs_off(self, monkeypatch, tmp_path): + tv = self._setup(monkeypatch, tmp_path, valid = False) + called = {"n": 0} + + def _fake_repair(): + called["n"] += 1 + return False + + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair) + # A failed repair must not route through the broken sidecar, neither on + # the failing attempt nor while the backoff suppresses the next attempt. + assert tv._overlay_transformers_dir("latest") is None + assert tv._overlay_transformers_dir("latest") is None + assert called["n"] == 1 + + +class TestStageAndSwapBeforeSwap: + """before_swap fires only when the staged install succeeded and the swap is next.""" + + def _setup(self, monkeypatch, tmp_path, build_ok): + import utils.transformers_version as tv + + live = tmp_path / "venv_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live)) + + def _fake_build(target, packages, label): + if build_ok: + Path(target).mkdir(parents = True, exist_ok = True) + return build_ok + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_build) + return tv, live + + def test_called_after_successful_staging(self, monkeypatch, tmp_path): + tv, live = self._setup(monkeypatch, tmp_path, build_ok = True) + calls = [] + assert tv._stage_and_swap_latest_venv( + "5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1) + ) + assert calls == [1] and live.is_dir() + + def test_not_called_when_staging_fails(self, monkeypatch, tmp_path): + tv, live = self._setup(monkeypatch, tmp_path, build_ok = False) + calls = [] + assert not tv._stage_and_swap_latest_venv( + "5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1) + ) + assert calls == [] and not live.exists() + + def test_failure_in_before_swap_keeps_previous_sidecar(self, monkeypatch, tmp_path): + tv, live = self._setup(monkeypatch, tmp_path, build_ok = True) + live.mkdir() + (live / "sentinel").write_text("old") + + def _boom(): + raise RuntimeError("worker teardown failed") + + assert not tv._stage_and_swap_latest_venv( + "5.99.0", ("transformers==5.99.0",), before_swap = _boom + ) + assert (live / "sentinel").read_text() == "old" + + +class TestKillSwitchBeatsMappingCache: + def test_cached_latest_probe_ignored_when_disabled(self, monkeypatch): + import utils.transformers_version as tv + + key = tv._probe_cache_key("some/model") + monkeypatch.setitem(tv._probe_tier_cache, key, "latest") + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + # With the switch set, the cached latest entry must not short-circuit; + # the probe re-resolves against the non-latest order (stub it to 530). + monkeypatch.setattr(tv, "_probe_tier_venvs", lambda: {}) + monkeypatch.setattr(tv, "_probe_tier_order", lambda: ()) + assert tv._probe_tier("some/model", None, "test") != "latest" + # Cached non-latest entries and the unset switch still short-circuit. + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS") + assert tv._probe_tier("some/model", None, "test") == "latest" + + def test_cached_latest_mapping_ignored_when_disabled(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setitem(tv._config_mapping_cache, "latest", frozenset({"brandnew"})) + # The cache is trusted only when the sidecar is intact; hold it intact so this + # test isolates the kill switch, not the sidecar-revalidation path. + monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True) + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + assert tv._config_model_types("latest") == frozenset() + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS") + assert tv._config_model_types("latest") == frozenset({"brandnew"}) + + +class TestRaiseTierForNested: + """_raise_tier_for_nested: a wrapper's nested model_type can raise a fast-path tier.""" + + def _patch_types(self, monkeypatch, per_tier): + import utils.transformers_version as tv + monkeypatch.setattr( + tv, "_config_model_types", lambda tier: frozenset(per_tier.get(tier, ())) + ) + + def test_nested_latest_only_type_raises(self, monkeypatch): + import utils.transformers_version as tv + + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}}) + cfg = {"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}} + assert tv._raise_tier_for_nested(cfg, "550") == "latest" + + def test_never_lowers_a_fast_path_tier(self, monkeypatch): + import utils.transformers_version as tv + + # Mapping alone would say 530, but the fast path (e.g. a name override) said 550. + self._patch_types(monkeypatch, {"530": {"qwen3_5"}, "550": {"qwen3_5"}}) + assert tv._raise_tier_for_nested({"model_type": "qwen3_5"}, "550") == "550" + + def test_no_config_keeps_tier(self): + import utils.transformers_version as tv + assert tv._raise_tier_for_nested(None, "550") == "550" + + def test_unknown_nested_type_never_vetoes(self, monkeypatch): + import utils.transformers_version as tv + + # A nested type unknown everywhere (not even latest) keeps the fast path. + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4"}}) + cfg = {"model_type": "gemma4", "text_config": {"model_type": "unreleased"}} + assert tv._raise_tier_for_nested(cfg, "550") == "550" + + def test_name_fast_path_folds_when_latest_pinned(self, monkeypatch): + """A fixed-tier name match with a latest-only model_type routes to latest + once the sidecar is pinned; without a pin the name tier stands (no I/O).""" + import utils.transformers_version as tv + + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"brandnew_arch"}}) + monkeypatch.setattr(tv, "_tier_from_name", lambda name: ("550", "gemma-4")) + monkeypatch.setattr( + tv, "_load_config_json", lambda name, tok = None: {"model_type": "brandnew_arch"} + ) + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0") + assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "latest" + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None) + monkeypatch.setattr( + tv, + "_load_config_json", + lambda name, tok = None: (_ for _ in ()).throw(AssertionError("no I/O without a pin")), + ) + assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "550" + + def test_fast_path_folds_nested_tier(self, monkeypatch, tmp_path): + """End to end: a local wrapper config on a fixed fast path routes to latest + when its nested type only exists in the installed latest sidecar.""" + import utils.transformers_version as tv + + ckpt = tmp_path / "wrapper" + ckpt.mkdir() + (ckpt / "config.json").write_text( + json.dumps({"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}}) + ) + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}}) + monkeypatch.setattr(tv, "_config_needs_510", lambda cfg: False) + monkeypatch.setattr(tv, "_config_needs_550", lambda cfg: True) + assert tv.get_transformers_tier(str(ckpt), probe = False) == "latest" diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 281ca24281..284bbb5745 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -698,6 +698,25 @@ if backend_dir not in sys.path: try: from transformers import AutoConfig + # Union the ACTIVE sidecar's registry into the inlined parent-process sets + # so architectures only the sidecar knows still classify correctly. + try: + from transformers.models.auto import modeling_auto as _ma + for _attr in ("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES", + "MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES"): + _d = dict(getattr(_ma, _attr, None) or {}) + _VLM_MODEL_TYPES |= set(_d) + _VLM_CLASS_NAMES |= set(_d.values()) + for _attr in ("MODEL_FOR_CTC_MAPPING_NAMES", + "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", + "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES", + "MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES"): + _AUDIO_ONLY_MODEL_TYPES |= set(dict(getattr(_ma, _attr, None) or {})) + except Exception: + pass + # Capability detection never executes model repo code. kwargs = {"trust_remote_code": False} if token: @@ -727,13 +746,23 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) """ token_arg = hf_token or "" + # Latest-only architectures need the latest sidecar for AutoConfig; + # other tiers keep the 5.5 sidecar. + sidecar_dir = _VENV_T5_DIR + try: + from utils.transformers_version import _VENV_T5_LATEST_DIR, get_transformers_tier + if get_transformers_tier(model_name, hf_token, probe = False) == "latest": + sidecar_dir = _VENV_T5_LATEST_DIR + except Exception: + pass + try: result = subprocess.run( [ sys.executable, "-c", _VISION_CHECK_SCRIPT, - _VENV_T5_DIR, + sidecar_dir, _BACKEND_DIR, model_name, token_arg, @@ -876,6 +905,17 @@ def _is_vision_model_uncached( model_name, hf_token = hf_token, local_files_only = local_files_only ) if raw is not None: + if raw is False and not local_files_only: + # Raw heuristics predate latest-only architectures; on the latest tier, + # trust that sidecar's AutoConfig probe over the heuristic False. An + # inconclusive probe (sidecar mid-repair, timeout) is transient: return + # None so the heuristic False is not cached and the model is re-probed. + try: + from utils.transformers_version import get_transformers_tier + if get_transformers_tier(model_name, hf_token, probe = False) == "latest": + return _is_vision_model_subprocess(model_name, hf_token = hf_token) + except Exception: + pass return raw # Raw read failed transiently: fall back to AutoConfig (remote code DISABLED), via a diff --git a/studio/backend/utils/transformers_latest.py b/studio/backend/utils/transformers_latest.py new file mode 100644 index 0000000000..40c8f729a5 --- /dev/null +++ b/studio/backend/utils/transformers_latest.py @@ -0,0 +1,607 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Latest-transformers support check for brand-new model architectures. + +When a model's ``model_type`` is absent from every installed transformers overlay +(base 4.57.x plus the .venv_t5_530/550/510 sidecars and, if provisioned, .venv_t5_latest), +Studio cannot load it today. This module answers, without authentication, code execution, +or trust_remote_code: + + 1. Does the LATEST transformers release on PyPI ship this ``model_type``? + 2. Does transformers ``main`` on GitHub ship it (dev-only, not yet installable)? + +Sources (all unauthenticated; raw.githubusercontent.com is not API rate-limited and +api.github.com is deliberately never used): + - https://pypi.org/pypi/transformers/json -> latest release version + - https://raw.githubusercontent.com/huggingface/transformers/{ref}/src/transformers/ + models/auto/configuration_auto.py + auto_mappings.py -> CONFIG_MAPPING_NAMES + +The fetched sources are parsed with the same AST extractor the static router uses +(:func:`utils.transformers_version._model_types_from_source`), so the remote answer is +computed exactly like the local overlay answer. + +Results are cached in memory and in a small JSON snapshot under ``studio_root()/cache`` +(ttl ~1 day) so repeated tier resolutions never re-fetch; failures are backed off in +memory. Every fetch is bounded (<=5s, one retry), so a hung network cannot block model +loading. Fully offline-safe: offline env vars or the kill switch +``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1`` make every check return None (current +behavior preserved). + +The consented install path (:func:`install_latest_transformers`) provisions the +persistent ``.venv_t5_latest`` sidecar via +:func:`utils.transformers_version.ensure_latest_transformers_venv`. +""" + +import json +import os +import threading +import time +from pathlib import Path + +from loggers import get_logger +from utils.paths.storage_roots import studio_root as _studio_root +from utils.transformers_version import ( + _env_offline, + _load_config_json, + _model_types_from_source, + _tier_from_config_mapping, + _config_model_types, + _NESTED_CONFIG_KEYS, + _TIER_RANK, + _model_types_from_config, + _TRANSFORMERS_510_MODEL_TYPES, + _TRANSFORMERS_530_MODEL_TYPES, + _TRANSFORMERS_550_MODEL_TYPES, + ensure_latest_transformers_venv, + latest_venv_pinned_version, +) + +logger = get_logger(__name__) + +_PYPI_JSON_URL = "https://pypi.org/pypi/transformers/json" +_RAW_URL = ( + "https://raw.githubusercontent.com/huggingface/transformers/{ref}" + "/src/transformers/models/auto/{name}" +) +_AUTO_FILES = ("configuration_auto.py", "auto_mappings.py") + +_FETCH_TIMEOUT_SECONDS = 5.0 +_FETCH_RETRIES = 1 +_CACHE_TTL_SECONDS = 24 * 60 * 60 +_FAILURE_BACKOFF_SECONDS = 300 + +_CACHE_FILE_NAME = "transformers_latest_check.json" +_SNAPSHOT_SCHEMA = 1 + +# Snapshot: {"schema", "fetched_at", "pypi_version", "pypi_model_types", "main_model_types"}. +# Install-in-progress state lives in utils.transformers_version (the sidecar swap reservation). +_lock = threading.Lock() +_memory_snapshot: dict | None = None +_last_failure_at: float = 0.0 +_is_fetching: bool = False + +_TRUE_VALUES = {"1", "true", "yes", "on"} + + +def _disabled() -> bool: + """True if the operator disabled the latest-transformers check entirely.""" + return ( + os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in _TRUE_VALUES + ) + + +def _cache_file() -> Path: + return _studio_root() / "cache" / _CACHE_FILE_NAME + + +# Sentinel for HTTP 404 (absent at ref), distinct from transient failures. +_FETCH_MISSING = "__unsloth_fetch_missing__" + + +def _fetch_text(url: str) -> str | None: + """GET *url* with a bounded timeout and one retry; None on any failure. + + Returns ``_FETCH_MISSING`` (without retrying) on HTTP 404 so callers can tell + "absent at this ref" apart from "network flaked". + """ + import urllib.error + import urllib.request + + for attempt in range(1 + _FETCH_RETRIES): + try: + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _FETCH_TIMEOUT_SECONDS) as resp: + return resp.read().decode("utf-8", "replace") + except urllib.error.HTTPError as exc: + if exc.code == 404: + return _FETCH_MISSING + logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc) + except Exception as exc: + logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc) + return None + + +def _fetch_latest_pypi_version() -> str | None: + """Latest transformers release version from PyPI's unauthenticated JSON API.""" + body = _fetch_text(_PYPI_JSON_URL) + if body is None or body == _FETCH_MISSING: + return None + try: + version = json.loads(body).get("info", {}).get("version") + except Exception as exc: + logger.debug("Could not parse PyPI JSON: %s", exc) + return None + return version if isinstance(version, str) and version else None + + +def _fetch_remote_model_types(ref: str) -> frozenset[str] | None: + """CONFIG_MAPPING_NAMES keys at *ref* (a release tag like ``v5.12.0`` or ``main``). + + Fetches configuration_auto.py plus auto_mappings.py (the 5.10+ split) from + raw.githubusercontent.com and parses them with the shared AST extractor. A file + that 404s (auto_mappings.py on pre-5.10 tags) is skipped, but a transient fetch + or parse failure of EITHER file fails the whole lookup: most model types live in + auto_mappings.py on current releases, so a partial map cached for the TTL would + make /validate skip the upgrade prompt for architectures the release does ship. + An empty result is likewise a failure so it is never cached as "supports nothing". + """ + keys: set[str] = set() + fetched_any = False + for name in _AUTO_FILES: + source = _fetch_text(_RAW_URL.format(ref = ref, name = name)) + if source is None: + return None + if source == _FETCH_MISSING: + continue + fetched_any = True + try: + keys |= _model_types_from_source(source) + except Exception as exc: + logger.debug("Could not parse %s at %s: %s", name, ref, exc) + return None + if not fetched_any or not keys: + return None + return frozenset(keys) + + +def _load_snapshot_file() -> dict | None: + """Persisted snapshot from disk, or None (missing/corrupt/old schema).""" + try: + with open(_cache_file(), encoding = "utf-8") as f: + data = json.load(f) + except Exception: + return None + if not isinstance(data, dict) or data.get("schema") != _SNAPSHOT_SCHEMA: + return None + if not isinstance(data.get("fetched_at"), (int, float)): + return None + if not isinstance(data.get("pypi_version"), str): + return None + for key in ("pypi_model_types", "main_model_types"): + value = data.get(key) + if not isinstance(value, list) or not all(isinstance(v, str) for v in value): + return None + return data + + +def _save_snapshot_file(snapshot: dict) -> None: + """Atomic best-effort write (tmp + os.replace, Windows-safe); failures only log.""" + path = _cache_file() + tmp = path.with_name(path.name + ".tmp") + try: + path.parent.mkdir(parents = True, exist_ok = True) + tmp.write_text(json.dumps(snapshot), encoding = "utf-8") + os.replace(tmp, path) + except Exception as exc: + logger.debug("Could not persist %s: %s", path, exc) + try: + tmp.unlink(missing_ok = True) + except Exception: + pass + + +def _snapshot_is_fresh(snapshot: dict | None) -> bool: + return ( + snapshot is not None + and (time.time() - float(snapshot.get("fetched_at", 0))) < _CACHE_TTL_SECONDS + ) + + +def _refresh_snapshot() -> dict | None: + """Fetch a fresh snapshot from PyPI + raw.githubusercontent.com; None on failure. + + The PyPI version and its tagged mapping are required; the ``main`` mapping is + best-effort (recorded as an empty list plus ``main_checked=False`` when unavailable, + so a dev-only architecture is reported as "unknown" rather than "unsupported"). + """ + version = _fetch_latest_pypi_version() + if version is None: + return None + pypi_types = _fetch_remote_model_types(f"v{version}") + if pypi_types is None: + return None + main_types = _fetch_remote_model_types("main") + return { + "schema": _SNAPSHOT_SCHEMA, + "fetched_at": time.time(), + "pypi_version": version, + "pypi_model_types": sorted(pypi_types), + "main_model_types": sorted(main_types) if main_types is not None else [], + "main_checked": main_types is not None, + } + + +def _get_snapshot() -> dict | None: + """Current support snapshot: memory -> disk -> network, with TTL and failure backoff. + + The network refresh runs outside the lock so a slow fetch cannot stall other + threads in the ASGI pool; _is_fetching deduplicates concurrent refreshes + (losers return None, the graceful fallthrough, rather than waiting). + """ + global _memory_snapshot, _last_failure_at, _is_fetching + with _lock: + if _snapshot_is_fresh(_memory_snapshot): + return _memory_snapshot + disk = _load_snapshot_file() + if _snapshot_is_fresh(disk): + _memory_snapshot = disk + return disk + if _disabled() or _env_offline(): + return None + if time.time() - _last_failure_at < _FAILURE_BACKOFF_SECONDS: + return None + if _is_fetching: + return None + _is_fetching = True + fresh = None + try: + fresh = _refresh_snapshot() + finally: + with _lock: + _is_fetching = False + if fresh is None: + _last_failure_at = time.time() + else: + _memory_snapshot = fresh + if fresh is None: + # A stale positive could offer a version PyPI no longer serves; be strict. + return None + _save_snapshot_file(fresh) + return fresh + + +def clear_caches() -> None: + """Test helper: drop the in-memory snapshot, failure backoff, and busy flags.""" + global _memory_snapshot, _last_failure_at, _is_fetching + with _lock: + _memory_snapshot = None + _last_failure_at = 0.0 + _is_fetching = False + from utils.transformers_version import end_sidecar_swap + + end_sidecar_swap() + + +def latest_transformers_supports(model_type: str) -> dict | None: + """Whether the newest transformers (PyPI release and/or GitHub main) ships *model_type*. + + Returns ``{"pypi_version": str, "supported_in_pypi": bool, "supported_in_main": bool}`` + or None when the answer is unavailable (offline, kill switch, network failure) — the + caller must then fall through to current behavior. Cached (memory + JSON snapshot on + disk, ttl ~1 day) so repeated tier resolutions never re-fetch. + """ + if not isinstance(model_type, str) or not model_type: + return None + if _disabled() or _env_offline(): + return None + snapshot = _get_snapshot() + if snapshot is None: + return None + return { + "pypi_version": snapshot["pypi_version"], + "supported_in_pypi": model_type in set(snapshot["pypi_model_types"]), + "supported_in_main": model_type in set(snapshot["main_model_types"]), + } + + +# model_types the hardcoded tier tables already route; never remote-check these. +def _hardcoded_model_types() -> frozenset[str]: + return frozenset( + _TRANSFORMERS_530_MODEL_TYPES + | _TRANSFORMERS_550_MODEL_TYPES + | _TRANSFORMERS_510_MODEL_TYPES + ) + + +def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dict | None: + """Upgrade signal for *model_name*, or None when current routing already handles it. + + The tier hook for the pre-load ``/validate`` path: fires ONLY when the model's + ``model_type`` is absent from every installed overlay (and from the hardcoded tier + tables), i.e. exactly when today's load would fail with an unrecognized-architecture + error. Returns ``{"model_type", "pypi_version", "supported_in_pypi", + "supported_in_main"}`` when the newest transformers knows the type, else None. + + Never raises; every network touch is bounded and cached. Offline or with the + ``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS`` kill switch it returns None immediately. + """ + try: + if _disabled() or _env_offline(): + return None + cfg = _load_config_json(model_name, hf_token) + if not isinstance(cfg, dict): + return None + candidates = _model_types_from_config(cfg) + if not candidates: + return None + # Without a readable base mapping every type looks brand new; bail out. + if not _config_model_types("default"): + return None + hardcoded = _hardcoded_model_types() + missing = [ + candidate + for candidate in candidates + if candidate not in hardcoded + and not any(candidate in _config_model_types(tier) for tier in _TIER_RANK) + ] + if not missing: + return None + # Latest must load EVERY missing type (wrappers build nested sub-configs + # through CONFIG_MAPPING) or the load still fails. + supports = [latest_transformers_supports(candidate) for candidate in missing] + if any( + s is None or not (s["supported_in_pypi"] or s["supported_in_main"]) for s in supports + ): + return None + # Offer the PyPI install only if the release ships every missing type; a + # main-only type in the mix surfaces as dev-only. + model_type = missing[0] + supported_in_pypi = all(s["supported_in_pypi"] for s in supports) + supported_in_main = all(s["supported_in_pypi"] or s["supported_in_main"] for s in supports) + logger.info( + "Model %s has model_type=%s unknown to every installed transformers " + "(latest PyPI %s: %s, main: %s)", + model_name, + model_type, + supports[0]["pypi_version"], + "supported" if supported_in_pypi else "unsupported", + "supported" if supported_in_main else "unsupported", + ) + return { + "model_type": model_type, + "pypi_version": supports[0]["pypi_version"], + "supported_in_pypi": supported_in_pypi, + "supported_in_main": supported_in_main, + } + except Exception as exc: + logger.debug("Latest-transformers check failed for '%s': %s", model_name, exc) + return None + + +# --- Dependency compatibility preflight ------------------------------------------------------ +# Sidecars install transformers --no-deps atop the base env. Before installing, compare +# requires_dist: unsatisfied shadowable deps become exact --target pins, anything else blocks. + +# Safe to shadow inside the sidecar dir (pure wheels, no torch coupling). +_SHADOWABLE_DEPS = frozenset({"tokenizers", "safetensors"}) +# Provided by the sidecar recipe; checked against its pin, not the base env. +_SIDECAR_PROVIDED = {"huggingface-hub": "1.8.0", "hf-xet": "1.4.2"} +# CLI-only; never imported at runtime in Studio's workers. +_IGNORED_DEPS = frozenset({"typer"}) + + +def _canonical_dep_name(name: str) -> str: + return name.lower().replace("_", "-") + + +def _fetch_requires_dist(version: str) -> list[str] | None: + """Core (marker-free, non-extra) requires_dist of transformers *version* from PyPI.""" + body = _fetch_text(f"https://pypi.org/pypi/transformers/{version}/json") + if body is None or body == _FETCH_MISSING: + return None + try: + reqs = json.loads(body).get("info", {}).get("requires_dist") + except Exception: + return None + if not isinstance(reqs, list): + return None + return [r for r in reqs if isinstance(r, str)] + + +def _resolve_exact_version(name: str, specifier) -> str | None: + """Newest PyPI release of *name* satisfying *specifier* (exact pin for the shadow).""" + body = _fetch_text(f"https://pypi.org/pypi/{name}/json") + if body is None or body == _FETCH_MISSING: + return None + try: + from packaging.version import InvalidVersion, Version + + releases = json.loads(body).get("releases", {}) + best = None + for candidate in releases: + try: + parsed = Version(candidate) + except InvalidVersion: + continue + if parsed.is_prerelease or not specifier.contains(candidate): + continue + if best is None or parsed > Version(best): + best = candidate + return best + except Exception as exc: + logger.debug("Could not resolve an exact %s version: %s", name, exc) + return None + + +def compat_plan(version: str) -> tuple[tuple[str, ...], list[str]]: + """(extra exact pins to shadow-install, blocking requirement strings) for *version*. + + Compares the release's core requires_dist against the running base env (the env the + workers overlay the sidecar onto). A requirement the base env satisfies needs nothing; + an unsatisfied shadowable dep becomes an exact pin inside the sidecar; any other + unsatisfied requirement is a blocker. An unavailable requires_dist BLOCKS the + install: proceeding unverified could pin a sidecar whose imports then crash the + workers, and the caller just reached PyPI for the version check so a retry is cheap. + """ + reqs = _fetch_requires_dist(version) + if reqs is None: + return (), ["dependency metadata for this release (could not be fetched from PyPI; retry)"] + try: + from importlib.metadata import PackageNotFoundError + from importlib.metadata import version as _installed_version + from packaging.requirements import InvalidRequirement, Requirement + except Exception: + return (), [] + extras: list[str] = [] + blockers: list[str] = [] + for raw in reqs: + try: + req = Requirement(raw) + except InvalidRequirement: + continue + if req.extras or (req.marker is not None and not req.marker.evaluate()): + continue + name = _canonical_dep_name(req.name) + if name in _IGNORED_DEPS: + continue + if name in _SIDECAR_PROVIDED: + if not req.specifier.contains(_SIDECAR_PROVIDED[name], prereleases = True): + blockers.append(raw) + continue + try: + installed = _installed_version(req.name) + except PackageNotFoundError: + installed = None + if installed is not None and req.specifier.contains(installed, prereleases = True): + continue + if name in _SHADOWABLE_DEPS: + exact = _resolve_exact_version(name, req.specifier) + if exact is None: + blockers.append(raw) + else: + extras.append(f"{name}=={exact}") + else: + blockers.append(raw) + return tuple(extras), blockers + + +def is_install_in_progress() -> bool: + """True while a latest-transformers install or lazy repair holds the sidecar swap + reservation. Training and export starts check this so a fresh worker never + activates the sidecar mid-swap.""" + from utils.transformers_version import sidecar_swap_in_progress + return sidecar_swap_in_progress() + + +def install_latest_transformers( + version: str, + before_swap = None, + reserved: bool = False, +) -> dict: + """Consented install of the latest transformers sidecar; returns a structured result. + + Guards: the requested *version* must match the current PyPI latest from the (cached) + snapshot, so a client cannot pin an arbitrary package version through this endpoint. + On success ``.venv_t5_latest`` is provisioned and pinned; routing then resolves the + new tier automatically on this and every future start. *before_swap* is forwarded + to the stage-and-swap: it runs only after the staged install succeeded, right + before the live sidecar is replaced. *reserved* means the caller already holds the + sidecar swap reservation (the install route takes it before waiting on the + inference lifecycle gate, so worker starts see it for the whole window). + """ + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + if not reserved and not try_begin_sidecar_swap(): + return { + "success": False, + "version": version, + "message": "A transformers installation is already in progress.", + } + try: + return _install_latest_transformers_locked(version, before_swap = before_swap) + finally: + if not reserved: + end_sidecar_swap() + + +def _install_latest_transformers_locked(version: str, before_swap = None) -> dict: + """Body of install_latest_transformers; runs with the in-progress flag held.""" + if _disabled(): + return { + "success": False, + "version": version, + "message": "Latest-transformers installs are disabled " + "(UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS).", + } + if _env_offline(): + return { + "success": False, + "version": version, + "message": "Cannot install: Studio is in offline mode.", + } + # Re-verify against a LIVE snapshot (a release may land inside the cache TTL); + # fall back to the cached one on fetch failure. + global _memory_snapshot + snapshot = _refresh_snapshot() + if snapshot is not None: + with _lock: + _memory_snapshot = snapshot + _save_snapshot_file(snapshot) + else: + snapshot = _get_snapshot() + if snapshot is None: + return { + "success": False, + "version": version, + "message": "Could not verify the latest transformers release on PyPI.", + } + if version != snapshot["pypi_version"]: + return { + "success": False, + "version": version, + "message": f"Requested version {version!r} is not the latest transformers " + f"release ({snapshot['pypi_version']}).", + # Lets the consent dialog retry with the release that superseded the + # one /validate saw, instead of re-sending the stale version forever. + "latest_version": snapshot["pypi_version"], + } + extra_packages, blockers = compat_plan(version) + if blockers: + return { + "success": False, + "version": version, + "message": "Cannot install transformers " + f"{version}: this environment does not satisfy {', '.join(blockers)}. " + "A Studio update is required first.", + } + if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap): + return { + "success": False, + "version": version, + "message": f"Installing transformers {version} failed; see the Studio logs.", + } + _invalidate_capability_caches() + return { + "success": True, + "version": version, + "message": f"Installed transformers {version} into the latest sidecar " + f"(pinned: {latest_venv_pinned_version()}).", + } + + +def _invalidate_capability_caches(): + """Drop caches computed before the new sidecar existed: tier probes and the + latest tier's model_type mapping (stale on upgrade) plus vision detection + (a raw-heuristic False may now defer to the sidecar AutoConfig probe).""" + try: + from utils import transformers_version as tv + tv._probe_tier_cache.clear() + tv._config_mapping_cache.pop("latest", None) + except Exception: + pass + try: + from utils.models import model_config as mc + mc._vision_detection_cache.clear() + except Exception: + pass diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index a69673f081..9f9f8aa3de 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -35,9 +35,12 @@ import json import structlog from loggers import get_logger import os +import re import shutil import subprocess import sys +import threading +import time from pathlib import Path from utils.native_path_leases import child_env_without_native_path_secret @@ -235,8 +238,12 @@ _VENV_T5_DIR = _VENV_T5_550_DIR # reuses the workspace torch (torch-agnostic). _VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") -# Tier precedence: higher rank wins in _higher_tier. -_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} +# User-consented "latest transformers" sidecar (utils/transformers_latest.py); pinned version in a marker file. +_VENV_T5_LATEST_DIR = str(_studio_root() / ".venv_t5_latest") +_LATEST_PIN_MARKER = ".unsloth_pinned_transformers" + +# Tier precedence: higher rank wins in _higher_tier. "latest" outranks every fixed tier. +_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3, "latest": 4} def _higher_tier(a: str, b: str) -> str: @@ -254,20 +261,40 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None = ``hf_token`` is forwarded to tier detection so a gated/private model whose only 5.x signal is an authenticated config/tokenizer reaches the right sidecar, not the default. """ - # Pre-resolve only LoRA adapters; full checkpoints go to get_transformers_tier so their - # local config.json drives the tier (a full checkpoint with a private/offline - # _name_or_path must not resolve to an unreachable HF id and skip its own config). + # Pre-resolve LoRA adapters (local dir or remote adapter repo); full checkpoints + # go to get_transformers_tier so their local config.json drives the tier (a full + # checkpoint with a private/offline _name_or_path must not resolve to an + # unreachable HF id and skip its own config). Remote adapters activate for their + # BASE model, matching latest_tier_active_for and the inference worker. if _is_lora_adapter_dir(Path(model_name)): resolved = _resolve_base_model(model_name) else: - resolved = model_name + resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name tier = get_transformers_tier(resolved, hf_token) if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): # Gate on a real local config.json: a checkpoint carries config the base may not # surface, but path names alone must not upgrade a plain adapter. tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token)) - if tier == "510": + if tier == "latest": + pinned = latest_venv_pinned_version() + if pinned is None or not _ensure_venv_t5_latest_exists(): + raise RuntimeError( + f"Cannot activate the latest-transformers sidecar: " + f".venv_t5_latest missing or unpinned at {_VENV_T5_LATEST_DIR}" + ) + if _VENV_T5_LATEST_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_LATEST_DIR) + logger.info( + "Prepended transformers %s venv to sys.path from %s " + "(path only; the loaded version is confirmed later by " + "'Subprocess loaded transformers ...' on first import)", + pinned, + _VENV_T5_LATEST_DIR, + ) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_LATEST_DIR + (os.pathsep + _pp if _pp else "") + elif tier == "510": if not _ensure_venv_t5_510_exists(): raise RuntimeError( f"Cannot activate transformers {TRANSFORMERS_510_VERSION}: " @@ -322,6 +349,34 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None = logger.info("Using default transformers (4.57.x) for %s", model_name) +def latest_tier_active_for(model_name: str, hf_token: str | None = None) -> bool: + """True when *model_name* routes to the consented latest-transformers sidecar. + + Mirrors the inference worker's pre-activation resolution (local adapter dir, + then a remote adapter's Hub adapter_config.json). ``latest`` only wins when + the sidecar exists with a valid pin, i.e. exactly the loads that will import + the newest release. Never raises: any resolution failure returns False so + callers treat the model as a known tier. + """ + try: + # No consented sidecar pin means nothing routes to latest; return before + # any resolution so the common case costs no config or network reads. + if latest_venv_pinned_version() is None: + return False + if _is_lora_adapter_dir(Path(model_name)): + resolved = _resolve_base_model(model_name) + else: + # A remote LoRA activates the sidecar for its BASE model; sizing and the + # worker's 4-bit guard must see that base too, not the adapter repo. + resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name + tier = get_transformers_tier(resolved, hf_token) + if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): + tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token)) + return tier == "latest" + except Exception: + return False + + def _has_adapter_weights(path: Path) -> bool: """True if *path* holds LoRA adapter weight files (``adapter_model.*``).""" try: @@ -881,17 +936,85 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None: _config_mapping_cache: dict[str, frozenset[str]] = {} +def _latest_tier_disabled() -> bool: + """Kill switch shared with utils.transformers_latest: lets operators roll + back a provisioned latest sidecar without deleting files.""" + return os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +# Failed lazy repairs back off so a broken sidecar can't turn every routing +# call into a pip install attempt. +_latest_repair_failed_at: float = 0.0 +_LATEST_REPAIR_BACKOFF_SECS = 5 * 60 + + +def _latest_sidecar_intact() -> bool: + """The pinned latest sidecar exists with its transformers dir and every pinned + package. False when the pin itself is gone: a cached 'latest' mapping must then be + dropped (routing re-resolves to no latest tier), not trusted, and a sidecar that kept + transformers/ but lost a pinned package must self-heal rather than route models to a + latest tier that fails activation in workers, which refuse parent-only repairs. + + (_overlay_transformers_dir only calls this after gating on a present pin, so the + pin-missing case here is the cache-revalidation caller whose pin was deleted after + the mapping was first cached.)""" + pin = _latest_pin_data() + if pin is None: + return False + return _venv_dir_is_valid(_VENV_T5_LATEST_DIR, tuple(pin["packages"])) + + def _overlay_transformers_dir(tier: str) -> str | None: """transformers source dir for a tier, located without importing it.""" + global _latest_repair_failed_at if tier != "default": - root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier) + # latest requires a valid pin and the kill switch off. + if tier == "latest" and (_latest_tier_disabled() or latest_venv_pinned_version() is None): + return None + root = { + "530": _VENV_T5_530_DIR, + "550": _VENV_T5_550_DIR, + "510": _VENV_T5_510_DIR, + "latest": _VENV_T5_LATEST_DIR, + }.get(tier) src = os.path.join(root, "transformers") if root else None + if src and tier == "latest" and not _latest_sidecar_intact(): + # A valid pin whose sidecar vanished or lost a pinned package (partial + # deletion, disk issue, interrupted external edits) must self-heal, or + # latest-only models either silently route to older tiers or reach a + # worker that cannot repair, failing every load until a manual + # reinstall. Repair under the swap reservation; back off after a + # failure so routing calls don't hammer pip. + repaired = False + if time.time() - _latest_repair_failed_at >= _LATEST_REPAIR_BACKOFF_SECS: + if _ensure_venv_t5_latest_exists(): + _latest_repair_failed_at = 0.0 + repaired = True + else: + _latest_repair_failed_at = time.time() + if not repaired: + # Still broken: treat the overlay as unavailable rather than route + # models to a tier whose worker activation is known to fail. Models + # an older tier supports keep loading there until a repair succeeds, + # matching the behavior when the sidecar dir is missing entirely. + return None return src if src and _safe_is_dir(Path(src)) else None # default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one # is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir. sidecars = tuple( os.path.abspath(d) + os.sep - for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR) + for d in ( + _VENV_T5_530_DIR, + _VENV_T5_550_DIR, + _VENV_T5_510_DIR, + _VENV_T5_LATEST_DIR, + _VENV_LLMCOMPRESSOR_DIR, + ) ) candidates = [] try: @@ -930,11 +1053,47 @@ def _mapping_first_keys(value: ast.AST) -> set[str]: return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)} +def _model_types_from_source(source: str) -> set[str]: + """model_type keys of CONFIG_MAPPING_NAMES in *source* (AST only, no execution). + + Handles the direct ``CONFIG_MAPPING_NAMES = ...`` binding (dict literal or + OrderedDict/dict call over 2-tuple lists and **{...} unpacking) and any + ``CONFIG_MAPPING_NAMES.update({...})`` mutation. Shared by the on-disk overlay + reader below and the remote latest-release checker (utils/transformers_latest.py). + """ + keys: set[str] = set() + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets + ): + keys |= _mapping_first_keys(node.value) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + fn = node.value.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "update" + and isinstance(fn.value, ast.Name) + and fn.value.id == "CONFIG_MAPPING_NAMES" + ): + keys |= _mapping_first_keys(node.value) + return keys + + def _config_model_types(tier: str) -> frozenset[str]: """model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py).""" + # Kill switch beats the cache: a stale mapping must not keep routing latest-only models until restart. + if tier == "latest" and _latest_tier_disabled(): + return frozenset() cached = _config_mapping_cache.get(tier) if cached is not None: - return cached + # A cached 'latest' mapping can outlive the sidecar it was parsed from: if the + # pinned sidecar was since deleted or lost a package in this process, drop the + # cache so routing re-resolves through _overlay_transformers_dir (which self-heals) + # instead of routing latest-only models to a broken tier until restart. + if tier != "latest" or _latest_sidecar_intact(): + return cached + _config_mapping_cache.pop("latest", None) tdir = _overlay_transformers_dir(tier) if tdir is None: return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads @@ -944,22 +1103,7 @@ def _config_model_types(tier: str) -> frozenset[str]: if not _safe_is_file(path): continue try: - tree = ast.parse(path.read_text(encoding = "utf-8")) - for node in ast.walk(tree): - # direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation - if isinstance(node, ast.Assign) and any( - isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets - ): - keys |= _mapping_first_keys(node.value) - elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): - fn = node.value.func - if ( - isinstance(fn, ast.Attribute) - and fn.attr == "update" - and isinstance(fn.value, ast.Name) - and fn.value.id == "CONFIG_MAPPING_NAMES" - ): - keys |= _mapping_first_keys(node.value) + keys |= _model_types_from_source(path.read_text(encoding = "utf-8")) except Exception: continue result = frozenset(keys) @@ -967,23 +1111,73 @@ def _config_model_types(tier: str) -> frozenset[str]: return result -def _tier_from_config_mapping(cfg: dict) -> str | None: - """Lowest tier whose transformers ships cfg's model_type, or None if unknown.""" - model_type = cfg.get("model_type") - if not isinstance(model_type, str): - for key in _NESTED_CONFIG_KEYS: - sub = cfg.get(key) - if isinstance(sub, dict) and isinstance(sub.get("model_type"), str): - model_type = sub["model_type"] - break - if not isinstance(model_type, str): - return None +def _model_types_from_config(cfg: dict) -> list[str]: + """All model_types in the config: the primary (top-level, else first nested) + first, then every other nested sub-config. Wrappers instantiate sub-configs + through CONFIG_MAPPING, so nested types matter for routing too.""" + seen: list[str] = [] + + def add(value): + if isinstance(value, str) and value and value not in seen: + seen.append(value) + + add(cfg.get("model_type")) + for key in _NESTED_CONFIG_KEYS: + sub = cfg.get(key) + if isinstance(sub, dict): + add(sub.get("model_type")) + for value in cfg.values(): + if isinstance(value, dict): + add(value.get("model_type")) + return seen + + +def _lowest_tier_for(model_type: str) -> str | None: for tier in sorted(_TIER_RANK, key = _TIER_RANK.get): if model_type in _config_model_types(tier): return tier return None +def _tier_from_config_mapping(cfg: dict) -> str | None: + """Lowest tier able to load every model_type in cfg, or None when the + primary type is unknown everywhere. A nested type can raise the tier (its + sub-config is built through CONFIG_MAPPING); an unknown nested type never + vetoes, since no installed tier could load it either way (the latest + checker handles surfacing the install prompt for it).""" + types = _model_types_from_config(cfg) + if not types: + return None + best = _lowest_tier_for(types[0]) + if best is None: + return None + for model_type in types[1:]: + tier = _lowest_tier_for(model_type) + if tier is not None and _TIER_RANK[tier] > _TIER_RANK[best]: + best = tier + return best + + +def _raise_tier_for_nested(cfg: dict | None, tier: str) -> str: + """Raise *tier* when the mapping resolver needs a higher one for *cfg*. + + A wrapper's top-level model_type can match a hardcoded fast path while a + nested text/vision config's type only exists in a newer sidecar (e.g. the + installed latest); its sub-config is built through CONFIG_MAPPING, so the + fast-path tier would fail to load it. Raise-only: never lowers a fast-path + match, so name overrides (Qwen3.6) keep their tier. Never raises an + exception: a resolution failure keeps the fast-path tier.""" + if not isinstance(cfg, dict): + return tier + try: + mapped = _tier_from_config_mapping(cfg) + if mapped is not None and _TIER_RANK.get(mapped, 0) > _TIER_RANK.get(tier, 0): + return mapped + except Exception: + pass + return tier + + # --- AutoConfig probe: general tier resolution for ambiguous models ---------- # When the cheap signals only say "needs some 5.x", parse config.json with the built-in # parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond @@ -1039,9 +1233,19 @@ def _probe_tier_venvs(): "530": (_VENV_T5_530_DIR, _ensure_venv_t5_530_exists), "550": (_VENV_T5_550_DIR, _ensure_venv_t5_550_exists), "510": (_VENV_T5_510_DIR, _ensure_venv_t5_510_exists), + "latest": (_VENV_T5_LATEST_DIR, _ensure_venv_t5_latest_exists), } +def _probe_tier_order() -> tuple[str, ...]: + """Sidecar probe order. The consented "latest" sidecar joins only once it is + provisioned (pin marker present): an absent optional tier must not flip the probe's + skipped-tier bookkeeping, keeping pre-latest behavior byte-identical.""" + if not _latest_tier_disabled() and latest_venv_pinned_version() is not None: + return _PROBE_TIER_ORDER + ("latest",) + return _PROBE_TIER_ORDER + + def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) -> bool | None: """Parse config.json with the built-in parser inside *target_dir*'s sidecar. True = parses, False = parse/version failure (escalate), None = transient @@ -1119,7 +1323,7 @@ def _probe_tier( stays on the default. Cached per _probe_cache_key (process lifetime). No Hub sha is resolved: that would import huggingface_hub before the sidecar is on sys.path. """ - if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes"): + if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes", "on"): return floor key = _probe_cache_key(model_name) # Key by probe mode: the default-first path can return 'default', which must not be @@ -1127,7 +1331,10 @@ def _probe_tier( if include_default or floor != "530": key = f"{key}\0floor={floor}:def={int(include_default)}" if key in _probe_tier_cache: - return _probe_tier_cache[key] + cached = _probe_tier_cache[key] + # Kill switch beats the cache (like _config_model_types): a stale 'latest' probe must not keep activating it. + if cached != "latest" or not _latest_tier_disabled(): + return cached def _cache(tier: str, *, skipped: bool) -> str: # Do not pin a result that depended on a skipped lower tier: once that sidecar is @@ -1137,7 +1344,8 @@ def _probe_tier( return tier venvs = _probe_tier_venvs() - order = (("default",) + _PROBE_TIER_ORDER) if include_default else _PROBE_TIER_ORDER + sidecar_order = _probe_tier_order() + order = (("default",) + sidecar_order) if include_default else sidecar_order probed_count = 0 skipped_any = False for tier in order: @@ -1264,17 +1472,21 @@ def get_transformers_tier( cfg = _load_config_json(model_name, hf_token) if cfg is not None: if _config_needs_510(cfg): + tier = _raise_tier_for_nested(cfg, "510") logger.info( - "Transformers tier 510 selected for %s (local config.json check)", + "Transformers tier %s selected for %s (local config.json check)", + tier, model_name, ) - return "510" + return tier if _config_needs_550(cfg): + tier = _raise_tier_for_nested(cfg, "550") logger.info( - "Transformers tier 550 selected for %s (local config.json check)", + "Transformers tier %s selected for %s (local config.json check)", + tier, model_name, ) - return "550" + return tier if _config_needs_530(cfg): # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name. Only a real # Hub id (or the folder basename) may override 530, so a stale local @@ -1287,17 +1499,20 @@ def get_transformers_tier( ) override = _higher_tier_name_override(hint_src) if override is not None: + override = _raise_tier_for_nested(cfg, override) logger.info( "Transformers tier %s selected for %s (name overrides 530 config)", override, model_name, ) return override + tier = _raise_tier_for_nested(cfg, "530") logger.info( - "Transformers tier 530 selected for %s (local config.json check)", + "Transformers tier %s selected for %s (local config.json check)", + tier, model_name, ) - return "530" + return tier # Unknown arch: resolve the base id from config. A resolved local dir # recurses (config check); a Hub id uses name rules only (no network). resolved = _resolve_base_model(model_name) @@ -1359,6 +1574,13 @@ def get_transformers_tier( result = _tier_from_name(model_name) if result is not None: tier, match = result + # With a consented latest sidecar pinned, a name that matches a fixed + # tier can still carry a latest-only model_type (e.g. a newer variant + # reusing a family name); consult the config so an accepted upgrade + # actually routes to the sidecar it installed. Costs a config read only + # in the pinned case, keeping the pre-latest path I/O-free. + if latest_venv_pinned_version() is not None: + tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), tier) logger.info( "Transformers tier %s selected for %s (substring match: %s)", tier, @@ -1369,11 +1591,13 @@ def get_transformers_tier( # --- Slow config fallbacks (network for HF IDs; authenticated with hf_token) -------- if _check_config_needs_510(model_name, hf_token): - logger.info("Transformers tier 510 selected for %s (config.json check)", model_name) - return "510" + tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "510") + logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name) + return tier if _check_config_needs_550(model_name, hf_token): - logger.info("Transformers tier 550 selected for %s (config.json check)", model_name) - return "550" + tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "550") + logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name) + return tier if _check_config_needs_530(model_name, hf_token): # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name; honor a real Hub-id name # hint from _name_or_path before selecting 530. @@ -1383,14 +1607,16 @@ def get_transformers_tier( base if isinstance(base, str) and base != model_name else None ) if override is not None: + override = _raise_tier_for_nested(remote_cfg, override) logger.info( "Transformers tier %s selected for %s (name overrides 530 config)", override, model_name, ) return override - logger.info("Transformers tier 530 selected for %s (config.json check)", model_name) - return "530" + tier = _raise_tier_for_nested(remote_cfg, "530") + logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name) + return tier # _load_config_json (not the cache-only reader) so a config served from the hub # cache during a transient outage still feeds the mapping resolver. remote_cfg = _load_config_json(model_name, hf_token) @@ -1657,6 +1883,471 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- User-consented "latest transformers" sidecar (.venv_t5_latest) -------------------------- +# Provisioned via ensure_latest_transformers_venv() after the user confirms the upgrade popup +# (utils/transformers_latest.py); pinned in a marker file so restarts revalidate and routing auto-picks it. + +# PEP 440-ish release strings only (guards the pip install spec against injection). +_LATEST_VERSION_RE = r"[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?" + + +def _is_valid_version_string(version: str) -> bool: + import re + return isinstance(version, str) and re.fullmatch(_LATEST_VERSION_RE, version) is not None + + +# Only the sidecar recipe's own packages, as plain (optionally ==pinned) specs, may +# come from the on-disk pin marker; anything else (URLs, extras, options) is rebuilt. +_PIN_SPEC_RE = re.compile(r"^[A-Za-z0-9_.-]+(==[A-Za-z0-9_.+-]+)?$") +_PIN_ALLOWED_NAMES = frozenset( + { + "transformers", + "huggingface_hub", + "huggingface-hub", + "hf_xet", + "hf-xet", + "tiktoken", + "tokenizers", + "safetensors", + } +) + + +def _is_safe_pin_spec(spec: str) -> bool: + if not _PIN_SPEC_RE.match(spec): + return False + name = spec.split("==", 1)[0].lower().replace("_", "-") + return name in {n.replace("_", "-") for n in _PIN_ALLOWED_NAMES} + + +def _recover_stranded_latest_sidecar() -> None: + """Restore a sidecar stranded at ``.old`` by a swap whose activation rename AND its + rollback both failed (e.g. a lingering worker file handle on Windows blocked both). + + That double failure leaves no live dir and the pin marker gone with it, so the + sidecar reads as unprovisioned and never self-heals. Recover only when no live dir + exists and no swap is in flight: the reservation is held throughout the swap, so the + transient live-absent window of a legitimate swap never triggers a restore.""" + live = Path(_VENV_T5_LATEST_DIR) + retired = Path(_VENV_T5_LATEST_DIR + ".old") + try: + if live.exists() or not retired.is_dir() or sidecar_swap_in_progress(): + return + os.rename(retired, live) + logger.info("Recovered .venv_t5_latest from a stranded .old after a failed swap") + except OSError: + pass + + +def _latest_pin_data() -> dict | None: + """Parsed pin marker: {"version": str, "packages": [specs...]}, or None. + + The marker is JSON; a plain version string (older/simpler writers) is tolerated and + expanded with the default package set. + """ + _recover_stranded_latest_sidecar() + marker = Path(_VENV_T5_LATEST_DIR) / _LATEST_PIN_MARKER + try: + if not marker.is_file(): + return None + raw = marker.read_text(encoding = "utf-8").strip() + except Exception: + return None + try: + data = json.loads(raw) + except ValueError: + data = raw + if isinstance(data, str): + if not _is_valid_version_string(data): + return None + return {"version": data, "packages": list(_venv_t5_latest_packages(data))} + if not isinstance(data, dict): + return None + version = data.get("version") + if not _is_valid_version_string(version): + return None + packages = data.get("packages") + if not ( + isinstance(packages, list) + and packages + and all(isinstance(p, str) and _is_safe_pin_spec(p) for p in packages) + ): + # Malformed or unexpected specs (the pin is user-writable on disk) never + # reach pip: rebuild the canonical set for the pinned version instead. + packages = list(_venv_t5_latest_packages(version)) + return {"version": version, "packages": packages} + + +def latest_venv_pinned_version() -> str | None: + """Exact transformers version pinned in .venv_t5_latest's marker, or None if the + sidecar was never provisioned (or the marker is unreadable/invalid).""" + data = _latest_pin_data() + return data["version"] if data else None + + +def _venv_t5_latest_packages(version: str, extra_packages: tuple[str, ...] = ()) -> tuple[str, ...]: + """Package set for the latest sidecar; mirrors the fixed .venv_t5_* sidecars. + *extra_packages* carries dep-compat shadows (e.g. a newer tokenizers) computed by + utils.transformers_latest before install.""" + return ( + f"transformers=={version}", + "huggingface_hub==1.8.0", + "hf_xet==1.4.2", + "tiktoken", + ) + tuple(extra_packages) + + +# Single reservation for ANY .venv_t5_latest replacement (consented install or lazy repair), +# checked by training/export starts so no worker spawns mid-swap. Backed by a lock FILE (not just +# this flag) so a lazy repair running in a worker subprocess stays visible to the parent's route +# checks; the in-process flag marks ownership (only the owner unlinks the file). +_sidecar_swap_lock = threading.Lock() +_sidecar_swap_active = False +_sidecar_swap_token: str | None = None +_sidecar_swap_kind: str | None = None +# An install is minutes; a lock this old is a crashed owner, not a live swap. +_SWAP_LOCK_STALE_SECS = 2 * 60 * 60 + + +def _swap_lock_path() -> Path: + return Path(_VENV_T5_LATEST_DIR + ".swaplock") + + +def _pid_alive(pid) -> bool: + if not isinstance(pid, int) or pid <= 0: + return False + try: + import psutil + return psutil.pid_exists(pid) + except Exception: + pass + if os.name == "nt": + # os.kill(pid, 0) is NOT a POSIX signal-0 liveness probe on Windows: signal 0 + # is CTRL_C_EVENT, so CPython routes it through GenerateConsoleCtrlEvent (a real + # Ctrl+C to that console group) rather than a harmless check. Probe via OpenProcess. + try: + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + # PROCESS_QUERY_LIMITED_INFORMATION: minimal right, granted across integrity levels. + handle = kernel32.OpenProcess(0x1000, False, pid) + if handle: + kernel32.CloseHandle(handle) + return True + # ERROR_ACCESS_DENIED means the process exists but we may not query it. + return ctypes.get_last_error() == 5 + except Exception: + return False + try: + os.kill(pid, 0) + return True + except OSError: + return False + except Exception: + return False + + +def _swap_lock_is_stale(path: Path) -> bool: + """Stale when the recorded owner is provably dead: a crashed installer is reclaimed + at once, not after the long cutoff, so `/load`, training, export, and repair are not + wedged for hours after a crash. A live but slow pip install keeps its lock (its PID + is alive), so breaking it and racing two swaps on the same staging dirs stays + impossible. Only a lock whose PID can't be read (mid-write or corrupt) falls back to + the age cutoff, so the create-before-metadata-write window is never mistaken for dead.""" + try: + age = time.time() - path.stat().st_mtime + except OSError: + return False + data = _read_swap_lock(path) or {} + pid = data.get("pid") + if not isinstance(pid, int) or pid <= 0: + return age > _SWAP_LOCK_STALE_SECS + return not _pid_alive(pid) + + +class SidecarSwapInProgress(RuntimeError): + """A worker start lost the race to a .venv_t5_latest install/repair; retryable.""" + + +def _read_swap_lock(path: Path) -> dict | None: + try: + data = json.loads(path.read_text(encoding = "utf-8")) + return data if isinstance(data, dict) else {} + except FileNotFoundError: + return None + except OSError: + return {} + except Exception: + return {} + + +def try_begin_sidecar_swap(kind: str = "install") -> bool: + """Reserve the sidecar swap window; False when one is already reserved + (in this process or, via the lock file, in any worker subprocess). + *kind* is "install" (consented route) or "repair" (lazy venv repair).""" + global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind + with _sidecar_swap_lock: + if _sidecar_swap_active: + return False + token = f"{os.getpid()}-{time.time_ns()}" + path = _swap_lock_path() + try: + path.parent.mkdir(parents = True, exist_ok = True) + except OSError: + pass + for attempt in range(2): + try: + fd = os.open(str(path), os.O_CREAT | os.O_EXCL | os.O_WRONLY) + break + except FileExistsError: + if attempt or not _swap_lock_is_stale(path): + return False + try: + path.unlink() + except OSError: + return False + except OSError: + # Lock file not creatable (odd filesystem): fall back to the process-local reservation. + fd = None + break + if fd is not None: + try: + with os.fdopen(fd, "w") as f: + f.write( + json.dumps( + {"pid": os.getpid(), "at": time.time(), "token": token, "kind": kind} + ) + ) + except OSError: + pass + _sidecar_swap_active = True + _sidecar_swap_token = token + _sidecar_swap_kind = kind + return True + + +def end_sidecar_swap() -> None: + """Release the reservation taken by :func:`try_begin_sidecar_swap`.""" + global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind + with _sidecar_swap_lock: + if _sidecar_swap_active: + # Only the file WE wrote is removed: if this reservation was declared + # stale and superseded, unlinking blindly would drop the new owner's + # live lock and unguard its in-flight swap. + path = _swap_lock_path() + data = _read_swap_lock(path) + if data is not None and data.get("token", _sidecar_swap_token) == _sidecar_swap_token: + try: + path.unlink() + except OSError: + pass + _sidecar_swap_active = False + _sidecar_swap_token = None + _sidecar_swap_kind = None + + +def sidecar_swap_in_progress() -> bool: + """True while a .venv_t5_latest install or repair holds the reservation, + in this process or any other Studio process (lock file).""" + return sidecar_swap_kind() is not None + + +def sidecar_swap_kind() -> str | None: + """The active reservation's kind ("install" / "repair"), or None when idle. + Lets guards that rely on the install route's own abort-on-active-worker + checks keep refusing for repairs, which have no such checks.""" + with _sidecar_swap_lock: + if _sidecar_swap_active: + return _sidecar_swap_kind or "install" + path = _swap_lock_path() + try: + if not path.is_file() or _swap_lock_is_stale(path): + return None + except OSError: + return None + data = _read_swap_lock(path) or {} + kind = data.get("kind") + return kind if kind in ("install", "repair") else "install" + + +def _stage_and_swap_latest_venv( + version: str, + packages: tuple[str, ...], + before_swap = None, +) -> bool: + """Stage-and-swap: build the new sidecar next to the live one and swap only + once complete, so a failed install or marker write never destroys a + previously working .venv_t5_latest or its pin. Shared by the consented + install and the lazy repair path. *before_swap* (optional callable) runs + after the staging build succeeds and immediately before the live dir is + replaced, so callers can tear down workers only when the swap is certain; + if it raises, the previous sidecar is left untouched.""" + staging = _VENV_T5_LATEST_DIR + ".staging" + retired = _VENV_T5_LATEST_DIR + ".old" + shutil.rmtree(staging, ignore_errors = True) + try: + if not _ensure_venv_dir(staging, packages, f"transformers {version} (latest)"): + # No exception, so the except cleanup below never runs; drop the partial dir. + shutil.rmtree(staging, ignore_errors = True) + return False + (Path(staging) / _LATEST_PIN_MARKER).write_text( + json.dumps({"version": version, "packages": list(packages)}), encoding = "utf-8" + ) + if before_swap is not None: + before_swap() + shutil.rmtree(retired, ignore_errors = True) + if os.path.isdir(_VENV_T5_LATEST_DIR): + os.rename(_VENV_T5_LATEST_DIR, retired) + try: + os.rename(staging, _VENV_T5_LATEST_DIR) + except OSError: + # Restore the previous sidecar if the final swap fails. + if not os.path.isdir(_VENV_T5_LATEST_DIR) and os.path.isdir(retired): + os.rename(retired, _VENV_T5_LATEST_DIR) + raise + except Exception as exc: + logger.error("Could not provision transformers %s into .venv_t5_latest: %s", version, exc) + shutil.rmtree(staging, ignore_errors = True) + return False + shutil.rmtree(retired, ignore_errors = True) + # CONFIG_MAPPING_NAMES may have changed: drop the cached key set. + _config_mapping_cache.pop("latest", None) + logger.info("Provisioned .venv_t5_latest with transformers %s", version) + return True + + +def _workers_active_for_repair() -> bool: + """Best-effort: any parent-visible chat/training/export worker alive. Never + raises; unavailable backends (worker subprocess, early startup) count idle.""" + try: + from core.training import get_training_backend + if get_training_backend().is_training_active(): + return True + except Exception: + pass + try: + from core.export import get_export_backend + + _export = get_export_backend() + if _export.is_export_active(): + return True + _alive = getattr(_export, "is_worker_alive", None) + if callable(_alive) and _alive(): + return True + except Exception: + pass + try: + from core.inference import get_inference_backend + + backend = get_inference_backend() + if getattr(backend, "active_model_name", None): + return True + # An in-flight load counts too: its worker spawns moments later. + if getattr(backend, "loading_models", None): + return True + _alive = getattr(backend, "is_worker_alive", None) + if callable(_alive) and _alive(): + return True + except Exception: + pass + return False + + +def _ensure_venv_t5_latest_exists() -> bool: + """Ensure .venv_t5_latest/ holds its pinned transformers version. + + Never installs without a pin: an unprovisioned sidecar (no marker) returns False so + routing and probing behave exactly as before the feature existed. With a pin present + it repairs a broken dir the same way the fixed sidecars do. + """ + pin = _latest_pin_data() + if pin is None: + return False + version = pin["version"] + packages = tuple(pin["packages"]) + if _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages): + return True + if _env_offline(): + logger.warning( + ".venv_t5_latest (transformers %s) is incomplete and offline mode is set; " + "cannot repair it.", + version, + ) + return False + # Repairs are a parent-process action: a worker child's backend singletons are + # empty, so it cannot see live siblings that may still lazy-import from the + # sidecar. Fail activation in the child instead; the parent's routing + # self-heal (guarded below) performs the actual repair. + try: + import multiprocessing as _mp + if _mp.parent_process() is not None: + logger.warning( + ".venv_t5_latest is incomplete; repairs run in the parent process. " + "Retry after the parent repairs the sidecar." + ) + return False + except Exception: + pass + # Same stage-and-swap as the install, under the same reservation so training/export starts + # (which check sidecar_swap_in_progress) wait out a lazy repair; a failed repair keeps the pin. + if not try_begin_sidecar_swap(kind = "repair"): + logger.warning( + "Cannot repair .venv_t5_latest: another sidecar install or repair is in progress." + ) + return False + try: + # Worker check UNDER the reservation (the install route quiesces workers; + # a repair has none): worker starts set their active markers BEFORE + # rechecking the reservation, so either this check sees them and aborts, + # or their recheck sees this reservation and aborts -- no interleaving + # lets a worker spawn against a mid-swap sidecar. + if _workers_active_for_repair(): + logger.warning( + "Cannot repair .venv_t5_latest: active chat/training/export workers " + "may be importing from it. Retry when they are idle." + ) + return False + return _stage_and_swap_latest_venv(version, packages) + finally: + end_sidecar_swap() + + +def ensure_latest_transformers_venv( + version: str, + extra_packages: tuple[str, ...] = (), + before_swap = None, +) -> bool: + """Provision .venv_t5_latest/ pinned to *version* (user-consented install path). + + Reuses the same --target/--no-deps installer as the fixed sidecars, then writes the pin + marker (version + full package set) so the venv persists across restarts and + :func:`latest_venv_pinned_version` / routing pick it up automatically. + *extra_packages* carries dep-compat shadows (see utils.transformers_latest). + Returns True on success. + """ + if not _is_valid_version_string(version): + logger.error("Refusing to install invalid transformers version %r", version) + return False + if _env_offline(): + logger.warning( + "Cannot install transformers %s: HF/transformers offline mode is set.", version + ) + return False + packages = _venv_t5_latest_packages(version, extra_packages) + pin = _latest_pin_data() + if ( + pin is not None + and pin["version"] == version + and tuple(pin["packages"]) == packages + and _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages) + ): + return True + return _stage_and_swap_latest_venv(version, packages, before_swap = before_swap) + + # --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- # Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize # Qwen3.5 / Gemma-4 / Llama. @@ -1819,7 +2510,7 @@ def _activate_venv(venv_dir: str, label: str) -> None: def _deactivate_5x() -> None: """Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport.""" - for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR): + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_T5_LATEST_DIR): while d in sys.path: sys.path.remove(d) logger.info("Removed venv_t5 dirs from sys.path") @@ -1853,14 +2544,25 @@ def ensure_transformers_version(model_name: str) -> None: if _is_lora_adapter_dir(Path(model_name)): resolved = _resolve_base_model(model_name) else: - resolved = model_name + # A remote adapter's tier is its BASE model's (see activation above). + resolved = _remote_lora_base(model_name) or model_name tier = get_transformers_tier(resolved) if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): # Gate on a real local config.json: a checkpoint carries config the base may not # surface, but path names alone must not upgrade a plain adapter. tier = _higher_tier(tier, get_transformers_tier(model_name)) - if tier == "510": + if tier == "latest": + pinned = latest_venv_pinned_version() + if pinned is None: + raise RuntimeError( + f"Cannot activate the latest-transformers sidecar: " + f"no pin marker at {_VENV_T5_LATEST_DIR}" + ) + target_version = pinned + venv_dir = _VENV_T5_LATEST_DIR + ensure_fn = _ensure_venv_t5_latest_exists + elif tier == "510": target_version = TRANSFORMERS_510_VERSION venv_dir = _VENV_T5_510_DIR ensure_fn = _ensure_venv_t5_510_exists diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 6f68917224..ba56ce7525 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -16,6 +16,7 @@ import { type ChatSearch, } from "@/features/chat"; import { RemoteCodeConsentDialog } from "@/features/security"; +import { TransformersUpgradeDialog } from "@/features/transformers-upgrade"; import { useTrainingUnloadGuard } from "@/features/training"; import { useExportRuntimeLifecycle } from "@/features/export"; import { hasAuthToken } from "@/features/auth"; @@ -230,6 +231,7 @@ function RootLayout() { {!isAuthFlowRoute && } + {hideNavbar ? (
}> diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 9ed11deccc..12c1c3b385 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1454,6 +1454,11 @@ async function autoLoadSmallestModel(): Promise<{ blockedByTrustRemoteCode = true; return false; } + // Never install packages from a background load; explicit loads raise the upgrade dialog. + if (validation.requires_transformers_upgrade) { + hadNonTrustFailure = true; + return false; + } return true; } diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index a659b7f83e..4646f67cb4 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -4,6 +4,10 @@ import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "@/lib/toast"; import { confirmRemoteCodeIfNeeded } from "@/features/security"; +import { + confirmTransformersUpgradeIfNeeded, + useTransformersUpgradeDialogStore, +} from "@/features/transformers-upgrade"; import { consumeNativePathToken } from "@/features/native-intents/api"; import { notifyNative, @@ -245,6 +249,10 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`; } +function getTransformersUpgradeRequiredMessage(modelName: string): string { + return `${modelName} was not loaded because it needs a newer transformers release that was not installed. Load it again to install it.`; +} + /** * Reconcile the chat runtime store against `/api/inference/status`: refresh the * models/loras catalogs and either re-pin the active checkpoint or clear the @@ -626,6 +634,30 @@ export function useChatModelRuntime() { is_lora: isLora, gguf_variant: ggufVariant ?? null, }); + // Upgrade consent runs before the security dialogs; Accept installs and the load continues. + if (validation.requires_transformers_upgrade) { + const upgraded = await confirmTransformersUpgradeIfNeeded({ + modelName: modelId, + upgrade: validation.transformers_upgrade, + // No installable release: custom-code models may fall back to the trust_remote_code gate below. + trustRemoteCodeFallback: validation.requires_trust_remote_code, + }); + // The install unloads the previous model before the swap (even when + // the swap then fails), so any exit after this point must roll back. + // False for the custom-code fallback, which resolves without installing. + if ( + useTransformersUpgradeDialogStore + .getState() + .consumeServerUnloadedChat() + && currentCheckpoint + ) { + previousWasUnloaded = true; + } + if (!upgraded) { + throw new Error(getTransformersUpgradeRequiredMessage(displayName)); + } + } + if (abortCtrl.signal.aborted) throw new Error("Cancelled"); // Open the consent dialog when the model needs custom-code consent or has a // flagged unsafe file. Fires even when trustRemoteCode is preset on, since the // worker requires a matching fingerprint that only the dialog produces. diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index a50af46e85..47a0720dac 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -67,6 +67,10 @@ import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge import { NewProjectDialog } from "./components/new-project-dialog"; import { useChatProjects } from "./hooks/use-chat-projects"; import { confirmRemoteCodeIfNeeded } from "@/features/security"; +import { + confirmTransformersUpgradeIfNeeded, + useTransformersUpgradeDialogStore, +} from "@/features/transformers-upgrade"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, @@ -929,6 +933,10 @@ export function SharedComposer({ return parts[parts.length - 1] || id; } + // Set when an accepted transformers install unloaded the active model + // server-side; a later failure must then clear the stale checkpoint. + let upgradeUnloadedActive = false; + // Helper: load a model and update store checkpoint async function ensureModelLoaded( sel: CompareModelSelection, @@ -955,6 +963,31 @@ export function SharedComposer({ trust_remote_code: loadTrustRemoteCode, chat_template_override: effectiveChatTemplateOverride, }); + // Upgrade dialog first (mirrors the primary load path). + if (validation.requires_transformers_upgrade) { + const upgraded = await confirmTransformersUpgradeIfNeeded({ + modelName: sel.id, + upgrade: validation.transformers_upgrade, + // No installable release: custom-code models may fall back to the trust_remote_code gate below. + trustRemoteCodeFallback: validation.requires_trust_remote_code, + }); + // The install unloads the active model before the swap (even when the + // swap then fails); if a later gate cancels or the load fails, the UI + // must stop pointing at that unloaded model. + if ( + useTransformersUpgradeDialogStore + .getState() + .consumeServerUnloadedChat() + && currentStore.params.checkpoint + ) { + upgradeUnloadedActive = true; + } + if (!upgraded) { + throw new Error( + `${modelDisplayName(sel.id)} needs a newer transformers release to load.`, + ); + } + } if ( validation.requires_trust_remote_code || validation.requires_security_review @@ -990,6 +1023,7 @@ export function SharedComposer({ tensor_parallel: currentStore.tensorParallel, }); saveSpeculativeType(specSettings.speculativeType); + upgradeUnloadedActive = false; const store = useChatRuntimeStore.getState(); store.setCheckpoint( resp.model, @@ -1097,6 +1131,11 @@ export function SharedComposer({ toast.success("Compare complete", { id: toastId, duration: 2000 }); } catch (err) { compareStepSucceededRef.current = false; + // The install already unloaded the previously active model; drop the + // checkpoint so the UI does not keep pointing at an unloaded model. + if (upgradeUnloadedActive) { + useChatRuntimeStore.getState().clearCheckpoint(); + } toast.error("Compare failed", { id: toastId, description: err instanceof Error ? err.message : "Unknown error", diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 954e88e86b..1a9a967263 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -1,6 +1,8 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import type { TransformersUpgradeInfo } from "@/features/transformers-upgrade"; + export interface BackendModelDetails { id: string; name?: string | null; @@ -78,6 +80,10 @@ export interface ValidateModelResponse { requires_security_review?: boolean; /** Native context length from the local GGUF header; null until downloaded. */ context_length?: number | null; + /** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */ + requires_transformers_upgrade?: boolean; + /** Set only when requires_transformers_upgrade. */ + transformers_upgrade?: TransformersUpgradeInfo | null; } export interface GgufVariantDetail { diff --git a/studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts b/studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts new file mode 100644 index 0000000000..2df9f712ff --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +interface InstallLatestTransformersResponse { + success: boolean; + version: string; + message: string; + /** The server unloaded the active chat model before the swap (set even on a + * structured failure, so callers can restore their model state). */ + model_unloaded?: boolean; + /** On a version-mismatch failure: the release that superseded the requested + * one, so Retry can use it. */ + latest_version?: string | null; +} + +/** Consented install of the latest transformers into the sidecar; synchronous, can take minutes. */ +export async function installLatestTransformers( + version: string, +): Promise { + const response = await authFetch("/api/inference/install-latest-transformers", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ version }), + }); + if (!response.ok) { + throw new Error(await readFastApiError(response)); + } + return (await response.json()) as InstallLatestTransformersResponse; +} diff --git a/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx b/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx new file mode 100644 index 0000000000..98c590e5e4 --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Spinner } from "@/components/ui/spinner"; +import { cn } from "@/lib/utils"; +import { PackageIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store"; + +function modelDisplayName(modelName: string | null): string { + if (!modelName) return "This model"; + return modelName.split("/").pop() || modelName; +} + +/** Root-mounted consent dialog for models needing a newer transformers; + * Install runs the sidecar install and resumes the paused load on success. */ +export function TransformersUpgradeDialog() { + const open = useTransformersUpgradeDialogStore((s) => s.open); + const modelName = useTransformersUpgradeDialogStore((s) => s.modelName); + const upgrade = useTransformersUpgradeDialogStore((s) => s.upgrade); + const phase = useTransformersUpgradeDialogStore((s) => s.phase); + const errorMessage = useTransformersUpgradeDialogStore((s) => s.errorMessage); + const trustRemoteCodeFallback = useTransformersUpgradeDialogStore( + (s) => s.trustRemoteCodeFallback, + ); + const install = useTransformersUpgradeDialogStore((s) => s.install); + const resolve = useTransformersUpgradeDialogStore((s) => s.resolve); + + const displayName = modelDisplayName(modelName); + const modelType = upgrade?.model_type ?? "unknown"; + const version = upgrade?.pypi_version ?? null; + // Only released PyPI versions are installable; dev (main) builds are never offered. + const installable = Boolean(upgrade?.supported_in_pypi && version); + const devOnly = !installable && Boolean(upgrade?.supported_in_main); + const installing = phase === "installing"; + + return ( + { + // Escape/overlay dismiss must not abandon an in-flight install. + if (!next && !installing) resolve(false); + }} + > + + +
+
+ +
+
+
+ New model architecture + + + {displayName} + {" "} + uses the{" "} + {modelType}{" "} + architecture, which your installed transformers does not + support yet.{" "} + {installable ? ( + <> + Install transformers{" "} + + {version} + {" "} + from PyPI to load it. The install runs once and can take + a minute; loading continues automatically afterwards. + + ) : devOnly ? ( + <> + Even the latest transformers release on PyPI does not + support it yet: the architecture is only available on the + transformers development branch (main), and Studio does + not install development builds. Support arrives with the + next transformers release on PyPI. + + ) : ( + <> + No released transformers version supports it yet, so it + cannot be loaded. + + )} + {!installable && trustRemoteCodeFallback ? ( + <> + {" "} + This model also ships its own modeling code; you can + continue and review enabling that custom code instead. + + ) : null} + +
+ + {phase === "error" && errorMessage ? ( +

+ {errorMessage} +

+ ) : null} + + {installing ? ( +

+ + Installing transformers {version}... This can take a minute. +

+ ) : null} +
+
+ + + + Cancel + {installable ? ( + <> + {phase === "error" && trustRemoteCodeFallback ? ( + // Install failed but the model ships custom code: offer the + // caller's trust_remote_code gate instead of forcing a retry. + resolve(true)} + > + Continue with custom code + + ) : null} + { + // Keep the dialog open; the store closes it on success. + event.preventDefault(); + void install(); + }} + > + {installing ? ( + <> + + Installing... + + ) : phase === "error" ? ( + "Retry install" + ) : ( + `Install transformers ${version}` + )} + + + ) : trustRemoteCodeFallback ? ( + // No installable release but the model ships custom code: continue + // into the caller's trust_remote_code gate as the last resort. + resolve(true)}> + Continue with custom code + + ) : null} + + + + ); +} diff --git a/studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts b/studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts new file mode 100644 index 0000000000..7d79d08d9c --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store"; +import type { TransformersUpgradeInfo } from "../types"; + +interface ConfirmArgs { + modelName: string; + /** validate's transformers_upgrade payload; null/undefined skips the dialog. */ + upgrade: TransformersUpgradeInfo | null | undefined; + /** When no release is installable, offer continuing into the caller's custom-code gate. */ + trustRemoteCodeFallback?: boolean; +} + +/** Pause a load needing a newer transformers on the consent dialog and run the install. + * Resolves true when the load can continue; false on cancel or not-installable with no fallback. */ +export async function confirmTransformersUpgradeIfNeeded({ + modelName, + upgrade, + trustRemoteCodeFallback, +}: ConfirmArgs): Promise { + if (!upgrade) return true; + return useTransformersUpgradeDialogStore + .getState() + .requestConsent(modelName, upgrade, { + trustRemoteCodeFallback: Boolean(trustRemoteCodeFallback), + }); +} diff --git a/studio/frontend/src/features/transformers-upgrade/index.ts b/studio/frontend/src/features/transformers-upgrade/index.ts new file mode 100644 index 0000000000..4318fd7830 --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/index.ts @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export { TransformersUpgradeDialog } from "./components/transformers-upgrade-dialog"; +export { confirmTransformersUpgradeIfNeeded } from "./hooks/use-transformers-upgrade-consent"; +export { installLatestTransformers } from "./api/transformers-upgrade-api"; +export { useTransformersUpgradeDialogStore } from "./stores/transformers-upgrade-dialog-store"; +export type { TransformersUpgradeInfo } from "./types"; diff --git a/studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts b/studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts new file mode 100644 index 0000000000..9e307fb1a1 --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { create } from "zustand"; +import { installLatestTransformers } from "../api/transformers-upgrade-api"; +import type { TransformersUpgradeInfo, TransformersUpgradePhase } from "../types"; + +type Resolver = (installed: boolean) => void; + +// One in-flight consent; a new request resolves any prior pending one as declined. +let pendingResolver: Resolver | null = null; + +interface TransformersUpgradeDialogStore { + open: boolean; + modelName: string | null; + upgrade: TransformersUpgradeInfo | null; + phase: TransformersUpgradePhase; + errorMessage: string | null; + /** Model ships custom code; without a PyPI install the load may fall back to trust_remote_code. */ + trustRemoteCodeFallback: boolean; + /** True once this consent's install completed. The install unloads the previous + * model before swapping, so the caller must treat it as already unloaded; the + * custom-code fallback resolves true without installing and leaves it loaded. */ + installRan: boolean; + /** True when the server unloaded the active chat model during this consent, + * including a swap that failed AFTER the unload: callers must then treat + * their previous model as gone and roll back on any later cancel. */ + serverUnloadedChat: boolean; + /** Read-and-clear serverUnloadedChat: each waiter consumes the signal once, + * so a superseding consent can neither erase it before the old waiter reads + * it nor leak it into an unrelated later load. */ + consumeServerUnloadedChat: () => boolean; + /** Open the dialog for a paused load; resolves true on install success or custom-code fallback. */ + requestConsent: ( + modelName: string, + upgrade: TransformersUpgradeInfo, + options?: { trustRemoteCodeFallback?: boolean }, + ) => Promise; + /** Accept/Retry: run the install; on success resolve(true) and close. */ + install: () => Promise; + resolve: (installed: boolean) => void; +} + +export const useTransformersUpgradeDialogStore = + create()((set, get) => ({ + open: false, + modelName: null, + upgrade: null, + phase: "consent", + errorMessage: null, + trustRemoteCodeFallback: false, + installRan: false, + serverUnloadedChat: false, + requestConsent: (modelName, upgrade, options) => + new Promise((resolve) => { + pendingResolver?.(false); + pendingResolver = resolve; + set({ + open: true, + modelName, + upgrade, + phase: "consent", + errorMessage: null, + trustRemoteCodeFallback: Boolean(options?.trustRemoteCodeFallback), + installRan: false, + }); + }), + consumeServerUnloadedChat: () => { + const value = get().serverUnloadedChat; + if (value) set({ serverUnloadedChat: false }); + return value; + }, + install: async () => { + const { upgrade, phase } = get(); + const version = upgrade?.pypi_version; + if (!version || phase === "installing") return; + const requestResolver = pendingResolver; + set({ phase: "installing", errorMessage: null }); + let result: Awaited>; + try { + result = await installLatestTransformers(version); + // Latch the server-side unload IMMEDIATELY, before any resolver-identity + // guard: even a superseded consent's install may have unloaded the chat + // model, and the signal must survive for whichever load consumes it next. + if (result.model_unloaded) { + set({ serverUnloadedChat: true }); + } + } catch (error) { + // Ignore the failure if a newer request superseded this consent. + if (pendingResolver === requestResolver) { + set({ + phase: "error", + errorMessage: + error instanceof Error && error.message + ? error.message + : "Failed to install transformers.", + }); + } + return; + } + if (pendingResolver === requestResolver) { + if (result.success) { + // serverUnloadedChat was latched above (and is never reset here): a + // retry after a failed-after-unload attempt reports false because the + // model is already gone, and a superseded install may have set it too. + set({ installRan: true }); + get().resolve(true); + return; + } + // Structured failure: the swap failed but may have already unloaded the + // chat model; record that so a later cancel still rolls the caller back. + // A version mismatch also carries the superseding release, so Retry + // re-requests a version that can actually succeed. + const { upgrade } = get(); + set({ + phase: "error", + errorMessage: result.message || "Failed to install transformers.", + serverUnloadedChat: + get().serverUnloadedChat || Boolean(result.model_unloaded), + ...(result.latest_version && upgrade + ? { upgrade: { ...upgrade, pypi_version: result.latest_version } } + : {}), + }); + } + }, + resolve: (installed) => { + const resolver = pendingResolver; + pendingResolver = null; + set({ + open: false, + modelName: null, + upgrade: null, + phase: "consent", + errorMessage: null, + trustRemoteCodeFallback: false, + }); + resolver?.(installed); + }, + })); diff --git a/studio/frontend/src/features/transformers-upgrade/types.ts b/studio/frontend/src/features/transformers-upgrade/types.ts new file mode 100644 index 0000000000..cf6bbe3d2f --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/types.ts @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** Wire shape of `transformers_upgrade` from /api/inference/validate. */ +export interface TransformersUpgradeInfo { + /** config.json model_type unknown to installed transformers. */ + model_type: string; + /** Latest transformers release on PyPI at check time. */ + pypi_version?: string | null; + /** Latest PyPI release ships this model_type (installable after consent). */ + supported_in_pypi?: boolean; + /** Only transformers main ships it (dev-only; not installable). */ + supported_in_main?: boolean; +} + +export type TransformersUpgradePhase = "consent" | "installing" | "error";