diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index d8b7eb383..4f528a689 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -1,46 +1,29 @@ # 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. +"""Boundary 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. +Reject only flags Studio manages (model identity, auth, network, +parallel slots). Everything else (sampling, ``-c``, ``-ngl``, +``--flash-attn``, ``--cache-type-*``, ``--spec-*``, ``--jinja``, ...) +is appended after Studio's auto-set flags so llama.cpp's last-wins +parser lets the user override. -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 +Ref: 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. +# Each group = every alias (short + long) of one hard-denied flag. +# Extend the matching group when llama.cpp adds a new alias. _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. + # Parallel slots: owned by typer --parallel; a pass-through would + # desync app.state.llama_parallel_slots from llama-server. + frozenset({"-np", "--parallel", "--n-parallel"}), + # Model identity: Studio resolves it from LoadRequest; a second + # -m would load a different model than Studio thinks it loaded. frozenset({"-m", "--model"}), frozenset({"-mu", "--model-url"}), frozenset({"-dr", "--docker-repo"}), @@ -51,28 +34,21 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( 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. + # Networking: Studio binds + proxies; retargeting orphans the proxy. 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. + # Auth / TLS: Studio terminates auth; upstream --api-key / TLS + # shadows Studio's key and breaks the 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. - # ``--webui``/``--no-webui`` are the legacy spelling; current - # upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions. - # Keep both so the denylist matches old and new llama-server - # binaries (Studio's prebuilt vs system-llama.cpp). + # Built-in web UI. --webui/--no-webui is the legacy spelling; + # upstream renamed to --ui/--no-ui + --ui-*. Keep both so prebuilt + # and system llama.cpp binaries both match. frozenset({"--webui", "--no-webui"}), frozenset({"--ui", "--no-ui"}), frozenset({"--ui-config"}), @@ -82,32 +58,46 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( frozenset({"--models-preset"}), frozenset({"--models-max"}), frozenset({"--models-autoload", "--no-models-autoload"}), + # Server-mode flips: --embedding / --rerank restrict llama-server to + # those endpoints, breaking Studio's /v1/chat/completions hop. + frozenset({"--embedding", "--embeddings"}), + frozenset({"--rerank", "--reranking"}), + # llama-server's own built-in tools flag would silently stack on top + # of Studio's --enable-tools / --disable-tools policy resolver. + frozenset({"--tools"}), ) _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. + """Flag name for ``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. + Peels `--key=value` to `--key`, treats `-1` / `-0.5` as values + (llama-server shorts always start with a letter), strips + whitespace, and normalises attached `-np8` / signed `-np-1` / + digit-prefix-junk `-np8x` to `-np`. Mirrors the CLI's + `_expand_attached_np_short`. """ + token = token.strip() 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] + name = token.split("=", 1)[0] + if len(name) > 3 and name.startswith("-np"): + suffix = name[3:] + if suffix[0].isdigit() or ( + len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit() + ): + return "-np" + return name 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. - """ + """Validate user-supplied llama-server args. Returns a flat list + ready to extend the llama-server command; raises ``ValueError`` + naming the offending flag on the first managed token.""" if not args: return [] out: list[str] = [] @@ -124,15 +114,15 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: def is_managed_flag(flag: str) -> bool: - """True if ``flag`` is a Studio-managed llama-server flag.""" - return flag in _DENYLIST + """True if ``flag`` is Studio-managed. Normalises via ``_flag_name`` + so `-np8` / `--parallel=8` classify like the canonical tokens.""" + normalised = _flag_name(flag) + return normalised is not None and normalised in _DENYLIST -# Pass-through flags that shadow first-class ``LoadRequest`` fields -# (max_seq_length, cache_type_kv, speculative_type, -# chat_template_override). Stripped from inherited extras so they -# can't last-wins-override an Apply that re-sets the same first-class -# field. +# Pass-through flags that shadow first-class LoadRequest fields; +# stripped from inherited extras so they can't last-wins-override an +# Apply that re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) _CACHE_FLAGS: frozenset[str] = frozenset( {"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"} @@ -169,9 +159,8 @@ _SHADOWING_FLAGS: frozenset[str] = ( _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS ) -# Boolean flags inside _SHADOWING_FLAGS that take no value. The -# value-consuming heuristic in strip_shadowing_flags must skip just the -# flag for these, never the following token. +# Shadowing flags that take no value -- strip the flag only, never the +# following token. _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( {"--spec-default", "--jinja", "--no-jinja"} ) @@ -187,14 +176,11 @@ def strip_shadowing_flags( ) -> list[str]: """Strip flags that shadow first-class Studio settings. - Used when the route inherits a previous load's ``llama_extra_args`` - so that an inherited ``-c 4096`` cannot override the current - request's ``max_seq_length`` (and equivalents for cache / - speculative / chat template). Each ``strip_*`` flag controls one - group; the route only strips groups whose corresponding first-class - field was actually supplied by the caller, so an inherited - ``--chat-template-file`` survives an Apply that omits both - ``llama_extra_args`` and ``chat_template_override``. + Used when inheriting a previous load's ``llama_extra_args`` so an + inherited `-c 4096` can't override the current `max_seq_length` + (same for cache / spec / template). Each ``strip_*`` toggle + controls one group; the route only strips groups whose first-class + field the caller actually supplied. """ shadowing: set[str] = set() if strip_context: @@ -216,9 +202,8 @@ def strip_shadowing_flags( out.append(tok) i += 1 continue - # Drop this token. Boolean shadowing flags never carry a value; - # other shadowing flags consume the next token when it isn't a - # flag and the value isn't already packed as ``--key=value``. + # Drop the flag; consume the next token too unless it's + # boolean, already inline (`-c=4096`), or another flag. if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok: i += 1 elif i + 1 < n and _flag_name(tokens[i + 1]) is None: diff --git a/studio/backend/run.py b/studio/backend/run.py index 3bde8abd3..e96e60965 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -846,11 +846,33 @@ if __name__ == "__main__": action = "store_true", help = "API server only, no frontend (for Tauri)", ) + # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 + # applies only to direct backend launches; `unsloth studio run` + # always passes its own value (4) explicitly. + _PARALLEL_MIN = 1 + _PARALLEL_MAX = 64 + _PARALLEL_DEFAULT_PLAIN = 1 + parser.add_argument( + "--parallel", + "--n-parallel", + type = int, + default = _PARALLEL_DEFAULT_PLAIN, + help = ( + f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). " + f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4." + ), + ) args = parser.parse_args() + if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX: + parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}") kwargs = dict( - host = args.host, port = args.port, silent = args.silent, api_only = args.api_only + host = args.host, + port = args.port, + silent = args.silent, + api_only = args.api_only, + llama_parallel_slots = args.parallel, ) if args.frontend is not None: kwargs["frontend_path"] = Path(args.frontend) diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 68a1c870f..02a272ba3 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -3,21 +3,35 @@ """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. +The validator is the boundary between user CLI/HTTP input and the +llama-server subprocess. These tests pin denylist behaviour so it +doesn't quietly regress when new managed flags are added. """ from __future__ import annotations +import importlib.util +import re +from pathlib import Path + import pytest -from core.inference.llama_server_args import ( - is_managed_flag, - strip_shadowing_flags, - validate_extra_args, +# Load llama_server_args.py directly so this test doesn't drag in the +# full backend chain (fastapi / structlog / loggers / utils.hardware) +# via core/inference/__init__.py. The validator is intentionally +# dependency-free and unit-tests should reflect that. +_LSA_PATH = ( + Path(__file__).resolve().parent.parent + / "core" + / "inference" + / "llama_server_args.py" ) +_spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH) +_lsa = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_lsa) +is_managed_flag = _lsa.is_managed_flag +strip_shadowing_flags = _lsa.strip_shadowing_flags +validate_extra_args = _lsa.validate_extra_args # ── Pass-through (allowed) ─────────────────────────────────────────── @@ -60,13 +74,12 @@ from core.inference.llama_server_args import ( # 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. + # Soft-managed: user-supplied flags last-wins-override Studio's + # auto-set version. --parallel / -np / --n-parallel are NOT + # here -- they're hard-denied (KV-cache + slot count would + # desync). Use `unsloth studio run --parallel N` instead. ["-c", "131072"], ["--ctx-size", "8192"], - ["--parallel", "1"], - ["-np", "8"], ["--flash-attn", "off"], ["-fa", "on"], ["--no-context-shift"], @@ -99,8 +112,7 @@ def test_value_with_equals_form_passes_through(): 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. + # Bare positionals are passed through; llama-server can reject them. assert validate_extra_args(["foo"]) == ["foo"] @@ -110,18 +122,33 @@ def test_non_flag_token_passes_through(): @pytest.mark.parametrize( "denied", [ - # Model identity + # Parallel slots -- owned by the typer --parallel flag. + "-np", + "--parallel", + "--n-parallel", + # Model identity (every alias; bumping llama.cpp must keep + # every form rejected, not just the long). "-m", "--model", + "-mu", + "--model-url", + "-dr", + "--docker-repo", "-hf", "-hfr", "--hf-repo", "-hff", "--hf-file", + "-hfv", + "-hfrv", + "--hf-repo-v", + "-hffv", + "--hf-file-v", "-hft", "--hf-token", "-mm", "--mmproj", + "-mmu", "--mmproj-url", # Networking (Studio binds + proxies) "--host", @@ -134,11 +161,28 @@ def test_non_flag_token_passes_through(): "--api-key-file", "--ssl-key-file", "--ssl-cert-file", - # Single-model server + # Single-model server (legacy --webui + current --ui group) "--webui", "--no-webui", + "--ui", + "--no-ui", + "--ui-config", + "--ui-config-file", + "--ui-mcp-proxy", + "--no-ui-mcp-proxy", "--models-dir", + "--models-preset", "--models-max", + "--models-autoload", + "--no-models-autoload", + # Server-mode flips: --embedding / --rerank would restrict + # llama-server to those endpoints and break Studio's chat hop. + "--embedding", + "--embeddings", + "--rerank", + "--reranking", + # llama-server's own --tools clashes with Studio's tool policy. + "--tools", ], ) def test_denylist_rejects_all_aliases(denied): @@ -146,14 +190,65 @@ def test_denylist_rejects_all_aliases(denied): validate_extra_args([denied, "value"]) +@pytest.mark.parametrize( + "args,offending", + [ + # Pass-through --parallel would last-wins-override the real + # slot count while Studio's KV-cache fit + llama_parallel_slots + # stay at the typer value -- plan vs. process disagree. + (["--parallel", "8"], "--parallel"), + (["--parallel=8"], "--parallel"), + (["--n-parallel", "16"], "--n-parallel"), + (["--n-parallel=16"], "--n-parallel"), + (["-np", "32"], "-np"), + # Attached short form: Click clusters it CLI-side; HTTP /load + # with `["-np8"]` must still resolve to managed. + (["-np8"], "-np"), + (["-np64"], "-np"), + # Out-of-range values that would bypass the typer 1..64 guard. + (["--parallel", "999"], "--parallel"), + (["-np", "0"], "-np"), + (["-np999"], "-np"), + # Signed attached forms; `-np-1` must not slip past. + (["-np-1"], "-np"), + (["-np+1"], "-np"), + ], +) +def test_parallel_flags_are_managed(args, offending): + with pytest.raises(ValueError, match = re.escape(offending)): + validate_extra_args(args) + + def test_denylist_rejects_equals_form(): with pytest.raises(ValueError, match = "--port"): validate_extra_args(["--port=9000"]) +@pytest.mark.parametrize( + "padded", + [" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"], +) +def test_denylist_rejects_whitespace_padded_forms(padded): + # `_flag_name` trims whitespace before lookup; otherwise a trailing + # space could slip a managed flag past the boundary. + with pytest.raises(ValueError, match = "parallel|np"): + validate_extra_args([padded, "8"]) + + +@pytest.mark.parametrize( + "attached", + ["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"], +) +def test_denylist_rejects_np_with_digit_prefix_and_junk(attached): + # Backend `_flag_name` must classify the same forms the CLI + # rewriter expands, else HTTP /load could smuggle `-np8x` through. + with pytest.raises(ValueError, match = "np"): + validate_extra_args([attached]) + + 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. + # `-m` is the short form of --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"]) @@ -165,9 +260,7 @@ def test_denylist_message_names_offending_flag(): 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. + # Validation stops at the first denied flag; the message names it. with pytest.raises(ValueError, match = "--port"): validate_extra_args(["--port", "1", "--host", "x"]) @@ -177,8 +270,7 @@ def test_first_denied_flag_short_circuits(): @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. + # `--seed -1`: the -1 is a value, not a flag. assert validate_extra_args(["--seed", value]) == ["--seed", value] @@ -190,6 +282,15 @@ def test_is_managed_flag_true_for_denied(): assert is_managed_flag("--api-key") is True assert is_managed_flag("-m") is True assert is_managed_flag("--model") is True + # Parallel slots owned by the typer --parallel flag. + assert is_managed_flag("--parallel") is True + assert is_managed_flag("--n-parallel") is True + assert is_managed_flag("-np") is True + # Normalised forms must classify like the canonical token so + # is_managed_flag filtering stays in sync with validate_extra_args. + assert is_managed_flag("-np8") is True + assert is_managed_flag("--parallel=8") is True + assert is_managed_flag("--port=9000") is True def test_is_managed_flag_false_for_pass_through(): @@ -199,7 +300,6 @@ def test_is_managed_flag_false_for_pass_through(): # 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 @@ -231,8 +331,8 @@ def test_strip_shadowing_flags_keeps_context_when_not_requested(): def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled(): - # Caller did not supply chat_template_override; the inherited - # --chat-template-file must survive the strip. + # No chat_template_override supplied; inherited + # --chat-template-file must survive. out = strip_shadowing_flags( ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"], strip_context = True, @@ -282,7 +382,7 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled(): def test_strip_shadowing_flags_drops_mtp_flags_when_requested(): - # MTP / draft-mtp flags must be stripped when speculative_type is re-applied. + # MTP / draft-mtp flags must drop when speculative_type re-applies. out = strip_shadowing_flags( [ "--spec-type", @@ -311,8 +411,7 @@ def test_is_managed_flag_false_for_mtp_pass_through(): def test_strip_shadowing_flags_boolean_does_not_consume_next_token(): - # --spec-default is a boolean shadowing flag; the value-skipping - # heuristic must skip just the flag, not the following positional. + # `--spec-default` is boolean; drop just the flag, keep the next token. out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True) assert out == ["ngram-mod"] @@ -343,8 +442,8 @@ def test_strip_shadowing_flags_handles_empty_input(): def test_strip_shadowing_flags_defaults_strip_everything(): - # The route's already-loaded comparator calls strip_shadowing_flags - # with no kwargs to detect ANY shadowing flag in stored extras. + # The route's already-loaded comparator calls with no kwargs to + # detect ANY shadowing flag in stored extras. out = strip_shadowing_flags( ["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"] ) diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py index a1af16d4f..1e2f0d0cf 100644 --- a/tests/studio/test_composer_rtl_bidi_attribute.py +++ b/tests/studio/test_composer_rtl_bidi_attribute.py @@ -26,7 +26,11 @@ def _block_around(src: str, anchor: str, radius: int = 600) -> str: def test_main_composer_has_dir_auto(): - block = _block_around(THREAD_TSX.read_text(), 'aria-label="Message input"') + # PR #5784 rewrote the literal attribute into a JSX conditional + # (`aria-label={overlay ? "Image edit instructions" : "Message input"}`), + # so anchor on the inner string literal instead -- it survives both + # the old and new spellings. + block = _block_around(THREAD_TSX.read_text(), '"Message input"') assert 'dir="auto"' in block, 'main composer is missing dir="auto"' diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 65834b3b9..c8ec4c66c 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import os.path as _osp +import sys as _sys + import typer from importlib.metadata import version as package_version, PackageNotFoundError @@ -8,7 +11,19 @@ from importlib.metadata import version as package_version, PackageNotFoundError 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 run as studio_run, studio_app +from unsloth_cli.commands.studio import ( + run as studio_run, + studio_app, + _expand_attached_np_short, +) + + +# Canonicalise `-np` only under the `unsloth` console-script; +# third-party scripts that import unsloth_cli keep their argv intact. +_entry_base = _osp.basename(_sys.argv[0]).lower() if _sys.argv else "" +if _entry_base in {"unsloth", "unsloth.exe"}: + _expand_attached_np_short() +del _entry_base def show_version(value: bool): @@ -47,9 +62,8 @@ 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. +# Top-level `unsloth run` aliases `unsloth studio run`; same context +# so unknown flags still pass through to llama-server. app.command( "run", context_settings = { diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index e37cd0a8d..edb8ea2cf 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -206,6 +206,14 @@ def _find_setup_script() -> Optional[Path]: return None +# Mirror in studio/backend/run.py argparse + backend denylist test; +# bumping the cap in one place only desyncs. +_PARALLEL_MIN = 1 +_PARALLEL_MAX = 64 +_PARALLEL_DEFAULT_RUN = 4 # pre-PR hardcoded for `unsloth studio run` +_PARALLEL_DEFAULT_PLAIN = 1 # pre-PR effective for plain `unsloth studio` + + def _iter_editable_studio_source_roots(venv_dir: Path): """Yield repo roots from setuptools `__editable___*_finder.py` files in *venv_dir*'s site-packages whose MAPPING includes a `studio` entry. @@ -587,14 +595,38 @@ def studio_default( "--api-only", help = "Run API server only, no frontend serving (for Tauri desktop app)", ), + parallel: int = typer.Option( + _PARALLEL_DEFAULT_PLAIN, + "--parallel", + "--n-parallel", + min = _PARALLEL_MIN, + max = _PARALLEL_MAX, + help = ( + f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). " + f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` " + f"defaults to {_PARALLEL_DEFAULT_RUN}." + ), + ), ): """Launch the Unsloth Studio server.""" - # Runs before any subcommand; covers run/setup/update/etc in one place. + # Runs before every subcommand (run/setup/update/...). _ensure_studio_env_exported() if ctx.invoked_subcommand is not None: + # Typer doesn't forward parent options to subcommands, so + # `unsloth studio --parallel N run ...` would silently drop N. + if parallel != _PARALLEL_DEFAULT_PLAIN: + typer.echo( + f"Error: --parallel on `unsloth studio` applies to the " + f"plain-server path only. For `unsloth studio " + f"{ctx.invoked_subcommand}`, put the flag after the " + f"subcommand: `unsloth studio {ctx.invoked_subcommand} " + f"--parallel {parallel} ...`", + err = True, + ) + raise typer.Exit(2) return - # Always use the studio venv if it exists and we're not already in it + # Use the studio venv if it exists and we aren't already in it. studio_venv_dir = STUDIO_HOME / "unsloth_studio" in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) @@ -611,6 +643,8 @@ def studio_default( host, "--port", str(port), + "--parallel", + str(parallel), ] # Resolve frontend explicitly so the spawned run.py uses a real # built dist regardless of where its __file__ lands. Skip in @@ -624,9 +658,8 @@ def studio_default( args.append("--silent") if api_only: args.append("--api-only") - # On Windows, os.execvp() spawns a child but the parent lingers, - # so Ctrl+C only kills the parent leaving the child orphaned. - # Use subprocess.run() on Windows so the parent waits for the child. + # On Windows os.execvp keeps the parent alive, so Ctrl+C + # would orphan the child; use Popen+wait instead. if sys.platform == "win32": import subprocess as _sp @@ -634,7 +667,7 @@ def studio_default( try: rc = proc.wait() except KeyboardInterrupt: - # Child has its own signal handler — let it finish + # Child handles its own signal; let it finish. rc = proc.wait() if rc != 0: typer.echo( @@ -661,7 +694,13 @@ def studio_default( display_host = _resolve_external_ip() if host == "0.0.0.0" else host typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}") - run_kwargs = dict(host = host, port = port, silent = silent, api_only = api_only) + run_kwargs = dict( + host = host, + port = port, + silent = silent, + api_only = api_only, + llama_parallel_slots = parallel, + ) if frontend is not None: run_kwargs["frontend_path"] = frontend run_server(**run_kwargs) @@ -670,8 +709,8 @@ def studio_default( try: if _shutdown_event is not None: - # NOTE: Event.wait() without a timeout blocks at the C level - # on Linux, preventing Python from delivering SIGINT (Ctrl+C). + # Event.wait() with no timeout blocks at C-level on Linux + # and swallows SIGINT; loop with a 1s timeout instead. while not _shutdown_event.is_set(): _shutdown_event.wait(timeout = 1) else: @@ -688,21 +727,15 @@ def studio_default( 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 :`` 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. - """ + """Split ``org/name:variant`` into ``(repo, variant)``; mirrors + llama.cpp's ``-hf :``. Local paths, Windows drives, + and ids without ``:`` pass through 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. + # Windows drive letter (e.g. "C:\path"): colon is a path separator. if len(s) >= 2 and s[1] == ":" and s[0].isalpha(): return s, None if ":" not in s: @@ -710,13 +743,80 @@ def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]: 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. + # Quant labels never contain a slash; `foo:bar/baz` isn't repo:variant. if "/" in variant: return s, None return repo, variant +def _expand_attached_np_short() -> None: + # Click clusters `-np8` as `-n -p 8` (-p = --port), dropping the + # parallel value. Split to `-np ` so typer's alias matches. + # Stops at `--`; accepts signed and digit-prefix-junk forms so + # typer can report a clean error against `-np`. Kept in lockstep + # with the backend `_flag_name` recogniser. + i = 0 + while i < len(sys.argv): + tok = sys.argv[i] + if tok == "--": + break + if len(tok) > 3 and tok.startswith("-np") and tok[3] != "=": + suffix = tok[3:] + first_numeric = suffix[0].isdigit() or ( + len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit() + ) + if first_numeric: + sys.argv[i : i + 1] = ["-np", suffix] + i += 2 + continue + i += 1 + + +def _consume_legacy_short_aliases( + args: List[str], + aliases: tuple[str, ...], + current: Optional[str], + canonical: str, +) -> tuple[Optional[str], List[str]]: + """Pop exact-match legacy shorts (`-m`/`-hfr`/`-f`) from args; + leave clusters (`-mg`/`-fa`/...) for the llama-server tail. Inline + `-x=value` form also accepted.""" + out: List[str] = [] + value = current + i, n = 0, len(args) + while i < n: + tok = args[i] + if tok == "--": # end of options; tail is raw payload. + out.extend(args[i:]) + break + name, sep, inline = tok.partition("=") + if name not in aliases: + out.append(tok) + i += 1 + continue + if value is not None: + raise typer.BadParameter( + f"{name} conflicts with {canonical} already provided" + ) + if sep: + if inline == "": # `-m=` would become --model '' (Path('')='.'). + raise typer.BadParameter(f"{name} requires a non-empty value") + value = inline + i += 1 + elif i + 1 < n: + nxt = args[i + 1] + # `--long` is unambiguously a flag; single-dash `-x` may be a path. + if nxt.startswith("--") and nxt != "--": + raise typer.BadParameter( + f"{name} expects a value but got the flag {nxt}" + ) + value = nxt + i += 2 + else: + raise typer.BadParameter(f"{name} requires a value") + return value, out + + @studio_app.command( context_settings = { "allow_extra_args": True, @@ -725,17 +825,18 @@ def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]: ) def run( ctx: typer.Context, - model: str = typer.Option( - ..., + model: Optional[str] = typer.Option( + None, "--model", - "-m", "-hf", - "-hfr", "--hf-repo", + # `-m` / `-hfr` removed (Click would cluster `-mg`/`-md`/...). + # Exact-match `-m`/`-hfr` still work via the legacy shim below. + # `-hf` stays (multi-char shorts don't cluster). help = ( "Model path or HF repo. Accepts llama.cpp-style " - "`org/repo:variant` syntax. The `-hf` / `--hf-repo` aliases " - "match llama-server's spelling." + "`org/repo:variant` syntax. `-hf` / `--hf-repo` match " + "llama-server's spelling." ), ), gguf_variant: Optional[str] = typer.Option( @@ -750,7 +851,8 @@ def run( ), port: int = typer.Option(8888, "--port", "-p"), host: str = typer.Option("127.0.0.1", "--host", "-H"), - frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"), + # `-f` removed (clustered `-fa`/`-fit*`); studio_default keeps it. + frontend: Optional[Path] = typer.Option(None, "--frontend"), silent: bool = typer.Option(False, "--silent", "-q"), enable_tools: Optional[bool] = typer.Option( None, @@ -766,26 +868,65 @@ def run( "-y", help = "Skip the 0.0.0.0 + --enable-tools confirmation prompt.", ), + parallel: int = typer.Option( + _PARALLEL_DEFAULT_RUN, + "--parallel", + "--n-parallel", + "-np", + min = _PARALLEL_MIN, + max = _PARALLEL_MAX, + help = ( + "llama-server parallel decode slots. N requests share one " + "loaded model; each slot gets ctx/N KV cache. Default " + f"{_PARALLEL_DEFAULT_RUN} (pre-PR hardcoded value)." + ), + ), ): - """Start Studio, load a model, and print an API key -- one-liner server. + """Start Studio, load a model, 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. + Unknown flags pass through to llama-server (GGUF only). Studio + rejects managed flags with HTTP 400: model identity, network + (--host/--port/--path/--api-prefix/--reuse-port), auth/TLS + (--api-key/--ssl-*), single-model UI (--ui/--models-*/--webui), + and parallel slots (use --parallel above). Full denylist in + studio/backend/core/inference/llama_server_args.py. Other knobs + (-c, -ngl, --jinja, --flash-attn, -t, ...) pass through and + last-wins-override Studio's auto-set value. 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 unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42 --parallel 8 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. + # Promote legacy exact `-m`/`-hfr`/`-f` back into typer params; + # clusters stay in extras. + model, extra_llama_args = _consume_legacy_short_aliases( + extra_llama_args, + ("-m", "-hfr"), + model, + "--model", + ) + legacy_frontend, extra_llama_args = _consume_legacy_short_aliases( + extra_llama_args, + ("-f",), + str(frontend) if frontend is not None else None, + "--frontend", + ) + if legacy_frontend is not None and frontend is None: + frontend = Path(legacy_frontend) + + if model is None: + typer.echo( + "Error: Missing option '--model' / '-hf' / '--hf-repo' " + "(legacy aliases '-m' / '-hfr' are still accepted).", + err = True, + ) + raise typer.Exit(2) + + # 0. Parse llama.cpp `repo:variant` in --model; error if also paired + # with --gguf-variant and they disagree. parsed_repo, embedded_variant = _split_repo_variant(model) if embedded_variant: if gguf_variant and gguf_variant != embedded_variant: @@ -798,8 +939,8 @@ def run( model = parsed_repo gguf_variant = gguf_variant or embedded_variant - # ── Resolve the server-side tool policy. The y/N prompt (if any) - # runs in the outer process so the re-exec'd child never re-prompts. + # Resolve tool policy here so the re-exec'd child inherits a + # concrete decision and never re-prompts. from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy enable_tools = resolve_tool_policy( @@ -809,7 +950,7 @@ def run( silent = silent, ) - # ── 1. Venv re-exec (same pattern as studio_default) ────────────── + # 1. Re-exec into the studio venv (same pattern as studio_default). studio_venv_dir = STUDIO_HOME / "unsloth_studio" in_studio_venv = sys.prefix.startswith(str(studio_venv_dir)) @@ -818,7 +959,7 @@ def run( if not studio_python: typer.echo("Studio not set up. Run install.sh first.") raise typer.Exit(1) - # Re-exec into the studio venv via its `unsloth` entry point + # Re-exec via the studio venv's `unsloth` console-script. studio_bin = studio_python.parent / "unsloth" if not studio_bin.is_file(): typer.echo( @@ -842,27 +983,26 @@ def run( ] if gguf_variant: args.extend(["--gguf-variant", gguf_variant]) - if not load_in_4bit: - args.append("--no-load-in-4bit") + # Forward the explicit polarity; a future default flip on one + # layer must not silently invert behaviour for the other. + args.append("--load-in-4bit" if load_in_4bit else "--no-load-in-4bit") if frontend: args.extend(["--frontend", str(frontend)]) if silent: args.append("--silent") - # Forward the resolved tool policy (always concrete True/False - # at this point — the resolver above ran before the re-exec). + # Forward the resolved tool policy so the child doesn't re-resolve. if enable_tools: args.append("--enable-tools") else: args.append("--disable-tools") - # Forward --yes whenever the parent already cleared the prompt - # (either operator passed --yes, or the parent's resolver - # accepted the network-bind confirmation). Otherwise the child - # re-runs the resolver and prompts a second time. + # Forward --yes if the parent already cleared the network-bind + # prompt, else the child re-prompts. if yes or (enable_tools and is_external_host(host)): args.append("--yes") - # 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. + # Typer claims --parallel outside ctx.args; without this the + # child reverts to its default and silently drops the value. + args.extend(["--parallel", str(parallel)]) + # llama-server pass-through extras → child ctx.args → load payload. if extra_llama_args: args.extend(extra_llama_args) @@ -879,34 +1019,31 @@ def run( # ── 2. Start server (always suppress built-in banner) ───────────── from studio.backend.run import run_server, _resolve_external_ip - run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = 4) + run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = parallel) if frontend is not None: run_kwargs["frontend_path"] = frontend app = run_server(**run_kwargs) actual_port = getattr(app.state, "server_port", port) or port - # ── Apply the resolved tool policy as a process-level override. - # Must use the same import path the route handlers use -- - # `studio/backend/run.py` adds `studio/backend/` to sys.path so the - # routes import this module as top-level `state.tool_policy`. If we - # imported via `studio.backend.state.tool_policy` instead, Python - # would cache two different module objects with two different - # `_tool_policy` globals, and the gates would never see our value. + # Match the route handlers' import path: run.py adds + # studio/backend/ to sys.path, so they import as `state.tool_policy`. + # Importing via `studio.backend.state.tool_policy` would cache a + # second module object whose flag the gates can't see. from state.tool_policy import set_tool_policy set_tool_policy(enable_tools) - # ── 3. Wait for server health ───────────────────────────────────── + # 3. Wait for server health. if not silent: typer.echo("Starting Unsloth Studio...") if not _wait_for_server(actual_port): typer.echo("Error: server did not become healthy within 30 seconds.", err = True) raise typer.Exit(1) - # ── 4. Create API key in-process ────────────────────────────────── + # 4. Create API key in-process. api_key = _create_api_key_inprocess(api_key_name) - # ── 5. Load model via HTTP ──────────────────────────────────────── + # 5. Load model via HTTP. if not silent: typer.echo(f"Loading model: {model}...") try: @@ -926,15 +1063,13 @@ def run( loaded_model = result.get("model", model) display_variant = f" ({gguf_variant})" if gguf_variant else "" - # ── 6. Print banner ─────────────────────────────────────────────── + # 6. Print banner. display_host = _resolve_external_ip() if host == "0.0.0.0" else host base_url = f"http://{display_host}:{actual_port}" sdk_base_url = f"{base_url}/v1" - # Claude orange (Claude Code's brand color) for tool-policy notices - # so they stand out from the surrounding banner. Always printed -- - # even under --silent / --yes -- so the operator never misses the - # current tool-execution status. + # Orange so the tool-policy notice stands out; printed under + # --silent / --yes too so the policy is never invisible. _tool_notice_fg = (217, 119, 87) _is_external = is_external_host(host) if _is_external and enable_tools: @@ -992,14 +1127,12 @@ def run( typer.echo(""" -d '{"input": "Hello", "stream": true}'""") typer.echo("") else: - # Silent mode still prints the essentials (URL, API key) plus - # the orange tool-status notice so the operator never loses - # visibility into the security-relevant policy. + # Silent still prints URL + API key + tool-status policy. typer.echo(f"URL: {base_url}") typer.echo(f"API Key: {api_key}") typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True) - # ── 7. Wait for Ctrl+C ──────────────────────────────────────────── + # 7. Wait for Ctrl+C. from studio.backend.run import _shutdown_event, _graceful_shutdown, _server try: diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py new file mode 100644 index 000000000..561600e43 --- /dev/null +++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py @@ -0,0 +1,416 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the `unsloth studio run --parallel` CLI flag. + +Pre-PR `llama_parallel_slots` was hardcoded to 4. These tests pin +the typer Option (aliases, default 4, 1..64 range), the +typer/denylist subset invariant, and re-exec forwarding. + +See ``test_studio_run_short_alias_clashes.py`` for the argv +canonicaliser and the legacy `-m` / `-hfr` / `-f` shim. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from typer.testing import CliRunner + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _load_run_command(): + """Import `studio` without triggering server start; backend imports + are lazy inside run().""" + from unsloth_cli.commands import studio as _studio + + return _studio + + +def test_parallel_option_is_registered(): + """The `--parallel` flag (with aliases) must be on the `run` command.""" + studio_mod = _load_run_command() + import inspect + + run_fn = studio_mod.run + sig = inspect.signature(run_fn) + assert "parallel" in sig.parameters, "missing `parallel` parameter on run()" + + param = sig.parameters["parallel"] + opt = param.default # typer.OptionInfo + flags = set() + decls = getattr(opt, "param_decls", None) or [] + for d in decls: + flags.add(d) + for required in ("--parallel", "--n-parallel", "-np"): + assert required in flags, f"flag {required!r} missing from --parallel option" + + +def test_parallel_default_is_four(): + """Default must stay at 4 so plain `unsloth studio run` is unchanged.""" + studio_mod = _load_run_command() + import inspect + + sig = inspect.signature(studio_mod.run) + opt = sig.parameters["parallel"].default + default = getattr(opt, "default", None) + assert ( + default == 4 + ), f"default changed to {default}; would silently alter existing deployments" + + +def test_parallel_range_guards_are_set(): + """Range guards: 1 <= N <= 64. Outside this is a hard reject.""" + studio_mod = _load_run_command() + import inspect + + sig = inspect.signature(studio_mod.run) + opt = sig.parameters["parallel"].default + assert getattr(opt, "min", None) == 1, "min must be 1 (0 = no decode possible)" + assert getattr(opt, "max", None) == 64, "max must be 64 (KV split sanity cap)" + + +def test_typer_parallel_aliases_are_subset_of_backend_denylist(): + """Every typer alias for --parallel must be denied on the backend + too; otherwise HTTP /load could smuggle the value via + `llama_extra_args` and desync llama_parallel_slots from the + running llama-server.""" + studio_mod = _load_run_command() + import inspect + import importlib.util + + # Load llama_server_args.py directly so the test doesn't need the + # backend's full runtime chain (fastapi / structlog / loggers / + # utils.hardware) installed -- the invariant is just about the + # _DENYLIST_GROUPS tuple. + lsa_path = ( + Path(__file__).resolve().parents[2] + / "studio" + / "backend" + / "core" + / "inference" + / "llama_server_args.py" + ) + spec = importlib.util.spec_from_file_location("_lsa_for_subset_test", lsa_path) + lsa = importlib.util.module_from_spec(spec) + spec.loader.exec_module(lsa) + _DENYLIST_GROUPS = lsa._DENYLIST_GROUPS + + parallel_group = next((g for g in _DENYLIST_GROUPS if "--parallel" in g), None) + assert parallel_group is not None, "denylist must include a --parallel group" + + sig = inspect.signature(studio_mod.run) + opt = sig.parameters["parallel"].default + typer_aliases = set(getattr(opt, "param_decls", []) or []) + missing = typer_aliases - parallel_group + assert not missing, ( + f"typer aliases {missing!r} are not in the backend denylist; " + f"add them to _DENYLIST_GROUPS to keep /load from desyncing " + f"llama_parallel_slots." + ) + + +# test_in_venv_path_passes_parallel_to_run_server (below) is the runtime +# equivalent of the retired source-text guard for hardcoded +# `llama_parallel_slots = 4`. + + +# Re-exec arg-builder coverage. run() re-execs into the studio venv +# (execvp on POSIX, Popen on Windows). Without explicit forwarding the +# child reverts to typer defaults and silently drops the user's value. + + +class _ExecCaptured(SystemExit): + def __init__(self, argv): + super().__init__(0) + self.argv = list(argv) + + +def _install_reexec_capture(monkeypatch, *, platform): + studio_mod = _load_run_command() + captured = [] + + monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") + + fake_venv = Path("/fake/studio/venv/unsloth_studio") + fake_python = fake_venv / "bin" / "python" + fake_bin = fake_venv / "bin" / "unsloth" + monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_python) + + real_is_file = Path.is_file + monkeypatch.setattr( + Path, + "is_file", + lambda self: True if str(self) == str(fake_bin) else real_is_file(self), + ) + + # resolve_tool_policy is imported lazily inside run(); patch the source. + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + + monkeypatch.setattr(sys, "platform", platform) + + def fake_execvp(file, argv): + captured.append({"kind": "execvp", "argv": list(argv)}) + raise _ExecCaptured(argv) + + class _FakePopen: + def __init__(self, argv, *a, **kw): + captured.append({"kind": "popen", "argv": list(argv)}) + self._argv = argv + + def wait(self): + raise _ExecCaptured(self._argv) + + monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp) + monkeypatch.setattr(studio_mod.subprocess, "Popen", _FakePopen) + + return captured + + +def _invoke_run(monkeypatch, args, *, platform = "linux"): + import typer as _typer + + studio_mod = _load_run_command() + captured = _install_reexec_capture(monkeypatch, platform = platform) + app = _typer.Typer() + app.command( + context_settings = { + "allow_extra_args": True, + "ignore_unknown_options": True, + }, + )(studio_mod.run) + result = CliRunner().invoke(app, args, catch_exceptions = True) + return result, captured + + +def _value_after(argv, flag): + for i, tok in enumerate(argv): + if tok == flag and i + 1 < len(argv): + return argv[i + 1] + return None + + +_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"] + + +@pytest.mark.parametrize( + "flag,value", + [("--parallel", "8"), ("--n-parallel", "16"), ("-np", "32")], +) +def test_reexec_forwards_parallel_all_aliases(monkeypatch, flag, value): + """Every alias the user can type must reach the re-exec'd child.""" + result, captured = _invoke_run(monkeypatch, _BASE + [flag, value]) + assert ( + len(captured) == 1 + ), f"expected one launch via re-exec, got {captured}; output={result.output!r}" + argv = captured[0]["argv"] + assert ( + _value_after(argv, "--parallel") == value + ), f"{flag} {value} was dropped on re-exec; argv = {argv}" + + +@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.""" + result, captured = _invoke_run( + monkeypatch, _BASE + ["--parallel", "12"], platform = platform + ) + assert len(captured) == 1 + expected_kind = "popen" if platform == "win32" else "execvp" + assert ( + captured[0]["kind"] == expected_kind + ), f"{platform}: expected launcher {expected_kind}, got {captured[0]['kind']}" + assert _value_after(captured[0]["argv"], "--parallel") == "12" + + +def test_reexec_np_is_first_class_alias(monkeypatch): + """`-np` must reach the child as --parallel . Pre-PR Click + clustered `-np 8` as `-p 8` (port=8) + stray `-n`; also pin that + --port is no longer collateral damage.""" + result, captured = _invoke_run(monkeypatch, _BASE + ["-np", "8"]) + assert len(captured) == 1 + argv = captured[0]["argv"] + assert ( + _value_after(argv, "--parallel") == "8" + ), f"-np 8 silently became 4 after re-exec; argv = {argv}" + # `-np 8` must not clobber --port (default 8888). + assert _value_after(argv, "--port") == "8888", argv + + +def test_reexec_mixed_parallel_with_passthrough(monkeypatch): + """--parallel + llama-server pass-through flags must all reach the child.""" + result, captured = _invoke_run( + monkeypatch, + _BASE + ["--parallel", "8", "--top-k", "20", "--temp", "0.7"], + ) + assert len(captured) == 1 + argv = captured[0]["argv"] + assert _value_after(argv, "--parallel") == "8", argv + assert _value_after(argv, "--top-k") == "20", argv + assert _value_after(argv, "--temp") == "0.7", argv + + +@pytest.mark.parametrize( + "user_flag,expected_in_child", + [ + ("--load-in-4bit", "--load-in-4bit"), + ("--no-load-in-4bit", "--no-load-in-4bit"), + (None, "--load-in-4bit"), # default True + ], +) +def test_reexec_forwards_load_in_4bit_in_both_directions( + monkeypatch, user_flag, expected_in_child +): + """Re-exec must emit the chosen polarity (or the typer default), + so a future default flip on one layer can't silently invert + behaviour for users who never typed the flag.""" + extras = [user_flag] if user_flag else [] + result, captured = _invoke_run(monkeypatch, _BASE + extras) + assert len(captured) == 1 + argv = captured[0]["argv"] + other_polarity = ( + "--no-load-in-4bit" + if expected_in_child == "--load-in-4bit" + else "--load-in-4bit" + ) + assert ( + expected_in_child in argv + ), f"expected {expected_in_child} in child argv; got {argv}" + assert ( + other_polarity not in argv + ), f"unexpected {other_polarity} in child argv; got {argv}" + + +# Runtime check: fake sys.prefix into the studio venv to bypass +# re-exec, then assert run_server receives --parallel as +# llama_parallel_slots. + + +class _RunServerCaptured(SystemExit): + def __init__(self, kwargs): + super().__init__(0) + self.kwargs = dict(kwargs) + + +def _types_module(name): + import types as _types + + return _types.ModuleType(name) + + +def test_studio_default_rejects_parallel_when_subcommand_invoked(): + """`unsloth studio --parallel 8 run ...` would silently drop the 8 + (typer doesn't forward parent options to subcommands). The + callback rejects with exit 2 and points at the subcommand flag.""" + studio_mod = _load_run_command() + import typer as _typer + + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + + runner = CliRunner() + result = runner.invoke(app, ["studio", "--parallel", "8", "run", "--model", "X"]) + assert result.exit_code == 2, ( + f"expected exit 2 when --parallel is on studio group with a " + f"subcommand invoked; got {result.exit_code}; output={result.output!r}" + ) + combined = (result.output or "") + (getattr(result, "stderr", "") or "") + assert "--parallel" in combined, combined + assert ( + "run --parallel 8" in combined + ), f"error message must show the corrected invocation; got: {combined}" + + +def test_studio_default_default_parallel_with_subcommand_does_not_error(): + """Omitting --parallel on the group must still let subcommands + run; the group's default 1 is benign.""" + studio_mod = _load_run_command() + import typer as _typer + + app = _typer.Typer() + app.add_typer(studio_mod.studio_app, name = "studio") + runner = CliRunner() + result = runner.invoke(app, ["studio", "--help"]) + assert result.exit_code == 0, result.output + + +def test_studio_default_exposes_parallel_option(): + """Plain `unsloth studio` exposes --parallel too so the API-only + path can raise concurrency without going through the denied + pass-through. Default stays at 1 (pre-PR); `run` keeps its 4.""" + studio_mod = _load_run_command() + import inspect + + sig = inspect.signature(studio_mod.studio_default) + assert "parallel" in sig.parameters, ( + "studio_default missing `parallel`; API-only path can't set " + "llama_parallel_slots" + ) + opt = sig.parameters["parallel"].default + decls = set(getattr(opt, "param_decls", []) or []) + assert "--parallel" in decls + assert "--n-parallel" in decls + assert ( + getattr(opt, "default", None) == 1 + ), "studio_default --parallel must default to 1 (pre-PR); `run` is 4" + assert getattr(opt, "min", None) == 1 + assert getattr(opt, "max", None) == 64 + + +@pytest.mark.parametrize("value", [1, 4, 8, 64]) +def test_in_venv_path_passes_parallel_to_run_server(monkeypatch, value): + """In-venv path must forward --parallel to + run_server(llama_parallel_slots=N), not the old hardcoded 4.""" + studio_mod = _load_run_command() + + fake_venv = Path("/fake/studio/venv/unsloth_studio") + monkeypatch.setattr(sys, "prefix", str(fake_venv)) + # Pin STUDIO_HOME so sys.prefix.startswith() picks the in-venv branch. + monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent) + + from unsloth_cli import _tool_policy as _tp_mod + + monkeypatch.setattr( + _tp_mod, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + + captured: dict = {} + + def fake_run_server(**kwargs): + captured.update(kwargs) + raise _RunServerCaptured(kwargs) + + fake_backend_run = sys.modules.setdefault( + "studio.backend.run", _types_module("studio.backend.run") + ) + fake_backend_run.run_server = fake_run_server + fake_backend_run._resolve_external_ip = lambda: "127.0.0.1" + + import typer as _typer + + app = _typer.Typer() + app.command( + context_settings = { + "allow_extra_args": True, + "ignore_unknown_options": True, + }, + )(studio_mod.run) + CliRunner().invoke(app, _BASE + ["--parallel", str(value)], catch_exceptions = True) + + assert ( + captured.get("llama_parallel_slots") == value + ), f"run_server got llama_parallel_slots={captured.get('llama_parallel_slots')!r}, expected {value}" diff --git a/unsloth_cli/tests/test_studio_run_short_alias_clashes.py b/unsloth_cli/tests/test_studio_run_short_alias_clashes.py new file mode 100644 index 000000000..8a3e94db4 --- /dev/null +++ b/unsloth_cli/tests/test_studio_run_short_alias_clashes.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for short-alias clashes with llama-server flags. + +`unsloth studio run` passes unknown flags through to llama-server. +Pre-cleanup it exposed 1-char shorts ``-m`` / ``-f`` plus ``-hfr``; +Click clustered llama-server tokens against them (``-fa`` -> ``-f a``, +``-mg 0`` -> ``-m g``, ``-fitt 1024`` -> ``-f itt``, ...), silently +breaking ~11 pass-through flags. + +The cleanup drops ``-m``, ``-f``, ``-hfr``. The 2-char ``-hf`` stays +(documented; multi-char shorts don't cluster). Long forms remain. +``studio_default`` keeps ``-f`` because it has no pass-through. + +See ``test_studio_run_parallel_flag.py`` for ``--parallel`` / +``-np`` coverage and re-exec forwarding. +""" + +from __future__ import annotations + +import inspect +import sys +from pathlib import Path + +import pytest +from typer.testing import CliRunner + + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _studio_mod(): + from unsloth_cli.commands import studio as _s + + return _s + + +def _decls_for(param_name): + sig = inspect.signature(_studio_mod().run) + opt = sig.parameters[param_name].default + return set(getattr(opt, "param_decls", []) or []) + + +# Surface checks: removed shorts must not reappear. + + +def test_model_short_aliases_removed(): + """`-m` / `-hfr` removed from --model; `-hf` kept (multi-char, + doesn't cluster).""" + decls = _decls_for("model") + assert "-m" not in decls, "`-m` re-added; brings back `-mg`/`-md` clustering" + assert "-hfr" not in decls, "`-hfr` was re-added; remove it" + assert "--model" in decls + assert "--hf-repo" in decls + assert "-hf" in decls, "`-hf` is documented and must keep working" + + +def test_frontend_short_alias_removed_from_run(): + """`-f` must not be on `run` (eats `-fa`/`-fit`/`-fitt`/`-fitc`).""" + decls = _decls_for("frontend") + assert "-f" not in decls, "`-f` re-added on run(); brings back `-fa` clustering" + assert "--frontend" in decls + + +def test_studio_default_keeps_dash_f(): + """`studio_default` keeps `-f`: no pass-through tail to clash with.""" + sig = inspect.signature(_studio_mod().studio_default) + opt = sig.parameters["frontend"].default + decls = set(getattr(opt, "param_decls", []) or []) + assert "-f" in decls + + +# Behaviour checks: llama-server shorts must reach the child verbatim. + + +class _ExecCaptured(SystemExit): + def __init__(self, argv): + super().__init__(0) + self.argv = list(argv) + + +def _install_capture(monkeypatch): + studio_mod = _studio_mod() + captured = [] + monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv") + fake_bin = Path("/fake/studio/venv/unsloth_studio/bin/unsloth") + monkeypatch.setattr( + studio_mod, "_studio_venv_python", lambda: fake_bin.parent / "python" + ) + real_is_file = Path.is_file + monkeypatch.setattr( + Path, + "is_file", + lambda self: True if str(self) == str(fake_bin) else real_is_file(self), + ) + from unsloth_cli import _tool_policy as _tp + + monkeypatch.setattr( + _tp, + "resolve_tool_policy", + lambda host, flag, yes, silent: False if flag is None else bool(flag), + ) + monkeypatch.setattr(sys, "platform", "linux") + + def fake_execvp(file, argv): + captured.append(list(argv)) + raise _ExecCaptured(argv) + + monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp) + return captured + + +def _invoke(monkeypatch, args): + import typer as _typer + + studio_mod = _studio_mod() + captured = _install_capture(monkeypatch) + app = _typer.Typer() + app.command( + context_settings = { + "allow_extra_args": True, + "ignore_unknown_options": True, + }, + )(studio_mod.run) + CliRunner().invoke(app, args, catch_exceptions = True) + return captured + + +# (short_flag, value, llama-server long name). All were silently +# mis-parsed pre-cleanup. +_PREVIOUSLY_BROKEN = [ + ("-fa", None, "--flash-attn"), + ("-fit", None, "--fit"), + ("-fitt", "1024", "--fit-target"), + ("-fitc", "4096", "--fit-ctx"), + ("-mg", "0", "--main-gpu"), + ("-md", "/path/draft.gguf", "--spec-draft-model"), + ("-hff", "Q4_K_M.gguf", "--hf-file"), + ("-cmoe", None, "--cpu-moe"), + ("-cram", "16384", "--cache-ram"), + ("-sm", "row", "--split-mode"), + ("-ncmoe", "8", "--n-cpu-moe"), +] + + +@pytest.mark.parametrize("flag,value,llama_long_name", _PREVIOUSLY_BROKEN) +def test_previously_broken_short_flag_now_passes_through( + monkeypatch, + flag, + value, + llama_long_name, +): + """Each of these was eaten by typer pre-cleanup; must pass through verbatim now.""" + extras = [flag] if value is None else [flag, value] + captured = _invoke(monkeypatch, ["--model", "X"] + extras) + assert len(captured) == 1, f"parent did not re-exec for {extras}" + argv = captured[0] + assert flag in argv, ( + f"llama-server short flag {flag!r} ({llama_long_name}) was eaten " + f"by typer; child argv = {argv}" + ) + if value is not None: + idx = argv.index(flag) + assert ( + idx + 1 < len(argv) and argv[idx + 1] == value + ), f"value for {flag!r} was lost or moved; argv = {argv}" + + +def test_dash_hf_documented_alias_still_works(monkeypatch): + """`-hf` is documented and must keep working (multi-char shorts + don't cluster in Click).""" + captured = _invoke( + monkeypatch, + ["-hf", "unsloth/gemma-4-26B-A4B-it-GGUF:UD-Q4_K_XL"], + ) + assert len(captured) == 1 + argv = captured[0] + # `_split_repo_variant` peels the `:variant` suffix before re-exec. + assert argv[argv.index("--model") + 1] == ("unsloth/gemma-4-26B-A4B-it-GGUF"), argv + assert argv[argv.index("--gguf-variant") + 1] == "UD-Q4_K_XL", argv + + +# Legacy `-m` / `-hfr` / `-f` were typer aliases pre-PR. The +# preprocessor promotes EXACT matches back to their typer params and +# leaves clustered tokens (`-mg`, `-fa`, ...) in the pass-through tail. + + +@pytest.mark.parametrize( + "legacy_args,expected_model", + [ + (["-m", "unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + (["-m=unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + (["-hfr", "unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + (["-hfr=unsloth/Qwen3-1.7B-GGUF"], "unsloth/Qwen3-1.7B-GGUF"), + ], +) +def test_legacy_model_aliases_still_promote_to_model( + monkeypatch, + legacy_args, + expected_model, +): + """Pre-PR `-m X` / `-hfr X` set --model X; preprocessor preserves that.""" + captured = _invoke(monkeypatch, legacy_args) + assert len(captured) == 1, f"parent did not re-exec for {legacy_args}" + argv = captured[0] + assert argv[argv.index("--model") + 1] == expected_model, argv + # Promoted alias must not also leak into the pass-through tail. + for alias in ("-m", "-hfr"): + if alias in legacy_args: + assert alias not in argv, f"legacy {alias} leaked into child argv: {argv}" + + +def test_legacy_frontend_alias_still_promotes_to_frontend(monkeypatch): + """Pre-PR `-f dist` set --frontend dist; preprocessor preserves it.""" + captured = _invoke(monkeypatch, ["--model", "X", "-f", "/tmp/dist"]) + assert len(captured) == 1 + argv = captured[0] + # Compare via Path so Windows's str(Path("/tmp/dist")) = "\tmp\dist" + # doesn't trip the assertion on the same logical path. + assert Path(argv[argv.index("--frontend") + 1]) == Path("/tmp/dist"), argv + assert "-f" not in argv, f"-f leaked into child argv: {argv}" + + +def test_legacy_model_alias_conflicts_with_long_form(monkeypatch): + """`--model X` plus `-m Y` is ambiguous; must error pre-re-exec.""" + captured = _invoke(monkeypatch, ["--model", "X", "-m", "Y"]) + assert ( + len(captured) == 0 + ), f"expected error before re-exec, got launch with argv = {captured}" + + +def test_clustered_tokens_are_not_promoted(monkeypatch): + """`-mg` / `-fa` / `-fitt` are llama-server flags and must survive + in the tail even though they start with `-m` / `-f`.""" + captured = _invoke( + monkeypatch, + ["--model", "X", "-mg", "0", "-fa", "-fitt", "1024"], + ) + assert len(captured) == 1 + argv = captured[0] + assert argv[argv.index("--model") + 1] == "X", argv + for flag in ("-mg", "-fa", "-fitt"): + assert flag in argv, f"{flag!r} was promoted instead of passed through: {argv}" + + +def test_legacy_m_with_repo_variant_syntax(monkeypatch): + """`-m repo:variant` must round-trip through preprocessor + + _split_repo_variant into --model + --gguf-variant.""" + captured = _invoke( + monkeypatch, + ["-m", "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL"], + ) + assert len(captured) == 1 + argv = captured[0] + assert argv[argv.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF", argv + assert argv[argv.index("--gguf-variant") + 1] == "UD-Q4_K_XL", argv + + +def test_missing_model_after_preprocessor_errors(monkeypatch): + """Neither --model nor a legacy alias → clean exit(2) before re-exec.""" + captured = _invoke(monkeypatch, ["--parallel", "8"]) + assert ( + len(captured) == 0 + ), f"expected exit before re-exec, got launch with argv = {captured}" + + +def test_legacy_m_inline_value_form(monkeypatch): + """`-m=foo` is promoted like `-m foo`.""" + captured = _invoke(monkeypatch, ["-m=unsloth/Qwen3-1.7B-GGUF"]) + assert len(captured) == 1 + argv = captured[0] + assert argv[argv.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF", argv + + +# Unit tests for _consume_legacy_short_aliases. + + +def test_consume_helper_exact_match_space_form(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-m", "FOO", "--top-k", "20"], + ("-m",), + None, + "--model", + ) + assert value == "FOO" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_exact_match_inline_form(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-m=FOO", "--top-k", "20"], + ("-m",), + None, + "--model", + ) + assert value == "FOO" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_leaves_clusters_alone(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-mg", "0", "-md", "/x"], + ("-m",), + None, + "--model", + ) + assert value is None + assert remaining == ["-mg", "0", "-md", "/x"] + + +def test_consume_helper_value_already_set_raises(): + helper = _studio_mod()._consume_legacy_short_aliases + import typer as _typer + + with pytest.raises(_typer.BadParameter): + helper(["-m", "Y"], ("-m",), "X", "--model") + + +def test_consume_helper_missing_value_raises(): + helper = _studio_mod()._consume_legacy_short_aliases + import typer as _typer + + with pytest.raises(_typer.BadParameter): + helper(["-m"], ("-m",), None, "--model") + + +def test_consume_helper_multiple_aliases_in_group(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["-hfr", "FOO", "--top-k", "20"], + ("-m", "-hfr"), + None, + "--model", + ) + assert value == "FOO" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_preserves_value_when_no_match(): + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["--top-k", "20"], + ("-m",), + "PRESET", + "--model", + ) + assert value == "PRESET" + assert remaining == ["--top-k", "20"] + + +# `-p` is typer short for --port, so Click clusters `-np8` as `-n -p 8` +# (port=8, parallel dropped). The rewrite splits to `-np 8` pre-parse. + + +def test_expand_np_rewrites_attached_form(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "studio", "run", "--model", "X", "-np8"], + ) + _studio_mod()._expand_attached_np_short() + assert sys.argv == [ + "unsloth", + "studio", + "run", + "--model", + "X", + "-np", + "8", + ] + + +@pytest.mark.parametrize("value", ["1", "8", "64", "999"]) +def test_expand_np_rewrites_all_digit_values(monkeypatch, value): + monkeypatch.setattr(sys, "argv", ["unsloth", "studio", "run", f"-np{value}"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "studio", "run", "-np", value] + + +def test_expand_np_leaves_space_form_alone(monkeypatch): + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-np", "8"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", "8"] + + +def test_expand_np_leaves_equals_form_alone(monkeypatch): + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-np=8"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np=8"] + + +def test_expand_np_leaves_non_digit_suffix_alone(monkeypatch): + # `-npfoo` isn't a numeric attached value; let typer reject it. + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-npfoo"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-npfoo"] + + +def test_expand_np_leaves_bare_np_alone(monkeypatch): + monkeypatch.setattr(sys, "argv", ["unsloth", "run", "-np"]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np"] + + +def test_expand_np_handles_multiple_occurrences(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "run", "-np8", "-np16"], + ) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", "8", "-np", "16"] + + +@pytest.mark.parametrize("attached,expected", [("-np-1", "-1"), ("-np+1", "+1")]) +def test_expand_np_handles_signed_attached_forms(monkeypatch, attached, expected): + """Signed `-np-1` / `-np+1` must split too, else Click sets port=-1.""" + monkeypatch.setattr(sys, "argv", ["unsloth", "run", attached]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", expected] + + +@pytest.mark.parametrize( + "attached,expected_suffix", + [("-np8x", "8x"), ("-np-1foo", "-1foo"), ("-np9bar", "9bar")], +) +def test_expand_np_rewrites_numeric_prefix_even_with_junk( + monkeypatch, attached, expected_suffix +): + """`-np8x` would surface as a baffling --port error; rewriting to + `-np 8x` makes typer report against `-np` where it was typed.""" + monkeypatch.setattr(sys, "argv", ["unsloth", "run", attached]) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "-np", expected_suffix] + + +def test_consume_helper_rejects_empty_inline_value(): + """`-m=` must error, not silently become --model ''.""" + import typer as _typer + + helper = _studio_mod()._consume_legacy_short_aliases + with pytest.raises(_typer.BadParameter, match = "non-empty"): + helper(["-m="], ("-m",), None, "--model") + + +# Gate isolation: importing unsloth_cli from a third-party script must +# leave its sys.argv intact. Pins the narrow basename set. + + +@pytest.mark.parametrize( + "third_party_argv0", + [ + "/home/user/myproj/cli.py", + "cli.py", + "/usr/bin/some-tool", + "pytest", + "/opt/wrapper/launch.py", + "unsloth-cli", + "unsloth-cli.py", + ], +) +def test_third_party_importers_do_not_trigger_np_rewrite( + monkeypatch, third_party_argv0 +): + """Only the `unsloth` / `unsloth.exe` console-script may run the + canonicaliser; third-party scripts must keep their argv intact.""" + import os as _os + import importlib + + starting_argv = [third_party_argv0, "subcmd", "-np8", "--input", "foo"] + monkeypatch.setattr(sys, "argv", list(starting_argv)) + # Force a fresh import so the import-time gate actually runs. + monkeypatch.delitem(sys.modules, "unsloth_cli", raising = False) + importlib.import_module("unsloth_cli") + assert sys.argv == starting_argv, ( + f"third-party argv[0]={third_party_argv0!r} triggered the " + f"unsloth -np canonicaliser; sys.argv was mutated to {sys.argv}" + ) + _ = _os # silence unused-import linters when monkeypatch lazy-binds + + +def test_attached_np8_no_longer_silently_sets_port(monkeypatch): + """After the gate runs, `-np8` produces --parallel=8 (not --port=8).""" + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "studio", "run", "--model", "X", "-np8"], + ) + _studio_mod()._expand_attached_np_short() + captured = _invoke(monkeypatch, sys.argv[2:]) # drop "unsloth studio" + assert len(captured) == 1, "parent did not re-exec" + argv = captured[0] + assert argv[argv.index("--parallel") + 1] == "8", argv + assert argv[argv.index("--port") + 1] == "8888", argv + + +def test_expand_np_stops_at_double_dash(monkeypatch): + """Tokens after `--` are positional; `-np8` stays raw.""" + monkeypatch.setattr( + sys, + "argv", + ["unsloth", "run", "--model", "X", "--", "-np8"], + ) + _studio_mod()._expand_attached_np_short() + assert sys.argv == ["unsloth", "run", "--model", "X", "--", "-np8"] + + +def test_consume_helper_stops_at_double_dash(): + """Alias promotion must not reach past `--`.""" + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper( + ["--top-k", "20", "--", "-m", "FOO"], + ("-m",), + None, + "--model", + ) + assert value is None + assert remaining == ["--top-k", "20", "--", "-m", "FOO"] + + +def test_consume_helper_rejects_long_flag_as_value(): + """`-m --flash-attn` errors; `--xxx` is unambiguously a flag.""" + import typer as _typer + + helper = _studio_mod()._consume_legacy_short_aliases + with pytest.raises(_typer.BadParameter, match = "--flash-attn"): + helper(["-m", "--flash-attn"], ("-m",), None, "--model") + + +def test_consume_helper_allows_bare_dash_as_value(): + """Lone `-` is a stdin/path sentinel, not a flag.""" + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper(["-m", "-", "--top-k", "20"], ("-m",), None, "--model") + assert value == "-" + assert remaining == ["--top-k", "20"] + + +def test_consume_helper_allows_short_dash_value(): + """`-foo` may be a path or a leading-dash model name; only `--long` + tokens are rejected as values.""" + helper = _studio_mod()._consume_legacy_short_aliases + value, remaining = helper(["-m", "-foo", "--top-k", "20"], ("-m",), None, "--model") + assert value == "-foo" + assert remaining == ["--top-k", "20"]