ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
* 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
Download patch file Download diff file Expand all files Collapse all files

156
.github/workflows/startup-profile-ci.yml vendored Normal file
View file

@ -0,0 +1,156 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Measures where Studio's startup time goes, on each platform.
#
# Nothing recorded a number before: main.py logs "lifespan startup completed in X ms"
# and studio_test_kit polls /healthz, but both throw the elapsed time away. A first
# local run (Linux, warm cache, 18-core server) put `import main` at 5.7-6.6s BEFORE
# the server can bind, dominated by eager module-level imports pulled in by routes:
# torch ~1.9s self, unsloth_zoo ~0.8s, routes ~0.6s, transformers ~0.5s.
#
# Not a gate yet: --max-healthz-seconds exists, but a budget should come from
# observed numbers rather than a guess.
name: Startup profile
on:
pull_request:
paths:
# The measured import graph is the whole backend tree: main.py imports auth,
# core, hub, loggers, models, picker, routes and utils at module scope.
- 'studio/backend/**'
- '!studio/backend/tests/**'
# The launch phase spawns `unsloth studio --api-only`, so the CLI counts too.
- 'unsloth_cli/**'
- 'studio/src-tauri/src/preflight**'
# The profiler hardcodes the desktop argv that process.rs::backend_args builds,
# so a change there must schedule a run or the two silently diverge.
- 'studio/src-tauri/src/process.rs'
- 'scripts/profile_startup.py'
- '.github/workflows/startup-profile-ci.yml'
# The job profiles whatever `install.sh --local` built: the installers pick the
# venv's Python and the dependency specs, and pyproject's include list is what
# makes --local overlay studio.backend*.
- 'install.sh'
- 'install.ps1'
- 'pyproject.toml'
# --local also runs the checkout's setup scripts (install.sh picks
# $_REPO_ROOT/studio/setup.sh, the editable install resolves setup.ps1 to the
# repo), and both call install_python_stack.py, which picks the dependencies.
- 'studio/setup.sh'
- 'studio/setup.ps1'
- 'studio/install_python_stack.py'
workflow_dispatch:
inputs:
repeats:
description: 'launch repeats per OS (median reported)'
type: string
default: '3'
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
profile:
name: startup ${{ matrix.os }}
runs-on: ${{ matrix.os }}
timeout-minutes: 60
continue-on-error: true
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, macos-14, windows-latest]
env:
UNSLOTH_STUDIO_HOME: ${{ github.workspace }}/.studio-home
# A wildcard bind calls ifconfig.me on the startup path; loopback times our code.
UNSLOTH_STUDIO_DISABLE_PUBLIC_CHECK: '1'
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- name: Install Studio
shell: bash
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
set -o pipefail
mkdir -p logs
# --local is load-bearing: it overlays the checkout, so the profiled server
# is this diff. Without it install.sh resolves unsloth from PyPI.
if [ "${{ runner.os }}" = "Windows" ]; then
pwsh -NoProfile -File ./install.ps1 --local 2>&1 | tee logs/install.log
else
bash install.sh --local 2>&1 | tee logs/install.log
fi
- name: Profile startup
shell: bash
run: |
BIN="$UNSLOTH_STUDIO_HOME/unsloth_studio/bin/unsloth"
[ -x "$BIN" ] || BIN="$UNSLOTH_STUDIO_HOME/unsloth_studio/Scripts/unsloth.exe"
[ -x "$BIN" ] || BIN=""
# Profile imports with the INSTALLED interpreter: that venv is what launches.
PY="$UNSLOTH_STUDIO_HOME/unsloth_studio/bin/python"
[ -x "$PY" ] || PY="$UNSLOTH_STUDIO_HOME/unsloth_studio/Scripts/python.exe"
[ -x "$PY" ] || PY="$(command -v python3 || command -v python)"
python3 scripts/profile_startup.py \
--python "$PY" \
${BIN:+--bin "$BIN"} \
--repeats "${{ inputs.repeats || '3' }}" \
--json "startup-${{ matrix.os }}.json" 2>&1 | tee logs/profile.log
- name: Summary
if: always()
shell: bash
run: |
f="startup-${{ matrix.os }}.json"
[ -f "$f" ] || { echo "no profile produced"; exit 0; }
python3 - "$f" >> "$GITHUB_STEP_SUMMARY" <<'PY'
import json, sys
d = json.load(open(sys.argv[1]))
print(f"### {d['platform']} / {d['machine']} (py {d['python']}, {d['cpu_count']} cpu)\n")
imp = d.get("imports", {})
# Gate on ok: a failed `import main` still leaves rows, so a total can lie.
if imp.get("ok"):
print(f"**`import main`: {imp['total_seconds']}s**\n")
print("| package | self ms |")
print("|---|---:|")
for k, v in list(imp.get("self_by_package_ms", {}).items())[:8]:
print(f"| {k} | {v} |")
print()
else:
print("**`import main` failed - no valid import profile**\n")
print("```\n" + (imp.get("error") or "")[-1500:] + "\n```\n")
lau = d.get("launch") or {}
runs = len(lau.get("runs") or [])
failed = lau.get("failed_runs") or 0
if lau.get("healthz_median_seconds") is not None:
# The aggregates cover only the runs that reached healthz, so flag the
# failures: bare numbers would read as a normal fast startup.
note = f" _({runs - failed} of {runs} launches; {failed} never became healthy)_" if failed else ""
print(f"**time to a healthy port: {lau['healthz_median_seconds']}s median, "
f"{lau['healthz_max_seconds']}s max**{note}\n")
elif lau.get("skipped"):
print(f"_launch phase skipped: {lau['skipped']}_\n")
elif runs:
print(f"**no launch measurement: all {runs} launches failed to become healthy**\n")
PY
- name: Upload profile
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: startup-profile-${{ matrix.os }}
path: |
startup-*.json
logs/
retention-days: 14
if-no-files-found: warn

377
scripts/profile_startup.py Normal file
View file

@ -0,0 +1,377 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Measure where Unsloth Studio's startup time goes, per platform.
Nothing measured this before: the backend logs "lifespan startup completed in X ms"
but no test or CI job asserted a budget, and studio_test_kit discards the elapsed
time of its /healthz poll. A first local run (Linux, warm cache, fast server CPU)
found `import main` alone costs 6.6s before the server can bind, dominated by eager
module-level imports pulled in by the `routes` package:
torch 1930 ms self
unsloth_zoo 914 ms self
routes 779 ms self
transformers 524 ms self
Phases measured:
import `python -X importtime -c "import main"`, top cumulative + per-package self
spawn process start -> first byte on stdout
healthz process start -> /api/health (or /healthz) answers 200
lifespan the backend's own "lifespan startup completed in X ms" log line
Usage:
python scripts/profile_startup.py --repeats 3 --json out.json
python scripts/profile_startup.py --import-only # no server, no port needed
Exit code is 0 unless --max-healthz-seconds is given and exceeded.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import platform
import re
import shutil
import socket
import statistics
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
BACKEND = REPO_ROOT / "studio" / "backend"
_IMPORTTIME_RE = re.compile(r"import time:\s+(\d+)\s+\|\s+(\d+)\s+\|(\s*)(\S.*)")
def _free_port() -> int:
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return int(s.getsockname()[1])
def profile_imports(python: str, top: int = 15) -> dict:
"""Cumulative and self import cost for the backend's module graph.
Run in a subprocess with -X importtime: the numbers are only meaningful for a
cold interpreter, and importing in-process would measure a warm sys.modules.
"""
proc = subprocess.run(
[python, "-X", "importtime", "-c", "import sys; sys.path.insert(0, '.'); import main"],
cwd = BACKEND,
capture_output = True,
text = True,
timeout = 900,
)
rows = []
for line in proc.stderr.splitlines():
m = _IMPORTTIME_RE.match(line)
if m:
rows.append((int(m.group(1)), int(m.group(2)), m.group(4).strip()))
if not rows:
return {"ok": False, "error": (proc.stderr or proc.stdout)[-2000:]}
if proc.returncode != 0:
# Rows survive up to the failure, so any total from a partial graph is wrong.
return {
"ok": False,
"error": (proc.stderr or proc.stdout)[-2000:],
"partial_rows": len(rows),
}
by_cum = sorted(rows, key = lambda r: -r[1])
# Total comes from the `main` row, not by_cum[0]: -X importtime also prints the
# interpreter's own startup graph (`site`), which can outrank a trivial main.
main_row = next((r for r in reversed(rows) if r[2] == "main"), None)
if main_row is None:
return {
"ok": False,
"error": "no `import main` row in -X importtime output\n"
+ (proc.stderr or proc.stdout)[-2000:],
}
self_by_pkg: dict[str, int] = {}
for self_us, _cum, name in rows:
pkg = name.split(".")[0]
self_by_pkg[pkg] = self_by_pkg.get(pkg, 0) + self_us
return {
"ok": True,
"total_seconds": round(main_row[1] / 1e6, 3),
"top_cumulative": [
{"module": n, "seconds": round(c / 1e6, 3)} for _s, c, n in by_cum[:top]
],
"self_by_package_ms": {
k: round(v / 1000) for k, v in sorted(self_by_pkg.items(), key = lambda x: -x[1])[:top]
},
}
def _terminate_tree(proc: subprocess.Popen) -> None:
"""Stop the server AND its children, which on Windows are a separate process.
CI profiles `Scripts/unsloth.exe`, a distlib launcher stub that CreateProcess's
the venv python and waits, so terminate() reaps the stub only: the real backend
keeps the inherited stdout handle, the reader thread never sees EOF, and
--repeats strands one server per iteration on the shared UNSLOTH_STUDIO_HOME.
taskkill /T walks the tree, as unsloth_cli/commands/start.py already does.
"""
if proc.poll() is not None:
return
if os.name == "nt":
try:
killed = subprocess.run(
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
capture_output = True,
timeout = 30,
check = False,
)
if killed.returncode == 0:
return
except Exception:
# taskkill missing or timed out; fall through so the stub still dies.
pass
# check=False: a nonzero taskkill does not raise, so fall through as well.
proc.terminate()
def profile_launch(
bin_path: str,
port: int,
timeout_s: int = 300,
) -> dict:
"""Spawn the backend the way the desktop app does and time it to first 200."""
log_lines: list[str] = []
first_byte: list[float] = []
t0 = time.perf_counter()
proc = subprocess.Popen(
[bin_path, "studio", "--api-only", "-H", "127.0.0.1", "-p", str(port)],
cwd = REPO_ROOT,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
bufsize = 1,
)
def _drain() -> None:
# Runs alongside the health polling: the first read timestamps the spawn
# phase, and an undrained pipe blocks the backend before it binds.
for line in proc.stdout:
if not first_byte:
first_byte.append(time.perf_counter() - t0)
log_lines.append(line.rstrip("\n"))
reader = threading.Thread(target = _drain, daemon = True)
reader.start()
t_healthz = None
deadline = t0 + timeout_s
try:
while time.perf_counter() < deadline:
if proc.poll() is not None:
break
if t_healthz is None:
for url in (
f"http://127.0.0.1:{port}/api/health",
f"http://127.0.0.1:{port}/healthz",
):
try:
with urllib.request.urlopen(url, timeout = 2) as r:
if r.status == 200:
t_healthz = time.perf_counter() - t0
break
except (urllib.error.URLError, OSError, TimeoutError):
pass
if t_healthz is not None:
break
time.sleep(0.25)
finally:
_terminate_tree(proc)
try:
# Safe: the reader drains the pipe, so the child cannot block on write().
proc.wait(timeout = 30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
reader.join(timeout = 10)
t_first_byte = first_byte[0] if first_byte else None
lifespan_ms = None
for line in log_lines:
m = re.search(r"lifespan startup completed in ([\d.]+)ms", line)
if m:
lifespan_ms = float(m.group(1))
return {
"spawn_seconds": round(t_first_byte, 3) if t_first_byte is not None else None,
"healthz_seconds": round(t_healthz, 3) if t_healthz is not None else None,
"lifespan_ms": lifespan_ms,
"reached_healthz": t_healthz is not None,
"log_tail": log_lines[-25:],
}
def python_version_of(python: str) -> str:
"""Version of the interpreter that runs the imports, not the one running us.
--python points at the installed Studio venv while this script runs under the
runner's system python, so platform.python_version() would label it wrong.
"""
if python == sys.executable:
return platform.python_version()
try:
proc = subprocess.run(
[python, "-c", "import platform; print(platform.python_version())"],
capture_output = True,
text = True,
timeout = 60,
)
if proc.returncode == 0 and proc.stdout.strip():
return proc.stdout.strip()
except (OSError, subprocess.SubprocessError):
pass
return "unknown"
def find_bin() -> str | None:
home = os.environ.get("UNSLOTH_STUDIO_HOME") or str(Path.home() / ".unsloth" / "studio")
names = ["unsloth.exe", "unsloth"] if platform.system() == "Windows" else ["unsloth"]
subdirs = ["unsloth_studio/Scripts", "unsloth_studio/bin", "bin", "Scripts"]
for sd in subdirs:
for n in names:
p = Path(home) / sd / n
if p.exists():
return str(p)
return shutil.which("unsloth")
def main(argv: list[str]) -> int:
ap = argparse.ArgumentParser(
description = __doc__, formatter_class = argparse.RawDescriptionHelpFormatter
)
ap.add_argument(
"--repeats",
type = int,
default = 1,
help = "launch repeats; the median is reported (imports are measured once)",
)
ap.add_argument(
"--python",
default = sys.executable,
help = "interpreter used for the import profile (default: this one)",
)
ap.add_argument("--bin", help = "path to the unsloth CLI (default: autodetect)")
ap.add_argument(
"--import-only",
action = "store_true",
help = "skip the server phases (no install needed beyond the deps)",
)
ap.add_argument(
"--max-healthz-seconds",
type = float,
help = "fail if the median time to a healthy port exceeds this",
)
ap.add_argument("--json", help = "write the full report here")
a = ap.parse_args(argv)
# range(0) launches nothing, leaving the budget check with nothing to fail on.
if a.repeats < 1:
ap.error("--repeats must be at least 1")
# Same reason: --import-only never launches anything.
if a.import_only and a.max_healthz_seconds is not None:
ap.error("--max-healthz-seconds cannot be combined with --import-only")
# nan and inf parse fine as floats but `med > budget` is then always False,
# so the gate would report success without ever bounding anything.
if a.max_healthz_seconds is not None and not math.isfinite(a.max_healthz_seconds):
ap.error("--max-healthz-seconds must be a finite number")
report: dict = {
"platform": platform.system().lower(),
"machine": platform.machine(),
"python": python_version_of(a.python),
"cpu_count": os.cpu_count(),
}
print("== import graph ==")
report["imports"] = profile_imports(a.python)
imp = report["imports"]
if imp.get("ok"):
print(f" import main: {imp['total_seconds']}s")
for row in imp["top_cumulative"][:8]:
print(f" {row['seconds']:7.3f}s {row['module']}")
print(" self time by package (ms):")
for k, v in list(imp["self_by_package_ms"].items())[:8]:
print(f" {v:8} ms {k}")
else:
print(f" FAILED: {imp.get('error', '')[:400]}")
if not a.import_only:
bin_path = a.bin or find_bin()
if not bin_path:
print(
"== launch == skipped: no unsloth CLI found "
"(set UNSLOTH_STUDIO_HOME or pass --bin)"
)
report["launch"] = {"skipped": "no unsloth CLI found"}
else:
print(f"== launch == {bin_path}")
runs = []
for i in range(a.repeats):
r = profile_launch(bin_path, _free_port())
runs.append(r)
print(
f" run {i + 1}: healthz={r['healthz_seconds']}s "
f"lifespan={r['lifespan_ms']}ms reached={r['reached_healthz']}"
)
got = [r["healthz_seconds"] for r in runs if r["healthz_seconds"] is not None]
report["launch"] = {
"runs": runs,
"failed_runs": sum(1 for r in runs if not r["reached_healthz"]),
"healthz_median_seconds": round(statistics.median(got), 3) if got else None,
"healthz_max_seconds": round(max(got), 3) if got else None,
}
if got:
print(
f" median time to healthy port: {report['launch']['healthz_median_seconds']}s"
)
if a.json:
Path(a.json).write_text(json.dumps(report, indent = 2), encoding = "utf-8")
print(f"\nwrote {a.json}")
if a.max_healthz_seconds is not None:
launch = report.get("launch") or {}
med = launch.get("healthz_median_seconds")
failed = launch.get("failed_runs") or 0
if failed:
# Failed launches fail the budget; dropping them would keep only the fast ones.
print(
f"::error::startup regression: {failed} of {len(launch.get('runs') or [])} "
f"launches never became healthy within the timeout"
)
return 1
if med is None:
# Nothing measured: exiting 0 would pass a requested budget without a
# single health request, so fail closed.
print(
"::error::startup regression: no healthz measurement, so the "
f"{a.max_healthz_seconds}s budget was never checked "
f"({launch.get('skipped') or 'launch phase produced no runs'})"
)
return 1
elif med > a.max_healthz_seconds:
print(
f"::error::startup regression: {med}s median to a healthy port "
f"exceeds the {a.max_healthz_seconds}s budget"
)
return 1
return 0
if __name__ == "__main__":
raise SystemExit(main(sys.argv[1:]))

243
tests/test_profile_startup_gate.py Normal file
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