Measure where Studio's startup time actually goes (#7553)

* Measure where Studio's startup time actually goes

Nothing measured this. studio/backend/main.py logs 'lifespan startup completed in
X ms' but no test or CI job ever asserted a budget, a repo-wide grep for
startup_ms or time_to_ready matches only that one file, and studio_test_kit polls
/healthz in a loop that discards the elapsed time it already computes. Its
default healthz_timeout_s of 180 was the only recorded expectation.

scripts/profile_startup.py breaks a launch into phases: import cost via
python -X importtime in a subprocess (top cumulative contributors), process spawn
to first output, and spawn to /healthz 200, over N repeats with median and p90.

First numbers on Linux: importing the backend module costs 5.7 to 6.6 seconds
before the server can even bind, and it dominates everything else. That is eager
module-level imports pulled in by the routes package, not the hardware detection
I first suspected: utils.hardware is 23ms and does not pull torch.

--max-healthz-seconds exists so a budget can be enforced once per-platform
numbers are agreed. It is not wired into a gate yet, deliberately: a threshold
picked before the data is in would either be meaningless or flaky.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Profile the code under test, and let the profile fail

Both installer calls omitted --local, so every phase measured the published PyPI
backend and could not move when a PR edits main.py, run.py or routes. t_first_byte
was a dead local, advertised in the docstring but never returned, and the reader
could deadlock once the child filled the pipe. A failed launch and an impossible
budget both produced a warning and exit 0, and the importtime parse reported the
largest cumulative row, which is site, not main, so a raising import published a
number as success. Pin the controller to the profiled venv's interpreter.

* Stop the startup summary hiding failed launches

The aggregates cover only the runs that reached healthz, so two dead launches
and one fast one rendered as a normal fast startup, and an all-failed phase
printed nothing at all. With continue-on-error and no budget wired, that summary
is the only thing anyone sees. Say how many launches the number is made of, and
say so explicitly when none came up.

* Reject --repeats below 1

range(0) launches nothing, so the empty runs list reached the budget check as
"no healthz measurement", warned and exited 0: a gate that cannot fail. The
value comes straight from a dispatch input, so reject it loudly instead.

* Run the startup profile when the imported startup tree changes

The path filter listed main.py, run.py and routes/**, but the graph the
profiler measures is far wider: main.py imports auth, core, hub, loggers,
models, picker and utils at module scope, and routes/models.py imports
utils.utils and utils.hidden_models. A change to any of those moved
`import main` without ever running this job, so the regressions the
workflow exists to catch went unmeasured.

Cover studio/backend/** (tests excluded) and unsloth_cli/**, since the
launch phase spawns `unsloth studio --api-only` and the CLI is on the
process-to-healthz path.

* Read the labelled main row and kill the Windows launcher tree

total_seconds took by_cum[0], the largest cumulative row in -X importtime
output. That output also carries the interpreter's own startup graph (site,
encodings, whatever a venv sitecustomize pulls in), which is not part of
import main, and the two are not ordered by construction. With a trivial main
the old code reported site's 0.027s as "import main" while main actually cost
0.000249s. Today's backend dwarfs site so the published figures are unchanged,
but the headline number must not silently become another module's cost once
the backend imports get optimized, so read the row named main.

profile_launch spawned Scripts/unsloth.exe on Windows. A pip console-script
.exe is a distlib launcher stub that CreateProcess's the venv python and waits,
so terminate() reaped the stub and left the backend holding the inherited
stdout handle: the reader thread never saw EOF and burned the full 10s join,
and with --repeats each iteration stranded another server on the shared
UNSLOTH_STUDIO_HOME. Walk the tree with taskkill /T, matching the cleanup in
unsloth_cli/commands/start.py and unsloth/dataprep/synthetic.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fail the startup budget when nothing was measured and fall back when taskkill fails

* Trigger on installer inputs and harden the startup gate tests

* Tighten comments in the startup profiler and its workflow

* Trigger the startup profile on the Studio setup scripts

install.sh --local runs the checkout's studio/setup.sh, install.ps1 reaches
studio/setup.ps1 through the editable install, and both call
install_python_stack.py, which decides the dependency set that gets imported.
Editing any of them could change startup time with no measurement taken.

* Shorten the startup profiler comments

Comments and docstrings only.

* Reject non-finite startup budgets and profile when the desktop argv changes

--max-healthz-seconds nan or inf parses as a float but compares False against
any median, so the gate reported success without bounding anything. Require a
finite value.

The profiler hardcodes the argv that process.rs::backend_args builds, but that
file was not in the trigger paths, so a change to the desktop launch command
scheduled no measurement. Add it, and anchor the two argv lists with a test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothai@gmail.com>
This commit is contained in:
Daniel Han 2026-07-28 22:24:34 -07:00 committed by GitHub
commit bd3972804d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 776 additions and 0 deletions

View file

@ -0,0 +1,243 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Regression coverage for the startup profiler's budget gate, teardown and triggers."""
from __future__ import annotations
import ast
import fnmatch
import importlib.util
import re
import subprocess
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
import yaml
REPO_ROOT = Path(__file__).resolve().parents[1]
SCRIPT = REPO_ROOT / "scripts" / "profile_startup.py"
WORKFLOW = REPO_ROOT / ".github" / "workflows" / "startup-profile-ci.yml"
PROCESS_RS = REPO_ROOT / "studio" / "src-tauri" / "src" / "process.rs"
# Checkout files that build the venv the workflow profiles.
INSTALLER_INPUTS = (
"studio/setup.sh",
"studio/setup.ps1",
"studio/install_python_stack.py",
)
# Checkout file that defines the argv the profiler reproduces.
LAUNCH_INPUTS = ("studio/src-tauri/src/process.rs",)
def _load():
spec = importlib.util.spec_from_file_location("profile_startup", SCRIPT)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
def _no_subprocesses(mod, monkeypatch):
# Keep the gate tests off the real interpreter and CLI.
monkeypatch.setattr(mod, "find_bin", lambda: None)
monkeypatch.setattr(mod, "profile_imports", lambda python, top = 15: {"ok": False, "error": ""})
monkeypatch.setattr(mod, "python_version_of", lambda python: "3.13.0")
class _Proc:
"""Stand-in for a still-running Popen."""
def __init__(self):
self.pid = 4321
self.terminated = False
def poll(self):
return None
def terminate(self):
self.terminated = True
def _nt(mod, monkeypatch, returncode):
calls: list[list[str]] = []
def _run(argv, **kwargs):
calls.append(argv)
return subprocess.CompletedProcess(argv, returncode, "", "")
# Patch the module's own references, not the real os/subprocess the session shares.
monkeypatch.setattr(mod, "os", SimpleNamespace(name = "nt"))
monkeypatch.setattr(mod, "subprocess", SimpleNamespace(run = _run))
return calls
def test_budget_fails_when_no_launch_was_measured(capsys, monkeypatch):
"""A requested budget must not pass just because the CLI was never found."""
mod = _load()
_no_subprocesses(mod, monkeypatch)
rc = mod.main(["--max-healthz-seconds", "30"])
out = capsys.readouterr().out
assert rc == 1
assert "::error::" in out and "no healthz measurement" in out
assert "no unsloth CLI found" in out
def _healthy_launch(
mod,
monkeypatch,
healthz = 1.5,
):
monkeypatch.setattr(mod, "find_bin", lambda: "unsloth")
monkeypatch.setattr(
mod,
"profile_launch",
lambda bin_path, port, **kw: {
"spawn_seconds": 0.1,
"healthz_seconds": healthz,
"lifespan_ms": 100.0,
"reached_healthz": True,
"log_tail": [],
},
)
def test_budget_still_passes_when_a_launch_was_measured(monkeypatch):
"""The fail-closed branch must not swallow a genuinely healthy run."""
mod = _load()
_no_subprocesses(mod, monkeypatch)
_healthy_launch(mod, monkeypatch)
assert mod.main(["--max-healthz-seconds", "30"]) == 0
assert mod.main(["--max-healthz-seconds", "1"]) == 1
# "=" form for -inf: a bare "-inf" is an option token to argparse, not a value.
@pytest.mark.parametrize(
"bad", ["--max-healthz-seconds=nan", "--max-healthz-seconds=inf", "--max-healthz-seconds=-inf"]
)
def test_budget_rejects_non_finite_values(bad, capsys, monkeypatch):
"""`med > nan` and `med > inf` are always False, so the gate would never bind."""
mod = _load()
_no_subprocesses(mod, monkeypatch)
_healthy_launch(mod, monkeypatch)
with pytest.raises(SystemExit) as exc:
mod.main([bad])
assert exc.value.code == 2
assert "finite" in capsys.readouterr().err
def test_budget_rejects_import_only(capsys):
"""--import-only launches nothing, so a budget on it could only ever pass."""
mod = _load()
with pytest.raises(SystemExit) as exc:
mod.main(["--import-only", "--max-healthz-seconds", "30"])
assert exc.value.code == 2
assert "--import-only" in capsys.readouterr().err
def test_terminate_tree_falls_back_when_taskkill_fails(monkeypatch):
"""A nonzero taskkill must still reach terminate(), not return silently."""
mod = _load()
calls = _nt(mod, monkeypatch, returncode = 1)
proc = _Proc()
mod._terminate_tree(proc)
assert calls == [["taskkill", "/PID", "4321", "/T", "/F"]]
assert proc.terminated
def test_terminate_tree_falls_back_when_taskkill_raises(monkeypatch):
"""A missing or hung taskkill must reach terminate() too."""
mod = _load()
monkeypatch.setattr(mod, "os", SimpleNamespace(name = "nt"))
def _boom(argv, **kwargs):
raise FileNotFoundError(argv)
monkeypatch.setattr(mod, "subprocess", SimpleNamespace(run = _boom))
proc = _Proc()
mod._terminate_tree(proc)
assert proc.terminated
def test_terminate_tree_returns_on_successful_taskkill(monkeypatch):
mod = _load()
_nt(mod, monkeypatch, returncode = 0)
proc = _Proc()
mod._terminate_tree(proc)
assert not proc.terminated
def test_terminate_tree_skips_an_exited_process(monkeypatch):
mod = _load()
calls = _nt(mod, monkeypatch, returncode = 0)
proc = _Proc()
proc.poll = lambda: 0
mod._terminate_tree(proc)
assert calls == [] and not proc.terminated
def _trigger_paths():
wf = yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8"))
# YAML 1.1 turns the bare `on:` key into True.
on = wf.get("on") or wf[True]
return [p for p in on["pull_request"]["paths"] if not p.startswith("!")]
@pytest.mark.parametrize("rel", INSTALLER_INPUTS)
def test_workflow_triggers_on_studio_installer_inputs(rel):
"""A setup script that changes the profiled venv must schedule a measurement."""
assert (REPO_ROOT / rel).is_file(), f"{rel} moved; revisit the trigger list"
paths = _trigger_paths()
assert any(fnmatch.fnmatch(rel, p) for p in paths), f"{rel} not covered by {paths}"
def test_studio_installer_inputs_are_on_the_local_install_path():
"""Anchor the list above: these files are what --local actually executes."""
# install.ps1 reaches setup.ps1 through the editable install, not by name.
assert "studio/setup.sh" in (REPO_ROOT / "install.sh").read_text(encoding = "utf-8")
for setup in ("studio/setup.sh", "studio/setup.ps1"):
text = (REPO_ROOT / setup).read_text(encoding = "utf-8", errors = "replace")
assert "install_python_stack.py" in text
@pytest.mark.parametrize("rel", LAUNCH_INPUTS)
def test_workflow_triggers_on_the_desktop_launch_command(rel):
"""The profiler copies process.rs's argv, so a change there must be measured."""
assert (REPO_ROOT / rel).is_file(), f"{rel} moved; revisit the trigger list"
paths = _trigger_paths()
assert any(fnmatch.fnmatch(rel, p) for p in paths), f"{rel} not covered by {paths}"
def _desktop_backend_argv():
body = re.search(
r"fn backend_args\(port: u16\) -> Vec<String> \{(.*?)\n\}",
PROCESS_RS.read_text(encoding = "utf-8"),
re.S,
)
assert body, "backend_args moved; revisit the trigger list"
return re.findall(r'"([^"]+)"', body.group(1))
def _profiler_argv():
tree = ast.parse(SCRIPT.read_text(encoding = "utf-8"))
fn = next(
n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "profile_launch"
)
call = next(
n for n in ast.walk(fn) if isinstance(n, ast.Call) and ast.unparse(n.func).endswith("Popen")
)
return [e.value for e in call.args[0].elts if isinstance(e, ast.Constant)]
def test_profiler_spawns_the_desktop_backend_argv():
"""Anchor the trigger above: these two argv lists must stay identical."""
assert _profiler_argv() == _desktop_backend_argv()
@pytest.mark.skipif(sys.platform == "win32", reason = "posix branch")
def test_terminate_tree_posix_uses_terminate():
mod = _load()
proc = _Proc()
mod._terminate_tree(proc)
assert proc.terminated