unsloth/tests/validate_studio_features.py
Daniel Han 9ca7be82c4 docker: trim redundant comments in the image build files
Comment-only pass over the PR's own files. No executable line changes.

- Dockerfile / Dockerfile.studio: drop the decorative stage banner rules, the
  stale "5)" / "6)" step numbering, and the entrypoint pre-flight list that
  restated (and had drifted from) entrypoint.sh's own accurate header. Cut the
  llama.cpp asset bullet list that repeats fetch_llama_prebuilt.py's docstring
  and the structlog rationale already spelled out at the install site.
- entrypoint.sh / studio_launch.sh: fold the section banners into the
  explanation lines that follow them.
- docker-publish.yml: remove the comment rule lines around the job headers.
- validate_studio_features.py: same for the numbered section headers.
- smoke_test.py: drop the stale "~125M params" note on a 1B model.
- unsloth_branding.py, unsloth_nb_view.py, unsloth_nb_pip_magic.py,
  colabTitle.ts: remove comments that restate the adjacent line.
2026-07-26 15:45:58 +00:00

269 lines
10 KiB
Python

#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0
"""Cross-platform validation of the Unsloth Docker JupyterLab/notebook features.
Runs WITHOUT Docker or a GPU, so it can execute on the Linux/macOS/Windows CI
lanes. It exercises the actual notebook-helper logic (not just py_compile) and
checks the shipped JupyterLab config + labextension source, so a regression in
the notebook organisation, Colab compatibility, Colab-intro/widget stripping,
sidecar-log gating, the labextension plugins, the JupyterLab defaults, or the
login branding fails CI on every device.
Usage: python tests/validate_studio_features.py
Exit 0 = all checks pass; non-zero = at least one failed.
"""
from __future__ import annotations
import importlib
import json
import os
import sys
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DOCKER = os.path.join(ROOT, "docker")
JUPYTER = os.path.join(DOCKER, "jupyter")
LABEXT = os.path.join(JUPYTER, "unsloth_labext")
sys.path.insert(0, DOCKER)
_failures: list[str] = []
def check(
name: str,
cond: bool,
detail: str = "",
) -> None:
status = "PASS" if cond else "FAIL"
print(f" [{status}] {name}" + (f" -- {detail}" if detail and not cond else ""))
if not cond:
_failures.append(name)
# 1. Colab cell-magic compatibility (#@title then %%capture)
def test_colab_compat() -> None:
print("colab cell-magic compat (unsloth_colab_compat):")
m = importlib.import_module("unsloth_colab_compat")
out = m.colab_cell_magic_fix(["#@title Setup\n", "%%capture\n", "!pip install x\n"])
check("magic hoisted above #@title", out[0] == "%%capture\n" and "#@title Setup\n" in out)
# idempotent / already on top
same = ["%%capture\n", "print(1)\n"]
check("no-op when magic already first", m.colab_cell_magic_fix(same) == same)
# non-magic cell untouched
plain = ["x = 1\n", "y = 2\n"]
check("plain cell untouched", m.colab_cell_magic_fix(plain) == plain)
# content magic (%%writefile) NOT hoisted into the written file body
wf = ["#@title Config\n", "%%writefile config.json\n", "{}\n"]
check("content magic (%%writefile) left untouched", m.colab_cell_magic_fix(wf) == wf)
# safe magic with arg still hoisted
bash = ["#@title Run\n", "%%bash\n", "echo hi\n"]
check("safe magic (%%bash) hoisted", m.colab_cell_magic_fix(bash)[0] == "%%bash\n")
# 2. Notebook categorisation (clean_section) + README parsing
def test_nb_view() -> None:
print("notebook view (unsloth_nb_view):")
v = importlib.import_module("unsloth_nb_view")
check(
"clean_section dash/slash -> space",
v.clean_section("### GRPO-Reinforcement/Learning Notebooks")
== "GRPO Reinforcement Learning Notebooks",
v.clean_section("### GRPO-Reinforcement/Learning Notebooks"),
)
check(
"clean_section strips hashes/space",
v.clean_section("## Main Notebooks ") == "Main Notebooks",
)
# 3. Colab-intro + stale-widget stripping
def test_strip() -> None:
print("notebook strip (unsloth_nb_strip_colab):")
s = importlib.import_module("unsloth_nb_strip_colab")
nb = {
"metadata": {"widgets": {"application/vnd.jupyter.widget-state+json": {"x": 1}}},
"cells": [
{
"cell_type": "markdown",
"source": [
'To run this, press "Runtime" ... Tesla T4 Google Colab instance!\n',
"\n",
"You will learn how to ...\n",
],
},
{
"cell_type": "code",
"source": ["print(1)\n"],
"outputs": [
{"output_type": "stream", "name": "stdout", "text": "ok\n"},
{
"output_type": "display_data",
"data": {
"application/vnd.jupyter.widget-view+json": {"model_id": "abc"},
"text/plain": "0%| | 0/10",
},
},
],
},
],
}
changed1 = s._strip_intro(nb)
changed2 = s._clean_widgets(nb)
check(
"intro line stripped",
changed1 and not any("to run this, press" in (l.lower()) for l in nb["cells"][0]["source"]),
)
check("intro body kept", any("You will learn" in l for l in nb["cells"][0]["source"]))
wv = sum(
1
for c in nb["cells"]
for o in (c.get("outputs", []) or [])
if "application/vnd.jupyter.widget-view+json" in (o.get("data", {}) or {})
)
check("widget-view outputs removed", changed2 and wv == 0)
check(
"non-widget outputs kept",
any(
o.get("output_type") == "stream"
for c in nb["cells"]
for o in (c.get("outputs", []) or [])
),
)
check("metadata.widgets removed", "widgets" not in nb["metadata"])
# idempotent
check("strip idempotent", not s._strip_intro(nb) and not s._clean_widgets(nb))
# 4. Sidecar-log gating
def test_sidecar_log_gate() -> None:
print("sidecar log gate (unsloth_nb_compat):")
c = importlib.import_module("unsloth_nb_compat")
old = os.environ.pop("UNSLOTH_ENABLE_LOGGING", None)
try:
check("logging off by default", c._logging_enabled() is False)
os.environ["UNSLOTH_ENABLE_LOGGING"] = "1"
check("logging on with env=1", c._logging_enabled() is True)
os.environ["UNSLOTH_ENABLE_LOGGING"] = "0"
check("logging off with env=0", c._logging_enabled() is False)
finally:
os.environ.pop("UNSLOTH_ENABLE_LOGGING", None)
if old is not None:
os.environ["UNSLOTH_ENABLE_LOGGING"] = old
# 5. JupyterLab defaults (overrides.json)
def test_overrides() -> None:
print("jupyterlab defaults (jupyter/overrides.json):")
path = os.path.join(JUPYTER, "overrides.json")
check("overrides.json exists", os.path.isfile(path))
if not os.path.isfile(path):
return
with open(path, encoding = "utf-8") as f:
d = json.load(f) # raises -> CI fails if invalid JSON
themes = d.get("@jupyterlab/apputils-extension:themes", {})
check(
"default theme = Unsloth Dark",
themes.get("theme") == "Unsloth Dark",
str(themes.get("theme")),
)
check("adaptive theme on", themes.get("adaptive-theme") is True)
check("preferred dark = Unsloth Dark", themes.get("preferred-dark-theme") == "Unsloth Dark")
tracker = d.get("@jupyterlab/notebook-extension:tracker", {})
check(
"windowingMode none",
tracker.get("windowingMode") == "none",
str(tracker.get("windowingMode")),
)
notif = d.get("@jupyterlab/apputils-extension:notification", {})
check(
"news prompt off",
str(notif.get("fetchNews")) == "false" and notif.get("checkForUpdates") is False,
)
panel = d.get("@jupyterlab/notebook-extension:panel", {})
labels = [t.get("label", "") for t in panel.get("toolbar", [])]
check(
"Restart & Run All label (single >>)",
any(l == "Restart & Run All" for l in labels) and not any(">>" in l for l in labels),
str(labels),
)
# 6. Labextension source (plugins) + login branding assets
def test_labext_and_branding() -> None:
print("labextension + branding assets:")
pkg = os.path.join(LABEXT, "package.json")
check("labext package.json exists", os.path.isfile(pkg))
if os.path.isfile(pkg):
with open(pkg, encoding = "utf-8") as f:
p = json.load(f)
check("labext name unsloth-jupyterlab", p.get("name") == "unsloth-jupyterlab")
check("labext themePath set", bool(p.get("jupyterlab", {}).get("themePath")))
# Concatenate every .ts module under src/ so plugins defined in their own
# files (cellNav, colabTitle, outputSelect, uiChrome) are all covered.
src_dir = os.path.join(LABEXT, "src")
all_src = ""
if os.path.isdir(src_dir):
for fn in sorted(os.listdir(src_dir)):
if fn.endswith(".ts"):
with open(os.path.join(src_dir, fn), encoding = "utf-8") as f:
all_src += f.read() + "\n"
for plug in [
"unsloth-jupyterlab:theme",
"unsloth-jupyterlab:cell-nav",
"unsloth-jupyterlab:logo",
"unsloth-jupyterlab:colab-title",
"unsloth-jupyterlab:output-select-all",
"unsloth-jupyterlab:ui-chrome",
]:
check(f"plugin present: {plug}", plug in all_src)
# The two newest plugins are also exported from index.ts (wired in).
index = os.path.join(src_dir, "index.ts")
index_src = open(index, encoding = "utf-8").read() if os.path.isfile(index) else ""
check("outputSelect wired in index.ts", "outputSelectPlugin" in index_src)
check("uiChrome wired in index.ts", "uiChromePlugin" in index_src)
# uiChrome hides the right activity bar; CTRL+A output-select selects nodes.
check("right activity bar hidden", "jp-mod-right" in all_src and "display: none" in all_src)
check("ctrl+A output select", "selectNodeContents" in all_src)
# branding assets
login = os.path.join(JUPYTER, "login.html")
login_src = open(login, encoding = "utf-8").read() if os.path.isfile(login) else ""
check("login.html branded", "unsloth-login-card" in login_src)
check(
"login.html uses sloth stickers",
'static_url("sloth/' in login_src or "static_url('sloth/" in login_src,
)
check("favicon.ico present", os.path.isfile(os.path.join(JUPYTER, "favicon.ico")))
check("logo.png present", os.path.isfile(os.path.join(JUPYTER, "logo.png")))
check(
"sloth sticker installer present",
os.path.isfile(os.path.join(JUPYTER, "install_sloth_stickers.py")),
)
def main() -> int:
print("=== Unsloth Studio/notebook feature validation ===")
for t in (
test_colab_compat,
test_nb_view,
test_strip,
test_sidecar_log_gate,
test_overrides,
test_labext_and_branding,
):
try:
t()
except Exception as e: # a thrown exception is a failure, not a crash
_failures.append(f"{t.__name__}: {e!r}")
print(f" [FAIL] {t.__name__} raised {e!r}")
print()
if _failures:
print(f"FAILED ({len(_failures)}): " + ", ".join(_failures))
return 1
print("ALL CHECKS PASSED")
return 0
if __name__ == "__main__":
sys.exit(main())