Polish model download progress output

This commit is contained in:
oobabooga 2026-07-21 15:15:10 -03:00
commit 715727b64b
2 changed files with 17 additions and 7 deletions

View file

@ -739,9 +739,7 @@ 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(f"Starting Unsloth server for {model}")
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
@ -795,6 +793,8 @@ def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess
if _studio_healthy(base) and ready_signal:
if progress is not None:
progress.complete()
progress.close()
progress = None
typer.echo(f"Unsloth server ready at {base}.")
return server
time.sleep(2.0)

View file

@ -1755,7 +1755,7 @@ 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 = {}
class FakePopen:
@ -1791,6 +1791,9 @@ def test_start_studio_server_builds_command_and_waits(monkeypatch):
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 for unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL…" in output
assert "No Unsloth server at" not in output
def test_start_studio_server_polls_progress_from_early_key(monkeypatch):
@ -1826,6 +1829,11 @@ def test_start_studio_server_polls_progress_from_early_key(monkeypatch):
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,
@ -1834,15 +1842,17 @@ def test_start_studio_server_polls_progress_from_early_key(monkeypatch):
)
assert server.pid == 4321
assert created[0] == (
assert (
BASE,
"sk-unsloth-early",
"owner/model-GGUF",
"Q4_K_M",
"created",
)
) in created
assert created.count("poll") == 2
assert created[-2:] == ["complete", "close"]
ready = ("echo", f"Unsloth server ready at {BASE}.")
assert created[-3:] == ["complete", "close", ready]
assert created.index("close") < created.index(ready)
def test_load_model_with_progress_uses_selected_gguf_size(monkeypatch, capsys):