From 8b3c37246c38579bc9525f28066d919e30880b8f Mon Sep 17 00:00:00 2001 From: oobabooga Date: Wed, 22 Jul 2026 06:36:24 -0300 Subject: [PATCH] Unsloth start improvements: download progress, server reuse, and safe model switching (#7313) * Improve unsloth start runtime lifecycle * Remove speculative Gemma prompt override * Polish model download progress output * Refine unsloth start status output * Clarify unsloth readiness banner * Clarify model reuse and switching output * Queue model switches behind active inference * Tighten unsloth start model switching * Reduce model switch bookkeeping * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Studio re-exec compatibility * Recheck sidecar reservation after inference drain * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pass start marker through child environment * Fix key redaction, switch-waiter ordering, and stop/messaging gaps for PR #7313 - Redact minted sk-unsloth keys from the startup-failure log tail: the early key marker lands in the server log before the model load finishes, so a load-phase crash printed a live key to the terminal - Deregister a finished switch waiter before releasing the swap gate so a swap on another event loop cannot count it as still queued and unload the model the finished request is about to generate against - Warn on same-repo quant switches: an explicit variant replaces the resident weights for every attached session, but the repo ids match so no switch warning was printed - Note the agent exit code when it is nonzero so the server keep-alive message does not read as a successful session - Use taskkill /T in unsloth studio stop so llama-server children stop too * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments in start, studio, and inference changes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/routes/inference.py | 136 +++--- .../backend/tests/test_openai_auto_switch.py | 131 ++++-- unsloth_cli/commands/start.py | 438 ++++++++++++++++-- unsloth_cli/commands/studio.py | 49 +- unsloth_cli/tests/test_start.py | 387 +++++++++++++++- .../tests/test_studio_run_parallel_flag.py | 39 +- 6 files changed, 1009 insertions(+), 171 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d3e588bb0b..41e1fc5589 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3406,9 +3406,8 @@ async def _acquire_swap_gate() -> None: await asyncio.sleep(0.02) -# Counts in-flight auto-switch requests per (target, variant). The busy guard -# subtracts same-target waiters so concurrent requests for one model load once -# instead of each 409-ing the other. +# Counts auto-switch requests queued to load each (target, variant). They are not +# generating, so the drain wait below excludes them from the active inference count. _auto_switch_waiters: dict[tuple[str, str], int] = {} _auto_switch_waiters_guard = threading.Lock() @@ -3426,35 +3425,31 @@ def _note_switch_waiter(key: tuple[str, str], delta: int) -> None: _auto_switch_waiters.pop(key, None) -def _same_target_waiters(key: tuple[str, str]) -> int: +def _switch_waiter_count() -> int: with _auto_switch_waiters_guard: - return _auto_switch_waiters.get(key, 0) + return sum(max(0, count) for count in _auto_switch_waiters.values()) -# A second waiter map keyed by the raw requested model, registered before the -# (slow) resolve. The middleware counts a concurrent same-model request as -# in-flight before it resolves and joins _auto_switch_waiters, so without this -# the first request would see it as an unrelated request and 409. -_auto_switch_request_waiters: dict[str, int] = {} -_auto_switch_request_waiters_guard = threading.Lock() +async def _wait_for_model_switch_idle(*, current_request_counted: bool) -> None: + """Wait until a model replacement cannot interrupt active inference. - -def _request_waiter_key(requested_model: str) -> str: - return requested_model.strip().lower() - - -def _note_request_waiter(key: str, delta: int) -> None: - with _auto_switch_request_waiters_guard: - n = _auto_switch_request_waiters.get(key, 0) + delta - if n > 0: - _auto_switch_request_waiters[key] = n - else: - _auto_switch_request_waiters.pop(key, None) - - -def _same_request_waiters(key: str) -> int: - with _auto_switch_request_waiters_guard: - return _auto_switch_request_waiters.get(key, 0) + The caller holds ``inference_lifecycle_gate``, which prevents new inference + from starting while existing requests drain. Auto-switch requests that have + resolved their targets are scheduler waiters, not active generations, so + exclude them to avoid a queue deadlock. + """ + from core.inference.llama_keepwarm import other_inference_request_count + while True: + queued_switches = _switch_waiter_count() + if current_request_counted and queued_switches > 0: + queued_switches -= 1 + active_others = other_inference_request_count( + current_request_counted = current_request_counted, + include_pending = False, + ) + if active_others <= queued_switches: + return + await asyncio.sleep(0.02) def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: @@ -3582,7 +3577,6 @@ async def _maybe_auto_switch_model( from core.inference.local_model_resolver import resolve_local_gguf from core.inference.llama_keepwarm import ( get_last_unloaded_model, - other_inference_request_count, inference_lifecycle_gate, ) @@ -3603,12 +3597,7 @@ async def _maybe_auto_switch_model( if not auto_switch_on and get_auto_unload_idle_seconds() <= 0: return - # Register by the raw requested model before resolving (which can be slow): - # the middleware already counts a concurrent same-model request as in-flight, - # so the busy guard must know it shares this target even while it resolves. - request_key = _request_waiter_key(requested_model) - _note_request_waiter(request_key, 1) - try: + async def _resolve_and_switch() -> None: # Off the loop: a cold-cache rebuild walks several model dirs + HF caches. # With auto-switch off (or an omitted-model reload-only request), skip the # resolve so only the reload-stash path runs and no name is ever matched. @@ -3706,6 +3695,7 @@ async def _maybe_auto_switch_model( ) key = _switch_key(override_id, variant) _note_switch_waiter(key, 1) + waiter_noted = True try: async with _auto_switch_lock(): # The asyncio lock is per loop; add a process-wide gate so a swap on @@ -3718,31 +3708,6 @@ async def _maybe_auto_switch_model( if _already_serving(): _record_serving_alias() return - # Single slot: refuse a cross-model swap while another inference - # request is active rather than killing its response. Requests - # heading to this same target (by resolved id or raw name) are - # excluded, so concurrent requests for one model load once. A - # pending request is still in the middleware, not generating, so - # it is not counted here. - same_others = max( - _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 - ) - others = other_inference_request_count( - current_request_counted = True, include_pending = False - ) - # Not gated on the GGUF being loaded: _load_model_impl also - # tears down an active Unsloth backend before loading a GGUF, - # so refuse whenever any other inference request is in flight. - if others > same_others: - raise HTTPException( - status_code = 409, - detail = openai_error_body( - "Cannot switch models while another inference request is in progress.", - status = 409, - code = "model_switch_busy", - param = "model", - ), - ) # Apply this model's saved launch flags so the swap honors the config. override = get_model_override(override_id) load_kwargs = {"model_path": target_id, "gguf_variant": variant} @@ -3757,16 +3722,22 @@ async def _maybe_auto_switch_model( LoadRequest(**load_kwargs), fastapi_request, current_subject, + current_request_counted = True, ) # Advertise the repo id (not the concrete load path) as the loaded # model's public id and override key for /v1/models and idle stash. get_llama_cpp_backend()._openai_advertised_id = override_id finally: + # Deregister before releasing the gate: otherwise a swap on another + # loop counts this finished request as queued and unloads its model. + _note_switch_waiter(key, -1) + waiter_noted = False _auto_switch_process_lock.release() finally: - _note_switch_waiter(key, -1) - finally: - _note_request_waiter(request_key, -1) + if waiter_noted: + _note_switch_waiter(key, -1) + + await _resolve_and_switch() async def _auto_switch_from_request_body(request: Request, current_subject: str): @@ -4186,6 +4157,15 @@ def _maybe_unsupported_message(msg: str) -> str: return msg +def _raise_if_sidecar_swap_in_progress() -> None: + from utils.transformers_version import sidecar_swap_in_progress + if sidecar_swap_in_progress(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + + @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, @@ -4206,24 +4186,23 @@ async def load_model( # 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 + _raise_if_sidecar_swap_in_progress() # 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. async with inference_lifecycle_gate(): - if sidecar_swap_in_progress(): - raise _swap_409 + _raise_if_sidecar_swap_in_progress() return await _load_model_impl(request, fastapi_request, current_subject) -async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): +async def _load_model_impl( + request: LoadRequest, + fastapi_request: Request, + current_subject: str, + *, + current_request_counted: bool = False, +): from core.inference.llama_cpp import LlamaServerNotFoundError # A new load starts here; arm the progress throttle so this load's first @@ -4557,6 +4536,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre ), ) + # Keep the resident model alive until every active generation finishes; + # the caller's lifecycle gate blocks new starts. + await _wait_for_model_switch_idle(current_request_counted = current_request_counted) + # A sidecar install can reserve the gate while inference drains, after the + # route-level checks above, so recheck before replacing either backend. + _raise_if_sidecar_swap_in_progress() + # Unload any active Unsloth model only after every hub conflict check. if unsloth_backend.active_model_name: logger.info( @@ -4767,6 +4753,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Unload any active GGUF model first llama_backend = get_llama_cpp_backend() + await _wait_for_model_switch_idle(current_request_counted = current_request_counted) + _raise_if_sidecar_swap_in_progress() if llama_backend.is_loaded: logger.info("Unloading GGUF model before loading Unsloth model") llama_backend.unload_model() @@ -7096,7 +7084,7 @@ async def openai_chat_completions( if payload.provider_id or payload.provider_type: # External provider: this request won't touch the local GGUF, so drop it # from the keep-warm count or its in-flight stream would falsely block a - # concurrent local auto-switch with model_switch_busy. + # concurrent local model switch from proceeding. from core.inference.llama_keepwarm import untrack_current_request untrack_current_request(request.scope) diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 1ee9ef36d3..9361db66bb 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -68,7 +68,13 @@ class _LoadRecorder: request, fastapi_request, current_subject = None, + *, + current_request_counted = False, ): + # Mirror the production load boundary before recording any replacement. + await inference_route._wait_for_model_switch_idle( + current_request_counted = current_request_counted + ) self.calls.append(request) if self.fail: from fastapi import HTTPException @@ -94,7 +100,6 @@ def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder): # gate that auto-switch already owns, so it calls the impl directly). monkeypatch.setattr(inference_route, "_load_model_impl", recorder) monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) - monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) def _run_hook(model = "some/model"): @@ -1205,10 +1210,9 @@ def test_middleware_ignores_non_post(monkeypatch): # ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ── -def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): - # A cross-model swap must 409 (not kill) while another inference request is in - # flight; the requesting call itself is excluded from the count. - from fastapi import HTTPException +def test_auto_switch_waits_for_another_inference_to_finish(monkeypatch): + # A cross-model swap queues while another request is generating, then loads + # after that request drains. The requesting call itself is excluded. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M") @@ -1222,10 +1226,18 @@ def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): ) monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one monkeypatch.setattr(kw, "_pending", 0) - with pytest.raises(HTTPException) as exc: - _run_hook("org/B-GGUF:Q8_0") - assert exc.value.status_code == 409 - assert rec.calls == [] + + async def _drive(): + task = asyncio.create_task( + inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester") + ) + await asyncio.sleep(0.05) + assert rec.calls == [] + kw._note_end() # the other generation finishes; this request remains counted + await asyncio.wait_for(task, timeout = 1) + + asyncio.run(_drive()) + assert len(rec.calls) == 1 def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch): @@ -1411,13 +1423,12 @@ def test_concurrent_same_target_requests_load_once(monkeypatch): monkeypatch.setattr(kw, "_pending", 0) inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) _run_hook("org/B-GGUF:Q8_0") - assert len(rec.calls) == 1 # loads once, no 409 + assert len(rec.calls) == 1 -def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch): - # A concurrent request heading to a different target still blocks the swap: the - # same-target exclusion must not swallow a genuinely conflicting request. - from fastapi import HTTPException +def test_queued_different_target_does_not_deadlock_current_swap(monkeypatch): + # A concurrent request already queued for another target is not generating, + # so it must not prevent the current serialized swap from proceeding. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF") @@ -1432,10 +1443,8 @@ def test_swap_still_refused_when_other_request_targets_different_model(monkeypat monkeypatch.setattr(kw, "_inflight", 2) monkeypatch.setattr(kw, "_pending", 0) inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) - with pytest.raises(HTTPException) as exc: - _run_hook("org/B-GGUF:Q8_0") - assert exc.value.status_code == 409 - assert rec.calls == [] + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): @@ -1481,6 +1490,37 @@ def test_load_route_holds_lifecycle_gate(monkeypatch): assert "_load_model_impl" in src +def test_model_replacements_recheck_sidecar_swap_before_either_backend_is_unloaded(): + # Both replacement directions drain active inference, then recheck whether a + # sidecar install reserved the lifecycle gate during that wait. Exact-model + # reuse exits earlier, so an already-loaded model never waits on unrelated inference. + import inspect + + src = inspect.getsource(inference_route._load_model_impl) + gguf_wait = src.index("await _wait_for_model_switch_idle", src.index("if config.is_gguf:")) + gguf_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", gguf_wait) + unload_unsloth = src.index("unsloth_backend.unload_model", gguf_wait) + standard_wait = src.index("await _wait_for_model_switch_idle", gguf_wait + 1) + standard_sidecar_check = src.index("_raise_if_sidecar_swap_in_progress()", standard_wait) + unload_gguf = src.index("llama_backend.unload_model()", standard_wait) + already_loaded = src.index('status = "already_loaded"') + + assert already_loaded < gguf_wait < gguf_sidecar_check < unload_unsloth + assert standard_wait < standard_sidecar_check < unload_gguf + + +def test_switch_waiter_deregisters_before_swap_gate_release(): + # A waiter left registered after the swap gate is released would let a swap on + # another event loop count the finished request as still queued, pass the drain + # early, and unload the model that request is about to generate against. + import inspect + + src = inspect.getsource(inference_route._maybe_auto_switch_model) + deregister = src.index("_note_switch_waiter(key, -1)") + release = src.index("_auto_switch_process_lock.release()") + assert deregister < release + + def _anthropic_payload(max_tokens = None): from models.inference import AnthropicMessagesRequest, AnthropicMessage return AnthropicMessagesRequest( @@ -1519,9 +1559,9 @@ def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): # ── review round 6: concurrency ordering, external untrack, unload gate, ids ── -def test_pending_same_target_request_does_not_force_409(monkeypatch): +def test_pending_same_target_request_does_not_block_swap(monkeypatch): # A second same-target request blocked in the middleware (pending, not yet - # generating) must not make the first request 409: pending is excluded. + # generating) must not block the first request: pending is excluded. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF") @@ -1536,13 +1576,13 @@ def test_pending_same_target_request_does_not_force_409(monkeypatch): monkeypatch.setattr(kw, "_inflight", 1) # just the caller monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware _run_hook("org/B-GGUF:Q8_0") - assert len(rec.calls) == 1 # loads once, no 409 + assert len(rec.calls) == 1 -def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch): +def test_swap_waits_until_concurrent_request_finishes_resolving(monkeypatch): # The real middleware counts a concurrent same-model request as in-flight - # before it resolves and registers a target waiter. The raw-request waiter, - # registered before resolve, must still exclude it so the first request loads. + # before it resolves and registers a target waiter. Treat it as active until + # its target is known, then recognize it as another queued switch request. from core.inference import llama_keepwarm as kw backend = _FakeBackend("org/A-GGUF") @@ -1556,10 +1596,20 @@ def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypat ) monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin monkeypatch.setattr(kw, "_pending", 0) - # The twin has only registered its raw requested model (not yet a target waiter). - inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) - _run_hook("org/B-GGUF:Q8_0") - assert len(rec.calls) == 1 # loads once, no 409 + # The twin is still resolving, so it is counted in-flight but has not joined + # the concrete target queue yet. + + async def _drive(): + task = asyncio.create_task( + inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester") + ) + await asyncio.sleep(0.05) + assert rec.calls == [] + inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + await asyncio.wait_for(task, timeout = 1) + + asyncio.run(_drive()) + assert len(rec.calls) == 1 def test_external_untrack_decrements_inflight_and_is_idempotent(): @@ -1595,11 +1645,9 @@ def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): assert not backend.is_loaded # torn down despite the active request -def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): +def test_auto_switch_waits_when_unsloth_stream_active(monkeypatch): # The GGUF slot is empty but an Unsloth model is streaming (counted in-flight). - # _load_model_impl would unload it, so auto-switch must 409, not only when a - # GGUF is loaded. - from fastapi import HTTPException + # The replacement waits for it just as it does for a GGUF generation. from core.inference import llama_keepwarm as kw backend = _FakeBackend(None) # no GGUF loaded @@ -1613,10 +1661,18 @@ def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): ) monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request monkeypatch.setattr(kw, "_pending", 0) - with pytest.raises(HTTPException) as exc: - _run_hook("org/B-GGUF:Q8_0") - assert exc.value.status_code == 409 - assert rec.calls == [] # the active Unsloth model is not torn down + + async def _drive(): + task = asyncio.create_task( + inference_route._maybe_auto_switch_model("org/B-GGUF:Q8_0", object(), "tester") + ) + await asyncio.sleep(0.05) + assert rec.calls == [] + kw._note_end() + await asyncio.wait_for(task, timeout = 1) + + asyncio.run(_drive()) + assert len(rec.calls) == 1 def test_public_model_id_prefers_advertised_over_path(): @@ -3097,6 +3153,8 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch): request, fastapi_request, current_subject = None, + *, + current_request_counted = False, ): with slock: state["cur"] += 1 @@ -3114,7 +3172,6 @@ def test_auto_switch_serializes_across_event_loops(monkeypatch): monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load) monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) - monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) barrier = threading.Barrier(2) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 3d73df65be..317d1f4f3e 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -14,12 +14,13 @@ import signal import subprocess import sys import tempfile +import threading import time import urllib.error import urllib.request from pathlib import Path from typing import NamedTuple, NoReturn, Optional -from urllib.parse import urlparse +from urllib.parse import urlencode, urlparse import click import typer @@ -105,8 +106,8 @@ _SERVE_OPTION = typer.Option( True, "--serve/--no-serve", help = ( - "If no Unsloth server is running, auto-start one for --model and stop it when the " - "agent exits. --no-serve keeps the old behavior of erroring out." + "If no Unsloth server is running, auto-start one for --model and keep it available " + "after the agent exits. --no-serve keeps the old behavior of erroring out." ), ) # Model-load knobs mirrored from `unsloth run`; only used when --model triggers a @@ -326,6 +327,13 @@ def _split_repo_variant(model: str) -> tuple: return repo, variant +def _display_model_spec(model: str, variant: Optional[str]) -> str: + """Return a user-facing model name that includes the selected GGUF variant.""" + repo, inline_variant = _split_repo_variant(model) + selected_variant = variant or inline_variant + return f"{repo}:{selected_variant}" if selected_variant else model + + def _fail(message: str) -> NoReturn: typer.echo(message, err = True) raise typer.Exit(code = 1) @@ -373,11 +381,265 @@ def _http_json( # A server that WE auto-started (never one we merely found). Kept at module scope so -# _run's finally and the atexit backstop can tear it down without threading a handle +# failure paths and the atexit backstop can tear it down without threading a handle # through all six agent commands. Only one agent runs per process, so one slot is enough. _auto_served_server: Optional[subprocess.Popen] = None # Model download + load can be slow; give the auto-started server room before giving up. _SERVER_START_TIMEOUT_S = 900 +_DOWNLOAD_POLL_INTERVAL_S = 1.0 +_START_API_KEY_PREFIX = "UNSLOTH_START_API_KEY: " +_START_API_KEY_MARKER_ENV = "_UNSLOTH_START_API_KEY_MARKER" + + +def _format_download_bytes(value: int) -> str: + value = max(0, int(value)) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if value < 1024 or unit == "TiB": + precision = 0 if unit in ("B", "KiB") else 1 + return f"{value:.{precision}f} {unit}" + value /= 1024 + return "0 B" + + +def _format_download_eta(seconds: float) -> str: + seconds = max(0, int(seconds)) + if seconds < 60: + return f"{seconds}s" + minutes, seconds = divmod(seconds, 60) + if minutes < 60: + return f"{minutes}m {seconds:02d}s" + hours, minutes = divmod(minutes, 60) + return f"{hours}h {minutes:02d}m" + + +class _DownloadProgressDisplay: + """Render download progress without making redirected output noisy.""" + + def __init__(self) -> None: + self._samples: list[tuple[float, int]] = [] + self._shown = False + self._last_bucket = -1 + self._last_line_length = 0 + self._last_expected = 0 + self._interactive = bool(getattr(sys.stdout, "isatty", lambda: False)()) + + def update(self, progress: dict) -> None: + downloaded = max(0, int(progress.get("downloaded_bytes") or 0)) + completed = max(0, int(progress.get("completed_bytes") or 0)) + expected = max(0, int(progress.get("expected_bytes") or 0)) + self._last_expected = max(self._last_expected, expected) + fraction = float(progress.get("progress") or 0) + if downloaded <= 0: + return + # A fully cached snapshot can report 99% with no incomplete bytes; that is + # not a transfer, so don't show it as a download. + if completed >= downloaded > 0: + return + + now = time.monotonic() + if self._samples and downloaded < self._samples[-1][1]: + self._samples.clear() + self._samples.append((now, downloaded)) + cutoff = now - 15.0 + while len(self._samples) > 2 and self._samples[0][0] < cutoff: + self._samples.pop(0) + + rate = 0.0 + if len(self._samples) >= 2: + elapsed = self._samples[-1][0] - self._samples[0][0] + delta = self._samples[-1][1] - self._samples[0][1] + if elapsed >= 1.0 and delta > 0: + rate = delta / elapsed + + if expected > 0: + # The endpoint caps at 99% while bytes remain in an incomplete file; trust it. + fraction = min(1.0, max(0.0, fraction)) + percent = min(100, max(0, int(fraction * 100))) + filled = min(24, int(fraction * 24)) + bar = "=" * filled + ">" + "." * max(0, 23 - filled) if filled < 24 else "=" * 24 + line = ( + f"Downloading model [{bar}] {percent:3d}% " + f"{_format_download_bytes(downloaded)} / {_format_download_bytes(expected)}" + ) + bucket = percent // 10 + if rate > 0: + line += f" | {_format_download_bytes(rate)}/s" + if downloaded < expected: + line += f" | ETA {_format_download_eta((expected - downloaded) / rate)}" + else: + line = f"Downloading model: {_format_download_bytes(downloaded)}" + bucket = downloaded // (1024**3) + if rate > 0: + line += f" | {_format_download_bytes(rate)}/s" + + if self._interactive: + padding = " " * max(0, self._last_line_length - len(line)) + typer.echo(f"\r{line}{padding}", nl = False) + sys.stdout.flush() + self._last_line_length = len(line) + elif not self._shown or bucket > self._last_bucket: + typer.echo(line) + self._last_bucket = bucket + self._shown = True + + def close(self) -> None: + if self._interactive and self._shown: + typer.echo() + self._last_line_length = 0 + + def complete(self) -> None: + """Finish a displayed transfer after the model load confirms success.""" + if not self._shown: + return + downloaded = self._samples[-1][1] if self._samples else 0 + expected = max(downloaded, getattr(self, "_last_expected", 0)) + self.update( + { + "downloaded_bytes": expected, + "expected_bytes": expected, + "progress": 1.0, + } + ) + + +def _normalized_variant(value: object) -> str: + return re.sub(r"[^a-z0-9]", "", str(value or "").lower()) + + +class _ModelDownloadProgress: + """Best-effort polling of the model download endpoints.""" + + def __init__(self, base: str, key: str, model: str, variant: Optional[str]) -> None: + self._base = base + self._key = key + self._model = model + self._variant = variant or "" + self._expected_bytes = 0 + self._display = _DownloadProgressDisplay() + self._configured = False + self._disabled = not _is_hub_model_id(model) + self._progress_prefix = "/api/hub" + + def _configure(self) -> None: + self._configured = True + if self._disabled: + return + # GGUF repos need the selected quant's size; the repo endpoint totals every + # quant. Resolve the variant first, otherwise show bytes only. + if self._variant or "gguf" in self._model.lower(): + try: + params = urlencode({"repo_id": self._model}) + try: + info = _http_json( + "GET", + f"{self._base}/api/hub/gguf-variants?{params}", + self._key, + timeout = 10, + ) + except urllib.error.HTTPError as exc: + if exc.code != 404: + raise + self._progress_prefix = "/api/models" + info = _http_json( + "GET", + f"{self._base}/api/models/gguf-variants?{params}", + self._key, + timeout = 10, + ) + self._variant = self._variant or str(info.get("default_variant") or "") + wanted = _normalized_variant(self._variant) + for item in info.get("variants") or []: + quant = _normalized_variant(item.get("quant")) + filename = _normalized_variant(item.get("filename")) + if wanted and (wanted == quant or wanted in filename): + self._expected_bytes = int( + item.get("download_size_bytes") or item.get("size_bytes") or 0 + ) + break + except Exception: + # Older servers lack this endpoint; byte progress is still useful. + pass + + def poll(self) -> None: + if not self._configured: + self._configure() + if self._disabled: + return + try: + if self._variant or "gguf" in self._model.lower(): + params = urlencode( + { + "repo_id": self._model, + "variant": self._variant, + "expected_bytes": self._expected_bytes, + } + ) + url = f"{self._base}{self._progress_prefix}/gguf-download-progress?{params}" + else: + url = ( + f"{self._base}{self._progress_prefix}/download-progress?" + f"{urlencode({'repo_id': self._model})}" + ) + try: + reading = _http_json("GET", url, self._key, timeout = 10) + except urllib.error.HTTPError as exc: + if exc.code != 404 or self._progress_prefix == "/api/models": + raise + self._progress_prefix = "/api/models" + self.poll() + return + self._display.update(reading) + except Exception: + # Progress is best-effort; never fail the load over a polling error. + self._disabled = True + + def close(self) -> None: + self._display.close() + + def complete(self) -> None: + self._display.complete() + + +def _load_model_with_progress( + base: str, key: str, model: str, load: LoadOptions, payload: dict +) -> dict: + """Run the blocking load request while polling its download progress.""" + result: list[tuple[bool, object]] = [] + done = threading.Event() + + def _load() -> None: + try: + value = _http_json( + "POST", + f"{base}/api/inference/load", + key, + payload, + timeout = 3600, + error = "Model load failed", + ) + result.append((True, value)) + except BaseException as exc: + result.append((False, exc)) + finally: + done.set() + + threading.Thread(target = _load, name = "unsloth-model-load", daemon = True).start() + progress = _ModelDownloadProgress(base, key, model, load.gguf_variant) + loading_announced = False + try: + while not done.wait(_DOWNLOAD_POLL_INTERVAL_S): + if not loading_announced: + typer.echo(f"Loading model: {_display_model_spec(model, load.gguf_variant)}") + loading_announced = True + progress.poll() + ok, value = result[0] + if not ok: + assert isinstance(value, BaseException) + raise value + progress.complete() + return value if isinstance(value, dict) else {} + finally: + progress.close() def _studio_healthy(base: str, timeout: float = 3.0) -> bool: @@ -396,6 +658,11 @@ def _log_tail(path: Path, lines: int = 20) -> str: return "(no server log)" +def _redacted_log_tail(path: Path, lines: int = 20) -> str: + """Tail with minted keys removed; only for tails shown on the terminal.""" + return re.sub(r"sk-unsloth-\S+", "sk-unsloth-[redacted]", _log_tail(path, lines)) + + def _shutdown_server(server: Optional[subprocess.Popen]) -> None: # Idempotent teardown of a server WE started, plus its own children (llama-server, # cloudflared). A no-op once the process is already gone. @@ -438,6 +705,14 @@ def _shutdown_auto_served() -> None: _shutdown_server(server) +def _keep_auto_served() -> bool: + """Release ownership so a successfully started server survives this CLI.""" + global _auto_served_server + server, _auto_served_server = _auto_served_server, None + atexit.unregister(_shutdown_auto_served) + return server is not None and server.poll() is None + + def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess.Popen: """Spawn `unsloth run` for `model`, wait until it is fully ready, and return it.""" global _auto_served_server @@ -467,9 +742,8 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess command += ["--tensor-parallel"] log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log" - typer.echo( - f"No Unsloth server at {base}. Starting one for {model} (loading the model can take a while)…" - ) + typer.echo("Starting Unsloth server") + typer.echo(f"Model: {_display_model_spec(model, load.gguf_variant)}") typer.echo(f"Server log: {log_path}") # 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and # the tempdir is world-traversable. Unlink first so a stale looser-mode file (pid @@ -477,8 +751,17 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess log_path.unlink(missing_ok = True) log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") # Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the - # server; we tear it down explicitly when the agent exits. - kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL} + # server. It survives a successful agent session; torn down on startup/launch failure. + child_env = os.environ.copy() + # Pass the marker via env so an older launcher ignores it instead of treating an + # unknown CLI flag as a llama-server arg; new launchers preserve it across re-exec. + child_env[_START_API_KEY_MARKER_ENV] = "1" + kwargs: dict = { + "stdout": log, + "stderr": subprocess.STDOUT, + "stdin": subprocess.DEVNULL, + "env": child_env, + } if os.name == "nt": kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP else: @@ -491,17 +774,45 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess atexit.register(_shutdown_auto_served) deadline = time.monotonic() + _SERVER_START_TIMEOUT_S - while time.monotonic() < deadline: - if server.poll() is not None: - tail = _log_tail(log_path) - _shutdown_auto_served() - _fail(f"The Unsloth server stopped before it was ready. Last log lines:\n{tail}") - # `unsloth run` prints the minted key only after the server is up AND the model is - # loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses). - if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400): - typer.echo(f"Unsloth server ready at {base}.") - return server - time.sleep(2.0) + progress: Optional[_ModelDownloadProgress] = None + early_key_seen = False + try: + while time.monotonic() < deadline: + if server.poll() is not None: + # The early key marker lands here before load finishes; redact it. + tail = _redacted_log_tail(log_path) + _shutdown_auto_served() + _fail(f"The Unsloth server stopped before it was ready. Last log lines:\n{tail}") + tail = _log_tail(log_path, lines = 400) + if progress is None: + marker = re.search( + rf"^{re.escape(_START_API_KEY_PREFIX)}(sk-unsloth-[^\s]+)$", + tail, + flags = re.MULTILINE, + ) + if marker: + early_key_seen = True + progress = _ModelDownloadProgress( + base, + marker.group(1), + model, + load.gguf_variant, + ) + if progress is not None: + progress.poll() + # New children emit an early key marker, so wait for the final model banner; + # older children only print the key after load, so fall back to that. + ready_signal = "Model loaded:" in tail if early_key_seen else "sk-unsloth-" in tail + if _studio_healthy(base) and ready_signal: + if progress is not None: + progress.complete() + progress.close() + progress = None + return server + time.sleep(2.0) + finally: + if progress is not None: + progress.close() _shutdown_auto_served() _fail( f"The Unsloth server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}." @@ -796,6 +1107,7 @@ def _resolve_model( load: LoadOptions = LoadOptions(), ) -> dict: models = _loaded_models(base, key) + load_requested = False # Only casefold-match ids against a loopback Unsloth, where _is_hub_model_id's # local existence probe can actually reject a server-side path; see the note there. allow_casefold = is_loopback_url(base) @@ -825,11 +1137,30 @@ def _resolve_model( ) ) if requested and match is None: - typer.echo( - f"Loading {requested} - please wait…" - if load_has_overrides - else f"Loading {requested} on the Unsloth server (this can take a while)…" - ) + load_requested = True + active = next((m for m in models if m.get("loaded") is not False), None) + active_id = active.get("id") if active else None + if active_id and not _model_id_matches( + active_id, + requested, + allow_casefold = allow_casefold, + ): + typer.echo(f"Switching the Unsloth server from {active_id} to {requested}.") + typer.echo("This unloads the current model for every attached session.") + elif active_id and load.gguf_variant: + # Same repo id but an explicit quant still replaces the resident + # weights; /v1/models has no variant, so ask the status endpoint. + try: + status = _http_json("GET", f"{base}/api/inference/status", key) + except Exception: + status = {} + resident = status.get("gguf_variant") if status.get("is_gguf") else None + if resident and _normalized_variant(resident) != _normalized_variant(load.gguf_variant): + typer.echo( + f"Switching the Unsloth server from {active_id}:{resident} " + f"to {requested}:{load.gguf_variant}." + ) + typer.echo("This unloads the current model for every attached session.") # Mirror `unsloth run`'s load knobs; keep the default payload as just # model_path so a bare `--model` load is unchanged. payload = {"model_path": requested} @@ -841,14 +1172,9 @@ def _resolve_model( payload["load_in_4bit"] = False if load.tensor_parallel: payload["tensor_parallel"] = True - loaded = _http_json( - "POST", - f"{base}/api/inference/load", - key, - payload, - timeout = 3600, - error = "Model load failed", - ) + loaded = _load_model_with_progress(base, key, requested, load, payload) + if loaded.get("status") == "already_loaded": + typer.echo(f"Reusing loaded model: {_display_model_spec(requested, load.gguf_variant)}") # Unsloth registers the model under a canonical id (resolved identifier, # casing) that /v1/models echoes but which may differ from the path we # passed; match on the id the load reports so we don't silently fall @@ -861,13 +1187,16 @@ def _resolve_model( ( m for m in models - if any( + if m.get("loaded") is not False + and any( _model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted ) ), None, ) if match is not None: + if requested and not load_requested: + typer.echo(f"Reusing loaded model: {_display_model_spec(requested, load.gguf_variant)}") return match if requested: # We asked Unsloth to load it and it didn't surface in /v1/models; don't @@ -881,7 +1210,13 @@ def _resolve_model( "No model is loaded in Unsloth. Load one from the model dropdown in " "the UI, or pass --model to load it from here." ) - return models[0] + resident = next((m for m in models if m.get("loaded") is not False), None) + if resident is None: + _fail( + "No model is currently resident in Unsloth. Pass --model " + "to reload one, or load it from the model dropdown in the UI." + ) + return resident def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: @@ -1356,7 +1691,7 @@ def _launch( env: dict, install_hint: str, unset_env: tuple = (), -) -> NoReturn: +) -> int: # Resolve well-known install dirs (e.g. ~/.local/bin) first, so an already-installed # agent not yet on PATH is found instead of prompting a needless reinstall. _augment_path_with_install_dirs() @@ -1382,7 +1717,7 @@ def _launch( finally: signal.signal(signal.SIGINT, previous) # Negative returncode means killed by signal N; shells expect 128+N. - raise typer.Exit(code = code if code >= 0 else 128 - code) + return code if code >= 0 else 128 - code def _connect( @@ -1434,16 +1769,35 @@ def _run( # --no-launch recipes stay intact. if launch and clear_screen: click.clear() - typer.echo(f"Unsloth {base} · model {entry['id']}") + typer.echo(f"Unsloth ready at {base} · model {entry['id']}") if not launch: env, wsl_env_bridge = _wsl_shim_env(command, env, unset_env) _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) + if _keep_auto_served(): + typer.echo(f"Unsloth Studio is still running at {base}.") + typer.echo("Stop it with: unsloth studio stop") return try: - _launch(command, env, install_hint = install_hint, unset_env = unset_env) - finally: - # Tear down a server we auto-started once the agent session ends (no-op otherwise). + code = _launch(command, env, install_hint = install_hint, unset_env = unset_env) + except BaseException: + # Startup succeeded but the agent failed to launch; tear the server down + # rather than orphan it. _shutdown_auto_served() + raise + auto_started = _auto_served_server is not None + kept = _keep_auto_served() + if auto_started and not kept: + typer.echo(f"The auto-started Unsloth server at {base} stopped during the session.") + raise typer.Exit(code = code) + if code: + # The server status below must not read as a successful agent session. + typer.echo(f"The agent exited with code {code}.") + if is_loopback_url(base): + typer.echo(f"Unsloth Studio is still running at {base}.") + typer.echo("Stop it with: unsloth studio stop") + else: + typer.echo(f"The remote Unsloth server is still running at {base}.") + raise typer.Exit(code = code) def _agents_config_root() -> Path: @@ -1893,7 +2247,7 @@ def codex( launch = launch, ) # This preflight runs after _connect may have auto-started a server but before _run - # installs its teardown finally, so tear the server down here if it rejects the model + # takes over its lifecycle, so tear the server down here if it rejects the model # (e.g. a transformers-backend model) rather than leaving it on the atexit backstop. try: _require_gguf_for_codex(base, key, entry["id"]) diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index f2f41fc583..e1924cce00 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -106,6 +106,13 @@ API_KEY_PBKDF2_SALT_KEY = "api_key_pbkdf2_salt" DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash" DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at" PBKDF2_ITERATIONS = 100_000 +_START_API_KEY_MARKER_ENV = "_UNSLOTH_START_API_KEY_MARKER" + + +def _consume_start_api_key_marker_env() -> bool: + """Consume the one-shot readiness marker passed across a Studio re-exec.""" + return os.environ.pop(_START_API_KEY_MARKER_ENV, None) == "1" + # __file__ is unsloth_cli/commands/studio.py -- two parents up is the package root # (either site-packages or the repo root for editable installs). @@ -1760,6 +1767,12 @@ def run( "decode speed, MoE usually don't." ), ), + start_api_key_marker: bool = typer.Option( + False, + "--start-api-key-marker", + hidden = True, + help = "Emit an early API key marker for the unsloth start parent process.", + ), password: str = typer.Option( "", "--password", @@ -1786,6 +1799,11 @@ def run( unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja unsloth studio run --model unsloth/Qwen3-27B-GGUF --gguf-variant Q8_0 --tensor-parallel """ + # A newer outer CLI can re-exec into an older Studio venv; pass this signal via + # env so an older child ignores it instead of treating it as a llama-server arg. + inherited_start_api_key_marker = _consume_start_api_key_marker_env() + start_api_key_marker = start_api_key_marker or inherited_start_api_key_marker + # Back-compat: --not-secure is a deprecated alias for --no-secure. secure = _resolve_secure(secure, not_secure) extra_llama_args: List[str] = list(ctx.args) if ctx.args else [] @@ -1991,15 +2009,21 @@ def run( if extra_llama_args: args.extend(extra_llama_args) - if sys.platform == "win32": - proc = subprocess.Popen(args) - try: - rc = proc.wait() - except KeyboardInterrupt: - rc = proc.wait() - raise typer.Exit(rc) - else: - os.execvp(str(studio_bin), args) + if start_api_key_marker: + os.environ[_START_API_KEY_MARKER_ENV] = "1" + try: + if sys.platform == "win32": + proc = subprocess.Popen(args) + try: + rc = proc.wait() + except KeyboardInterrupt: + rc = proc.wait() + raise typer.Exit(rc) + else: + os.execvp(str(studio_bin), args) + finally: + # execvp doesn't return on success; restore env after a Windows wait or a failed launch. + os.environ.pop(_START_API_KEY_MARKER_ENV, None) # ── 2. Start server (always suppress built-in banner) ───────────── run_mod = _load_run_module() @@ -2045,6 +2069,10 @@ def run( # 4. Create API key in-process. api_key = _create_api_key_inprocess(api_key_name) + if start_api_key_marker: + # `unsloth start` reads this key from a private 0600 log to authenticate + # download-progress polling; the normal `unsloth run` output is unchanged. + typer.echo(f"UNSLOTH_START_API_KEY: {api_key}") # 5. Load model via HTTP. if not silent: @@ -2236,7 +2264,8 @@ def stop(): # Send SIGTERM (graceful shutdown) or TerminateProcess on Windows try: if sys.platform == "win32": - subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True) + # /T also stops llama-server children, which otherwise keep GPU and port. + subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], check = True) else: os.kill(pid, _signal.SIGTERM) typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).") diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 1e03d390d1..7c070fa5f4 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -20,6 +20,7 @@ if str(_REPO_ROOT) not in sys.path: import pytest +import typer from typer.testing import CliRunner import unsloth_cli.commands.start as start @@ -639,8 +640,13 @@ def fake_studio(tmp_path, monkeypatch): if url.endswith("/api/auth/api-keys"): return {"key": "sk-unsloth-feedfacefeedface"} if url.endswith("/api/inference/load"): + already_loaded = state["models"][0]["id"] == payload["model_path"] state["models"] = [{"id": payload["model_path"], "context_length": 4096}] - return {} + return { + "status": "already_loaded" if already_loaded else "loaded", + "model": payload["model_path"], + "display_name": payload["model_path"], + } raise AssertionError(f"unexpected request: {method} {url}") monkeypatch.setattr(start, "find_studio_server", lambda: BASE) @@ -824,7 +830,7 @@ def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, t assert profile["model"] == MODEL["id"] -def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch): +def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch, capsys): calls = [] state = {"loaded": False} @@ -862,6 +868,8 @@ def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch): assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF" assert any(c[1].endswith("/api/inference/load") for c in calls) + output = capsys.readouterr().out + assert "please wait" not in output def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch): @@ -903,6 +911,35 @@ def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch): assert any(u.endswith("/api/inference/load") for _, u in calls) +def test_resolve_model_does_not_attach_if_catalog_stays_unloaded(monkeypatch): + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/Gemma-4-GGUF", + "loaded": False, + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + return {"status": "loaded", "model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + with pytest.raises(typer.Exit): + start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch): # The mirror case: a loaded entry (loaded == True) that case-matches attaches with # no /api/inference/load call. @@ -931,6 +968,25 @@ def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch assert not any(u.endswith("/api/inference/load") for _, u in calls) +def test_resolve_model_without_request_rejects_unloaded_catalog(monkeypatch): + monkeypatch.setattr( + start, + "_http_json", + lambda *a, **k: { + "data": [ + { + "id": "unsloth/Gemma-4-GGUF", + "loaded": False, + "context_length": 131072, + } + ] + }, + ) + + with pytest.raises(typer.Exit): + start._resolve_model(BASE, "sk-test", None) + + def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch): # Against a remote Unsloth the local existence probe cannot see server-side paths, # so a case-variant loaded id must NOT attach without a load: it could be a distinct @@ -1213,6 +1269,9 @@ def test_connect_model_flag_loads_on_server(fake_studio): assert loads == [ ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) ] + assert result.output.index( + f"Switching the Unsloth server from {MODEL['id']} to unsloth/Qwen3.5-35B-A3B.\n" + ) < result.output.index("This unloads the current model for every attached session.\n") _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") @@ -1303,6 +1362,7 @@ def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio): assert result.exit_code == 0, result.output loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] assert loads == [] + assert f"Reusing loaded model: {MODEL['id']}\n" in result.output _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) @@ -1324,6 +1384,7 @@ def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio): {"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"}, ) ] + assert f"Reusing loaded model: {MODEL['id']}:UD-Q4_K_XL\n" in result.output _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) @@ -1730,8 +1791,9 @@ def _reset_auto_served(): start._auto_served_server = None -def test_start_studio_server_builds_command_and_waits(monkeypatch): +def test_start_studio_server_builds_command_and_waits(monkeypatch, capsys): captured = {} + monkeypatch.setenv(start._START_API_KEY_MARKER_ENV, "parent") class FakePopen: def __init__(self, command, **kwargs): @@ -1761,13 +1823,200 @@ def test_start_studio_server_builds_command_and_waits(monkeypatch): assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL" assert cmd[cmd.index("--context-length") + 1] == "8192" assert "--tensor-parallel" in cmd + assert "--start-api-key-marker" not in cmd + assert captured["kwargs"]["env"][start._START_API_KEY_MARKER_ENV] == "1" + assert start.os.environ[start._START_API_KEY_MARKER_ENV] == "parent" assert cmd[cmd.index("-p") + 1] == "8888" assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd assert captured["kwargs"].get("start_new_session") is True # own process group assert server.pid == 4321 + output = capsys.readouterr().out + assert "Starting Unsloth server\n" in output + assert "Model: unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL\n" in output + assert "No Unsloth server at" not in output + assert "server ready" not in output -def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch): +def test_start_studio_server_polls_progress_from_early_key(monkeypatch): + class FakePopen: + pid = 4321 + + def poll(self): + return None + + tails = iter( + [ + "UNSLOTH_START_API_KEY: sk-unsloth-early\nLoading model...", + "UNSLOTH_START_API_KEY: sk-unsloth-early\nModel loaded: owner/model", + ] + ) + created = [] + + class FakeProgress: + def __init__(self, base, key, model, variant): + created.append((base, key, model, variant, "created")) + + def poll(self): + created.append("poll") + + def close(self): + created.append("close") + + def complete(self): + created.append("complete") + + monkeypatch.setattr(start.subprocess, "Popen", lambda *a, **k: FakePopen()) + monkeypatch.setattr(start, "_studio_healthy", lambda *a, **k: True) + monkeypatch.setattr(start, "_log_tail", lambda *a, **k: next(tails)) + monkeypatch.setattr(start, "_ModelDownloadProgress", FakeProgress) + monkeypatch.setattr(start.time, "sleep", lambda _s: None) + monkeypatch.setattr( + start.typer, + "echo", + lambda message = "", **_kwargs: created.append(("echo", message)), + ) + + server = start._start_studio_server( + BASE, + "owner/model-GGUF", + start.LoadOptions(gguf_variant = "Q4_K_M"), + ) + + assert server.pid == 4321 + assert (BASE, "sk-unsloth-early", "owner/model-GGUF", "Q4_K_M", "created") in created + assert created.count("poll") == 2 + assert created[-2:] == ["complete", "close"] + assert not any(isinstance(event, tuple) and "server ready" in event[-1] for event in created) + + +def test_load_model_with_progress_uses_selected_gguf_size(monkeypatch, capsys): + release = start.threading.Event() + calls = [] + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/api/inference/load"): + assert release.wait(timeout = 2) + return {"model": "owner/model-GGUF"} + if "/api/hub/gguf-variants?" in url: + return { + "default_variant": "Q8_0", + "variants": [ + { + "quant": "UD-Q4_K_XL", + "filename": "model-UD-Q4_K_XL.gguf", + "size_bytes": 4 * 1024**3, + "download_size_bytes": 4 * 1024**3, + } + ], + } + if "/api/hub/gguf-download-progress?" in url: + release.set() + return { + "downloaded_bytes": 2 * 1024**3, + "expected_bytes": 4 * 1024**3, + "progress": 0.5, + } + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + monkeypatch.setattr(start, "_DOWNLOAD_POLL_INTERVAL_S", 0.001) + result = start._load_model_with_progress( + BASE, + "sk-test", + "owner/model-GGUF", + start.LoadOptions(gguf_variant = "UD-Q4_K_XL"), + {"model_path": "owner/model-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + + assert result == {"model": "owner/model-GGUF"} + output = capsys.readouterr().out + assert "Downloading model" in output + assert "100%" in output + progress_url = next(url for method, url, _ in calls if "gguf-download-progress" in url) + assert "variant=UD-Q4_K_XL" in progress_url + assert f"expected_bytes={4 * 1024**3}" in progress_url + + +def test_download_progress_ignores_fully_cached_bytes(capsys): + display = start._DownloadProgressDisplay() + display.update( + { + "downloaded_bytes": 4 * 1024**3, + "completed_bytes": 4 * 1024**3, + "expected_bytes": 4 * 1024**3, + "progress": 0.99, + } + ) + display.close() + + assert capsys.readouterr().out == "" + + +def test_resolve_model_warns_on_same_repo_quant_switch(monkeypatch, capsys): + models = [{"id": "owner/model-GGUF", "loaded": True}] + + def http_json( + method, + url, + key, + payload = None, + timeout = 30, + error = None, + ): + assert url.endswith("/api/inference/status"), url + return {"is_gguf": True, "gguf_variant": "Q4_K_M"} + + monkeypatch.setattr(start, "_loaded_models", lambda base, key: models) + monkeypatch.setattr(start, "_http_json", http_json) + monkeypatch.setattr( + start, + "_load_model_with_progress", + lambda base, key, model, load, payload: {"status": "loaded", "model": "owner/model-GGUF"}, + ) + + start._resolve_model(BASE, "key", "owner/model-GGUF", start.LoadOptions(gguf_variant = "Q8_0")) + + out = capsys.readouterr().out + assert ( + "Switching the Unsloth server from owner/model-GGUF:Q4_K_M to owner/model-GGUF:Q8_0." in out + ) + assert "every attached session" in out + + +def test_resolve_model_same_quant_prints_no_switch_warning(monkeypatch, capsys): + models = [{"id": "owner/model-GGUF", "loaded": True}] + + monkeypatch.setattr(start, "_loaded_models", lambda base, key: models) + monkeypatch.setattr( + start, + "_http_json", + lambda *a, **k: {"is_gguf": True, "gguf_variant": "Q8_0"}, + ) + monkeypatch.setattr( + start, + "_load_model_with_progress", + lambda base, key, model, load, payload: { + "status": "already_loaded", + "model": "owner/model-GGUF", + }, + ) + + start._resolve_model(BASE, "key", "owner/model-GGUF", start.LoadOptions(gguf_variant = "Q8_0")) + + out = capsys.readouterr().out + assert "Switching" not in out + assert "Reusing loaded model: owner/model-GGUF:Q8_0" in out + + +def test_auto_serves_when_no_server_then_keeps_server(fake_studio, monkeypatch): monkeypatch.setattr(start, "find_studio_server", lambda: None) started = {} fake = SimpleNamespace(pid = 999, poll = lambda: None) @@ -1793,8 +2042,134 @@ def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch): assert started["model"] == "unsloth/Qwen3-1.7B-GGUF" assert started["load"].gguf_variant == "UD-Q4_K_XL" assert started["base"] == BASE - # Torn down after the agent session ended. - assert started.get("down") is fake + # A successful agent exit releases ownership and leaves the server available + # for another terminal. Explicit startup failures still use the cleanup path. + assert "down" not in started + assert start._auto_served_server is None + assert "is still running" in result.output + assert "unsloth studio stop" in result.output + + +def test_auto_served_agent_launch_failure_stops_server(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + stopped = [] + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(*_args): + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr(start, "_shutdown_server", stopped.append) + monkeypatch.setattr( + start, + "_launch", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("agent launch failed")), + ) + + result = CliRunner().invoke( + start.start_app, + ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"], + ) + + assert result.exit_code == 1 + assert stopped == [fake] + assert "is still running" not in result.output + + +def test_auto_served_server_exit_is_not_reported_as_running(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + fake = SimpleNamespace(pid = 999, poll = lambda: 1) + + def fake_start(*_args): + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr(start, "_launch", lambda *a, **k: 0) + + result = CliRunner().invoke( + start.start_app, + ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"], + ) + + assert result.exit_code == 0, result.output + assert "stopped during the session" in result.output + assert "is still running" not in result.output + + +def test_attached_server_prints_stop_hint_after_agent_exits(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: []) + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, env: SimpleNamespace(returncode = 0), + ) + + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert f"Unsloth ready at {BASE} · model {MODEL['id']}\n" in result.output + assert f"Unsloth Studio is still running at {BASE}." in result.output + assert "Stop it with: unsloth studio stop\n" in result.output + + +def test_no_launch_recipe_does_not_print_stop_hint(fake_studio): + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "is still running" not in result.output + + +def test_nonzero_agent_exit_notes_code_before_stop_hint(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda *a, **k: []) + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, env: SimpleNamespace(returncode = 3), + ) + + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 3 + assert "The agent exited with code 3." in result.output + assert f"Unsloth Studio is still running at {BASE}." in result.output + + +def test_redacted_log_tail_strips_minted_keys(tmp_path): + log = tmp_path / "server.log" + log.write_text( + "booting\nUNSLOTH_START_API_KEY: sk-unsloth-feedfacefeedface\nerror: load failed\n", + encoding = "utf-8", + ) + + tail = start._redacted_log_tail(log) + + assert "sk-unsloth-feedfacefeedface" not in tail + assert "sk-unsloth-[redacted]" in tail + assert "error: load failed" in tail + + +def test_startup_failure_output_redacts_minted_key(monkeypatch, tmp_path, capsys): + monkeypatch.setattr(start.tempfile, "gettempdir", lambda: str(tmp_path)) + fake = SimpleNamespace(pid = 4242, poll = lambda: 1) + + def fake_popen(command, **kwargs): + # The child prints the early key marker, then dies before it is ready. + kwargs["stdout"].write(b"UNSLOTH_START_API_KEY: sk-unsloth-secretsecret\nload failed\n") + kwargs["stdout"].flush() + return fake + + monkeypatch.setattr(start.subprocess, "Popen", fake_popen) + + with pytest.raises(start.typer.Exit): + start._start_studio_server(BASE, "owner/model-GGUF", start.LoadOptions()) + + err = capsys.readouterr().err + assert "stopped before it was ready" in err + assert "sk-unsloth-secretsecret" not in err + assert "sk-unsloth-[redacted]" in err def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch): diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py index 558b268a4d..74ea607753 100644 --- a/unsloth_cli/tests/test_studio_run_parallel_flag.py +++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py @@ -170,13 +170,24 @@ def _install_reexec_capture(monkeypatch, *, platform): monkeypatch.setattr(sys, "platform", platform) + def capture(kind, argv): + captured.append( + { + "kind": kind, + "argv": list(argv), + "start_api_key_marker": studio_mod.os.environ.get( + studio_mod._START_API_KEY_MARKER_ENV + ), + } + ) + def fake_execvp(file, argv): - captured.append({"kind": "execvp", "argv": list(argv)}) + capture("execvp", argv) raise _ExecCaptured(argv) class _FakePopen: def __init__(self, argv, *a, **kw): - captured.append({"kind": "popen", "argv": list(argv)}) + capture("popen", argv) self._argv = argv def wait(self): @@ -235,6 +246,30 @@ def test_reexec_forwards_parallel_all_aliases(monkeypatch, flag, value): ), f"{flag} {value} was dropped on re-exec; argv = {argv}" +@pytest.mark.parametrize("platform", ["linux", "darwin", "win32"]) +def test_reexec_hands_off_start_api_key_marker_out_of_band(monkeypatch, platform): + """A new child receives the marker while an old child sees no unknown flag.""" + result, captured = _invoke_run( + monkeypatch, + _BASE + ["--start-api-key-marker"], + platform = platform, + ) + assert len(captured) == 1, result.output + assert "--start-api-key-marker" not in captured[0]["argv"] + assert captured[0]["start_api_key_marker"] == "1" + + +def test_reexeced_child_consumes_start_api_key_marker_env(monkeypatch): + """A supported child consumes the handoff before starting descendants.""" + studio_mod = _load_run_command() + monkeypatch.setenv(studio_mod._START_API_KEY_MARKER_ENV, "1") + + inherited = studio_mod._consume_start_api_key_marker_env() + + assert inherited is True + assert studio_mod._START_API_KEY_MARKER_ENV not in studio_mod.os.environ + + @pytest.mark.parametrize("platform", ["linux", "darwin", "win32"]) def test_reexec_argv_is_consistent_across_platforms(monkeypatch, platform): """Linux/Darwin (execvp) and Windows (Popen) must build the same argv."""