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 <danielhanchen@gmail.com>
This commit is contained in:
oobabooga
GitHub
parent
59bda2e1f7
commit
8b3c37246c
6 changed files with 1009 additions and 171 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <hf-id-or-path> 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 <hf-id-or-path> "
|
||||
"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"])
|
||||
|
|
|
|||
|
|
@ -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}).")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue