mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-16 11:18:43 +00:00
Studio: forward llama-server args from unsloth studio run , activate unsloth run , and allow passing model:quant to load models (#5271)
* Studio: forward unknown CLI args directly to llama-server `unsloth studio run --model X --top-k 20 --chat-template-file foo.jinja` now passes the unknown flags through to the llama-server subprocess. Adds a denylist for flags Studio manages (port, -m, -c, --api-key, -ngl, --flash-attn, --no-context-shift, --jinja, GPU-fit, model-identity, ...) that returns HTTP 400 on collision. HTTP callers can supply the same list via LoadRequest.llama_extra_args. * Studio: accept `--model org/repo:variant` shorthand in `unsloth studio run` Mirrors llama.cpp's `-hf <repo>:<quant>` and ollama's pull syntax so `unsloth studio run --model unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL` is equivalent to `--model unsloth/... --gguf-variant UD-Q4_K_XL`. Local paths and Windows drive letters are preserved verbatim. If both an embedded variant and an explicit `--gguf-variant` are given and they disagree, the command fails with a clear error. * Studio: register `unsloth run` as alias for `unsloth studio run` Top-level `unsloth run --model ...` is now equivalent to `unsloth studio run --model ...`. Same context_settings, so unknown flags continue to pass through to llama-server. * Studio: let users override soft-managed llama-server flags from CLI Trims the denylist to flags Studio fundamentally cannot share with the user (model identity, --host/--port/--path/--api-prefix, --api-key, --ssl-*, --webui, --models-*). Soft-managed flags -- -c/--ctx-size, --parallel, --flash-attn, --no-context-shift, --jinja, -ngl, -t/--threads, --fit* -- now pass through and override Studio's auto-set version via llama.cpp's last-wins CLI parsing. Lets users tune their run on the spot: unsloth run --model X -c 131072 --parallel 1 --threads 32 * Studio: accept `-hf` / `-hfr` / `--hf-repo` as aliases for `--model` Matches llama-server's `-hf <repo>:<quant>` spelling so users coming from llama.cpp can use the same flag. Typer claims the aliases before the pass-through validator runs, so the HTTP-API denylist on those flags is unaffected. unsloth run -hf unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL
This commit is contained in:
Roland Tannous
GitHub
parent
e35cbfb454
commit
dbea77e347
9 changed files with 656 additions and 5 deletions
|
|
@ -22,7 +22,7 @@ import sys
|
|||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Generator, Optional
|
||||
from typing import Generator, List, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
|
@ -1381,6 +1381,7 @@ class LlamaCppBackend:
|
|||
n_threads: Optional[int] = None,
|
||||
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
|
||||
n_parallel: int = 1,
|
||||
extra_args: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Start llama-server with a GGUF model.
|
||||
|
|
@ -1767,6 +1768,17 @@ class LlamaCppBackend:
|
|||
else:
|
||||
self._api_key = None
|
||||
|
||||
# User-supplied pass-through args go last so llama.cpp's
|
||||
# last-wins flag parsing lets the user override Studio's
|
||||
# auto-set tier-2 flags (e.g. --cache-type-k, --spec-type).
|
||||
# The route layer has already validated this list against
|
||||
# the managed-flag denylist via validate_extra_args().
|
||||
if extra_args:
|
||||
cmd.extend(str(a) for a in extra_args)
|
||||
logger.info(
|
||||
f"Appending user extra args to llama-server: {list(extra_args)}"
|
||||
)
|
||||
|
||||
_log_cmd = list(cmd)
|
||||
if "--api-key" in _log_cmd:
|
||||
_ki = _log_cmd.index("--api-key") + 1
|
||||
|
|
|
|||
120
studio/backend/core/inference/llama_server_args.py
Normal file
120
studio/backend/core/inference/llama_server_args.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Validator for user-supplied llama-server pass-through args.
|
||||
|
||||
Studio runs llama-server as a managed subprocess and lets callers pass
|
||||
extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP:
|
||||
``LoadRequest.llama_extra_args``). This module is the boundary that
|
||||
rejects only flags Studio fundamentally cannot share with the user --
|
||||
model identity, the auth key, and the network endpoint Studio's HTTP
|
||||
proxy targets. Anything else passes through.
|
||||
|
||||
User-supplied args are appended to ``cmd`` after Studio's auto-set
|
||||
flags, so llama.cpp's last-wins CLI parsing makes the user's value
|
||||
override the auto-set one. That covers tunable knobs the user might
|
||||
reasonably want to override -- ``-c``/``--ctx-size``,
|
||||
``-np``/``--parallel``, ``-fa``/``--flash-attn``,
|
||||
``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``,
|
||||
``--cache-type-k/v``, ``--chat-template-file/-kwargs``,
|
||||
``--spec-*``, ``--jinja``/``--no-jinja``,
|
||||
``--no-context-shift``/``--context-shift``, sampling params, etc.
|
||||
|
||||
Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, Optional
|
||||
|
||||
# Each group is the full set of aliases (short + long) for one
|
||||
# hard-denied flag, taken from the llama-server README. If llama.cpp
|
||||
# adds a new alias for an existing denied flag, extend the relevant
|
||||
# group.
|
||||
#
|
||||
# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl,
|
||||
# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*,
|
||||
# --chat-template-*, --spec-*) pass through and override Studio's
|
||||
# auto-set version via llama.cpp's last-wins CLI parsing.
|
||||
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
|
||||
# Model identity -- Studio resolves the model from LoadRequest and
|
||||
# passes -m / mmproj after downloading from HF if needed. A second
|
||||
# -m would point at a different model than the one Studio thinks
|
||||
# is loaded.
|
||||
frozenset({"-m", "--model"}),
|
||||
frozenset({"-mu", "--model-url"}),
|
||||
frozenset({"-dr", "--docker-repo"}),
|
||||
frozenset({"-hf", "-hfr", "--hf-repo"}),
|
||||
frozenset({"-hff", "--hf-file"}),
|
||||
frozenset({"-hfv", "-hfrv", "--hf-repo-v"}),
|
||||
frozenset({"-hffv", "--hf-file-v"}),
|
||||
frozenset({"-hft", "--hf-token"}),
|
||||
frozenset({"-mm", "--mmproj"}),
|
||||
frozenset({"-mmu", "--mmproj-url"}),
|
||||
# Networking -- Studio binds llama-server's port and reverse-proxies
|
||||
# HTTP traffic to it. Retargeting host/port/path/prefix would
|
||||
# orphan Studio's proxy and the UI would lose the server.
|
||||
frozenset({"--host"}),
|
||||
frozenset({"--port"}),
|
||||
frozenset({"--path"}),
|
||||
frozenset({"--api-prefix"}),
|
||||
frozenset({"--reuse-port"}),
|
||||
# Auth / TLS -- Studio terminates auth at its own layer; an
|
||||
# upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM
|
||||
# key, and TLS on llama-server would break the local proxy hop.
|
||||
frozenset({"--api-key"}),
|
||||
frozenset({"--api-key-file"}),
|
||||
frozenset({"--ssl-key-file"}),
|
||||
frozenset({"--ssl-cert-file"}),
|
||||
# Single-model server -- Studio runs one model per llama-server
|
||||
# process and serves its own UI. Enabling multi-model loading or
|
||||
# llama-server's built-in web UI changes the surface clients see.
|
||||
frozenset({"--webui", "--no-webui"}),
|
||||
frozenset({"--models-dir"}),
|
||||
frozenset({"--models-preset"}),
|
||||
frozenset({"--models-max"}),
|
||||
frozenset({"--models-autoload", "--no-models-autoload"}),
|
||||
)
|
||||
|
||||
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
|
||||
|
||||
|
||||
def _flag_name(token: str) -> Optional[str]:
|
||||
"""Return the flag name for a token, or None if it isn't a flag.
|
||||
|
||||
Peels ``--key=value`` to the bare ``--key``. Plain numeric values
|
||||
like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags;
|
||||
llama-server short-form flags always start with a letter.
|
||||
"""
|
||||
if not token.startswith("-") or token in {"-", "--"}:
|
||||
return None
|
||||
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
|
||||
return None
|
||||
return token.split("=", 1)[0]
|
||||
|
||||
|
||||
def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
|
||||
"""Validate user-supplied llama-server args.
|
||||
|
||||
Returns the args as a flat list ready to extend the llama-server
|
||||
command. Raises ``ValueError`` (with the offending flag in the
|
||||
message) the moment a token resolves to a Studio-managed flag.
|
||||
"""
|
||||
if not args:
|
||||
return []
|
||||
out: list[str] = []
|
||||
for raw in args:
|
||||
token = str(raw)
|
||||
flag = _flag_name(token)
|
||||
if flag is not None and flag in _DENYLIST:
|
||||
raise ValueError(
|
||||
f"llama-server flag '{flag}' is managed by Unsloth Studio "
|
||||
f"and cannot be passed as an extra arg"
|
||||
)
|
||||
out.append(token)
|
||||
return out
|
||||
|
||||
|
||||
def is_managed_flag(flag: str) -> bool:
|
||||
"""True if ``flag`` is a Studio-managed llama-server flag."""
|
||||
return flag in _DENYLIST
|
||||
|
|
@ -55,6 +55,16 @@ class LoadRequest(BaseModel):
|
|||
None,
|
||||
description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.",
|
||||
)
|
||||
llama_extra_args: Optional[List[str]] = Field(
|
||||
None,
|
||||
description = (
|
||||
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
|
||||
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
|
||||
"Studio-managed flags (model identity, port, context length, GPU placement, "
|
||||
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
|
||||
"non-GGUF models."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class UnloadRequest(BaseModel):
|
||||
|
|
|
|||
|
|
@ -119,6 +119,7 @@ try:
|
|||
_DEFAULT_T_MAX_PREDICT_MS,
|
||||
detect_reasoning_flags,
|
||||
)
|
||||
from core.inference.llama_server_args import validate_extra_args
|
||||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
from utils.models.model_config import load_model_defaults
|
||||
|
|
@ -140,6 +141,7 @@ except ImportError:
|
|||
_DEFAULT_T_MAX_PREDICT_MS,
|
||||
detect_reasoning_flags,
|
||||
)
|
||||
from core.inference.llama_server_args import validate_extra_args
|
||||
from utils.models import ModelConfig
|
||||
from utils.inference import load_inference_config
|
||||
from utils.models.model_config import load_model_defaults
|
||||
|
|
@ -431,6 +433,13 @@ async def load_model(
|
|||
native_grant_backed = False
|
||||
model_log_label = request.model_path
|
||||
try:
|
||||
# Validate user-supplied llama-server pass-through args up front
|
||||
# so a managed-flag collision returns 400 before any model work.
|
||||
try:
|
||||
extra_llama_args = validate_extra_args(request.llama_extra_args)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
|
||||
model_identifier, model_log_label, native_grant_backed = (
|
||||
_resolve_model_identifier_for_request(request, operation = "load-model")
|
||||
)
|
||||
|
|
@ -605,6 +614,7 @@ async def load_model(
|
|||
cache_type_kv = request.cache_type_kv,
|
||||
speculative_type = request.speculative_type,
|
||||
n_parallel = _n_parallel,
|
||||
extra_args = extra_llama_args,
|
||||
)
|
||||
else:
|
||||
# Local mode: llama-server loads via -m <path>
|
||||
|
|
@ -623,6 +633,7 @@ async def load_model(
|
|||
cache_type_kv = request.cache_type_kv,
|
||||
speculative_type = request.speculative_type,
|
||||
n_parallel = _n_parallel,
|
||||
extra_args = extra_llama_args,
|
||||
)
|
||||
|
||||
if not success:
|
||||
|
|
|
|||
189
studio/backend/tests/test_llama_server_args.py
Normal file
189
studio/backend/tests/test_llama_server_args.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for the llama-server pass-through args validator.
|
||||
|
||||
The validator is the security boundary between user-supplied CLI / HTTP
|
||||
input and the llama-server subprocess command. These tests pin the
|
||||
denylist behavior so the boundary doesn't quietly regress when new
|
||||
managed flags are added.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference.llama_server_args import (
|
||||
is_managed_flag,
|
||||
validate_extra_args,
|
||||
)
|
||||
|
||||
|
||||
# ── Pass-through (allowed) ───────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
# Sampling
|
||||
["--top-k", "20"],
|
||||
["--top-p", "0.9", "--min-p", "0.05"],
|
||||
["--seed", "-1"], # negative value, not a flag
|
||||
["--temp", "0.0"],
|
||||
["--repeat-penalty", "1.05"],
|
||||
["--mirostat", "2", "--mirostat-lr", "0.1"],
|
||||
["--xtc-probability", "0.05", "--xtc-threshold", "0.1"],
|
||||
["--dry-multiplier", "0.5"],
|
||||
# Tier-2 knobs that map to LoadRequest fields
|
||||
["--cache-type-k", "q8_0"],
|
||||
["--cache-type-v", "q8_0"],
|
||||
["--chat-template-file", "/tmp/tpl.jinja"],
|
||||
["--chat-template-kwargs", '{"reasoning_effort":"high"}'],
|
||||
["--spec-type", "ngram-mod"],
|
||||
["--spec-default"],
|
||||
# Reasoning controls
|
||||
["--reasoning-format", "deepseek"],
|
||||
["-rea", "auto"],
|
||||
# Soft-managed flags the user may want to override on the CLI;
|
||||
# llama.cpp's last-wins parsing means these win over Studio's
|
||||
# auto-set version.
|
||||
["-c", "131072"],
|
||||
["--ctx-size", "8192"],
|
||||
["--parallel", "1"],
|
||||
["-np", "8"],
|
||||
["--flash-attn", "off"],
|
||||
["-fa", "on"],
|
||||
["--no-context-shift"],
|
||||
["--context-shift"],
|
||||
["--jinja"],
|
||||
["--no-jinja"],
|
||||
["-ngl", "-1"],
|
||||
["--gpu-layers", "32"],
|
||||
["-t", "16"],
|
||||
["--threads", "32"],
|
||||
["-fit", "off"],
|
||||
["--fit", "on"],
|
||||
["--fit-ctx", "8192"],
|
||||
],
|
||||
)
|
||||
def test_pass_through_allowed(args):
|
||||
assert validate_extra_args(args) == args
|
||||
|
||||
|
||||
def test_none_returns_empty_list():
|
||||
assert validate_extra_args(None) == []
|
||||
|
||||
|
||||
def test_empty_list_returns_empty_list():
|
||||
assert validate_extra_args([]) == []
|
||||
|
||||
|
||||
def test_value_with_equals_form_passes_through():
|
||||
assert validate_extra_args(["--top-k=20"]) == ["--top-k=20"]
|
||||
|
||||
|
||||
def test_non_flag_token_passes_through():
|
||||
# A bare positional value (not preceded by a flag) is preserved
|
||||
# verbatim. llama-server may reject it, but that's not our job.
|
||||
assert validate_extra_args(["foo"]) == ["foo"]
|
||||
|
||||
|
||||
# ── Denylist (rejected) ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"denied",
|
||||
[
|
||||
# Model identity
|
||||
"-m",
|
||||
"--model",
|
||||
"-hf",
|
||||
"-hfr",
|
||||
"--hf-repo",
|
||||
"-hff",
|
||||
"--hf-file",
|
||||
"-hft",
|
||||
"--hf-token",
|
||||
"-mm",
|
||||
"--mmproj",
|
||||
"--mmproj-url",
|
||||
# Networking (Studio binds + proxies)
|
||||
"--host",
|
||||
"--port",
|
||||
"--path",
|
||||
"--api-prefix",
|
||||
"--reuse-port",
|
||||
# Auth / TLS
|
||||
"--api-key",
|
||||
"--api-key-file",
|
||||
"--ssl-key-file",
|
||||
"--ssl-cert-file",
|
||||
# Single-model server
|
||||
"--webui",
|
||||
"--no-webui",
|
||||
"--models-dir",
|
||||
"--models-max",
|
||||
],
|
||||
)
|
||||
def test_denylist_rejects_all_aliases(denied):
|
||||
with pytest.raises(ValueError, match = denied):
|
||||
validate_extra_args([denied, "value"])
|
||||
|
||||
|
||||
def test_denylist_rejects_equals_form():
|
||||
with pytest.raises(ValueError, match = "--port"):
|
||||
validate_extra_args(["--port=9000"])
|
||||
|
||||
|
||||
def test_denylist_rejects_short_form_when_long_is_denied():
|
||||
# -m is the short form of the hard-denied --model; rejecting only
|
||||
# the long form would leave a trivial bypass.
|
||||
with pytest.raises(ValueError, match = "-m"):
|
||||
validate_extra_args(["-m", "/some/other/path.gguf"])
|
||||
|
||||
|
||||
def test_denylist_message_names_offending_flag():
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
validate_extra_args(["--top-k", "20", "--api-key", "secret"])
|
||||
assert "--api-key" in str(excinfo.value)
|
||||
|
||||
|
||||
def test_first_denied_flag_short_circuits():
|
||||
# Validation stops at the first denied flag; later denied flags
|
||||
# in the same call don't matter for behaviour, but the message
|
||||
# should name the first one we hit.
|
||||
with pytest.raises(ValueError, match = "--port"):
|
||||
validate_extra_args(["--port", "1", "--host", "x"])
|
||||
|
||||
|
||||
# ── Numeric values that look flag-ish ─────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["-1", "-0.5", "-42", "-.5"])
|
||||
def test_negative_number_value_is_not_flag(value):
|
||||
# ``--seed -1`` is a value, not a flag. Validator must not try
|
||||
# to look up "-1" in the denylist.
|
||||
assert validate_extra_args(["--seed", value]) == ["--seed", value]
|
||||
|
||||
|
||||
# ── is_managed_flag helper ───────────────────────────────────────────
|
||||
|
||||
|
||||
def test_is_managed_flag_true_for_denied():
|
||||
assert is_managed_flag("--port") is True
|
||||
assert is_managed_flag("--api-key") is True
|
||||
assert is_managed_flag("-m") is True
|
||||
assert is_managed_flag("--model") is True
|
||||
|
||||
|
||||
def test_is_managed_flag_false_for_pass_through():
|
||||
assert is_managed_flag("--top-k") is False
|
||||
assert is_managed_flag("--cache-type-k") is False
|
||||
assert is_managed_flag("--chat-template-file") is False
|
||||
# Soft-managed flags pass through (last-wins override)
|
||||
assert is_managed_flag("-c") is False
|
||||
assert is_managed_flag("--ctx-size") is False
|
||||
assert is_managed_flag("--parallel") is False
|
||||
assert is_managed_flag("--flash-attn") is False
|
||||
assert is_managed_flag("-ngl") is False
|
||||
assert is_managed_flag("--threads") is False
|
||||
145
tests/studio/test_cli_repo_variant.py
Normal file
145
tests/studio/test_cli_repo_variant.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
"""Tests for the ``repo:variant`` shorthand parser used by ``unsloth studio run``.
|
||||
|
||||
Loads ``unsloth_cli/commands/studio.py`` directly via ``importlib`` with a
|
||||
minimal ``typer`` stub so the test doesn't drag in the rest of
|
||||
``unsloth_cli`` (which transitively imports the unsloth training stack).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_split_repo_variant():
|
||||
"""Load ``_split_repo_variant`` from studio.py with typer stubbed.
|
||||
|
||||
studio.py decorates Typer commands at import time, so a stub that
|
||||
accepts (and discards) those calls is enough to let module
|
||||
execution complete and expose the helper we want to test.
|
||||
"""
|
||||
if "typer" not in sys.modules:
|
||||
typer_stub = types.ModuleType("typer")
|
||||
|
||||
class _Typer:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def callback(self, *args, **kwargs):
|
||||
return lambda fn: fn
|
||||
|
||||
def command(self, *args, **kwargs):
|
||||
return lambda fn: fn
|
||||
|
||||
typer_stub.Typer = _Typer
|
||||
typer_stub.Option = lambda *args, **kwargs: (args[0] if args else None)
|
||||
typer_stub.Context = type("Context", (), {})
|
||||
typer_stub.Exit = type("Exit", (Exception,), {})
|
||||
typer_stub.echo = lambda *args, **kwargs: None
|
||||
sys.modules["typer"] = typer_stub
|
||||
|
||||
studio_py = (
|
||||
Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands" / "studio.py"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"_studio_for_repo_variant_test", studio_py
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module._split_repo_variant
|
||||
|
||||
|
||||
_split = _load_split_repo_variant()
|
||||
|
||||
|
||||
# ── HF-style repo:variant inputs -------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_arg, expected",
|
||||
[
|
||||
(
|
||||
"unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL",
|
||||
("unsloth/gpt-oss-20b-GGUF", "UD-Q4_K_XL"),
|
||||
),
|
||||
("unsloth/gpt-oss-120b-GGUF:Q4_K_XL", ("unsloth/gpt-oss-120b-GGUF", "Q4_K_XL")),
|
||||
("unsloth/Qwen3-0.6B-GGUF:Q4_K_M", ("unsloth/Qwen3-0.6B-GGUF", "Q4_K_M")),
|
||||
# Variants commonly contain dashes, dots, and underscores.
|
||||
("org/repo:UD-Q5_K_M", ("org/repo", "UD-Q5_K_M")),
|
||||
("org/repo:F16", ("org/repo", "F16")),
|
||||
],
|
||||
)
|
||||
def test_repo_variant_split(model_arg, expected):
|
||||
assert _split(model_arg) == expected
|
||||
|
||||
|
||||
# ── No variant suffix ------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_arg",
|
||||
[
|
||||
"unsloth/gpt-oss-20b-GGUF",
|
||||
"unsloth/Qwen3-0.6B-GGUF",
|
||||
"shorthand-no-org-no-colon",
|
||||
],
|
||||
)
|
||||
def test_no_colon_returns_none_variant(model_arg):
|
||||
repo, variant = _split(model_arg)
|
||||
assert repo == model_arg
|
||||
assert variant is None
|
||||
|
||||
|
||||
# ── Local paths must NOT be split ------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"local_path",
|
||||
[
|
||||
"/abs/path/to/model.gguf",
|
||||
"/abs/path:with-colon-in-name",
|
||||
"./relative/model",
|
||||
"../parent/model",
|
||||
"~/home/model",
|
||||
".",
|
||||
"C:\\Users\\me\\model.gguf",
|
||||
"C:/Users/me/model.gguf",
|
||||
"D:/data/model:Q4", # Windows drive + colon-suffixed filename: drive wins
|
||||
],
|
||||
)
|
||||
def test_local_path_passthrough(local_path):
|
||||
repo, variant = _split(local_path)
|
||||
assert repo == local_path
|
||||
assert variant is None
|
||||
|
||||
|
||||
# ── Edge cases -------------------------------------------------------
|
||||
|
||||
|
||||
def test_empty_string():
|
||||
assert _split("") == ("", None)
|
||||
|
||||
|
||||
def test_trailing_colon_no_variant():
|
||||
# "org/repo:" -- no quant label after the colon. Pass through
|
||||
# unchanged so the backend's existing validation surfaces a
|
||||
# clearer error than "variant ''".
|
||||
repo, variant = _split("org/repo:")
|
||||
assert repo == "org/repo:"
|
||||
assert variant is None
|
||||
|
||||
|
||||
def test_slash_in_variant_disqualifies_split():
|
||||
# "foo:bar/baz" -- the suffix has a slash, so this isn't a quant
|
||||
# label; treat the whole thing as opaque.
|
||||
repo, variant = _split("foo:bar/baz")
|
||||
assert repo == "foo:bar/baz"
|
||||
assert variant is None
|
||||
|
||||
|
||||
def test_whitespace_stripped():
|
||||
assert _split(" org/repo:Q4 ") == ("org/repo", "Q4")
|
||||
69
tests/studio/test_cli_run_alias.py
Normal file
69
tests/studio/test_cli_run_alias.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
"""Tests that ``unsloth run`` is registered as a top-level alias for
|
||||
``unsloth studio run``.
|
||||
|
||||
AST-based to avoid importing ``unsloth_cli`` (which pulls in the heavy
|
||||
training stack) at test-collection time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
_CLI_INIT = Path(__file__).resolve().parents[2] / "unsloth_cli" / "__init__.py"
|
||||
|
||||
|
||||
def _module_calls(source: str):
|
||||
tree = ast.parse(source)
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Call):
|
||||
yield node
|
||||
|
||||
|
||||
def test_top_level_run_alias_registered():
|
||||
"""`app.command("run", ...)` must be invoked with studio_run as its target."""
|
||||
source = _CLI_INIT.read_text()
|
||||
|
||||
# Find ``app.command("run", ...)`` call -- the decorator-call form.
|
||||
found_decorator_call = False
|
||||
for call in _module_calls(source):
|
||||
# Match ``app.command(...)`` syntactically.
|
||||
if not (
|
||||
isinstance(call.func, ast.Attribute)
|
||||
and call.func.attr == "command"
|
||||
and isinstance(call.func.value, ast.Name)
|
||||
and call.func.value.id == "app"
|
||||
):
|
||||
continue
|
||||
# Decorator-call form has a string literal "run" as the first
|
||||
# positional or as keyword ``name="run"``.
|
||||
first_pos = call.args[0] if call.args else None
|
||||
keyword_name = next(
|
||||
(kw.value for kw in call.keywords if kw.arg == "name"), None
|
||||
)
|
||||
is_run = (isinstance(first_pos, ast.Constant) and first_pos.value == "run") or (
|
||||
isinstance(keyword_name, ast.Constant) and keyword_name.value == "run"
|
||||
)
|
||||
if is_run:
|
||||
found_decorator_call = True
|
||||
break
|
||||
assert (
|
||||
found_decorator_call
|
||||
), 'Expected `app.command("run", ...)` registration in unsloth_cli/__init__.py'
|
||||
|
||||
|
||||
def test_studio_run_imported_for_alias():
|
||||
"""The alias must wire up to the studio.run function, not redefine it."""
|
||||
source = _CLI_INIT.read_text()
|
||||
tree = ast.parse(source)
|
||||
has_import = False
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.ImportFrom):
|
||||
continue
|
||||
if node.module != "unsloth_cli.commands.studio":
|
||||
continue
|
||||
for alias in node.names:
|
||||
if alias.name == "run":
|
||||
has_import = True
|
||||
break
|
||||
assert has_import, "Expected `from unsloth_cli.commands.studio import run` in unsloth_cli/__init__.py"
|
||||
|
|
@ -6,7 +6,7 @@ import typer
|
|||
from unsloth_cli.commands.train import train
|
||||
from unsloth_cli.commands.inference import inference
|
||||
from unsloth_cli.commands.export import export, list_checkpoints
|
||||
from unsloth_cli.commands.studio import studio_app
|
||||
from unsloth_cli.commands.studio import run as studio_run, studio_app
|
||||
|
||||
app = typer.Typer(
|
||||
help = "Command-line interface for Unsloth training, inference, and export.",
|
||||
|
|
@ -18,3 +18,15 @@ app.command()(inference)
|
|||
app.command()(export)
|
||||
app.command("list-checkpoints")(list_checkpoints)
|
||||
app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")
|
||||
|
||||
# Top-level alias: `unsloth run ...` is equivalent to `unsloth studio run ...`.
|
||||
# Same context_settings as the studio_app registration so unknown flags
|
||||
# still pass through to llama-server.
|
||||
app.command(
|
||||
"run",
|
||||
context_settings = {
|
||||
"allow_extra_args": True,
|
||||
"ignore_unknown_options": True,
|
||||
},
|
||||
help = "Alias for `unsloth studio run`.",
|
||||
)(studio_run)
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import time
|
|||
import types
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
import typer
|
||||
|
||||
studio_app = typer.Typer(help = "Unsloth Studio commands.")
|
||||
|
|
@ -374,6 +374,7 @@ def _load_model_via_http(
|
|||
gguf_variant: Optional[str],
|
||||
max_seq_length: int,
|
||||
load_in_4bit: bool,
|
||||
llama_extra_args: Optional[List[str]] = None,
|
||||
timeout: int = 600,
|
||||
) -> dict:
|
||||
"""POST to ``/api/inference/load`` using the API key for auth."""
|
||||
|
|
@ -388,6 +389,8 @@ def _load_model_via_http(
|
|||
}
|
||||
if gguf_variant:
|
||||
payload["gguf_variant"] = gguf_variant
|
||||
if llama_extra_args:
|
||||
payload["llama_extra_args"] = list(llama_extra_args)
|
||||
|
||||
data = json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
|
|
@ -514,9 +517,57 @@ def studio_default(
|
|||
# ── unsloth studio run ───────────────────────────────────────────────
|
||||
|
||||
|
||||
@studio_app.command()
|
||||
def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]:
|
||||
"""Split ``org/name:variant`` HF-style identifiers into (repo, variant).
|
||||
|
||||
Mirrors llama.cpp's ``-hf <repo>:<quant>`` convention so users can
|
||||
write ``unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL`` instead of passing
|
||||
``--gguf-variant`` separately. Local paths (absolute, ``./``,
|
||||
``~/``, Windows drive letters) and identifiers without a ``:``
|
||||
suffix are returned verbatim.
|
||||
"""
|
||||
s = model_arg.strip()
|
||||
if not s:
|
||||
return s, None
|
||||
if s.startswith(("/", "./", "../", "~")) or s == ".":
|
||||
return s, None
|
||||
# Windows drive letter (e.g. "C:\\path" or "C:/path") -- the colon
|
||||
# here is a path separator, not a variant suffix.
|
||||
if len(s) >= 2 and s[1] == ":" and s[0].isalpha():
|
||||
return s, None
|
||||
if ":" not in s:
|
||||
return s, None
|
||||
repo, _, variant = s.rpartition(":")
|
||||
if not repo or not variant:
|
||||
return s, None
|
||||
# A real quant label has no slashes; ``foo:bar/baz`` is not
|
||||
# ``repo:variant`` syntax.
|
||||
if "/" in variant:
|
||||
return s, None
|
||||
return repo, variant
|
||||
|
||||
|
||||
@studio_app.command(
|
||||
context_settings = {
|
||||
"allow_extra_args": True,
|
||||
"ignore_unknown_options": True,
|
||||
},
|
||||
)
|
||||
def run(
|
||||
model: str = typer.Option(..., "--model", "-m", help = "Model path or HF repo"),
|
||||
ctx: typer.Context,
|
||||
model: str = typer.Option(
|
||||
...,
|
||||
"--model",
|
||||
"-m",
|
||||
"-hf",
|
||||
"-hfr",
|
||||
"--hf-repo",
|
||||
help = (
|
||||
"Model path or HF repo. Accepts llama.cpp-style "
|
||||
"`org/repo:variant` syntax. The `-hf` / `--hf-repo` aliases "
|
||||
"match llama-server's spelling."
|
||||
),
|
||||
),
|
||||
gguf_variant: Optional[str] = typer.Option(
|
||||
None, "--gguf-variant", help = "GGUF quant variant (e.g. UD-Q4_K_XL)"
|
||||
),
|
||||
|
|
@ -534,9 +585,35 @@ def run(
|
|||
):
|
||||
"""Start Studio, load a model, and print an API key -- one-liner server.
|
||||
|
||||
Any flag this command does not recognize is forwarded verbatim to
|
||||
the underlying llama-server (GGUF only). Studio-managed flags
|
||||
(--port, -c / --ctx-size, --api-key, -ngl, --jinja, --flash-attn,
|
||||
--no-context-shift, model-identity flags, ...) are rejected with
|
||||
HTTP 400.
|
||||
|
||||
Example:
|
||||
unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL
|
||||
unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42
|
||||
unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja
|
||||
"""
|
||||
extra_llama_args: List[str] = list(ctx.args) if ctx.args else []
|
||||
|
||||
# ── 0. Parse llama.cpp-style ``repo:variant`` syntax in --model. ───
|
||||
# Lets users write ``--model unsloth/foo-GGUF:UD-Q4_K_XL`` instead
|
||||
# of pairing ``--model`` with ``--gguf-variant``. If both are given
|
||||
# and disagree, fail loudly instead of silently picking one.
|
||||
parsed_repo, embedded_variant = _split_repo_variant(model)
|
||||
if embedded_variant:
|
||||
if gguf_variant and gguf_variant != embedded_variant:
|
||||
typer.echo(
|
||||
f"Error: --model embeds variant '{embedded_variant}' but "
|
||||
f"--gguf-variant '{gguf_variant}' was also provided.",
|
||||
err = True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
model = parsed_repo
|
||||
gguf_variant = gguf_variant or embedded_variant
|
||||
|
||||
# ── 1. Venv re-exec (same pattern as studio_default) ──────────────
|
||||
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
|
||||
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
|
||||
|
|
@ -576,6 +653,11 @@ def run(
|
|||
args.extend(["--frontend", str(frontend)])
|
||||
if silent:
|
||||
args.append("--silent")
|
||||
# Forward unknown args (llama-server pass-through) to the
|
||||
# re-exec'd command so the studio venv sees them in ctx.args
|
||||
# and the re-execed run() can include them in the load payload.
|
||||
if extra_llama_args:
|
||||
args.extend(extra_llama_args)
|
||||
|
||||
if sys.platform == "win32":
|
||||
proc = subprocess.Popen(args)
|
||||
|
|
@ -617,6 +699,7 @@ def run(
|
|||
gguf_variant = gguf_variant,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = load_in_4bit,
|
||||
llama_extra_args = extra_llama_args,
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
typer.echo(f"Error: {exc}", err = True)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue