mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Studio: allow --secure with --api-only (headless secure API server) and add --api-only to unsloth studio run (#6591)
* Studio: start the Cloudflare tunnel for --secure even in --api-only, and add --api-only to `unsloth studio run` --secure exposes ONLY the Cloudflare link (it forces a loopback bind), but _cloudflare_tunnel_should_start gated the tunnel on `not api_only`, so `run.py --secure --api-only` started no tunnel and then fail-closed with "A secure Cloudflare link is not allowed". That blocked the natural headless use: serve just the API (no web UI) over the authenticated tunnel. Make --secure start the tunnel regardless of api_only (the non-secure path is unchanged: tunnel only a 0.0.0.0 bind, never api-only Tauri or Colab). Then expose --api-only on `unsloth studio run` and forward it through both the re-exec args and the in-venv run_server call, so `unsloth studio run --secure --api-only --model ...` is a one-liner secure API server. Verified end to end: `run.py --secure --api-only` now brings up the tunnel and serves /api/health over it (200), with / returning 404 (no UI). Tests: update the tunnel-gate truth table (secure+api-only now tunnels; secure+colab still does not) and add --api-only registration + re-exec/in-venv forwarding coverage to the run CLI tests. * Trim comments to be succinct (no behavior change) * studio: address review on parent --api-only and secure api-only CORS - Reject --api-only on the parent `unsloth studio` group when a subcommand is invoked, with the same redirect guidance used for --parallel/--secure; otherwise the flag was silently dropped and the UI served anyway. - Keep CORS any-origin for secure api-only serving: that mode publishes the API over Cloudflare for remote browser clients, so the Tauri-only lockdown (still applied to plain local api-only) would break preflight. Factored the decision into cors_origins_for_mode() and gate it on api_only and not secure; run_server exports UNSLOTH_SECURE before importing main. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: suppress TAURI_PORT and de-dup test for headless run --api-only - run_server gains emit_tauri_port (default True, unchanged for the Tauri/ desktop path). The new headless `run --api-only` path passes False so the Tauri-only TAURI_PORT= line no longer prepends the documented URL/API key banner (it ran even under --silent and could break one-liner parsers). - Remove a duplicate test_reexec_forwards_api_only that shadowed the parametrized one; fold the --secure --api-only case into it so the secure headless path is actually collected. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
7bac461471
commit
76cbddb859
6 changed files with 216 additions and 21 deletions
|
|
@ -855,24 +855,16 @@ async def _recipes_redirect(rest: str = ""):
|
|||
return _RedirectResponse(url = target, status_code = 308)
|
||||
|
||||
|
||||
_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1"
|
||||
_cors_origins = ["*"]
|
||||
if _api_only:
|
||||
_cors_origins = [
|
||||
"tauri://localhost", # Linux/macOS Tauri webview
|
||||
"http://tauri.localhost", # Windows Tauri webview
|
||||
"http://localhost", # dev fallback
|
||||
"http://localhost:5173", # Tauri dev/Vite
|
||||
"http://127.0.0.1:5173", # Tauri dev/Vite fallback
|
||||
]
|
||||
_cors_origin_regex = None
|
||||
else:
|
||||
_cors_origin_regex = None
|
||||
from utils.host_policy import cors_origins_for_mode # noqa: E402
|
||||
|
||||
_cors_origins = cors_origins_for_mode(
|
||||
api_only = os.environ.get("UNSLOTH_API_ONLY") == "1",
|
||||
secure = os.environ.get("UNSLOTH_SECURE") == "1",
|
||||
)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins = _cors_origins,
|
||||
allow_origin_regex = _cors_origin_regex,
|
||||
allow_credentials = True,
|
||||
allow_methods = ["*"],
|
||||
allow_headers = ["*"],
|
||||
|
|
|
|||
|
|
@ -893,9 +893,14 @@ def _setup_server_disk_logging():
|
|||
def _cloudflare_tunnel_should_start(
|
||||
*, cloudflare: bool, host: str, secure: bool, api_only: bool, is_colab: bool
|
||||
) -> bool:
|
||||
"""Whether to start the Cloudflare tunnel. --secure tunnels a loopback bind too;
|
||||
non-secure keeps the 0.0.0.0-only rule. Colab/api-only never tunnel."""
|
||||
return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab
|
||||
"""Whether to start the Cloudflare tunnel. --secure exposes only the tunnel
|
||||
(loopback bind), so it tunnels even api-only (headless secure API serving);
|
||||
otherwise tunnel only a 0.0.0.0 bind, never api-only (Tauri) or Colab."""
|
||||
if is_colab or not cloudflare:
|
||||
return False
|
||||
if secure:
|
||||
return True
|
||||
return host == "0.0.0.0" and not api_only
|
||||
|
||||
|
||||
def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
|
||||
|
|
@ -919,6 +924,7 @@ def run_server(
|
|||
cloudflare: bool = True,
|
||||
secure: bool = False,
|
||||
enable_tools: "Optional[bool]" = None,
|
||||
emit_tauri_port: bool = True,
|
||||
):
|
||||
"""
|
||||
Start the FastAPI server.
|
||||
|
|
@ -932,6 +938,9 @@ def run_server(
|
|||
llama_parallel_slots: parallel slots for llama-server
|
||||
enable_tools: explicit --enable-tools/--disable-tools policy; None leaves
|
||||
the default (tools on, per-request enable_tools honored)
|
||||
emit_tauri_port: print the machine-readable TAURI_PORT line the desktop
|
||||
app parses from stdout; the headless `run --api-only` path turns it
|
||||
off so it does not pollute the documented URL/API-key banner
|
||||
|
||||
Note:
|
||||
Signal handlers are NOT registered here so embedders (e.g. Colab) keep
|
||||
|
|
@ -974,9 +983,13 @@ def run_server(
|
|||
if _session_log is not None and not silent:
|
||||
print(f"Session log: {_session_log}")
|
||||
|
||||
# Set env var BEFORE importing main so CORS middleware picks it up.
|
||||
# Set env vars BEFORE importing main so CORS middleware picks them up.
|
||||
# secure api-only is a remote server behind Cloudflare, so it keeps the
|
||||
# any-origin CORS profile; plain api-only stays locked to the Tauri app.
|
||||
if api_only:
|
||||
os.environ["UNSLOTH_API_ONLY"] = "1"
|
||||
if secure:
|
||||
os.environ["UNSLOTH_SECURE"] = "1"
|
||||
|
||||
import nest_asyncio
|
||||
|
||||
|
|
@ -1158,7 +1171,8 @@ def run_server(
|
|||
atexit.register(terminate_all)
|
||||
|
||||
# Output port for Tauri (api-only), only after sockets bind and startup done.
|
||||
if api_only:
|
||||
# The headless `run --api-only` path opts out so it does not leak this line.
|
||||
if api_only and emit_tauri_port:
|
||||
print(f"TAURI_PORT={port}", flush = True)
|
||||
|
||||
# Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often
|
||||
|
|
|
|||
|
|
@ -31,11 +31,14 @@ from run import _cloudflare_tunnel_should_start as should_start # noqa: E402
|
|||
# --no-cloudflare always wins.
|
||||
(False, "0.0.0.0", False, False, False, False),
|
||||
(False, "127.0.0.1", True, False, False, False),
|
||||
# api-only and Colab never tunnel.
|
||||
# Non-secure api-only never tunnels (Tauri).
|
||||
(True, "0.0.0.0", False, True, False, False),
|
||||
(True, "127.0.0.1", True, True, False, False),
|
||||
# --secure tunnels even api-only (headless secure API server).
|
||||
(True, "127.0.0.1", True, True, False, True),
|
||||
# Colab never tunnels, even --secure.
|
||||
(True, "0.0.0.0", False, False, True, False),
|
||||
(True, "127.0.0.1", True, False, True, False),
|
||||
(True, "127.0.0.1", True, True, True, False),
|
||||
],
|
||||
)
|
||||
def test_cloudflare_gate(cloudflare, host, secure, api_only, is_colab, expected):
|
||||
|
|
@ -162,3 +165,46 @@ def test_failclosed_message_present_in_source():
|
|||
"A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link"
|
||||
in src
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_only,secure,expected",
|
||||
[
|
||||
(False, False, ["*"]), # plain server: any origin
|
||||
(False, True, ["*"]), # secure UI server: any origin
|
||||
(True, True, ["*"]), # secure api-only: remote browsers need any origin
|
||||
(True, False, "tauri"), # local api-only: locked to the Tauri app
|
||||
],
|
||||
)
|
||||
def test_cors_origins_for_mode(api_only, secure, expected):
|
||||
from utils.host_policy import cors_origins_for_mode
|
||||
origins = cors_origins_for_mode(api_only = api_only, secure = secure)
|
||||
if expected == "tauri":
|
||||
assert origins != ["*"] and any(o.startswith("tauri://") for o in origins)
|
||||
else:
|
||||
assert origins == expected
|
||||
|
||||
|
||||
def test_run_server_exports_secure_env_for_cors():
|
||||
# run_server must export UNSLOTH_SECURE before importing main so the CORS
|
||||
# profile can tell remote secure serving from local Tauri use.
|
||||
src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
|
||||
assert 'os.environ["UNSLOTH_SECURE"] = "1"' in src
|
||||
|
||||
|
||||
def test_run_server_emit_tauri_port_defaults_on():
|
||||
# Default on keeps the desktop app's stdout contract; the headless
|
||||
# `run --api-only` path opts out explicitly.
|
||||
import inspect
|
||||
|
||||
import run
|
||||
|
||||
params = inspect.signature(run.run_server).parameters
|
||||
assert "emit_tauri_port" in params
|
||||
assert params["emit_tauri_port"].default is True
|
||||
|
||||
|
||||
def test_tauri_port_print_is_gated_in_source():
|
||||
# The TAURI_PORT line must depend on emit_tauri_port, not api_only alone.
|
||||
src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
|
||||
assert "if api_only and emit_tauri_port:" in src
|
||||
|
|
|
|||
|
|
@ -34,6 +34,26 @@ def is_external_host(host: str) -> bool:
|
|||
return host.lower() not in _LOOPBACK_HOSTS
|
||||
|
||||
|
||||
# Tauri desktop webview origins. api-only serving (the desktop app calling a
|
||||
# local backend) locks CORS to these.
|
||||
_TAURI_CORS_ORIGINS = (
|
||||
"tauri://localhost", # Linux/macOS Tauri webview
|
||||
"http://tauri.localhost", # Windows Tauri webview
|
||||
"http://localhost", # dev fallback
|
||||
"http://localhost:5173", # Tauri dev/Vite
|
||||
"http://127.0.0.1:5173", # Tauri dev/Vite fallback
|
||||
)
|
||||
|
||||
|
||||
def cors_origins_for_mode(*, api_only: bool, secure: bool) -> list[str]:
|
||||
"""Allowed CORS origins. Default is any-origin (["*"]); api-only locks down
|
||||
to the Tauri desktop app, except in secure mode where the API is published
|
||||
over Cloudflare and must stay reachable from remote browser origins."""
|
||||
if api_only and not secure:
|
||||
return list(_TAURI_CORS_ORIGINS)
|
||||
return ["*"]
|
||||
|
||||
|
||||
def apply_stdio_mcp_loopback_default(host: str, *, is_colab: bool = False) -> None:
|
||||
"""Default stdio MCP servers on when bound to loopback.
|
||||
|
||||
|
|
|
|||
|
|
@ -776,6 +776,16 @@ def studio_default(
|
|||
err = True,
|
||||
)
|
||||
raise typer.Exit(2)
|
||||
# Same for --api-only: dropping it here would silently serve the UI.
|
||||
if api_only:
|
||||
typer.echo(
|
||||
f"Error: --api-only on `unsloth studio` applies to the "
|
||||
f"plain-server path only. For `unsloth studio "
|
||||
f"{ctx.invoked_subcommand}`, put it after the subcommand: "
|
||||
f"`unsloth studio {ctx.invoked_subcommand} --api-only ...`",
|
||||
err = True,
|
||||
)
|
||||
raise typer.Exit(2)
|
||||
return
|
||||
|
||||
# --secure requires the tunnel; force a loopback bind.
|
||||
|
|
@ -1031,6 +1041,12 @@ def run(
|
|||
host: str = typer.Option("127.0.0.1", "--host", "-H"),
|
||||
# `-f` removed (clustered `-fa`/`-fit*`); studio_default keeps it.
|
||||
frontend: Optional[Path] = typer.Option(None, "--frontend"),
|
||||
api_only: bool = typer.Option(
|
||||
False,
|
||||
"--api-only",
|
||||
help = "Serve only the API (no web UI), for a headless model server. "
|
||||
"Pairs with --secure to expose the API over the Cloudflare link alone.",
|
||||
),
|
||||
silent: bool = typer.Option(False, "--silent", "-q"),
|
||||
enable_tools: Optional[bool] = typer.Option(
|
||||
None,
|
||||
|
|
@ -1214,6 +1230,8 @@ def run(
|
|||
args.append("--load-in-4bit" if load_in_4bit else "--no-load-in-4bit")
|
||||
if frontend:
|
||||
args.extend(["--frontend", str(frontend)])
|
||||
if api_only:
|
||||
args.append("--api-only")
|
||||
if silent:
|
||||
args.append("--silent")
|
||||
# Forward the resolved tool policy so the child doesn't re-resolve.
|
||||
|
|
@ -1262,9 +1280,13 @@ def run(
|
|||
host = host,
|
||||
port = port,
|
||||
silent = True,
|
||||
api_only = api_only,
|
||||
llama_parallel_slots = parallel,
|
||||
cloudflare = cloudflare,
|
||||
secure = secure,
|
||||
# Headless serving prints its own URL/API-key banner; the Tauri-only
|
||||
# TAURI_PORT line would corrupt that machine-parseable output.
|
||||
emit_tauri_port = False,
|
||||
)
|
||||
if frontend is not None:
|
||||
run_kwargs["frontend_path"] = frontend
|
||||
|
|
|
|||
|
|
@ -361,6 +361,29 @@ def test_studio_default_rejects_parallel_when_subcommand_invoked():
|
|||
), f"error message must show the corrected invocation; got: {combined}"
|
||||
|
||||
|
||||
def test_studio_default_rejects_api_only_when_subcommand_invoked():
|
||||
"""`unsloth studio --api-only run ...` would silently serve the UI (the
|
||||
parent's --api-only never reaches run). 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", "--api-only", "run", "--model", "X"])
|
||||
assert result.exit_code == 2, (
|
||||
f"expected exit 2 when --api-only 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 "--api-only" in combined, combined
|
||||
assert (
|
||||
"run --api-only" 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."""
|
||||
|
|
@ -445,3 +468,81 @@ def test_in_venv_path_passes_parallel_to_run_server(monkeypatch, value):
|
|||
assert (
|
||||
captured.get("llama_parallel_slots") == value
|
||||
), f"run_server got llama_parallel_slots={captured.get('llama_parallel_slots')!r}, expected {value}"
|
||||
|
||||
|
||||
# --api-only: serve API only (no UI). Both re-exec and in-venv paths must carry it.
|
||||
|
||||
|
||||
def test_api_only_option_is_registered():
|
||||
studio_mod = _load_run_command()
|
||||
import inspect
|
||||
|
||||
opt = inspect.signature(studio_mod.run).parameters["api_only"].default
|
||||
assert "--api-only" in set(getattr(opt, "param_decls", []) or [])
|
||||
assert getattr(opt, "default", None) is False # opt-in; plain run keeps the UI
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra,present",
|
||||
[
|
||||
(["--api-only"], True),
|
||||
(["--secure", "--api-only"], True), # secure headless path
|
||||
([], False),
|
||||
],
|
||||
)
|
||||
def test_reexec_forwards_api_only(monkeypatch, extra, present):
|
||||
"""`--api-only` (and only when typed) must reach the re-exec'd child."""
|
||||
result, captured = _invoke_run(monkeypatch, _BASE + extra)
|
||||
assert len(captured) == 1, result.output
|
||||
argv = captured[0]["argv"]
|
||||
assert ("--api-only" in argv) is present, argv
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extra,expected", [(["--api-only"], True), ([], False)])
|
||||
def test_in_venv_path_passes_api_only_to_run_server(monkeypatch, extra, expected):
|
||||
"""In-venv path must forward --api-only to run_server(api_only=...)."""
|
||||
studio_mod = _load_run_command()
|
||||
|
||||
fake_venv = Path("/fake/studio/venv/unsloth_studio")
|
||||
monkeypatch.setattr(sys, "prefix", str(fake_venv))
|
||||
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"
|
||||
monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run)
|
||||
|
||||
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 + extra, catch_exceptions = True)
|
||||
|
||||
assert (
|
||||
captured.get("api_only") is expected
|
||||
), f"run_server got api_only={captured.get('api_only')!r}, expected {expected}"
|
||||
# Headless serving must suppress the Tauri-only TAURI_PORT line.
|
||||
assert (
|
||||
captured.get("emit_tauri_port") is False
|
||||
), f"run_server got emit_tauri_port={captured.get('emit_tauri_port')!r}, expected False"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue