diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8cfa7f998b..e018803a5f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -14,5 +14,8 @@ repos:
entry: scripts/run_ruff_format.py
language: python
types: [python]
+ # Mirror ruff's [tool.ruff] extend-exclude so this hook does not
+ # half-process files ruff itself skips (which produced churn).
+ exclude: '(chat_templates|ollama_template_mappers|_auto_install|mapper)\.py$'
additional_dependencies:
- ruff==0.6.9
diff --git a/pyproject.toml b/pyproject.toml
index 24fd66dbbb..713146ff32 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1308,6 +1308,7 @@ repository = "https://github.com/unslothai/unsloth"
[tool.ruff]
target-version = "py311"
+line-length = 100
force-exclude = true
extend-exclude = [
"*chat_templates.py",
diff --git a/scripts/check_frontend_dep_removal.py b/scripts/check_frontend_dep_removal.py
index 260ad5215a..3ec4e9037f 100644
--- a/scripts/check_frontend_dep_removal.py
+++ b/scripts/check_frontend_dep_removal.py
@@ -52,9 +52,7 @@ EXPECTED_NOISE_FILES = {
}
# Only quoted-string occurrences in these file types can be module specifiers.
-JS_LIKE_EXT = re.compile(
- r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$"
-)
+JS_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$")
# Files where JS-syntactic import patterns (static/dynamic/require/re-export)
# could be a real module reference. Markdown gets a separate gate (.mdx is
# real ESM; .md code fences are not).
@@ -273,9 +271,7 @@ def classify(pkg: str, file: str, content: str) -> str | None:
if is_script and re.search(rf"\bimport\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "dynamic_import"
# require / require.resolve
- if is_script and re.search(
- rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content
- ):
+ if is_script and re.search(rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "require"
# Re-exports: `export * from "pkg"`, `export { x } from "pkg"`,
# `export type { Foo } from "pkg"`. Multi-line supported.
@@ -289,16 +285,12 @@ def classify(pkg: str, file: str, content: str) -> str | None:
# segment bounded by a quote / `#` / `?` or a subpath `/`, so
# `/node_modules/foo-extra/...` is NOT treated as usage of `foo`.
html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])"
- if is_html and re.search(
- rf""
- )
+ req = _build_request("127.0.0.1:8902", origin = "data:text/html,")
assert _is_same_origin_request(req) is False
@@ -150,7 +148,6 @@ def test_is_same_origin_request_localhost_vs_127_is_cross_origin():
def test_is_same_origin_request_127_vs_localhost_is_cross_origin():
from main import _is_same_origin_request
-
req = _build_request("localhost:8902", origin = "http://127.0.0.1:8902")
assert _is_same_origin_request(req) is False
diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py
index ebd9c6c722..faf5a67873 100644
--- a/studio/backend/tests/test_inference_model_validation.py
+++ b/studio/backend/tests/test_inference_model_validation.py
@@ -202,10 +202,7 @@ def test_walkback_skips_explicitly_consumed_tool_call_id():
{"role": "tool", "content": "second result"},
]
)
- assert [m.tool_call_id for m in req.messages if m.role == "tool"] == [
- "call_a",
- "call_b",
- ]
+ assert [m.tool_call_id for m in req.messages if m.role == "tool"] == ["call_a", "call_b"]
def test_walkback_handles_malformed_function_string():
diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py
index d52a58a25c..001b5f1bee 100644
--- a/studio/backend/tests/test_kv_cache_estimation.py
+++ b/studio/backend/tests/test_kv_cache_estimation.py
@@ -128,7 +128,9 @@ def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
def _backend_from_gguf(
- arch: str, fields: dict, general: dict | None = None
+ arch: str,
+ fields: dict,
+ general: dict | None = None,
) -> LlamaCppBackend:
"""Create a LlamaCppBackend with parsed GGUF metadata from given fields.
@@ -346,8 +348,7 @@ class TestArchSwaPatternDefaults:
assert kv_default > 0
assert kv_legacy > 0
assert kv_default < kv_legacy, (
- f"arch fallback should under-shoot legacy estimate: "
- f"{kv_default} >= {kv_legacy}"
+ f"arch fallback should under-shoot legacy estimate: " f"{kv_default} >= {kv_legacy}"
)
def test_scalar_sliding_window_pattern_expanded(self):
@@ -430,25 +431,12 @@ class TestDynamicSwaResolver:
from core.inference.llama_cpp import _period_from_layer_types
# gemma3 (1 global per 6), gpt-oss (alternating), gemma3n (1 per 5).
- assert (
- _period_from_layer_types(
- (["sliding_attention"] * 5 + ["full_attention"]) * 4
- )
- == 6
- )
- assert (
- _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2
- )
- assert (
- _period_from_layer_types(
- (["sliding_attention"] * 4 + ["full_attention"]) * 7
- )
- == 5
- )
+ assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6
+ assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2
+ assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5
def test_period_from_layer_types_returns_none_for_aperiodic(self):
from core.inference.llama_cpp import _period_from_layer_types
-
lt = [
"sliding_attention",
"full_attention",
@@ -469,9 +457,7 @@ class TestDynamicSwaResolver:
== "google/gemma-3-1b-it"
)
assert (
- _hf_repo_from_url(
- "https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json"
- )
+ _hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json")
== "google/gemma-3-1b-it"
)
for bad in [
@@ -524,9 +510,7 @@ class TestDynamicSwaResolver:
b = _backend_from_gguf(
"newmodel",
_SWA_FIELDS,
- general = {
- "general.source.huggingface.repository": "vendor/newmodel-1b-instruct"
- },
+ general = {"general.source.huggingface.repository": "vendor/newmodel-1b-instruct"},
)
assert b._sliding_window_pattern == [(i + 1) % 4 != 0 for i in range(12)]
assert calls == ["vendor/newmodel-1b-instruct"]
@@ -573,9 +557,7 @@ class TestDynamicSwaResolver:
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None)
# Force the failure into the Tier 3 path; bypass Tier 2.5.
- monkeypatch.setattr(
- lc, "_resolve_swa_entry_from_transformers", lambda arch: None
- )
+ monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None)
b = _backend_from_gguf(
"newmodel",
_SWA_FIELDS,
@@ -621,18 +603,14 @@ class TestTransformersIntrospection:
class _FakeLazyMapping(dict):
def __getitem__(self, k):
- return (
- _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k)
- )
+ return _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k)
import sys, types as _types
fake_auto = _types.ModuleType("transformers.models.auto.configuration_auto")
fake_auto.CONFIG_MAPPING_NAMES = {"brokenarch": "FakeBroken"}
fake_auto.CONFIG_MAPPING = _FakeLazyMapping({"brokenarch": "FakeBroken"})
- monkeypatch.setitem(
- sys.modules, "transformers.models.auto.configuration_auto", fake_auto
- )
+ monkeypatch.setitem(sys.modules, "transformers.models.auto.configuration_auto", fake_auto)
assert lc._resolve_swa_entry_from_transformers("brokenarch") == 7
def test_returns_none_when_transformers_unavailable(self, monkeypatch):
@@ -658,12 +636,9 @@ class TestTransformersIntrospection:
def test_returns_none_for_arch_unknown_to_transformers(self):
from core.inference.llama_cpp import _resolve_swa_entry_from_transformers
-
assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None
- def test_full_resolver_uses_transformers_before_hf_fetch(
- self, monkeypatch, tmp_path
- ):
+ def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path):
# With bootstrap empty, Tier 2.5 must answer before Tier 3 fires.
self._isolate_cache(monkeypatch, tmp_path)
from core.inference import llama_cpp as lc
@@ -1328,8 +1303,7 @@ class TestServerFlags:
"_kv_key_length": 256,
"_kv_value_length": 256,
"_sliding_window": 512,
- "_sliding_window_pattern": [True, True, True, True, True, False] * 4
- + [True, True],
+ "_sliding_window_pattern": [True, True, True, True, True, False] * 4 + [True, True],
}
defaults.update(overrides)
b = LlamaCppBackend()
@@ -1385,9 +1359,7 @@ class TestServerFlags:
def test_swa_full_suppresses_checkpoint_term(self):
b = self._swa_backend()
with_cp = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8)
- with_cp_full = b._estimate_kv_cache_bytes(
- 8192, "f16", ctx_checkpoints = 8, swa_full = True
- )
+ with_cp_full = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8, swa_full = True)
no_cp_full = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True)
# Checkpoints only matter when SWA layers don't already keep n_ctx.
assert with_cp_full == no_cp_full
@@ -1405,9 +1377,7 @@ class TestServerFlags:
for slots in (1, 2, 4, 8):
for unified in (True, False):
assert (
- b._estimate_kv_cache_bytes(
- 4096, "f16", n_parallel = slots, kv_unified = unified
- )
+ b._estimate_kv_cache_bytes(4096, "f16", n_parallel = slots, kv_unified = unified)
== baseline
)
@@ -1416,9 +1386,7 @@ class TestServerFlags:
baseline = b._estimate_kv_cache_bytes(4096, "f16")
for unified in (True, False):
assert (
- b._estimate_kv_cache_bytes(
- 4096, "f16", n_parallel = 0, kv_unified = unified
- )
+ b._estimate_kv_cache_bytes(4096, "f16", n_parallel = 0, kv_unified = unified)
== baseline
)
@@ -1432,9 +1400,7 @@ class TestServerFlags:
per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back
per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1
global_bytes = sum(
- ctx * per_token_global
- for f in b._sliding_window_pattern[: b._n_layers]
- if not f
+ ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
)
swa_bytes_per_slot = sum(
per_slot_swa_cells * per_token_swa
@@ -1445,16 +1411,12 @@ class TestServerFlags:
assert global_bytes + swa_bytes_per_slot == baseline
# Only SWA portion scales by parallel
for slots in (1, 2, 3, 4):
- scaled = b._estimate_kv_cache_bytes(
- ctx, "f16", n_parallel = slots, kv_unified = False
- )
+ scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
# SWA cells get clamped to per_slot_ctx when ctx/slots < 2*swa
per_slot_ctx = max(1, ctx // slots)
cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bps = sum(
- cells * per_token_swa
- for f in b._sliding_window_pattern[: b._n_layers]
- if f
+ cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
)
assert scaled == global_bytes + slots * swa_bps
@@ -1469,9 +1431,7 @@ class TestServerFlags:
for slots in (1, 2, 4, 8):
for unified in (True, False):
assert (
- b._estimate_kv_cache_bytes(
- 8192, "f16", n_parallel = slots, kv_unified = unified
- )
+ b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified)
== baseline
)
@@ -1493,9 +1453,7 @@ class TestServerFlags:
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
# 22 SWA layers * 4 checkpoints * 512 cells * 4 heads * (256+256) * 2 bytes
- n_swa_layers = sum(
- 1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f
- )
+ n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f)
per_layer = 4 * 512 * 4 * (256 + 256) * 2
assert flagged == baseline + n_swa_layers * per_layer
@@ -1529,9 +1487,7 @@ class TestServerFlags:
flagged = b._estimate_kv_cache_bytes(
ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False
)
- assert flagged == global_bytes + slots * (
- swa_bytes_per_slot + cp_extra_per_slot
- )
+ assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot)
# ── --kv-offload (kv_on_gpu) ───────────────────────────────────
@@ -1655,9 +1611,7 @@ class TestParallelSWAScaling:
"_kv_value_length": 256,
"_sliding_window": 512,
# 15 SWA + 3 global, mirrors gemma-3-270m
- "_sliding_window_pattern": [
- t == "swa" for t in (["swa"] * 5 + ["global"]) * 3
- ],
+ "_sliding_window_pattern": [t == "swa" for t in (["swa"] * 5 + ["global"]) * 3],
}
defaults.update(overrides)
b = LlamaCppBackend()
@@ -1673,9 +1627,7 @@ class TestParallelSWAScaling:
for slots in (1, 2, 4, 8):
for unified in (True, False):
assert (
- b._estimate_kv_cache_bytes(
- 8192, "f16", n_parallel = slots, kv_unified = unified
- )
+ b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified)
== baseline
)
@@ -1729,9 +1681,7 @@ class TestParallelSWAScaling:
cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bps = n_swa * cells * per_token
for unified in (True, False):
- got = b._estimate_kv_cache_bytes(
- ctx, "f16", n_parallel = slots, kv_unified = unified
- )
+ got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
assert got == global_bytes + slots * swa_bps
def test_swa_fallback_scales_only_swa_portion(self):
@@ -1776,8 +1726,7 @@ class TestParallelSWAScaling:
baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
for slots in (1, 2, 4, 8):
assert (
- b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots)
- == baseline
+ b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline
)
# ── kv_unified: no-op for memory math ──────────────────────────
@@ -1791,12 +1740,8 @@ class TestParallelSWAScaling:
]
for label, b in backends:
for slots in (1, 2, 4, 8):
- u = b._estimate_kv_cache_bytes(
- 8192, "f16", n_parallel = slots, kv_unified = True
- )
- nu = b._estimate_kv_cache_bytes(
- 8192, "f16", n_parallel = slots, kv_unified = False
- )
+ u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
+ nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
assert u == nu, f"{label} parallel={slots} unified-mismatch"
# ── Empirical Gemma-3 270m formula ─────────────────────────────
@@ -1937,9 +1882,7 @@ class TestSharedKVLayers:
assert full_in_unshared == 4
kv_per = 4 * (256 + 256) * 2
swa_cells = min(ctx, 2 * 1024)
- expected = (
- full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per
- )
+ expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
def test_shared_layers_reduces_estimate(self):
@@ -1995,9 +1938,7 @@ class TestSharedKVLayers:
per_slot_ctx = max(1, ctx // slots)
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token
- flagged = b._estimate_kv_cache_bytes(
- ctx, "f16", n_parallel = slots, kv_unified = False
- )
+ flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
assert flagged == global_bytes + slots * swa_bytes_per_slot
def test_composes_with_ctx_checkpoints(self):
diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
index 5d2d672890..f887f7747c 100644
--- a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
+++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
@@ -134,12 +134,8 @@ def test_unknown_gpu_not_in_families():
def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows):
host = _make_rocm_host(gfx, windows = windows)
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
- result = resolve_lemonade_rocm_choice(
- host, os_prefix, "default", llama_tag = "latest"
- )
- assert (
- result is not None
- ), f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})"
+ result = resolve_lemonade_rocm_choice(host, os_prefix, "default", llama_tag = "latest")
+ assert result is not None, f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})"
assert _lookup_family(gfx) in result.name
assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm")
@@ -213,9 +209,7 @@ def test_simple_policy_plans_lemonade_for_windows_hip_host():
"assets": [],
}
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
- plan = direct_upstream_release_plan(
- release, host, "ggml-org/llama.cpp", "latest"
- )
+ plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt"
kinds = [a.install_kind for a in plan.attempts]
assert (
@@ -245,9 +239,7 @@ def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavaila
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None
kinds = [a.install_kind for a in plan.attempts]
- assert (
- "windows-hip" in kinds
- ), f"upstream HIP asset not included as fallback; got {kinds}"
+ assert "windows-hip" in kinds, f"upstream HIP asset not included as fallback; got {kinds}"
hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip")
assert hip_attempt.source_label == "upstream"
@@ -294,9 +286,7 @@ def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = bad_release):
- res = resolve_lemonade_rocm_choice(
- host, "ubuntu", "linux-rocm", llama_tag = "latest"
- )
+ res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
@@ -349,9 +339,7 @@ def test_lemonade_resolver_rejects_empty_browser_download_url():
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = release):
- res = resolve_lemonade_rocm_choice(
- host, "ubuntu", "linux-rocm", llama_tag = "latest"
- )
+ res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py
index 6fe5372147..ee8d54443a 100644
--- a/studio/backend/tests/test_llama_cpp_context_fit.py
+++ b/studio/backend/tests/test_llama_cpp_context_fit.py
@@ -143,7 +143,11 @@ def _drive(
model_size = int(model_gib * GIB)
cache_type_kv = None
- def fake_estimate(n_ctx_, _type = None, **_kwargs):
+ def fake_estimate(
+ n_ctx_,
+ _type = None,
+ **_kwargs,
+ ):
return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes
inst._estimate_kv_cache_bytes = fake_estimate
@@ -233,9 +237,7 @@ def _drive(
elif gpus:
gpu_indices, use_fit = inst._select_gpus(model_size, gpus)
if use_fit and not explicit_ctx:
- effective_ctx = (
- min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX
- )
+ effective_ctx = min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX
return {
"c_arg": effective_ctx if effective_ctx > 0 else 0,
diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py
index b32aeefcdb..c7dd111ee3 100644
--- a/studio/backend/tests/test_llama_cpp_freshness.py
+++ b/studio/backend/tests/test_llama_cpp_freshness.py
@@ -172,9 +172,7 @@ def test_latest_published_release_returns_none_on_network_failure(monkeypatch):
assert fr.latest_published_release("unslothai/llama.cpp") is None
-def test_latest_published_release_keeps_old_cache_on_transient_failure(
- monkeypatch, tmp_path
-):
+def test_latest_published_release_keeps_old_cache_on_transient_failure(monkeypatch, tmp_path):
# Disk entry older than TTL + network fail -> return cached value.
cache_dir = tmp_path / ".freshness"
cache_dir.mkdir()
@@ -188,9 +186,7 @@ def test_latest_published_release_keeps_old_cache_on_transient_failure(
# check_prebuilt_freshness end-to-end.
-def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(
- monkeypatch, tmp_path
-):
+def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
_write_marker(
install_dir,
@@ -200,9 +196,7 @@ def test_check_prebuilt_freshness_reports_stale_when_old_and_behind(
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
- monkeypatch.setattr(
- fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
- )
+ monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["has_marker"] is True
assert info["stale"] is True
@@ -222,9 +216,7 @@ def test_check_prebuilt_freshness_not_stale_when_tag_matches(monkeypatch, tmp_pa
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
- monkeypatch.setattr(
- fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
- )
+ monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["stale"] is False
assert info["installed_tag"] == "b9300"
@@ -242,9 +234,7 @@ def test_check_prebuilt_freshness_not_stale_within_threshold(monkeypatch, tmp_pa
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
- monkeypatch.setattr(
- fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
- )
+ monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["stale"] is False
assert info["age_days"] == 1
@@ -257,9 +247,7 @@ def test_check_prebuilt_freshness_fails_open_without_marker(tmp_path):
assert info["stale"] is False
-def test_check_prebuilt_freshness_fails_open_when_github_unreachable(
- monkeypatch, tmp_path
-):
+def test_check_prebuilt_freshness_fails_open_when_github_unreachable(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
_write_marker(
install_dir,
@@ -276,15 +264,11 @@ def test_check_prebuilt_freshness_fails_open_when_github_unreachable(
assert info["latest_tag"] is None
-def test_check_prebuilt_freshness_handles_unparseable_install_timestamp(
- monkeypatch, tmp_path
-):
+def test_check_prebuilt_freshness_handles_unparseable_install_timestamp(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
_write_marker(install_dir, tag = "b9190", installed_at_utc = "not-a-date")
bin_path = _fake_binary(install_dir, layout = "root")
- monkeypatch.setattr(
- fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
- )
+ monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300")
info = fr.check_prebuilt_freshness(str(bin_path))
assert info["stale"] is False
assert info["age_days"] is None
@@ -300,9 +284,7 @@ def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_pat
.replace("+00:00", "Z"),
)
bin_path = _fake_binary(install_dir, layout = "root")
- monkeypatch.setattr(
- fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300"
- )
+ monkeypatch.setattr(fr, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9300")
info = fr.check_prebuilt_freshness(str(bin_path), threshold_days = 1)
assert info["stale"] is True
@@ -311,9 +293,7 @@ def test_check_prebuilt_freshness_respects_custom_threshold(monkeypatch, tmp_pat
def test_format_stale_warning_contains_actionable_command():
- msg = fr.format_stale_warning(
- {"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5}
- )
+ msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 5})
assert "b9190" in msg
assert "b9300" in msg
assert "5 days" in msg
@@ -321,8 +301,6 @@ def test_format_stale_warning_contains_actionable_command():
def test_format_stale_warning_singular_day():
- msg = fr.format_stale_warning(
- {"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1}
- )
+ msg = fr.format_stale_warning({"installed_tag": "b9190", "latest_tag": "b9300", "age_days": 1})
assert "1 day" in msg
assert "1 days" not in msg
diff --git a/studio/backend/tests/test_llama_cpp_load_progress.py b/studio/backend/tests/test_llama_cpp_load_progress.py
index f46751b798..f95d8bf1a4 100644
--- a/studio/backend/tests/test_llama_cpp_load_progress.py
+++ b/studio/backend/tests/test_llama_cpp_load_progress.py
@@ -150,7 +150,6 @@ class TestLoadProgressSingleShard:
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
-
return io.StringIO(f"Name:\ttest\nVmRSS:\t{10 * 1024 ** 2}\tkB\n")
return open(path, *args, **kwargs) # fall through
@@ -175,7 +174,6 @@ class TestLoadProgressSingleShard:
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
-
return io.StringIO(f"VmRSS:\t{8 * 1024 ** 2}\tkB\n")
return open(path, *args, **kwargs)
@@ -210,7 +208,6 @@ class TestLoadProgressMultiShard:
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
-
return io.StringIO("VmRSS:\t0\tkB\n")
return open(path, *args, **kwargs)
@@ -233,7 +230,6 @@ class TestLoadProgressDegradation:
def fake_open(path, *args, **kwargs):
if str(path).startswith("/proc/"):
import io
-
return io.StringIO("VmRSS:\t1024\tkB\n")
return open(path, *args, **kwargs)
diff --git a/studio/backend/tests/test_llama_cpp_load_progress_live.py b/studio/backend/tests/test_llama_cpp_load_progress_live.py
index beed8713c1..44a8f00834 100644
--- a/studio/backend/tests/test_llama_cpp_load_progress_live.py
+++ b/studio/backend/tests/test_llama_cpp_load_progress_live.py
@@ -75,7 +75,11 @@ pytestmark = pytest.mark.skipif(
)
-def _make_backend(pid: int, gguf_path: str, healthy: bool = False):
+def _make_backend(
+ pid: int,
+ gguf_path: str,
+ healthy: bool = False,
+):
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._process = type("P", (), {"pid": pid})()
inst._gguf_path = gguf_path
diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py
index 22e4cda7d1..aa0198892f 100644
--- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py
+++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py
@@ -109,7 +109,12 @@ def _make_backend(native_ctx = 131072):
return inst
-def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes = 325_000):
+def _compute_max_available_ctx(
+ native_ctx,
+ model_gib,
+ gpus,
+ kv_per_token_bytes = 325_000,
+):
"""Run the ceiling-probe block from load_model and return the final
``max_available_ctx`` value the backend would assign to
``_max_context_length``.
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 4a8276adc0..9e7944913a 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -77,9 +77,7 @@ def _enc_kv_string(key: str, value: str) -> bytes:
def _enc_kv_uint32(key: str, value: int) -> bytes:
- return (
- _enc_string(key) + struct.pack(" Path:
"""Bash stub that prints `help_text` on --help."""
- path.write_text("#!/usr/bin/env bash\n" f"cat <<'EOF'\n{help_text}\nEOF\n")
+ path.write_text(f"#!/usr/bin/env bash\ncat <<'EOF'\n{help_text}\nEOF\n")
path.chmod(0o755)
return path
@@ -532,8 +530,7 @@ def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
# Renamed upstream: draft-mtp -> mtp.
fake = _make_fake_llama_server(
tmp_path / "llama-server",
- "--spec-type [none|mtp|ngram-cache|ngram-simple|ngram-map-k|"
- "ngram-map-k4v|ngram-mod]",
+ "--spec-type [none|mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]",
)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
@@ -655,14 +652,7 @@ def test_build_ngram_mod_flags_new():
def test_build_ngram_mod_flags_legacy():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "legacy"})
- assert flags == [
- "--spec-ngram-size-n",
- "24",
- "--draft-min",
- "48",
- "--draft-max",
- "64",
- ]
+ assert flags == ["--spec-ngram-size-n", "24", "--draft-min", "48", "--draft-max", "64"]
def test_build_ngram_mod_flags_empty_when_unsupported():
@@ -672,9 +662,7 @@ def test_build_ngram_mod_flags_empty_when_unsupported():
def test_build_ngram_mod_flags_respects_custom_values():
- flags = _build_ngram_mod_flags(
- {"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32
- )
+ flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"}, n_match = 16, n_min = 24, n_max = 32)
assert flags == [
"--spec-ngram-mod-n-match",
"16",
@@ -826,9 +814,7 @@ def _patch_probe(monkeypatch, ngram_supported):
)
-def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported(
- monkeypatch,
-):
+def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported(monkeypatch):
# 0.8B MTP request -- load_model would have promoted to ngram-mod
# (no MTP head); reload check must match a ngram-mod backend.
_patch_probe(monkeypatch, ngram_supported = True)
@@ -1011,7 +997,12 @@ def test_canonicalize_spec_mode(value, expected):
# ── _build_speculative_flags resolver matrix ──────────────────────
-def _resolver_backend(monkeypatch, *, ngram_supported = True, mtp_token = "draft-mtp"):
+def _resolver_backend(
+ monkeypatch,
+ *,
+ ngram_supported = True,
+ mtp_token = "draft-mtp",
+):
"""Backend with a deterministic probe so the resolver is hermetic."""
fake = {
"found": True,
@@ -1093,13 +1084,7 @@ _SUB_3B_MTP_MODEL = "unsloth/Qwen3.5-0.8B-MTP-GGUF"
],
)
def test_build_speculative_flags_matrix(
- monkeypatch,
- requested,
- gpus,
- model,
- expect_spec_type,
- expect_n_max,
- expect_ngram_knobs,
+ monkeypatch, requested, gpus, model, expect_spec_type, expect_n_max, expect_ngram_knobs
):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
index e647ff2c7d..202fd36c86 100644
--- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py
+++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
@@ -30,9 +30,7 @@ sys.modules.setdefault("loggers", _loggers_stub)
# Give the structlog stub a real get_logger: a bare ModuleType poisons
# sys.modules for later tests that call structlog.get_logger at import time.
_structlog_stub = _types.ModuleType("structlog")
-_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger(
- "structlog"
-)
+_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
if not hasattr(sys.modules["structlog"], "get_logger"):
sys.modules["structlog"].get_logger = _structlog_stub.get_logger
@@ -109,8 +107,7 @@ class TestUnsupportedNonDiffusionArchitecture:
class TestOllamaAndFallback:
_OLLAMA_GGUF = (
- f"/home/u/.ollama{__import__('os').sep}ollama_links"
- f"{__import__('os').sep}m.gguf"
+ f"/home/u/.ollama{__import__('os').sep}ollama_links" f"{__import__('os').sep}m.gguf"
)
def test_ollama_compat_message_still_works(self):
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index 00295d6283..125b13782a 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -131,9 +131,7 @@ def test_stale_kill_skips_wait():
LlamaCppBackend._wait_for_vram_settle(
**_kw(since_kill = long_ago, max_wait = 2.0, interval = 0.25)
)
- assert (
- state["calls"] == 0
- ), "kill older than _VRAM_SETTLE_WINDOW_S must skip the wait"
+ assert state["calls"] == 0, "kill older than _VRAM_SETTLE_WINDOW_S must skip the wait"
def test_empty_first_sample_returns_immediately():
@@ -221,9 +219,7 @@ def test_max_wait_respected_when_probe_is_slow():
elapsed = time.monotonic() - start
# First probe (0.30 s) + at most one short clipped sleep + bail.
# Hard cap well below the old behaviour of 0.30 + 0.25 + 0.30 = 0.85.
- assert (
- elapsed < 0.85
- ), f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s"
+ assert elapsed < 0.85, f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s"
def test_gpu_index_set_change_returns():
diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
index 7d4719c0e7..a5c9d2255f 100644
--- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
+++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
@@ -185,9 +185,7 @@ class TestWindowsPipNvidiaDllDirs:
# If sys.prefix points to a path that doesn't exist (unusual,
# but possible during test setup), the resolver must just
# return [] rather than raising.
- result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(
- "/this/path/does/not/exist/anywhere"
- )
+ result = LlamaCppBackend._windows_pip_nvidia_dll_dirs("/this/path/does/not/exist/anywhere")
assert result == []
def test_picks_up_cu13_bin_x86_64_layout(self, tmp_path):
@@ -196,9 +194,7 @@ class TestWindowsPipNvidiaDllDirs:
# ``nvidia/cu13/bin/x86_64/`` instead of ``nvidia//bin/``.
# Without this, users on the new CUDA 13 wheel generation hit
# the original #5106 failure mode.
- dll_dir = (
- tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
- )
+ dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
dll_dir.mkdir(parents = True)
for name in ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"):
(dll_dir / name).write_bytes(b"")
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index 2f7431497d..2c7362a9ff 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -20,12 +20,7 @@ import pytest
# 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"
-)
+_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)
@@ -359,14 +354,7 @@ def test_strip_shadowing_flags_keeps_cache_when_cache_disabled():
["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"],
strip_cache = False,
)
- assert out == [
- "--cache-type-k",
- "q8_0",
- "--cache-type-v",
- "q8_0",
- "--top-k",
- "20",
- ]
+ assert out == ["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"]
def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
@@ -374,14 +362,7 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"],
strip_spec = False,
)
- assert out == [
- "--spec-type",
- "ngram-mod",
- "--draft-min",
- "48",
- "--top-k",
- "20",
- ]
+ assert out == ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"]
def test_strip_shadowing_flags_drops_mtp_flags_when_requested():
@@ -505,9 +486,7 @@ def test_strip_shadowing_flags_jinja_boolean_preserves_positional():
def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional():
- out = strip_shadowing_flags(
- ["--no-jinja", "trailing-positional"], strip_template = True
- )
+ out = strip_shadowing_flags(["--no-jinja", "trailing-positional"], strip_template = True)
assert out == ["trailing-positional"]
diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py
index c8498d4857..1b084cf436 100644
--- a/studio/backend/tests/test_login_rate_limit.py
+++ b/studio/backend/tests/test_login_rate_limit.py
@@ -45,9 +45,12 @@ def env_trust_proxy(monkeypatch):
class _FakeRequest:
- def __init__(self, client_host = "127.0.0.1", headers = None):
+ def __init__(
+ self,
+ client_host = "127.0.0.1",
+ headers = None,
+ ):
from starlette.datastructures import Headers
-
self.client = type("Client", (), {"host": client_host})()
self.headers = Headers(headers or {})
@@ -58,12 +61,10 @@ class _FakeRequest:
class TestClientIp:
def test_uses_request_client_host_by_default(self, env_no_proxy):
from routes.auth import _client_ip
-
assert _client_ip(_FakeRequest("203.0.113.5")) == "203.0.113.5"
def test_ignores_xff_when_trust_off(self, env_no_proxy):
from routes.auth import _client_ip
-
req = _FakeRequest(
"127.0.0.1",
{"x-forwarded-for": "198.51.100.7, 10.0.0.1"},
@@ -74,7 +75,6 @@ class TestClientIp:
def test_honours_first_xff_when_trust_on(self, env_trust_proxy):
from routes.auth import _client_ip
-
req = _FakeRequest(
"127.0.0.1",
{"x-forwarded-for": "198.51.100.7, 10.0.0.1"},
@@ -83,12 +83,10 @@ class TestClientIp:
def test_falls_back_to_client_host_when_xff_missing(self, env_trust_proxy):
from routes.auth import _client_ip
-
assert _client_ip(_FakeRequest("203.0.113.9")) == "203.0.113.9"
def test_honours_forwarded_header_when_trust_on(self, env_trust_proxy):
from routes.auth import _client_ip
-
req = _FakeRequest(
"127.0.0.1",
{"forwarded": 'for="198.51.100.42";proto=https'},
@@ -104,34 +102,22 @@ class TestClientIp:
def test_xff_strips_ipv4_port(self, env_trust_proxy):
from routes.auth import _client_ip
-
- req = _FakeRequest(
- "127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"}
- )
+ req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"})
assert _client_ip(req) == "198.51.100.7"
def test_xff_strips_bracketed_ipv6_port(self, env_trust_proxy):
from routes.auth import _client_ip
-
- req = _FakeRequest(
- "127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"}
- )
+ req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"})
assert _client_ip(req) == "2001:db8::1"
def test_forwarded_strips_ipv4_port(self, env_trust_proxy):
from routes.auth import _client_ip
-
- req = _FakeRequest(
- "127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'}
- )
+ req = _FakeRequest("127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'})
assert _client_ip(req) == "198.51.100.7"
def test_forwarded_strips_bracketed_ipv6_port(self, env_trust_proxy):
from routes.auth import _client_ip
-
- req = _FakeRequest(
- "127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'}
- )
+ req = _FakeRequest("127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'})
assert _client_ip(req) == "2001:db8::1"
def test_forwarded_isolates_first_element(self, env_trust_proxy):
@@ -247,9 +233,7 @@ class TestLogin429Body:
import secrets as _secrets
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
- monkeypatch.setattr(
- storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password"
- )
+ monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password")
monkeypatch.setattr(storage, "_bootstrap_password", None)
storage.create_initial_user(
username = storage.DEFAULT_ADMIN_USERNAME,
diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py
index 10a6eb012b..0d263d7630 100644
--- a/studio/backend/tests/test_mcp_servers.py
+++ b/studio/backend/tests/test_mcp_servers.py
@@ -45,9 +45,7 @@ def test_list_servers_ordered_by_created_at(tmp_path, monkeypatch):
def test_update_server_coerces_bools(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
- assert mcp_servers_db.update_server(
- "srv1", {"is_enabled": False, "use_oauth": True}
- )
+ assert mcp_servers_db.update_server("srv1", {"is_enabled": False, "use_oauth": True})
row = mcp_servers_db.get_server("srv1")
assert row["is_enabled"] == 0
assert row["use_oauth"] == 1
@@ -81,7 +79,6 @@ def test_validate_url_accepts_http_and_https():
@pytest.mark.parametrize("bad", ["", " ", "ftp://x", "http://", "noscheme.com"])
def test_validate_url_rejects_bad(bad):
from routes.mcp_servers import _validate_url
-
with pytest.raises(HTTPException) as exc:
_validate_url(bad)
assert exc.value.status_code == 400
@@ -90,9 +87,7 @@ def test_validate_url_rejects_bad(bad):
def test_normalize_headers():
from routes.mcp_servers import _normalize_headers
- assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == {
- "Auth": "Bearer x"
- }
+ assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == {"Auth": "Bearer x"}
assert _normalize_headers({"X": 42}) == {"X": "42"}
assert _normalize_headers({}) is None
assert _normalize_headers(None) is None
@@ -104,15 +99,12 @@ def test_changes_from_payload_tristate_headers():
from models.mcp_servers import McpServerUpdate
# omitted → key absent
- assert "headers_json" not in _changes_from_payload(
- McpServerUpdate(display_name = "x")
- )
+ assert "headers_json" not in _changes_from_payload(McpServerUpdate(display_name = "x"))
# null → stored as None (clear all headers)
assert _changes_from_payload(McpServerUpdate(headers = None))["headers_json"] is None
# dict → serialised JSON
assert (
- _changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"]
- == '{"a": "1"}'
+ _changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"] == '{"a": "1"}'
)
@@ -135,7 +127,6 @@ def test_mcp_specs_skip_oversized_names():
def test_execute_tool_malformed_mcp_name():
from core.inference.tools import execute_tool
-
out = execute_tool("mcp__no_double_underscore", {})
assert out.startswith("Error: malformed MCP tool name")
@@ -143,11 +134,7 @@ def test_execute_tool_malformed_mcp_name():
def test_execute_tool_unknown_server(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
from core.inference.tools import execute_tool
-
- assert (
- execute_tool("mcp__missing__do_thing", {})
- == "Error: MCP server 'missing' not found"
- )
+ assert execute_tool("mcp__missing__do_thing", {}) == "Error: MCP server 'missing' not found"
def test_execute_tool_disabled_server(tmp_path, monkeypatch):
@@ -160,10 +147,7 @@ def test_execute_tool_disabled_server(tmp_path, monkeypatch):
)
from core.inference.tools import execute_tool
- assert (
- execute_tool("mcp__srv1__do_thing", {})
- == "Error: MCP server 'srv1' is disabled"
- )
+ assert execute_tool("mcp__srv1__do_thing", {}) == "Error: MCP server 'srv1' is disabled"
def test_mcp_specs_skip_invalid_openai_function_names():
@@ -219,7 +203,6 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
async def call_tool(self, name, args):
import asyncio as _asyncio
-
await _asyncio.sleep(30) # never finishes within the test
monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
@@ -439,8 +422,7 @@ def test_tool_healing_strip_handles_hyphenated_function_names():
from core.tool_healing import strip_tool_call_markup
out = strip_tool_call_markup(
- "before "
- "x after"
+ "before x after"
)
assert out == "before after"
@@ -558,9 +540,7 @@ def test_tool_xml_parser_handles_hyphenated_function_names():
from core.inference.tool_call_parser import parse_tool_calls_from_text
calls = parse_tool_calls_from_text(
- ""
- "octocat/hello"
- ""
+ "octocat/hello"
)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "mcp__srv__list-issues"
@@ -584,8 +564,7 @@ def test_tool_xml_strip_handles_hyphenated_function_names():
rx = ns["_TOOL_XML_RE"]
stripped = rx.sub(
"",
- "before "
- "x after",
+ "before x after",
)
assert stripped == "before after"
diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py
index e980e6a057..515e39d5b6 100644
--- a/studio/backend/tests/test_mcp_stdio_improvements.py
+++ b/studio/backend/tests/test_mcp_stdio_improvements.py
@@ -62,9 +62,7 @@ def test_create_forces_oauth_off_for_stdio(tmp_path, monkeypatch):
_enable(monkeypatch)
resp = asyncio.run(
routes_mcp.create_mcp_server(
- McpServerCreate(
- display_name = "FS", url = "npx -y server /tmp", use_oauth = True
- ),
+ McpServerCreate(display_name = "FS", url = "npx -y server /tmp", use_oauth = True),
current_subject = "u",
)
)
@@ -94,12 +92,8 @@ def test_update_url_to_stdio_clears_oauth(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
- monkeypatch.setattr(
- routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0)
- )
- mcp_servers_db.create_server(
- id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True
- )
+ monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0))
+ mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True)
resp = asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "npx -y server /tmp"), current_subject = "u"
@@ -148,9 +142,7 @@ def test_switch_keeps_explicitly_supplied_headers(tmp_path, monkeypatch):
resp = asyncio.run(
routes_mcp.update_mcp_server(
"s1",
- McpServerUpdate(
- url = "https://remote/mcp", headers = {"Authorization": "Bearer new"}
- ),
+ McpServerUpdate(url = "https://remote/mcp", headers = {"Authorization": "Bearer new"}),
current_subject = "u",
)
)
@@ -171,9 +163,7 @@ def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch):
)
# editing only the display name (still stdio) must not wipe env vars
resp = asyncio.run(
- routes_mcp.update_mcp_server(
- "s1", McpServerUpdate(display_name = "B"), current_subject = "u"
- )
+ routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u")
)
assert resp.headers == {"API_KEY": "secret"}
@@ -183,7 +173,6 @@ def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch):
def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch):
from routes.mcp_servers import _validate_url
-
_enable(monkeypatch)
for bad in ["ftp://host/x", "file:///etc/passwd", "ws://h/y"]:
with pytest.raises(HTTPException) as exc:
@@ -193,12 +182,9 @@ def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch):
def test_validate_url_allows_url_in_argument(monkeypatch):
from routes.mcp_servers import _validate_url
-
_enable(monkeypatch)
# :// inside an ARGUMENT (not the first token) is still a valid command
- assert _validate_url("npx server --url https://x/mcp") == (
- "npx server --url https://x/mcp"
- )
+ assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp")
# ── P6: Data Recipe stdio path obeys the same host gate ─────────────
diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py
index f49edcdaae..277306836b 100644
--- a/studio/backend/tests/test_mcp_stdio_pr5863.py
+++ b/studio/backend/tests/test_mcp_stdio_pr5863.py
@@ -83,9 +83,7 @@ def transport(monkeypatch):
monkeypatch.setattr(
mcp_client,
"_client",
- lambda url, headers, use_oauth = False: _RecordingClient(
- url, headers, use_oauth, recorder
- ),
+ lambda url, headers, use_oauth = False: _RecordingClient(url, headers, use_oauth, recorder),
)
return recorder
@@ -130,9 +128,12 @@ def test_parse_basic_argv():
def test_parse_keeps_url_argument_as_one_command():
# gemini "high": a :// inside an ARGUMENT must not break the command.
- assert mcp_client.parse_stdio_command(
- "npx server --endpoint https://example.com/mcp"
- ) == ["npx", "server", "--endpoint", "https://example.com/mcp"]
+ assert mcp_client.parse_stdio_command("npx server --endpoint https://example.com/mcp") == [
+ "npx",
+ "server",
+ "--endpoint",
+ "https://example.com/mcp",
+ ]
def test_parse_quoted_arg():
@@ -158,9 +159,7 @@ def test_parse_windows_strips_wrapping_quotes(monkeypatch):
# gemini "medium": posix=False keeps backslash paths but also the wrapping
# quotes; the PR strips a matched pair so argv[0] reaches the OS clean.
monkeypatch.setattr(sys, "platform", "win32")
- parts = mcp_client.parse_stdio_command(
- r'"C:\Program Files\node\node.exe" server.js'
- )
+ parts = mcp_client.parse_stdio_command(r'"C:\Program Files\node\node.exe" server.js')
assert parts[0] == r"C:\Program Files\node\node.exe"
assert parts[1] == "server.js"
@@ -242,14 +241,10 @@ def test_validate_url_gate_on_accepts_stdio(monkeypatch):
# http still works when stdio is on
assert _validate_url("https://x/mcp") == "https://x/mcp"
# url-bearing argument accepted as a command
- assert _validate_url("npx server --url https://x/mcp") == (
- "npx server --url https://x/mcp"
- )
+ assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp")
# A lone token is ambiguous; keep the prior behaviour and accept it as a
# command rather than guessing it's a URL (no regression for single binaries).
- assert (
- _validate_url("/usr/local/bin/my-mcp-server") == "/usr/local/bin/my-mcp-server"
- )
+ assert _validate_url("/usr/local/bin/my-mcp-server") == "/usr/local/bin/my-mcp-server"
assert _validate_url("mcp-server-sqlite") == "mcp-server-sqlite"
# empty / unparseable still rejected
for bad in [" ", '"unclosed']:
@@ -336,9 +331,7 @@ def test_refresh_route_gate(tmp_path, monkeypatch, transport):
assert transport == []
_enable(monkeypatch)
- res = asyncio.run(
- routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")
- )
+ res = asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u"))
assert res.ok and res.tool_count == 2
assert len(transport) == 1
@@ -349,9 +342,7 @@ def test_discovery_gate(tmp_path, monkeypatch, transport):
from core.inference.tools import get_enabled_mcp_tools
_reset_db(tmp_path, monkeypatch)
- mcp_servers_db.create_server(
- id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True
- )
+ mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True)
_disable(monkeypatch)
assert asyncio.run(get_enabled_mcp_tools()) == []
@@ -367,9 +358,7 @@ def test_execute_gate(tmp_path, monkeypatch, transport):
from core.inference.tools import execute_tool
_reset_db(tmp_path, monkeypatch)
- mcp_servers_db.create_server(
- id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True
- )
+ mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True)
_disable(monkeypatch)
out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"})
diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py
index 5e5b5a3c3b..2d6efa9f8b 100644
--- a/studio/backend/tests/test_middleware.py
+++ b/studio/backend/tests/test_middleware.py
@@ -24,7 +24,6 @@ if str(_BACKEND_ROOT) not in sys.path:
@pytest.fixture(scope = "module")
def main_module():
import main as _main # noqa: F401
-
return _main
@@ -168,9 +167,7 @@ class TestMaxBodyMiddleware:
assert r.status_code == 200
assert r.json()["total"] == 512
- def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(
- self, main_module
- ):
+ def test_upload_passthrough_rejects_declared_body_over_dedicated_cap(self, main_module):
app = _make_protected_app(
128,
main_module,
@@ -274,9 +271,7 @@ class TestSecurityHeadersMiddleware:
csp = r.headers["content-security-policy"]
assert f"'nonce-{nonce}'" in csp
# Internal handoff header must not leak to clients.
- assert main_module._CSP_SCRIPT_NONCE_HEADER not in {
- k.lower() for k in r.headers.keys()
- }
+ assert main_module._CSP_SCRIPT_NONCE_HEADER not in {k.lower() for k in r.headers.keys()}
def test_build_csp_helper_shape(self, main_module):
plain = main_module._build_csp()
@@ -290,9 +285,7 @@ class TestSecurityHeadersMiddleware:
# this allowlist entry citation favicons fall back to gray initials.
csp = main_module._build_csp()
img_directive = next(
- chunk.strip()
- for chunk in csp.split(";")
- if chunk.strip().startswith("img-src ")
+ chunk.strip() for chunk in csp.split(";") if chunk.strip().startswith("img-src ")
)
# Tokenise and compare with `==` so CodeQL's URL-substring rule does
# not read directive-string `in` membership as URL sanitisation.
diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py
index 16cca7dd40..49afab048d 100644
--- a/studio/backend/tests/test_mlx_inference_backend.py
+++ b/studio/backend/tests/test_mlx_inference_backend.py
@@ -103,8 +103,7 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):
def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite(
- monkeypatch,
- tmp_path,
+ monkeypatch, tmp_path
):
_install_fake_mlx(monkeypatch)
calls = []
@@ -199,7 +198,7 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
return ""
monkeypatch.setattr(
- "core.inference.chat_template_helpers." "apply_chat_template_for_generation",
+ "core.inference.chat_template_helpers.apply_chat_template_for_generation",
_fake_apply,
raising = True,
)
@@ -228,7 +227,11 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
class _Tok:
chat_template = "x"
- def decode(self, ids, skip_special_tokens = False):
+ def decode(
+ self,
+ ids,
+ skip_special_tokens = False,
+ ):
return "hi"
backend = MLXInferenceBackend()
diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py
index c36363b1ae..811971ec36 100644
--- a/studio/backend/tests/test_mlx_training_worker_config.py
+++ b/studio/backend/tests/test_mlx_training_worker_config.py
@@ -45,12 +45,8 @@ def _load_worker_module():
setattr(wheel_utils, name, lambda *_args, **_kwargs: None)
sys.modules["utils.wheel_utils"] = wheel_utils
- worker_path = (
- Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py"
- )
- spec = importlib.util.spec_from_file_location(
- "mlx_training_worker_under_test", worker_path
- )
+ worker_path = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py"
+ spec = importlib.util.spec_from_file_location("mlx_training_worker_under_test", worker_path)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py
index 3481e29948..417ec74d17 100644
--- a/studio/backend/tests/test_models_get_model_config_case_resolution.py
+++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py
@@ -50,9 +50,7 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch):
return _DummyModelConfig()
monkeypatch.setattr(models_route, "is_local_path", lambda _: False)
- monkeypatch.setattr(
- models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model"
- )
+ monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model")
monkeypatch.setattr(models_route, "load_model_defaults", _record_load)
monkeypatch.setattr(models_route, "is_vision_model", _record_vision)
monkeypatch.setattr(models_route, "is_embedding_model", _record_embedding)
diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py
index 4d7528d238..b2986f2c3a 100644
--- a/studio/backend/tests/test_multimodal_document.py
+++ b/studio/backend/tests/test_multimodal_document.py
@@ -317,11 +317,7 @@ def test_openai_base64_pdf_becomes_input_file(monkeypatch):
user_msg = captured["body"]["input"][0]
parts = user_msg["content"]
fileblk = next(p for p in parts if p.get("type") == "input_file")
- assert fileblk == {
- "type": "input_file",
- "file_data": _PDF_DATA_URI,
- "filename": "paper.pdf",
- }
+ assert fileblk == {"type": "input_file", "file_data": _PDF_DATA_URI, "filename": "paper.pdf"}
def test_openai_url_pdf_becomes_input_file(monkeypatch):
@@ -344,10 +340,7 @@ def test_openai_url_pdf_becomes_input_file(monkeypatch):
)
parts = captured["body"]["input"][0]["content"]
fileblk = next(p for p in parts if p.get("type") == "input_file")
- assert fileblk == {
- "type": "input_file",
- "file_url": "https://example.com/doc.pdf",
- }
+ assert fileblk == {"type": "input_file", "file_url": "https://example.com/doc.pdf"}
def test_openai_empty_data_uri_falls_back_to_file_url(monkeypatch):
@@ -512,9 +505,7 @@ def test_build_external_messages_passes_input_document_for_anthropic_and_openai(
)
]
for provider in ("anthropic", "openai"):
- out = _build_external_messages(
- msgs, supports_vision = True, provider_type = provider
- )
+ out = _build_external_messages(msgs, supports_vision = True, provider_type = provider)
assert len(out) == 1, (provider, out)
parts = out[0]["content"]
assert parts[0] == {"type": "text", "text": "summarise"}, provider
@@ -550,9 +541,7 @@ def test_build_external_messages_strips_input_document_for_unmapped_providers():
)
]
for provider in ("gemini", "mistral", "kimi", "openrouter", "deepseek", "qwen"):
- out = _build_external_messages(
- msgs, supports_vision = True, provider_type = provider
- )
+ out = _build_external_messages(msgs, supports_vision = True, provider_type = provider)
assert len(out) == 1, (provider, out)
parts = out[0]["content"]
types = [p.get("type") for p in parts if isinstance(p, dict)]
diff --git a/studio/backend/tests/test_native_context_length.py b/studio/backend/tests/test_native_context_length.py
index 60622c776d..01c05d3ec6 100644
--- a/studio/backend/tests/test_native_context_length.py
+++ b/studio/backend/tests/test_native_context_length.py
@@ -332,9 +332,7 @@ class TestPydanticModels:
def test_status_response_chat_template_roundtrip(self):
"""chat_template serializes and validates as part of status."""
resp = InferenceStatusResponse(chat_template = "{{ messages }}")
- roundtripped = InferenceStatusResponse.model_validate_json(
- resp.model_dump_json()
- )
+ roundtripped = InferenceStatusResponse.model_validate_json(resp.model_dump_json())
assert roundtripped.chat_template == "{{ messages }}"
def test_roundtrip_preserves_value(self):
@@ -390,9 +388,7 @@ class TestRouteCompleteness:
def test_gguf_load_responses_have_field(self):
"""Every GGUF LoadResponse (is_gguf = True) includes native_context_length."""
blocks = self._find_construction_blocks("LoadResponse")
- gguf_blocks = [
- b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b
- ]
+ gguf_blocks = [b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b]
assert (
len(gguf_blocks) >= 2
), f"Expected at least 2 GGUF LoadResponse blocks, found {len(gguf_blocks)}"
@@ -404,9 +400,7 @@ class TestRouteCompleteness:
def test_non_gguf_load_responses_omit_field(self):
"""Non-GGUF LoadResponse blocks do not set native_context_length (defaults to None)."""
blocks = self._find_construction_blocks("LoadResponse")
- non_gguf = [
- b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b
- ]
+ non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b]
# Non-GGUF paths should not reference native_context_length
# (Pydantic defaults it to None, so not setting it is correct)
for block in non_gguf:
@@ -422,7 +416,9 @@ class TestRouteCompleteness:
if "llama_backend" in block and "native_context_length" in block:
found = True
break
- assert found, "No InferenceStatusResponse block with llama_backend has native_context_length"
+ assert (
+ found
+ ), "No InferenceStatusResponse block with llama_backend has native_context_length"
# =====================================================================
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index d3b2f553a2..df7652a590 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -149,8 +149,7 @@ def _siblings(items: dict[str, int]):
"""Mock ``hf_model_info(...).siblings`` payload."""
return _types.SimpleNamespace(
siblings = [
- _types.SimpleNamespace(rfilename = name, size = size)
- for name, size in items.items()
+ _types.SimpleNamespace(rfilename = name, size = size) for name, size in items.items()
],
)
@@ -174,12 +173,8 @@ class TestIterHfCacheSnapshots:
assert list(_iter_hf_cache_snapshots("unsloth/bare")) == []
def test_yields_newest_first(self, hf_cache):
- old = _build_cache(
- hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40
- )
- new = _build_cache(
- hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40
- )
+ old = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40)
+ new = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40)
os.utime(old, (1000, 1000))
os.utime(new, (2000, 2000))
out = list(_iter_hf_cache_snapshots("unsloth/multi"))
@@ -218,9 +213,7 @@ class TestListGgufVariantsFromCache:
class TestListGgufVariantsOffline:
- def test_offline_env_short_circuits_api(
- self, hf_cache, clean_offline_env, monkeypatch
- ):
+ def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch):
_build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1})
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
@@ -232,11 +225,7 @@ class TestListGgufVariantsOffline:
assert len(variants) == 1
assert variants[0].quant == "UD-Q4_K_XL"
- def test_api_exception_falls_back_to_cache(
- self,
- hf_cache,
- clean_offline_env,
- ):
+ def test_api_exception_falls_back_to_cache(self, hf_cache, clean_offline_env):
_build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1})
def boom(*a, **k):
@@ -302,12 +291,7 @@ class TestDetectGgufFromCache:
class TestDetectGgufModelRemoteOffline:
- def test_offline_env_short_circuits_retries(
- self,
- hf_cache,
- clean_offline_env,
- monkeypatch,
- ):
+ def test_offline_env_short_circuits_retries(self, hf_cache, clean_offline_env, monkeypatch):
_build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1})
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
@@ -331,11 +315,7 @@ class TestDetectGgufModelRemoteOffline:
out = detect_gguf_model_remote("unsloth/a")
assert out == "a-Q4_K_M.gguf"
- def test_repository_not_found_does_not_consult_cache(
- self,
- hf_cache,
- clean_offline_env,
- ):
+ def test_repository_not_found_does_not_consult_cache(self, hf_cache, clean_offline_env):
# Cache has a file but the API explicitly says repo is gone.
_build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1})
@@ -430,12 +410,7 @@ class TestHfOfflineIfDnsDead:
assert did_set is False
assert "HF_HUB_OFFLINE" not in os.environ
- def test_user_set_hf_hub_offline_is_preserved(
- self,
- dns,
- clean_offline_env,
- monkeypatch,
- ):
+ def test_user_set_hf_hub_offline_is_preserved(self, dns, clean_offline_env, monkeypatch):
# User explicitly set offline before launching Studio.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
dns.fail()
@@ -445,12 +420,7 @@ class TestHfOfflineIfDnsDead:
# Helper must not pop a variable it did not set.
assert os.environ.get("HF_HUB_OFFLINE") == "1"
- def test_user_set_transformers_offline_is_preserved(
- self,
- dns,
- clean_offline_env,
- monkeypatch,
- ):
+ def test_user_set_transformers_offline_is_preserved(self, dns, clean_offline_env, monkeypatch):
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
dns.fail()
with _hf_offline_if_dns_dead():
@@ -461,11 +431,7 @@ class TestHfOfflineIfDnsDead:
# TRANSFORMERS_OFFLINE pre-existed -> preserved.
assert os.environ.get("TRANSFORMERS_OFFLINE") == "1"
- def test_exception_inside_block_still_restores_env(
- self,
- dns,
- clean_offline_env,
- ):
+ def test_exception_inside_block_still_restores_env(self, dns, clean_offline_env):
dns.fail()
with pytest.raises(RuntimeError, match = "boom"):
with _hf_offline_if_dns_dead():
@@ -503,10 +469,7 @@ class TestDownloadMmprojOfflineCacheFallback:
offline vision GGUF load path returns ``None`` even when the mmproj
is present in cache."""
- def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails(
- self,
- hf_cache,
- ):
+ def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails(self, hf_cache):
_build_cache(
hf_cache,
"unsloth/vision-GGUF",
@@ -520,7 +483,12 @@ class TestDownloadMmprojOfflineCacheFallback:
def boom_list(*a, **k):
raise OSError("offline")
- def fake_download(*, repo_id, filename, token = None):
+ def fake_download(
+ *,
+ repo_id,
+ filename,
+ token = None,
+ ):
# Echo back so the test can verify the cache-resolved filename
return f"/fake/cache/{repo_id}/{filename}"
@@ -551,7 +519,12 @@ class TestDownloadMmprojOfflineCacheFallback:
captured = {}
- def fake_download(*, repo_id, filename, token = None):
+ def fake_download(
+ *,
+ repo_id,
+ filename,
+ token = None,
+ ):
captured["filename"] = filename
return f"/fake/{filename}"
@@ -655,9 +628,7 @@ class TestListGgufVariantsPermanentErrors:
list_gguf_variants("u/gated-gguf")
assert type(exc_info.value).__name__ == "GatedRepoError"
- def test_transient_error_still_falls_back_to_cache(
- self, hf_cache, clean_offline_env
- ):
+ def test_transient_error_still_falls_back_to_cache(self, hf_cache, clean_offline_env):
from utils.models.model_config import list_gguf_variants
_build_cache(hf_cache, "u/transient-gguf", {"foo-Q4_K_M.gguf": 1})
@@ -676,7 +647,6 @@ class TestDetectGgufFromCacheExcludesMmproj:
def test_mmproj_only_returns_none(self, hf_cache):
from utils.models.model_config import _detect_gguf_from_hf_cache
-
_build_cache(
hf_cache,
"u/vision-only-mmproj",
@@ -739,7 +709,6 @@ class TestProbeDnsDeadNoGlobalTimeoutMutation:
# Simulate a wedged resolver: thread blocks forever.
def wedged(host):
import threading
-
threading.Event().wait()
monkeypatch.setattr(_socket, "gethostbyname", wedged)
diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py
index 088be4fcd5..fbb0aa8999 100644
--- a/studio/backend/tests/test_offline_inference_parent.py
+++ b/studio/backend/tests/test_offline_inference_parent.py
@@ -103,10 +103,7 @@ class TestEnvOffline:
class TestTransformersVersionOfflineShortCircuits:
def test_tokenizer_config_skips_urllib_when_offline(
- self,
- monkeypatch,
- clean_offline_env,
- tmp_path,
+ self, monkeypatch, clean_offline_env, tmp_path
):
# No local config + offline env -> must NOT call urlopen.
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
@@ -118,12 +115,7 @@ class TestTransformersVersionOfflineShortCircuits:
with patch("urllib.request.urlopen", boom):
assert _check_tokenizer_config_needs_v5(unique) is False
- def test_config_550_skips_urllib_when_offline(
- self,
- monkeypatch,
- clean_offline_env,
- tmp_path,
- ):
+ def test_config_550_skips_urllib_when_offline(self, monkeypatch, clean_offline_env, tmp_path):
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
unique = f"unsloth/never-cached-{tmp_path.name}-cfg"
@@ -139,9 +131,7 @@ class TestLoraDetectOffline:
OfflineModeIsEnabled; cached adapter_config.json wins."""
def test_hf_model_info_short_circuits_with_OfflineModeIsEnabled(
- self,
- monkeypatch,
- clean_offline_env,
+ self, monkeypatch, clean_offline_env
):
from unittest.mock import MagicMock
@@ -171,10 +161,7 @@ class TestLoraDetectOffline:
)
def test_cached_lora_detected_when_api_unreachable(
- self,
- monkeypatch,
- clean_offline_env,
- tmp_path,
+ self, monkeypatch, clean_offline_env, tmp_path
):
"""A cached adapter_config.json must still mark the repo as a
LoRA when the HF API is unreachable."""
diff --git a/studio/backend/tests/test_openai_citation_markers_edge.py b/studio/backend/tests/test_openai_citation_markers_edge.py
index ffe8c6b6eb..e8d0be6246 100644
--- a/studio/backend/tests/test_openai_citation_markers_edge.py
+++ b/studio/backend/tests/test_openai_citation_markers_edge.py
@@ -320,9 +320,7 @@ def test_split_helper_buffers_only_after_last_open_byte():
assert head == f"pre {complete} mid "
assert tail == partial
# And the head, once rewritten, drops every private-use byte.
- rewritten = _replace_openai_citation_markers(
- head, [{"source_id": "done", "url": "https://d"}]
- )
+ rewritten = _replace_openai_citation_markers(head, [{"source_id": "done", "url": "https://d"}])
assert rewritten == "pre [[1]](https://d) mid "
diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py
index 0b1a65c69e..d4f814856b 100644
--- a/studio/backend/tests/test_openai_code_execution.py
+++ b/studio/backend/tests/test_openai_code_execution.py
@@ -117,10 +117,7 @@ def test_shell_tool_added_on_cloud_with_container_auto(monkeypatch):
_drive(run())
tools = captured["body"].get("tools") or []
- assert {
- "type": "shell",
- "environment": {"type": "container_auto"},
- } in tools
+ assert {"type": "shell", "environment": {"type": "container_auto"}} in tools
def test_shell_tool_uses_container_reference_when_id_supplied(monkeypatch):
@@ -273,11 +270,7 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch):
# backend stamps onto every provider-side tool_start so the
# frontend serializer can distinguish hosted tools from
# user-declared functions on history replay.
- assert starts[0]["arguments"] == {
- "kind": "bash",
- "command": "ls -la",
- "_server_tool": True,
- }
+ assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True}
assert ends[0]["tool_call_id"] == "scall_1"
assert "total 24" in ends[0]["result"]
@@ -539,7 +532,5 @@ def test_expired_container_retries_only_once(monkeypatch):
# infinite loop.
assert call_count["n"] == 2
# The second failure surfaces normally as an error SSE line.
- error_lines = [
- line for line in lines if '"error"' in line and "_toolEvent" not in line
- ]
+ error_lines = [line for line in lines if '"error"' in line and "_toolEvent" not in line]
assert len(error_lines) >= 1
diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py
index 161a6fab83..48acf97e1f 100644
--- a/studio/backend/tests/test_openai_container_crud.py
+++ b/studio/backend/tests/test_openai_container_crud.py
@@ -80,17 +80,12 @@ def test_create_sends_openai_beta_header(monkeypatch):
return httpx.Response(200, json = {"id": "cntr_new", "name": "analysis"})
_mock_http_client(monkeypatch, handler)
- result = _drive(
- _make_client().create_openai_container(name = "analysis", ttl_minutes = 30)
- )
+ result = _drive(_make_client().create_openai_container(name = "analysis", ttl_minutes = 30))
assert result == {"id": "cntr_new", "name": "analysis"}
assert seen["headers"].get("openai-beta") == "containers=v1"
assert seen["body"]["name"] == "analysis"
- assert seen["body"]["expires_after"] == {
- "anchor": "last_active_at",
- "minutes": 30,
- }
+ assert seen["body"]["expires_after"] == {"anchor": "last_active_at", "minutes": 30}
def test_delete_sends_openai_beta_header_and_accepts_confirmation(monkeypatch):
diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py
index f177ed5ef3..95e9b2d63c 100644
--- a/studio/backend/tests/test_openai_responses_translation.py
+++ b/studio/backend/tests/test_openai_responses_translation.py
@@ -154,10 +154,7 @@ def test_responses_translates_image_parts(monkeypatch):
parts = captured["body"]["input"][0]["content"]
assert parts[0] == {"type": "input_text", "text": "What is this?"}
- assert parts[1] == {
- "type": "input_image",
- "image_url": "data:image/png;base64,AAA",
- }
+ assert parts[1] == {"type": "input_image", "image_url": "data:image/png;base64,AAA"}
# No max_output_tokens key when caller passes max_tokens=None.
assert "max_output_tokens" not in captured["body"]
@@ -497,8 +494,7 @@ def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch)
finish_reasons = [
json.loads(line[len("data:") :].strip())["choices"][0]["finish_reason"]
for line in lines
- if line.startswith("data:")
- and line[len("data:") :].strip() not in ("", "[DONE]")
+ if line.startswith("data:") and line[len("data:") :].strip() not in ("", "[DONE]")
]
assert "length" in finish_reasons
@@ -730,8 +726,7 @@ def test_responses_reasoning_summary_wrapped_in_think_tags(monkeypatch):
data_lines = [
line[len("data:") :].strip()
for line in lines
- if line.startswith("data:")
- and line[len("data:") :].strip() not in ("", "[DONE]")
+ if line.startswith("data:") and line[len("data:") :].strip() not in ("", "[DONE]")
]
payloads = [json.loads(raw) for raw in data_lines]
combined = "".join(
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index 05efee1595..7d488ae4c9 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -271,10 +271,7 @@ class TestChatCompletionRequestToolFields:
assert self._make(stop = "\nUser:").stop == "\nUser:"
def test_stop_list(self):
- assert self._make(stop = ["\nUser:", "\nAssistant:"]).stop == [
- "\nUser:",
- "\nAssistant:",
- ]
+ assert self._make(stop = ["\nUser:", "\nAssistant:"]).stop == ["\nUser:", "\nAssistant:"]
def test_tools_default_none(self):
req = self._make()
@@ -316,9 +313,7 @@ class TestChatCompletionRequestToolFields:
req = self._make()
assert req.stream is False
- def test_post_without_stream_field_decodes_to_stream_false_over_http(
- self, monkeypatch
- ):
+ def test_post_without_stream_field_decodes_to_stream_false_over_http(self, monkeypatch):
# Wire-level guard for the same default: a POST body that omits
# `stream` entirely (the exact shape naive curl / .NET clients
# send) must deserialise into stream=False *and* the response
@@ -419,13 +414,8 @@ class TestAnthropicToolChoiceToOpenAI:
assert anthropic_tool_choice_to_openai({"type": "none"}) == "none"
def test_tool_named(self):
- result = anthropic_tool_choice_to_openai(
- {"type": "tool", "name": "get_weather"}
- )
- assert result == {
- "type": "function",
- "function": {"name": "get_weather"},
- }
+ result = anthropic_tool_choice_to_openai({"type": "tool", "name": "get_weather"})
+ assert result == {"type": "function", "function": {"name": "get_weather"}}
def test_tool_missing_name_returns_none(self):
assert anthropic_tool_choice_to_openai({"type": "tool"}) is None
@@ -528,15 +518,11 @@ class TestFriendlyErrorHttpx:
# Non-httpx exceptions still fall through to the existing substring
# heuristics — a context-size message must still produce the
# "Message too long" path.
- ctx_msg = (
- "request (4096 tokens) exceeds the available context size (2048 tokens)"
- )
+ ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)"
assert "Message too long" in _friendly_error(ValueError(ctx_msg))
def test_generic_exception_returns_generic_message(self):
- assert (
- _friendly_error(RuntimeError("unrelated")) == "An internal error occurred"
- )
+ assert _friendly_error(RuntimeError("unrelated")) == "An internal error occurred"
from routes.inference import ( # noqa: E402
@@ -554,10 +540,7 @@ class TestDropEmptyAssistantSentinels:
{"role": "user", "content": "again"},
]
out = _drop_empty_assistant_sentinels(msgs)
- assert out == [
- {"role": "user", "content": "hi"},
- {"role": "user", "content": "again"},
- ]
+ assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}]
def test_drops_assistant_with_no_content_key(self):
# exclude_none=True strips the content key entirely; filter must catch this.
@@ -567,10 +550,7 @@ class TestDropEmptyAssistantSentinels:
{"role": "user", "content": "ok"},
]
out = _drop_empty_assistant_sentinels(msgs)
- assert out == [
- {"role": "user", "content": "hi"},
- {"role": "user", "content": "ok"},
- ]
+ assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "ok"}]
def test_preserves_assistant_with_text(self):
msgs = [
@@ -670,16 +650,10 @@ class TestGgufVisionMessages:
messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True)
assert has_image is True
- assert messages[0]["content"][0] == {
- "type": "text",
- "text": "describe image one",
- }
+ assert messages[0]["content"][0] == {"type": "text", "text": "describe image one"}
assert messages[0]["content"][1]["type"] == "image_url"
assert len(messages[0]["content"]) == 2
- assert messages[2]["content"][0] == {
- "type": "text",
- "text": "describe image two",
- }
+ assert messages[2]["content"][0] == {"type": "text", "text": "describe image two"}
assert messages[2]["content"][1]["type"] == "image_url"
assert len(messages[2]["content"]) == 2
assert isinstance(messages[1]["content"], str)
@@ -702,14 +676,9 @@ class TestGgufVisionMessages:
messages, has_image = _openai_messages_for_gguf_chat(req, is_vision = True)
assert has_image is True
- assert messages[0]["content"][0] == {
- "type": "text",
- "text": "describe this image",
- }
+ assert messages[0]["content"][0] == {"type": "text", "text": "describe this image"}
assert messages[0]["content"][1]["type"] == "image_url"
- assert messages[0]["content"][1]["image_url"]["url"].startswith(
- "data:image/png;base64,"
- )
+ assert messages[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
def test_rejects_image_parts_for_text_only_gguf(self):
req = ChatCompletionRequest(
@@ -775,9 +744,7 @@ class TestGgufVisionMessages:
{"role": "user", "content": "now"},
]
- updated = _set_or_prepend_system_message(
- messages, "Mid instructions.\n\nUse tools."
- )
+ updated = _set_or_prepend_system_message(messages, "Mid instructions.\n\nUse tools.")
assert [m["role"] for m in updated] == ["system", "user", "user"]
assert updated[0]["content"] == "Mid instructions.\n\nUse tools."
@@ -837,10 +804,7 @@ class TestGgufVisionToolRouting:
{
"type": "image_url",
"image_url": {
- "url": (
- "data:image/png;base64,"
- f"{TestGgufVisionMessages._PNG_B64}"
- ),
+ "url": (f"data:image/png;base64,{TestGgufVisionMessages._PNG_B64}"),
},
},
],
@@ -849,9 +813,7 @@ class TestGgufVisionToolRouting:
)
response = self._drive(
- openai_chat_completions(
- payload, request = self._Request(), current_subject = "test"
- )
+ openai_chat_completions(payload, request = self._Request(), current_subject = "test")
)
self._consume_response(response)
diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py
index 313c15a441..d669747b1a 100644
--- a/studio/backend/tests/test_pricing.py
+++ b/studio/backend/tests/test_pricing.py
@@ -21,7 +21,11 @@ from core.inference.pricing import (
)
-def _isclose(a, b, tol = 1e-6):
+def _isclose(
+ a,
+ b,
+ tol = 1e-6,
+):
return math.isclose(a, b, rel_tol = tol, abs_tol = tol)
@@ -241,9 +245,7 @@ def test_openai_cache_read_subtracted_from_input_at_discount():
)
# 20k charged at full price, 80k charged at 0.1x
assert _isclose(out["input_usd"], 20_000 / 1_000_000.0 * base)
- assert _isclose(
- out["cache_read_usd"], 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT
- )
+ assert _isclose(out["cache_read_usd"], 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT)
def test_openai_billable_input_tokens_does_not_double_count_cache_read():
@@ -390,9 +392,7 @@ def test_openai_web_search_charged_per_thousand():
"openai_tool_use": {"web_search_requests": 250},
},
)
- assert _isclose(
- out["server_tools_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K
- )
+ assert _isclose(out["server_tools_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K)
assert _isclose(out["total_usd"], 250 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K)
@@ -426,8 +426,7 @@ def test_openai_tool_surcharges_added_to_total():
expected_input = 100_000 / 1_000_000.0 * 5.0
expected_output = 5_000 / 1_000_000.0 * 30.0
expected_tools = (
- 3 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K
- + 0.25 * OPENAI_CONTAINER_USD_PER_HOUR
+ 3 / 1_000.0 * OPENAI_WEB_SEARCH_USD_PER_1K + 0.25 * OPENAI_CONTAINER_USD_PER_HOUR
)
assert _isclose(
out["total_usd"],
@@ -600,10 +599,7 @@ def test_openai_chat_style_envelope_reads_cache_from_prompt_tokens_details():
)
# Both envelopes must price identically.
assert _isclose(chat_style["input_usd"], raw["input_usd"]), (chat_style, raw)
- assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), (
- chat_style,
- raw,
- )
+ assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), (chat_style, raw)
# 80k at 0.1x base, 20k at full.
assert _isclose(
chat_style["cache_read_usd"],
diff --git a/studio/backend/tests/test_pricing_edge.py b/studio/backend/tests/test_pricing_edge.py
index ca4be258e0..5818b42e8a 100644
--- a/studio/backend/tests/test_pricing_edge.py
+++ b/studio/backend/tests/test_pricing_edge.py
@@ -18,7 +18,11 @@ from core.inference.pricing import (
)
-def _isclose(a, b, tol = 1e-6):
+def _isclose(
+ a,
+ b,
+ tol = 1e-6,
+):
return math.isclose(a, b, rel_tol = tol, abs_tol = tol)
@@ -188,9 +192,7 @@ def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable():
assert out["billable_input_tokens"] == 500 # 0 uncached + 500 cache_read
# cache_read still priced at the discount rate.
base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"]
- assert _isclose(
- out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT
- )
+ assert _isclose(out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT)
def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached():
@@ -207,9 +209,7 @@ def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached():
)
assert out["input_usd"] == 0.0
# Cache read still priced (the 0.1x bucket).
- assert _isclose(
- out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT
- )
+ assert _isclose(out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT)
# ── long-context tier crosses on billable, including cache_creation ──
diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py
index 0e668944f4..88df886b6c 100644
--- a/studio/backend/tests/test_providers_api.py
+++ b/studio/backend/tests/test_providers_api.py
@@ -225,9 +225,7 @@ class TestAuth:
json = {"username": USERNAME, "password": PASSWORD},
timeout = 10,
)
- assert (
- resp.status_code == 200
- ), f"Login failed ({resp.status_code}): {resp.text}"
+ assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}"
body = resp.json()
assert body.get("access_token"), "access_token is missing or empty"
assert body.get("token_type") == "bearer"
@@ -237,9 +235,7 @@ class TestAuth:
class TestPublicKey:
- def test_public_key_is_valid_pem(
- self, auth_headers: dict[str, str], public_key_pem: str
- ):
+ def test_public_key_is_valid_pem(self, auth_headers: dict[str, str], public_key_pem: str):
"""GET /api/providers/public-key returns an importable RSA PEM key."""
pem_bytes = public_key_pem.encode("utf-8")
key = serialization.load_pem_public_key(pem_bytes)
@@ -261,9 +257,7 @@ class TestRegistry:
)
assert resp.status_code == 200, f"Registry failed: {resp.text}"
providers = resp.json()
- assert (
- len(providers) == 9
- ), f"Expected 9 providers, got {len(providers)}: {providers}"
+ assert len(providers) == 9, f"Expected 9 providers, got {len(providers)}: {providers}"
print(f"\n {'Provider':<12} {'Base URL'}")
print(f" {'-'*12} {'-'*45}")
for p in providers:
@@ -283,9 +277,7 @@ class TestRegistry:
def test_registry_entries_have_required_fields(self, auth_headers: dict[str, str]):
"""Each registry entry has provider_type, display_name, base_url, default_models."""
- resp = requests.get(
- _url("/api/providers/registry"), headers = auth_headers, timeout = 10
- )
+ resp = requests.get(_url("/api/providers/registry"), headers = auth_headers, timeout = 10)
assert resp.status_code == 200
for entry in resp.json():
for field in (
@@ -320,9 +312,7 @@ class TestProviderCRUD:
json = {"provider_type": "openai", "display_name": "Test OpenAI (pytest)"},
timeout = 10,
)
- assert (
- resp.status_code == 201
- ), f"Create failed ({resp.status_code}): {resp.text}"
+ assert resp.status_code == 201, f"Create failed ({resp.status_code}): {resp.text}"
body = resp.json()
assert body.get("id"), "No id in response"
assert body["provider_type"] == "openai"
@@ -333,9 +323,7 @@ class TestProviderCRUD:
def test_list_includes_created(self, auth_headers: dict[str, str]):
"""GET /api/providers/ includes the newly created config."""
- assert (
- TestProviderCRUD._created_id
- ), "No created_id (run test_create_provider first)"
+ assert TestProviderCRUD._created_id, "No created_id (run test_create_provider first)"
resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10)
assert resp.status_code == 200
ids = [p["id"] for p in resp.json()]
@@ -354,9 +342,7 @@ class TestProviderCRUD:
json = {"display_name": new_name},
timeout = 10,
)
- assert (
- resp.status_code == 200
- ), f"Update failed ({resp.status_code}): {resp.text}"
+ assert resp.status_code == 200, f"Update failed ({resp.status_code}): {resp.text}"
assert resp.json()["display_name"] == new_name
print(f"\n updated display_name to '{new_name}'")
@@ -368,14 +354,10 @@ class TestProviderCRUD:
headers = auth_headers,
timeout = 10,
)
- assert (
- resp.status_code == 204
- ), f"Delete failed ({resp.status_code}): {resp.text}"
+ assert resp.status_code == 204, f"Delete failed ({resp.status_code}): {resp.text}"
# Confirm gone from list
- list_resp = requests.get(
- _url("/api/providers/"), headers = auth_headers, timeout = 10
- )
+ list_resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10)
ids = [p["id"] for p in list_resp.json()]
assert TestProviderCRUD._created_id not in ids, "Deleted provider still in list"
print(f"\n deleted id={TestProviderCRUD._created_id} confirmed gone")
@@ -423,9 +405,7 @@ class TestProviderInference:
json = {"provider_type": provider_type, "encrypted_api_key": encrypted},
timeout = 30,
)
- assert (
- resp.status_code == 200
- ), f"Request failed ({resp.status_code}): {resp.text}"
+ assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}"
body = resp.json()
assert (
body["success"] is True
@@ -449,9 +429,7 @@ class TestProviderInference:
json = {"provider_type": provider_type, "encrypted_api_key": encrypted},
timeout = 30,
)
- assert (
- resp.status_code == 200
- ), f"Request failed ({resp.status_code}): {resp.text}"
+ assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}"
models = resp.json()
assert isinstance(models, list), f"Expected list, got {type(models)}"
assert len(models) > 0, f"No models returned for {provider_type}"
@@ -498,9 +476,7 @@ class TestProviderInference:
# ── TestVisionInference ─────────────────────────────────────────────
# Sloth photo — used to test vision routing across providers
-_VISION_IMAGE_URL = (
- "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg"
-)
+_VISION_IMAGE_URL = "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg"
_VISION_PARAMS = [
pytest.param(
@@ -602,8 +578,6 @@ class TestLocalInferenceUnaffected:
f"This likely means the provider fields broke the base request schema."
)
status_label = (
- "local model responded"
- if resp.status_code == 200
- else "no model loaded (expected)"
+ "local model responded" if resp.status_code == 200 else "no model loaded (expected)"
)
print(f"\n status={resp.status_code} ({status_label}) — local path unaffected")
diff --git a/studio/backend/tests/test_responses_api.py b/studio/backend/tests/test_responses_api.py
index 5b55f87259..c0ec876ad0 100644
--- a/studio/backend/tests/test_responses_api.py
+++ b/studio/backend/tests/test_responses_api.py
@@ -173,9 +173,7 @@ class TestResponsesResponse:
resp = ResponsesResponse(
model = "test-model",
output = [
- ResponsesOutputMessage(
- content = [ResponsesOutputTextContent(text = "Hello!")]
- ),
+ ResponsesOutputMessage(content = [ResponsesOutputTextContent(text = "Hello!")]),
],
usage = ResponsesUsage(input_tokens = 10, output_tokens = 5, total_tokens = 15),
)
@@ -324,5 +322,4 @@ class TestNormaliseResponsesInput:
if __name__ == "__main__":
import pytest
-
pytest.main([__file__, "-v"])
diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py
index 2f1161c329..146d6017d0 100644
--- a/studio/backend/tests/test_responses_tool_passthrough.py
+++ b/studio/backend/tests/test_responses_tool_passthrough.py
@@ -160,9 +160,7 @@ class TestResponsesMultiTurnInput:
def test_function_call_output_missing_call_id_rejected(self):
with pytest.raises(ValidationError):
- ResponsesFunctionCallOutputInputItem(
- type = "function_call_output", output = "x"
- )
+ ResponsesFunctionCallOutputInputItem(type = "function_call_output", output = "x")
def test_function_call_output_accepts_content_array(self):
item = ResponsesFunctionCallOutputInputItem(
@@ -222,9 +220,7 @@ class TestToolsTranslation:
assert _translate_responses_tools_to_chat([]) is None
def test_only_builtin_tools_returns_none(self):
- assert (
- _translate_responses_tools_to_chat([{"type": "web_search_preview"}]) is None
- )
+ assert _translate_responses_tools_to_chat([{"type": "web_search_preview"}]) is None
def test_description_optional(self):
out = _translate_responses_tools_to_chat(
@@ -256,9 +252,7 @@ class TestToolChoiceTranslation:
"""If a client happens to send the Chat Completions nested shape,
we don't double-wrap it."""
already_nested = {"type": "function", "function": {"name": "get_weather"}}
- assert (
- _translate_responses_tool_choice_to_chat(already_nested) == already_nested
- )
+ assert _translate_responses_tool_choice_to_chat(already_nested) == already_nested
def test_unknown_shape_passes_through(self):
obj = {"type": "allowed_tools", "tools": [{"type": "function", "name": "x"}]}
@@ -625,9 +619,7 @@ class TestCodexStyleRequestShapes:
input = [
{
"role": "assistant",
- "content": [
- {"type": "output_text", "text": "ok", "annotations": []}
- ],
+ "content": [{"type": "output_text", "text": "ok", "annotations": []}],
},
{"role": "user", "content": "next"},
],
diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py
index 2ce9b55789..d27ee83f1d 100644
--- a/studio/backend/tests/test_rocm_oom_guard.py
+++ b/studio/backend/tests/test_rocm_oom_guard.py
@@ -135,9 +135,7 @@ class TestDeviceNameFallback:
props = _props(name = device_name)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "", f"expected empty gcn_arch, got {gcn!r}"
- assert (
- is_unified is True
- ), f"device {device_name!r} should be classified as unified-memory"
+ assert is_unified is True, f"device {device_name!r} should be classified as unified-memory"
# --- discrete devices that must NOT be mis-classified ---
diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py
index c3ee5b9ff1..5e3d2cc9d1 100644
--- a/studio/backend/tests/test_safetensors_capability_advertise.py
+++ b/studio/backend/tests/test_safetensors_capability_advertise.py
@@ -369,11 +369,7 @@ def test_worker_load_reply_payload_includes_chat_template_info():
"is_gguf": False,
}
_bm = getattr(backend, "models", {}) or {}
- _entry = (
- _bm.get(mc.identifier)
- or _bm.get(getattr(backend, "active_model_name", None))
- or {}
- )
+ _entry = _bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {}
_tpl_info = _entry.get("chat_template_info")
if isinstance(_tpl_info, dict):
model_info["chat_template_info"] = {
diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py
index 50f6011959..ff5653c742 100644
--- a/studio/backend/tests/test_safetensors_tool_loop.py
+++ b/studio/backend/tests/test_safetensors_tool_loop.py
@@ -53,9 +53,7 @@ from utils.datasets import is_gpt_oss_model_name
class TestParser:
def test_json_tool_call(self):
- text = (
- '{"name":"web_search","arguments":{"query":"hello"}}'
- )
+ text = '{"name":"web_search","arguments":{"query":"hello"}}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
tc = result[0]
@@ -214,7 +212,12 @@ def _collect_events(generator, max_events = 200):
return events
-def _make_loop(*, turns, exec_results = None, **kwargs):
+def _make_loop(
+ *,
+ turns,
+ exec_results = None,
+ **kwargs,
+):
"""Build a configured loop with a multi-turn fake generator.
``turns`` is a list of chunk-lists; iteration N yields chunks from
@@ -342,9 +345,7 @@ class TestLoopBasic:
assert exec_fn.calls[0][0] == "render_html"
assert "" in exec_fn.calls[0][1]["code"]
- def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start(
- self,
- ):
+ def test_python_tool_containing_render_html_signal_does_not_emit_provisional_start(self):
loop, exec_fn = _make_loop(
turns = [
[
@@ -361,9 +362,7 @@ class TestLoopBasic:
assert len(tool_starts) == 1
assert tool_starts[0]["tool_name"] == "python"
- assert exec_fn.calls == [
- ("python", {"code": "print('')"})
- ]
+ assert exec_fn.calls == [("python", {"code": "print('')"})]
def test_render_html_success_blocks_second_artifact_call(self):
exec_fn = FakeExecuteTool(["Rendered HTML artifact."])
@@ -398,10 +397,7 @@ class TestLoopBasic:
tool_starts = [e for e in events if e["type"] == "tool_start"]
assert exec_fn.calls == [("render_html", {"code": "one"})]
- assert [e["arguments"] for e in tool_starts] == [
- {},
- {"code": "one"},
- ]
+ assert [e["arguments"] for e in tool_starts] == [{}, {"code": "one"}]
def test_truncated_unclosed_tool_call(self):
loop, exec_fn = _make_loop(
@@ -425,9 +421,7 @@ class TestLoopBasic:
# ``arguments`` is a string that is not itself valid
# JSON for ``_coerce_arguments`` to parse, so the
# heal path runs.
- [
- '{"name":"web_search","arguments":"hello world"}'
- ],
+ ['{"name":"web_search","arguments":"hello world"}'],
["ok"],
],
exec_results = ["..."],
@@ -444,12 +438,8 @@ class TestLoopBehaviour:
# called only once.
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
["final"],
],
exec_results = ["search-result-1"],
@@ -467,9 +457,7 @@ class TestLoopBehaviour:
# tool_end event still carries the raw result for the UI.
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"python","arguments":{"code":"plot()"}}'
- ],
+ ['{"name":"python","arguments":{"code":"plot()"}}'],
["see chart"],
],
exec_results = ["chart\n__IMAGES__:/tmp/chart.png"],
@@ -507,9 +495,7 @@ class TestLoopBehaviour:
tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
assert tool_msgs, "no tool message reached the model"
for tm in tool_msgs:
- assert (
- "__IMAGES__" not in tm["content"]
- ), f"sentinel leaked to model: {tm['content']!r}"
+ assert "__IMAGES__" not in tm["content"], f"sentinel leaked to model: {tm['content']!r}"
def test_image_sentinel_stripped_with_multiple_markers(self):
# Consecutive sentinels: cut at the first, nothing leaks.
@@ -539,19 +525,13 @@ class TestLoopBehaviour:
tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
assert tool_msgs
for tm in tool_msgs:
- assert (
- "__IMAGES__" not in tm["content"]
- ), f"second sentinel leaked: {tm['content']!r}"
- assert (
- tm["content"] == "panel"
- ), f"expected payload-only 'panel', got {tm['content']!r}"
+ assert "__IMAGES__" not in tm["content"], f"second sentinel leaked: {tm['content']!r}"
+ assert tm["content"] == "panel", f"expected payload-only 'panel', got {tm['content']!r}"
def test_tool_execution_error_is_emitted_but_loop_continues(self):
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
["sorry, that failed"],
],
exec_results = ["Error: network unreachable"],
@@ -566,9 +546,7 @@ class TestLoopBehaviour:
def test_exception_in_executor_does_not_raise(self):
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
["recovered"],
],
exec_results = [RuntimeError("boom")],
@@ -588,8 +566,7 @@ class TestLoopControl:
events = list(
run_safetensors_tool_loop(
single_turn = _const_stream(
- '{"name":"web_search",'
- '"arguments":{"query":"x"}}'
+ '{"name":"web_search","arguments":{"query":"x"}}'
),
messages = [{"role": "user", "content": "hi"}],
tools = [],
@@ -606,9 +583,7 @@ class TestLoopControl:
loop, exec_fn = _make_loop(
turns = [
# : tool call (executes once)
- [
- '{"name":"web_search","arguments":{"query":"a"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"a"}}'],
# : model gives a final answer when nudged.
["here is the final answer"],
],
@@ -625,24 +600,19 @@ class TestStatusFormatting:
def test_status_for_known_tools(self):
# Use the private helper directly to verify status formatting.
assert (
- safetensors_agentic._status_for_tool("web_search", {"query": "abc"})
- == "Searching: abc"
+ safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) == "Searching: abc"
)
assert (
- safetensors_agentic._status_for_tool(
- "web_search", {"url": "https://www.example.com/x"}
- )
+ safetensors_agentic._status_for_tool("web_search", {"url": "https://www.example.com/x"})
== "Reading: example.com"
)
- assert safetensors_agentic._status_for_tool(
- "python", {"code": "x = 1"}
- ).startswith("Running Python:")
- assert safetensors_agentic._status_for_tool(
- "terminal", {"command": "ls"}
- ).startswith("Running:")
- assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith(
- "Calling:"
+ assert safetensors_agentic._status_for_tool("python", {"code": "x = 1"}).startswith(
+ "Running Python:"
)
+ assert safetensors_agentic._status_for_tool("terminal", {"command": "ls"}).startswith(
+ "Running:"
+ )
+ assert safetensors_agentic._status_for_tool("unknown_tool", {}).startswith("Calling:")
class TestProseMentioningToolCall:
@@ -655,9 +625,7 @@ class TestProseMentioningToolCall:
turns = [
# : a real tool call so the loop moves to
# .
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
# : prose that mentions the literal text.
["the docs say means an LLM tool call wrapper"],
],
@@ -677,9 +645,7 @@ class TestProseMentioningToolCall:
# result, so we should see exactly one call.
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
["the docs mention wrappers"],
],
exec_results = ["Page text: appears here in the docs"],
@@ -695,7 +661,6 @@ class TestChatTemplateHelper:
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
)
-
self.apply = apply_chat_template_for_generation
class _Tok:
@@ -705,7 +670,12 @@ class TestChatTemplateHelper:
self.last_kwargs = None
def apply_chat_template(
- self, messages, *, tokenize = False, add_generation_prompt = True, **kw
+ self,
+ messages,
+ *,
+ tokenize = False,
+ add_generation_prompt = True,
+ **kw,
):
self.call_count += 1
unknown = set(kw) - self.accepted
@@ -758,9 +728,7 @@ class TestGuardrails:
exec_fn = FakeExecuteTool([])
loop = run_safetensors_tool_loop(
single_turn = _fake_stream(
- [
- '{"name":"terminal","arguments":{"command":"echo bypass"}}'
- ]
+ ['{"name":"terminal","arguments":{"command":"echo bypass"}}']
),
messages = [{"role": "user", "content": "hi"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
@@ -776,9 +744,7 @@ class TestGuardrails:
exec_fn = FakeExecuteTool(["OK"])
loop = run_safetensors_tool_loop(
single_turn = _fake_stream(
- [
- '{"name":"python","arguments":{"code":"print(1)"}}'
- ]
+ ['{"name":"python","arguments":{"code":"print(1)"}}']
),
messages = [{"role": "user", "content": "hi"}],
tools = [],
@@ -790,11 +756,7 @@ class TestGuardrails:
def test_max_iterations_zero_executes_no_tools(self):
loop, exec_fn = _make_loop(
- turns = [
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ]
- ],
+ turns = [['{"name":"web_search","arguments":{"query":"x"}}']],
exec_results = ["OK"],
max_tool_iterations = 0,
)
@@ -825,9 +787,7 @@ class TestGuardrails:
def test_auto_heal_disabled_still_parses_valid_tool_call(self):
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"x"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
["done"],
],
exec_results = ["OK"],
@@ -840,47 +800,30 @@ class TestGuardrails:
def test_non_consecutive_duplicate_is_short_circuited(self):
loop, exec_fn = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"A"}}'
- ],
- [
- '{"name":"web_search","arguments":{"query":"B"}}'
- ],
- [
- '{"name":"web_search","arguments":{"query":"A"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"A"}}'],
+ ['{"name":"web_search","arguments":{"query":"B"}}'],
+ ['{"name":"web_search","arguments":{"query":"A"}}'],
["final"],
],
exec_results = ["res-A", "res-B"],
max_tool_iterations = 4,
)
events = _collect_events(loop)
- assert exec_fn.calls == [
- ("web_search", {"query": "A"}),
- ("web_search", {"query": "B"}),
- ]
+ assert exec_fn.calls == [("web_search", {"query": "A"}), ("web_search", {"query": "B"})]
tool_ends = [e for e in events if e["type"] == "tool_end"]
assert "already made this exact call" in tool_ends[-1]["result"]
def test_coerce_string_args_python_uses_code_key(self):
- assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {
- "code": "print(1)"
- }
+ assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"}
def test_coerce_string_args_terminal_uses_command_key(self):
- assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == {
- "command": "ls -la"
- }
+ assert _coerce_arguments("ls -la", heal = True, tool_name = "terminal") == {"command": "ls -la"}
def test_tool_call_ids_unique_across_loop_iterations(self):
loop, _exec = _make_loop(
turns = [
- [
- '{"name":"web_search","arguments":{"query":"A"}}'
- ],
- [
- '{"name":"web_search","arguments":{"query":"B"}}'
- ],
+ ['{"name":"web_search","arguments":{"query":"A"}}'],
+ ['{"name":"web_search","arguments":{"query":"B"}}'],
["done"],
],
exec_results = ["A", "B"],
diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py
index cd8957c3dd..92fee2e8e5 100644
--- a/studio/backend/tests/test_sandbox_tools.py
+++ b/studio/backend/tests/test_sandbox_tools.py
@@ -88,9 +88,7 @@ class TestTrustedHostAllowlist:
_ok(f"import requests; requests.get({url!r})")
def test_wikipedia_subdomain_passes(self):
- _ok(
- 'import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")'
- )
+ _ok('import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")')
def test_hf_co_short_form_passes(self):
_ok('import requests; requests.get("https://hf.co/unsloth/Qwen3.5-4B-GGUF")')
@@ -221,10 +219,7 @@ class TestUploadDenylist:
)
def test_plain_post_json_not_blocked(self):
- _ok(
- "import requests\n"
- 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})'
- )
+ _ok("import requests\n" 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})')
class TestSandboxEnvIsolation:
@@ -361,7 +356,6 @@ class TestBashBlocklistPosition:
@staticmethod
def _find():
from core.inference.tools import _find_blocked_commands
-
return _find_blocked_commands
# ---- argument-position: must NOT be blocked ----
@@ -526,7 +520,7 @@ class TestHfUploadImportGate:
def test_bare_name_upload_file_without_hf_import_allowed(self):
# No HF import -- local helper named upload_file should pass.
- _ok("def upload_file(*a, **k):\n pass\n" "upload_file('x', 'y', 'z')")
+ _ok("def upload_file(*a, **k):\n pass\nupload_file('x', 'y', 'z')")
class TestHfUploadSandboxLocalPaths:
diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py
index 521c99e126..70a103a2f8 100644
--- a/studio/backend/tests/test_studio_api.py
+++ b/studio/backend/tests/test_studio_api.py
@@ -72,11 +72,7 @@ DEFAULT_VARIANT = "UD-Q4_K_XL"
PORT = 18222 # high port unlikely to collide
HOST = "127.0.0.1"
STARTUP_TIMEOUT = 120 # seconds to wait for banner
-LOG_FILE = (
- Path(__file__).resolve().parent.parent.parent.parent
- / "temp"
- / "test_studio_api.log"
-)
+LOG_FILE = Path(__file__).resolve().parent.parent.parent.parent / "temp" / "test_studio_api.log"
# ── Helpers ──────────────────────────────────────────────────────────
@@ -219,9 +215,7 @@ def test_openai_sdk(base_url: str, api_key: str):
client = OpenAI(base_url = f"{base_url}/v1", api_key = api_key)
response = client.chat.completions.create(
model = "current",
- messages = [
- {"role": "user", "content": "What is 2+2? Answer with just the number."}
- ],
+ messages = [{"role": "user", "content": "What is 2+2? Answer with just the number."}],
stream = True,
)
content_parts = []
@@ -390,9 +384,7 @@ def test_openai_tools_nonstream(base_url: str, api_key: str):
assert "city" in parsed, f"Tool call missing required 'city' arg: {parsed}"
# Usage must be non-zero (was 0 before the fix)
usage = data.get("usage") or {}
- assert (
- usage.get("prompt_tokens", 0) > 0
- ), f"Expected non-zero prompt_tokens; got {usage}"
+ assert usage.get("prompt_tokens", 0) > 0, f"Expected non-zero prompt_tokens; got {usage}"
assert data.get("id"), "Missing response id"
print(
f" PASS openai tools non-stream: "
@@ -417,8 +409,7 @@ def test_openai_tools_stream(base_url: str, api_key: str):
assert status == 200, f"Expected 200, got {status}"
assert len(chunks) > 0, "No SSE chunks received"
assert _final_finish_reason(chunks) == "tool_calls", (
- f"Expected final finish_reason='tool_calls', got "
- f"{_final_finish_reason(chunks)!r}"
+ f"Expected final finish_reason='tool_calls', got " f"{_final_finish_reason(chunks)!r}"
)
assembled = _collect_streamed_tool_calls(chunks)
assert len(assembled) >= 1, "No tool_calls reassembled from stream"
@@ -501,8 +492,7 @@ def test_openai_sdk_tool_calling(base_url: str, api_key: str):
stream = False,
)
assert resp.choices[0].finish_reason == "tool_calls", (
- f"Expected finish_reason='tool_calls', got "
- f"{resp.choices[0].finish_reason!r}"
+ f"Expected finish_reason='tool_calls', got " f"{resp.choices[0].finish_reason!r}"
)
tool_calls = resp.choices[0].message.tool_calls
assert tool_calls and len(tool_calls) >= 1, "No tool_calls from SDK"
@@ -510,9 +500,7 @@ def test_openai_sdk_tool_calling(base_url: str, api_key: str):
assert tc.function.name == "get_weather"
parsed = json.loads(tc.function.arguments)
assert "city" in parsed
- print(
- f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}"
- )
+ print(f" PASS openai SDK tool calling: " f"tool={tc.function.name}, args={parsed}")
def test_invalid_key_rejected(base_url: str):
@@ -655,9 +643,7 @@ def test_anthropic_sdk(base_url: str, api_key: str):
message = client.messages.create(
model = "default",
max_tokens = 100,
- messages = [
- {"role": "user", "content": "What is 2+2? Answer with just the number."}
- ],
+ messages = [{"role": "user", "content": "What is 2+2? Answer with just the number."}],
)
assert message.role == "assistant"
assert len(message.content) > 0, "Empty content"
@@ -708,9 +694,7 @@ def test_anthropic_with_tools(base_url: str, api_key: str):
assert "message_stop" in event_types, "Missing message_stop"
full = _collect_anthropic_text(events)
- print(
- f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content"
- )
+ print(f" PASS anthropic with tools: {len(events)} events, {len(full)} chars content")
def test_anthropic_tool_choice_any(base_url: str, api_key: str):
@@ -770,8 +754,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str):
tool_use_starts = [
e
for e in events
- if e[0] == "content_block_start"
- and e[1].get("content_block", {}).get("type") == "tool_use"
+ if e[0] == "content_block_start" and e[1].get("content_block", {}).get("type") == "tool_use"
]
assert len(tool_use_starts) >= 1, "No tool_use content block emitted"
print(
@@ -821,9 +804,7 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st
if proc.poll() is not None:
log_fh.flush()
log_text = LOG_FILE.read_text()
- raise RuntimeError(
- f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}"
- )
+ raise RuntimeError(f"Server exited early (code {proc.returncode}):\n{log_text[-2000:]}")
log_text = LOG_FILE.read_text()
m = re.search(r"API Key:\s+(sk-unsloth-[a-f0-9]+)", log_text)
if m:
@@ -833,9 +814,7 @@ def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, st
if not api_key:
log_text = LOG_FILE.read_text()
_kill_server(proc)
- raise RuntimeError(
- f"Timed out waiting for API key in server output:\n{log_text[-2000:]}"
- )
+ raise RuntimeError(f"Timed out waiting for API key in server output:\n{log_text[-2000:]}")
# Wait a moment for the model to be fully loaded
time.sleep(2)
@@ -862,9 +841,7 @@ def _kill_server(proc: subprocess.Popen):
def main():
- parser = argparse.ArgumentParser(
- description = "End-to-end tests for unsloth studio run"
- )
+ parser = argparse.ArgumentParser(description = "End-to-end tests for unsloth studio run")
parser.add_argument(
"--model",
default = DEFAULT_MODEL,
@@ -898,9 +875,7 @@ def main():
run_test(test_help_output)
# ── 2-16. Start server and run API tests ─────────────────────────
- print(
- f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}..."
- )
+ print(f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}...")
proc = None
try:
proc, api_key = _start_server(args.model, args.gguf_variant)
diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py
index 8b90a46d5a..857302d543 100644
--- a/studio/backend/tests/test_tool_xml_strip.py
+++ b/studio/backend/tests/test_tool_xml_strip.py
@@ -77,9 +77,7 @@ def test_strips_orphan_tool_call_no_close():
def test_strips_orphan_function_no_close():
- text = (
- "I'll call python:\n\n\nprint(1)\n"
- )
+ text = "I'll call python:\n\n\nprint(1)\n"
cleaned = _TOOL_XML_RE.sub("", text)
assert "")
self.assertEqual(result.dataset[1]["text"], "world")
self.assertTrue(
- any(
- "null or non-string 'text' values" in notice.message
- for notice in result.notices
- )
+ any("null or non-string 'text' values" in notice.message for notice in result.notices)
)
diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py
index 94279c28b4..edc47c705e 100644
--- a/studio/backend/tests/test_training_worker_flash_attn.py
+++ b/studio/backend/tests/test_training_worker_flash_attn.py
@@ -15,7 +15,13 @@ from core.training import worker
def _missing_flash_attn_import():
real_import = builtins.__import__
- def fake_import(name, globals = None, locals = None, fromlist = (), level = 0):
+ def fake_import(
+ name,
+ globals = None,
+ locals = None,
+ fromlist = (),
+ level = 0,
+ ):
if name == "flash_attn":
raise ImportError
return real_import(name, globals, locals, fromlist, level)
@@ -26,7 +32,13 @@ def _missing_flash_attn_import():
def _missing_module_import(missing: str):
real_import = builtins.__import__
- def fake_import(name, globals = None, locals = None, fromlist = (), level = 0):
+ def fake_import(
+ name,
+ globals = None,
+ locals = None,
+ fromlist = (),
+ level = 0,
+ ):
if name == missing:
raise ImportError
return real_import(name, globals, locals, fromlist, level)
@@ -37,9 +49,7 @@ def _missing_module_import(missing: str):
def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch):
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
assert worker._should_try_runtime_flash_attn_install(32767) is False
- assert worker._should_try_runtime_flash_attn_install(
- 32768
- ) is sys.platform.startswith("linux")
+ assert worker._should_try_runtime_flash_attn_install(32768) is sys.platform.startswith("linux")
monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1")
assert worker._should_try_runtime_flash_attn_install(32768) is False
@@ -105,7 +115,12 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch):
)
monkeypatch.setattr(worker, "install_wheel", mock.Mock())
- def fake_run(cmd, stdout = None, stderr = None, text = None):
+ def fake_run(
+ cmd,
+ stdout = None,
+ stderr = None,
+ text = None,
+ ):
calls.append(list(cmd))
return subprocess.CompletedProcess(cmd, 0, "")
@@ -131,9 +146,7 @@ def test_runtime_flash_attn_skips_on_blackwell(monkeypatch):
install_mock = mock.Mock()
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
- monkeypatch.setattr(
- worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True
- )
+ monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True)
monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True)
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
monkeypatch.setattr(
@@ -504,14 +517,10 @@ def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch):
# Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY (no tilelang).
assert "--force-reinstall" in repair_args
- assert (
- "--no-deps" in repair_args
- ), "Repair MUST use --no-deps to avoid replacing torch / CUDA"
+ assert "--no-deps" in repair_args, "Repair MUST use --no-deps to avoid replacing torch / CUDA"
assert "--only-binary=:all:" in repair_args
assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in repair_args
- assert all(
- "tilelang" not in a for a in repair_args
- ), "Repair MUST only touch apache-tvm-ffi"
+ assert all("tilelang" not in a for a in repair_args), "Repair MUST only touch apache-tvm-ffi"
# Install: regular dep-resolving install, NO --force-reinstall.
assert "--force-reinstall" not in install_args
@@ -689,16 +698,12 @@ def test_hook_installs_when_gate_returns_false(monkeypatch):
conv_install = mock.Mock(side_effect = _conv_install_side_effect)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -722,9 +727,7 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch):
fla_install = mock.Mock()
tile_install = mock.Mock()
conv_install = mock.Mock()
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install)
# Tilelang healthy so the post_available path is a no-op (otherwise
@@ -734,9 +737,7 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch):
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.9")
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -764,16 +765,12 @@ def test_hook_idempotent_on_repeat_call(monkeypatch):
return True
conv_install = mock.Mock(side_effect = _conv_install_side_effect)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -794,18 +791,12 @@ def test_hook_handles_install_failure_gracefully(monkeypatch):
def raising_install(eq):
raise RuntimeError("pip failed to fetch wheel")
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", raising_install
- )
- monkeypatch.setattr(
- worker, "_ensure_tilelang_backend_unconditional", lambda eq: None
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", raising_install)
+ monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: None)
monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -819,14 +810,10 @@ def test_hook_can_be_disabled_via_env(monkeypatch):
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
fla_install = mock.Mock()
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1")
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -841,18 +828,12 @@ def test_hook_clears_lru_cache_before_first_check(monkeypatch):
conv_gate = _make_fake_gate(initial_return = True)
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None
- )
- monkeypatch.setattr(
- worker, "_ensure_tilelang_backend_unconditional", lambda eq: None
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: None)
+ monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: None)
monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: None)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
_iu.is_flash_linear_attention_available()
@@ -880,18 +861,12 @@ def test_hook_rewrites_previously_imported_module_bindings(monkeypatch):
fla_gate.next_return = True
return True
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fake_install
- )
- monkeypatch.setattr(
- worker, "_ensure_tilelang_backend_unconditional", lambda eq: True
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fake_install)
+ monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
# The fake module's local binding has been rewritten to the wrapper.
assert fake_mod.is_flash_linear_attention_available is not fla_gate
@@ -915,30 +890,20 @@ def test_hook_skips_when_import_utils_unavailable(monkeypatch):
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
# Should not raise.
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch):
"""Hook disabled -> legacy gate falls back to auto-discovered model types."""
install_mock = mock.Mock()
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", install_mock
- )
- monkeypatch.setattr(
- worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"})
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", install_mock)
+ monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"}))
monkeypatch.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1")
- worker._ensure_flash_linear_attention(
- event_queue = [], model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._ensure_flash_linear_attention(event_queue = [], model_name = "unsloth/Qwen3.5-2B")
assert install_mock.call_count == 1
- worker._ensure_flash_linear_attention(
- event_queue = [], model_name = "meta-llama/Llama-3.1-8B"
- )
+ worker._ensure_flash_linear_attention(event_queue = [], model_name = "meta-llama/Llama-3.1-8B")
assert install_mock.call_count == 1
@@ -968,13 +933,9 @@ def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch)
fla_install = mock.Mock(side_effect = _fla_install)
tile_install = mock.Mock(return_value = True)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
- monkeypatch.setattr(
- worker, "_install_package_wheel_first", mock.Mock(return_value = True)
- )
+ monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
# Hermetize the auto-discovered set so the test stays valid as new
# transformers releases add FLA-using model_types (eg olmo_hybrid in
@@ -1009,18 +970,12 @@ def test_hook_does_install_tilelang_for_qwen35(monkeypatch):
fla_install = mock.Mock(side_effect = _fla_install)
tile_install = mock.Mock(return_value = True)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
- monkeypatch.setattr(
- worker, "_install_package_wheel_first", mock.Mock(return_value = True)
- )
+ monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -1079,20 +1034,14 @@ def test_hook_trusts_installer_bool_not_metadata(monkeypatch):
return False # but deep import is broken
fake_fla_install = mock.Mock(side_effect = _bad_install)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fake_fla_install)
monkeypatch.setattr(
worker, "_ensure_tilelang_backend_unconditional", mock.Mock(return_value = True)
)
- monkeypatch.setattr(
- worker, "_install_package_wheel_first", mock.Mock(return_value = True)
- )
+ monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -1145,14 +1094,10 @@ def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch):
monkeypatch.setenv(worker._FLA_SKIP_ENV, "1")
tile_install = mock.Mock(return_value = True)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
- monkeypatch.setattr(
- worker, "_install_package_wheel_first", mock.Mock(return_value = True)
- )
+ monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -1172,21 +1117,15 @@ def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch):
fla_install = mock.Mock(return_value = True)
tile_install = mock.Mock(return_value = True)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", fla_install
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
- monkeypatch.setattr(
- worker, "_install_package_wheel_first", mock.Mock(return_value = True)
- )
+ monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
# tilelang missing AND tvm-ffi is on broken list — both trigger repair.
monkeypatch.setattr(worker, "_tilelang_importable", lambda: False)
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11")
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
from transformers.utils import import_utils as _iu
@@ -1290,17 +1229,11 @@ def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch):
monkeypatch.delenv("FLA_TILELANG", raising = False)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True
- )
- monkeypatch.setattr(
- worker, "_ensure_tilelang_backend_unconditional", lambda eq: True
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True)
+ monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
assert _os.environ.get("FLA_TILELANG") == "0"
@@ -1314,17 +1247,11 @@ def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch
monkeypatch.setenv("FLA_TILELANG", "1")
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True
- )
- monkeypatch.setattr(
- worker, "_ensure_tilelang_backend_unconditional", lambda eq: True
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True)
+ monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
assert _os.environ["FLA_TILELANG"] == "1"
@@ -1336,17 +1263,11 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch):
monkeypatch.delenv("FLA_TILELANG", raising = False)
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
monkeypatch.setattr(worker, "_torch_has_hip", lambda: False)
- monkeypatch.setattr(
- worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True
- )
- monkeypatch.setattr(
- worker, "_ensure_tilelang_backend_unconditional", lambda eq: True
- )
+ monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True)
+ monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True)
monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True)
- worker._install_fast_path_hooks(
- event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B"
- )
+ worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
assert _os.environ.get("FLA_TILELANG") is None
@@ -1356,9 +1277,7 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch):
# ───────────────────────────────────────────────────────────────────
-def _make_fake_transformers_tree(
- tmp_path, fla_types: list[str], non_fla_types: list[str]
-):
+def _make_fake_transformers_tree(tmp_path, fla_types: list[str], non_fla_types: list[str]):
"""Lay out a tmp dir as `transformers/models/{type}/modeling_{type}.py`."""
pkg = tmp_path / "transformers"
models = pkg / "models"
@@ -1401,9 +1320,7 @@ def test_discover_fla_model_types_returns_only_fla_users(tmp_path, monkeypatch):
def test_discover_fla_model_types_caches_across_calls(tmp_path, monkeypatch):
- pkg = _make_fake_transformers_tree(
- tmp_path, fla_types = ["qwen3_5"], non_fla_types = []
- )
+ pkg = _make_fake_transformers_tree(tmp_path, fla_types = ["qwen3_5"], non_fla_types = [])
fake = mock.MagicMock(__file__ = str(pkg / "__init__.py"))
monkeypatch.setitem(sys.modules, "transformers", fake)
_reset_fla_cache(monkeypatch)
@@ -1432,7 +1349,13 @@ def test_discover_fla_model_types_handles_missing_transformers(monkeypatch):
real_import = builtins.__import__
- def fake_import(name, globals = None, locals = None, fromlist = (), level = 0):
+ def fake_import(
+ name,
+ globals = None,
+ locals = None,
+ fromlist = (),
+ level = 0,
+ ):
if name == "transformers":
raise ImportError("transformers not installed")
return real_import(name, globals, locals, fromlist, level)
@@ -1443,9 +1366,7 @@ def test_discover_fla_model_types_handles_missing_transformers(monkeypatch):
def test_discover_fla_model_types_handles_unreadable_file(tmp_path, monkeypatch):
- pkg = _make_fake_transformers_tree(
- tmp_path, fla_types = ["qwen3_5"], non_fla_types = []
- )
+ pkg = _make_fake_transformers_tree(tmp_path, fla_types = ["qwen3_5"], non_fla_types = [])
fake = mock.MagicMock(__file__ = str(pkg / "__init__.py"))
monkeypatch.setitem(sys.modules, "transformers", fake)
_reset_fla_cache(monkeypatch)
@@ -1491,9 +1412,7 @@ def test_model_wants_tilelang_empty_when_transformers_has_no_fla(monkeypatch):
def test_model_wants_tilelang_normalizes_separators(monkeypatch):
- monkeypatch.setattr(
- worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"})
- )
+ monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"}))
for variant in (
"qwen3-next",
"Qwen3.Next",
diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py
index c031c2fea3..ff5e1f1381 100644
--- a/studio/backend/tests/test_transformers_version.py
+++ b/studio/backend/tests/test_transformers_version.py
@@ -295,9 +295,7 @@ class TestGetTransformersTier:
"utils.transformers_version._check_config_needs_550",
return_value = False,
):
- assert (
- get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530"
- )
+ assert get_transformers_tier("mistralai/Ministral-3-8B-Instruct-2512") == "530"
def test_llama_returns_default(self):
with (
diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py
index 64c9907119..bdb1cd2ce8 100644
--- a/studio/backend/tests/test_utils.py
+++ b/studio/backend/tests/test_utils.py
@@ -25,14 +25,12 @@ import pytest
# --- Conditional framework imports ---
try:
import torch
-
HAS_TORCH = True
except ImportError:
HAS_TORCH = False
try:
import mlx.core as mx
-
HAS_MLX = True
except ImportError:
HAS_MLX = False
@@ -196,15 +194,12 @@ class TestGetGpuMemoryInfo:
# can render the correct label. On CUDA / XPU / MLX / CPU hosts
# it is equivalent to `get_device().value`.
from utils.hardware.hardware import _backend_label
-
result = get_gpu_memory_info()
assert result["backend"] == _backend_label(get_device())
# --- When a GPU IS available ---
- @pytest.mark.skipif(
- _actual_device() == "cpu", reason = "No GPU available on this machine"
- )
+ @pytest.mark.skipif(_actual_device() == "cpu", reason = "No GPU available on this machine")
def test_gpu_available_fields(self):
result = get_gpu_memory_info()
assert result["available"] is True
@@ -302,9 +297,7 @@ class TestLogGpuMemory:
"free_gb": 14.0,
}
- with patch(
- "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
- ):
+ with patch("utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info):
log_gpu_memory("unit-test")
captured = capfd.readouterr()
@@ -315,9 +308,7 @@ class TestLogGpuMemory:
def test_logs_cpu_fallback_when_no_gpu(self, capfd):
fake_info = {"available": False, "backend": "cpu"}
- with patch(
- "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
- ):
+ with patch("utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info):
log_gpu_memory("cpu-test")
captured = capfd.readouterr()
diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py
index 9e7bbdd1fb..2af64dac91 100644
--- a/studio/backend/tests/test_vision_cache.py
+++ b/studio/backend/tests/test_vision_cache.py
@@ -209,9 +209,7 @@ class TestVisionCacheDirectPath:
@patch("utils.transformers_version.needs_transformers_5", return_value = False)
@patch("utils.models.model_config.load_model_config")
- def test_vision_config_attr_detected_and_cached(
- self, mock_load_config, mock_needs_t5
- ):
+ def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5):
"""Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True."""
cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist
cfg.model_type = "qwen2_vl"
diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py
index e54ae6dcf8..65964908d7 100644
--- a/studio/backend/tests/test_vram_estimation.py
+++ b/studio/backend/tests/test_vram_estimation.py
@@ -316,12 +316,8 @@ class TestLoraParams(unittest.TestCase):
self.assertLess(qv_only, all_mods)
def test_moe_mlp_modules_scale_with_experts(self):
- dense_lora = compute_lora_params(
- LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"]
- )
- moe_lora = compute_lora_params(
- MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"]
- )
+ dense_lora = compute_lora_params(LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"])
+ moe_lora = compute_lora_params(MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"])
ratio = moe_lora / dense_lora
self.assertAlmostEqual(ratio, 8.0, delta = 0.5)
@@ -338,12 +334,8 @@ class TestLoraParams(unittest.TestCase):
self.assertGreater(moe_lora, dense_lora * 20)
def test_attention_modules_same_for_moe(self):
- dense_attn = compute_lora_params(
- LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]
- )
- moe_attn = compute_lora_params(
- MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]
- )
+ dense_attn = compute_lora_params(LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"])
+ moe_attn = compute_lora_params(MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"])
self.assertEqual(dense_attn, moe_attn)
def test_all_linear_uses_default_text_modules(self):
@@ -466,9 +458,7 @@ class TestActivationBytes(unittest.TestCase):
def test_non_flash_attention_uses_quadratic_path(self):
seq_len = 4096
- expected_quadratic = (
- 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0
- )
+ expected_quadratic = 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0
for attention_implementation in ("eager", "unknown_impl", None):
with self.subTest(attention_implementation = attention_implementation):
non_flash = compute_activation_bytes(
@@ -483,9 +473,7 @@ class TestActivationBytes(unittest.TestCase):
def test_non_flash_attention_without_gc_scales_quadratic_path_by_layers(self):
seq_len = 4096
- one_layer = (
- 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0
- )
+ one_layer = 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0
non_flash = compute_activation_bytes(
STRUCTURED_MIXED,
1,
@@ -717,9 +705,7 @@ class TestEstimateTrainingVram(unittest.TestCase):
)
v8 = estimate_training_vram(LLAMA_8B, opt8)
v32 = estimate_training_vram(LLAMA_8B, opt32)
- self.assertAlmostEqual(
- v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1
- )
+ self.assertAlmostEqual(v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1)
def test_min_gpu_vram_treats_activations_as_per_gpu_fixed(self):
config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True)
@@ -769,9 +755,7 @@ class TestEstimateTrainingVram(unittest.TestCase):
optimizer = "adamw_8bit",
load_in_4bit = False,
)
- expected_floor = int(
- compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15
- )
+ expected_floor = int(compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15)
with patch(
"utils.hardware.vram_estimation.compute_gradient_bytes",
return_value = 1,
@@ -1059,7 +1043,6 @@ class TestDenseLayerIndices(unittest.TestCase):
class TestKvSharedLayer(unittest.TestCase):
def test_fully_shared_kv_returns_false_matching_upstream(self):
from utils.hardware.vram_estimation import _is_kv_shared_layer
-
arch = ModelArchConfig(
hidden_size = 512,
num_hidden_layers = 4,
@@ -1293,9 +1276,7 @@ class TestSharedExperts(unittest.TestCase):
delta_per_layer = 4096 * 1407 * 3 * 2
expected_delta = delta_per_layer * 32 * 2
actual_delta = w_yes - w_no
- self.assertAlmostEqual(
- actual_delta, expected_delta, delta = expected_delta * 0.01
- )
+ self.assertAlmostEqual(actual_delta, expected_delta, delta = expected_delta * 0.01)
def test_deepseek_v3_params_in_range(self):
total = compute_total_params(DEEPSEEK_V3)
@@ -1411,9 +1392,7 @@ class TestDenseMoEMix(unittest.TestCase):
moe_intermediate_size = 1024,
num_dense_layers = 5,
)
- lora_all = compute_lora_params(
- all_moe, 16, ["gate_proj", "up_proj", "down_proj"]
- )
+ lora_all = compute_lora_params(all_moe, 16, ["gate_proj", "up_proj", "down_proj"])
lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"])
self.assertNotEqual(lora_all, lora_mix)
@@ -1497,9 +1476,7 @@ class TestPerLayerInputSkipAlias(unittest.TestCase):
delta = _compute_skipped_quantizable_elements(arch)
self.assertEqual(
delta,
- arch.hidden_size
- * arch.num_hidden_layers
- * arch.hidden_size_per_layer_input,
+ arch.hidden_size * arch.num_hidden_layers * arch.hidden_size_per_layer_input,
)
def test_layer_aggregate_skip_includes_per_layer_input_modules(self):
@@ -1578,9 +1555,7 @@ class TestSharedExpertVariants(unittest.TestCase):
def test_shared_expert_size_separate_from_routed_changes_weight_count(self):
from utils.hardware.vram_estimation import _compute_moe_mlp_elements
- arch_separate = extract_arch_config(
- self._hf(shared_expert_intermediate_size = 64)
- )
+ arch_separate = extract_arch_config(self._hf(shared_expert_intermediate_size = 64))
arch_implicit = extract_arch_config(self._hf(n_shared_experts = 1))
# Different shared sizes (64 vs default moe_intermediate_size=128) must
# produce different MoE element counts.
@@ -1624,9 +1599,7 @@ class TestSharedExpertActivation(unittest.TestCase):
moe_intermediate_size = 64,
**fields,
)
- return extract_arch_config(
- SimpleNamespace(text_config = text_config, quantization_config = {})
- )
+ return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {}))
def test_shared_expert_increases_activation_bytes(self):
with_shared = self._make(shared_expert_intermediate_size = 64)
@@ -1678,9 +1651,7 @@ class TestPerLayerInputActivation(unittest.TestCase):
tie_word_embeddings = False,
**fields,
)
- return extract_arch_config(
- SimpleNamespace(text_config = text_config, quantization_config = {})
- )
+ return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {}))
def test_ple_increases_activation_bytes(self):
with_ple = self._make(
@@ -1744,9 +1715,7 @@ class TestKvSharedActivation(unittest.TestCase):
num_kv_shared_layers = kv_shared,
layer_types = ["full_attention"] * 4,
)
- return extract_arch_config(
- SimpleNamespace(text_config = text_config, quantization_config = {})
- )
+ return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {}))
def test_kv_shared_layers_keep_activation_bytes(self):
shared = self._make(kv_shared = 2)
@@ -1792,10 +1761,7 @@ class TestSparseMoeSkipAliases(unittest.TestCase):
def test_gemma4_layers_experts_alias_pulls_routed(self):
from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements
-
- arch = extract_arch_config(
- self._hf(["model.layers.0.experts"], enable_moe_block = True)
- )
+ arch = extract_arch_config(self._hf(["model.layers.0.experts"], enable_moe_block = True))
self.assertGreater(_compute_skipped_quantizable_elements(arch), 0)
def test_qwen_shared_expert_skip_pulls_only_shared(self):
@@ -1823,7 +1789,6 @@ class TestSparseMoeSkipAliases(unittest.TestCase):
def test_exaone_shared_experts_plural_alias(self):
from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements
-
arch = extract_arch_config(
self._hf(
["model.layers.0.mlp.shared_experts"],
@@ -1847,9 +1812,7 @@ class TestAllLinearMoELoraExclusion(unittest.TestCase):
moe_intermediate_size = 64,
**fields,
)
- return extract_arch_config(
- SimpleNamespace(text_config = text_config, quantization_config = {})
- )
+ return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {}))
def test_all_linear_drops_routed_moe_expert_lora(self):
arch = self._arch()
@@ -1867,9 +1830,7 @@ class TestAllLinearMoELoraExclusion(unittest.TestCase):
def test_all_linear_includes_attention_lora(self):
arch = self._arch()
all_linear = compute_lora_params(arch, 8, "all-linear")
- attn_only = compute_lora_params(
- arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"]
- )
+ attn_only = compute_lora_params(arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"])
# all-linear still attaches to attention nn.Linear modules.
self.assertGreaterEqual(all_linear, attn_only)
@@ -1887,9 +1848,7 @@ class TestExplicitPerLayerInputLora(unittest.TestCase):
hidden_size_per_layer_input = 32,
vocab_size_per_layer_input = 128,
)
- return extract_arch_config(
- SimpleNamespace(text_config = text_config, quantization_config = {})
- )
+ return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {}))
def test_explicit_per_layer_input_gate_returns_nonzero(self):
arch = self._arch()
@@ -1928,9 +1887,7 @@ class TestTopKExpertActivation(unittest.TestCase):
moe_intermediate_size = 64,
**fields,
)
- return extract_arch_config(
- SimpleNamespace(text_config = text_config, quantization_config = {})
- )
+ return extract_arch_config(SimpleNamespace(text_config = text_config, quantization_config = {}))
def test_num_experts_per_tok_extracted(self):
arch = self._make(num_experts_per_tok = 4)
@@ -2180,13 +2137,11 @@ class TestLlama4ArchExtraction(unittest.TestCase):
def test_llama4_moe_layers_dispatch_uses_explicit_indices(self):
from utils.hardware.vram_estimation import _compute_dense_layer_indices
-
cfg = SimpleNamespace(num_hidden_layers = 4, moe_layers = [1, 3])
self.assertEqual(_compute_dense_layer_indices(cfg, 4), (0, 2))
def test_llama4_moe_layers_takes_priority_over_first_k_dense_replace(self):
from utils.hardware.vram_estimation import _compute_dense_layer_indices
-
cfg = SimpleNamespace(
num_hidden_layers = 6,
moe_layers = [2, 4],
@@ -2288,7 +2243,6 @@ class TestDbrxFfnConfigExtraction(unittest.TestCase):
class TestErniePhaseModuloDispatch(unittest.TestCase):
def test_phase_modulo_with_interval_two_matches_decoder(self):
from utils.hardware.vram_estimation import _compute_dense_layer_indices
-
cfg = SimpleNamespace(
num_hidden_layers = 10,
moe_layer_start_index = 2,
@@ -2300,7 +2254,6 @@ class TestErniePhaseModuloDispatch(unittest.TestCase):
def test_phase_modulo_with_interval_three(self):
from utils.hardware.vram_estimation import _compute_dense_layer_indices
-
cfg = SimpleNamespace(
num_hidden_layers = 9,
moe_layer_start_index = 0,
diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py
index 023630fb9a..bc887f8cc5 100644
--- a/studio/backend/tests/test_windows_gpu_detection_mock.py
+++ b/studio/backend/tests/test_windows_gpu_detection_mock.py
@@ -169,14 +169,14 @@ def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None:
def _build_path_dirs_like_start_llama_server(
- binary_dir: Path, prefix: Path, cuda_path: str = ""
+ binary_dir: Path,
+ prefix: Path,
+ cuda_path: str = "",
) -> list[str]:
"""Path-friendly wrapper around LlamaCppBackend._build_windows_path_dirs.
Asserting against the staticmethod (not a hand-copy) is the point:
if the win32 PATH order drops _windows_pip_nvidia_dll_dirs, tests fail."""
- return LlamaCppBackend._build_windows_path_dirs(
- str(binary_dir), str(prefix), cuda_path
- )
+ return LlamaCppBackend._build_windows_path_dirs(str(binary_dir), str(prefix), cuda_path)
def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch":
@@ -210,9 +210,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
fake_csv = "0, 22805\n"
with _mock_nvidia_smi_run(fake_csv):
gpus = LlamaCppBackend._get_gpu_free_memory()
- assert gpus == [
- (0, 22805)
- ], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}"
+ assert gpus == [(0, 22805)], f"GPU probe failed to parse mocked nvidia-smi output: {gpus}"
def test_nvidia_smi_probe_respects_cuda_visible_devices(self, monkeypatch):
"""CUDA_VISIBLE_DEVICES=1 -> only GPU 1 visible."""
@@ -247,9 +245,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
site / "nvidia" / "cu13" / "bin" / "x86_64",
site / "torch" / "lib",
):
- assert (
- str(expected) in out
- ), f"resolver missed {expected.relative_to(prefix)}: {out}"
+ assert str(expected) in out, f"resolver missed {expected.relative_to(prefix)}: {out}"
def test_path_assembly_makes_cudart_reachable_without_toolkit(self, tmp_path):
"""The #5106 scenario: GPU detected, pip nvidia wheels present,
@@ -260,9 +256,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
_populate_studio_venv(prefix)
_populate_studio_install(install, runtime = "13.1")
binary_dir = install / "build" / "bin" / "Release"
- path_dirs = _build_path_dirs_like_start_llama_server(
- binary_dir, prefix, cuda_path = ""
- )
+ path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix, cuda_path = "")
# binary_dir first -- Windows DLL search step 1.
assert path_dirs[0] == str(
binary_dir
@@ -278,9 +272,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
)
# Defence in depth: both fix paths contribute cudart.
sources = {Path(e).relative_to(tmp_path).parts[0] for e, _ in cudart_locations}
- assert (
- "studio_install" in sources
- ), f"#5322's cudart drop not reachable: {cudart_locations}"
+ assert "studio_install" in sources, f"#5322's cudart drop not reachable: {cudart_locations}"
assert (
"studio_venv" in sources
), f"#5324's pip nvidia dir not contributing cudart: {cudart_locations}"
@@ -297,8 +289,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]:
reachable = any((Path(d) / required).exists() for d in path_dirs)
assert reachable, (
- f"{required} unreachable from PATH; #5106 not fixed.\n"
- f"PATH entries: {path_dirs}"
+ f"{required} unreachable from PATH; #5106 not fixed.\n" f"PATH entries: {path_dirs}"
)
def test_no_pip_nvidia_wheels_still_works_via_install_dir(self, tmp_path):
@@ -310,9 +301,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
_populate_studio_install(install, runtime = "13.1")
binary_dir = install / "build" / "bin" / "Release"
path_dirs = _build_path_dirs_like_start_llama_server(binary_dir, prefix)
- assert path_dirs == [
- str(binary_dir)
- ], f"bare venv produced unexpected PATH: {path_dirs}"
+ assert path_dirs == [str(binary_dir)], f"bare venv produced unexpected PATH: {path_dirs}"
for required in REAL_UPSTREAM_CUDART_BUNDLE["13.1"]:
assert (
binary_dir / required
@@ -336,8 +325,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
(rel / fn).write_bytes(b"PE-stub")
path_dirs = _build_path_dirs_like_start_llama_server(rel, prefix)
cudart_reachable = any(
- (Path(d) / "cudart64_12.dll").exists()
- or (Path(d) / "cudart64_13.dll").exists()
+ (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists()
for d in path_dirs
)
assert cudart_reachable, (
@@ -345,8 +333,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
f"on cudart-less install. PATH entries: {path_dirs}"
)
cublas_reachable = any(
- (Path(d) / "cublas64_12.dll").exists()
- or (Path(d) / "cublas64_13.dll").exists()
+ (Path(d) / "cublas64_12.dll").exists() or (Path(d) / "cublas64_13.dll").exists()
for d in path_dirs
)
assert cublas_reachable, "cublas unreachable on cudart-less install"
@@ -365,8 +352,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
# Pre-PR PATH: binary_dir only. No pip nvidia dirs, no toolkit.
pre_pr_path_dirs = [str(rel)]
cudart_reachable_pre = any(
- (Path(d) / "cudart64_12.dll").exists()
- or (Path(d) / "cudart64_13.dll").exists()
+ (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists()
for d in pre_pr_path_dirs
)
assert not cudart_reachable_pre, (
@@ -387,7 +373,5 @@ class TestWindowsSysPlatformMocked:
out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix))
assert out, f"resolver returned empty under sys.platform=win32: {out}"
# cu13 arch dir must be in the output.
- cu13_arch = (
- prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
- )
+ cu13_arch = prefix / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
assert str(cu13_arch) in out
diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py
index 4c8e6239a0..9d01b40add 100644
--- a/studio/backend/utils/cache_cleanup.py
+++ b/studio/backend/utils/cache_cleanup.py
@@ -75,8 +75,7 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None)
if preserve_patterns:
logger.info(
- f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): "
- f"{cache_dir}"
+ f"Cleaning unsloth compiled cache (preserving {preserve_patterns}): " f"{cache_dir}"
)
for item in cache_dir.iterdir():
diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py
index 687da74c21..2955ec9023 100644
--- a/studio/backend/utils/datasets/data_collators.py
+++ b/studio/backend/utils/datasets/data_collators.py
@@ -28,19 +28,13 @@ class DataCollatorSpeechSeq2SeqWithPadding:
processor: Any
def __call__(self, features: List[dict]) -> dict:
- input_features = [
- {"input_features": feature["input_features"]} for feature in features
- ]
- batch = self.processor.feature_extractor.pad(
- input_features, return_tensors = "pt"
- )
+ input_features = [{"input_features": feature["input_features"]} for feature in features]
+ batch = self.processor.feature_extractor.pad(input_features, return_tensors = "pt")
label_features = [{"input_ids": feature["labels"]} for feature in features]
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt")
- labels = labels_batch["input_ids"].masked_fill(
- labels_batch.attention_mask.ne(1), -100
- )
+ labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
labels = labels[:, 1:]
@@ -169,9 +163,7 @@ class VLMDataCollator:
# Apply chat template
texts = [
- self.processor.apply_chat_template(
- msgs, tokenize = False, add_generation_prompt = False
- )
+ self.processor.apply_chat_template(msgs, tokenize = False, add_generation_prompt = False)
for msgs in all_messages
]
diff --git a/studio/backend/utils/datasets/dataset_none_detect.py b/studio/backend/utils/datasets/dataset_none_detect.py
index 1e884271cc..e48e153fb8 100644
--- a/studio/backend/utils/datasets/dataset_none_detect.py
+++ b/studio/backend/utils/datasets/dataset_none_detect.py
@@ -75,9 +75,7 @@ def _probe_conversation(dataset: Dataset, candidates = None):
# a list holding a dict/None turn); scalars and list-of-strings must
# not look like chatml. Upgrade a non-plausible fallback when a later
# candidate is plausible, so probe order keeps the best match.
- if all_corrupt_fallback is None or not all_corrupt_fallback.get(
- "has_plausible_turns"
- ):
+ if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"):
has_plausible_turns = False
for i in range(min(len(dataset), 100)):
cell = dataset[i][col]
@@ -122,9 +120,7 @@ def _probe_conversation(dataset: Dataset, candidates = None):
_CONV_KEYS = {"role", "from", "content", "value"}
if not any(keys <= turn_keys for keys in _CHAT_KEY_SETS):
schema_less_plausible = bool(turn_keys & _CONV_KEYS)
- if all_corrupt_fallback is None or not all_corrupt_fallback.get(
- "has_plausible_turns"
- ):
+ if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"):
all_corrupt_fallback = {
"column": col,
"turn_keys": turn_keys,
@@ -167,14 +163,11 @@ def is_none_or_empty(value) -> bool:
non_text_blocks = [item for item in dict_blocks if item.get("type") != "text"]
if non_text_blocks:
return False
- text_values = [
- item.get("text") for item in dict_blocks if item.get("type") == "text"
- ]
+ text_values = [item.get("text") for item in dict_blocks if item.get("type") == "text"]
if text_values and all(
t is None
or (
- isinstance(t, str)
- and not t.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip()
+ isinstance(t, str) and not t.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip()
)
for t in text_values
):
@@ -285,9 +278,7 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
stats["rows_with_none_turns"] += 1
stats["total_none_turns"] += 1
stats["rows_all_none"] += 1
- stats["none_by_role"]["unknown"] = (
- stats["none_by_role"].get("unknown", 0) + 1
- )
+ stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1
stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1
stats["findings"].append(
{
@@ -306,9 +297,7 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
stats["rows_with_none_turns"] += 1
stats["total_none_turns"] += 1
stats["rows_all_none"] += 1
- stats["none_by_role"]["unknown"] = (
- stats["none_by_role"].get("unknown", 0) + 1
- )
+ stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1
stats["none_by_type"]["empty_conversation"] = (
stats["none_by_type"].get("empty_conversation", 0) + 1
)
@@ -336,9 +325,7 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
"raw_value": repr(turn),
}
)
- stats["none_by_role"]["unknown"] = (
- stats["none_by_role"].get("unknown", 0) + 1
- )
+ stats["none_by_role"]["unknown"] = stats["none_by_role"].get("unknown", 0) + 1
vtype = "None" if turn is None else "invalid_type"
stats["none_by_type"][vtype] = stats["none_by_type"].get(vtype, 0) + 1
continue
@@ -359,20 +346,14 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
if "from" in turn and "value" in turn:
content = turn.get("value")
elif "role" in turn:
- content = (
- turn.get("content") if "content" in turn else turn.get("value")
- )
+ content = turn.get("content") if "content" in turn else turn.get("value")
elif "from" in turn:
content = turn.get("value")
else:
- content = (
- turn.get("content") if "content" in turn else turn.get("value")
- )
+ content = turn.get("content") if "content" in turn else turn.get("value")
# Assistant tool-call turns carry empty content + tool_calls and are
# valid; the exemption is assistant-only.
- if is_none_or_empty(content) and not (
- role == "assistant" and turn.get("tool_calls")
- ):
+ if is_none_or_empty(content) and not (role == "assistant" and turn.get("tool_calls")):
vtype = _classify_empty(content)
row_findings.append(
{
@@ -469,9 +450,7 @@ FORMAT_REGISTRY = [
},
{
"name": "sharegpt",
- "match": lambda ds, conv: (
- conv is not None and {"from", "value"} <= conv["turn_keys"]
- ),
+ "match": lambda ds, conv: (conv is not None and {"from", "value"} <= conv["turn_keys"]),
"scan": find_none_sharegpt,
},
{
@@ -532,13 +511,11 @@ def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict:
_dict_types = []
try:
from datasets import DatasetDict as _DatasetDict
-
_dict_types.append(_DatasetDict)
except ImportError:
pass
try:
from datasets import IterableDatasetDict as _IterableDatasetDict
-
_dict_types.append(_IterableDatasetDict)
except ImportError:
pass
@@ -552,7 +529,6 @@ def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict:
# instead of a confusing TypeError downstream.
try:
from datasets import IterableDataset as _IterableDataset
-
if isinstance(dataset, _IterableDataset):
raise ValueError(
"scan_dataset requires a materialized Dataset, not an IterableDataset. "
@@ -658,7 +634,11 @@ def _print_summary_header(stats: dict, fmt: str) -> bool:
return True
-def print_report(stats: dict, fmt: str, summary_only: bool = False):
+def print_report(
+ stats: dict,
+ fmt: str,
+ summary_only: bool = False,
+):
"""Print a human-readable summary, optionally with full findings list."""
has_findings = _print_summary_header(stats, fmt)
if not has_findings or summary_only:
@@ -689,7 +669,12 @@ def print_report(stats: dict, fmt: str, summary_only: bool = False):
print(f"{'=' * 64}")
-def show_row(dataset: Dataset, row_indices: list[int], fmt: str, col: str = None):
+def show_row(
+ dataset: Dataset,
+ row_indices: list[int],
+ fmt: str,
+ col: str = None,
+):
"""Print the full contents of specific rows for inspection.
Used by test_codex_fixes.py to verify row rendering behaviour.
@@ -750,9 +735,7 @@ def show_row(dataset: Dataset, row_indices: list[int], fmt: str, col: str = None
# Mirror scanner logic: tool_calls exemption is assistant-only;
# other roles with empty content + tool_calls are still bad.
r = t.get("role") if t.get("role") is not None else t.get("from")
- if is_none_or_empty(c) and not (
- str(r) == "assistant" and t.get("tool_calls")
- ):
+ if is_none_or_empty(c) and not (str(r) == "assistant" and t.get("tool_calls")):
return True
return False
@@ -773,19 +756,11 @@ def show_row(dataset: Dataset, row_indices: list[int], fmt: str, col: str = None
if "from" in turn and "value" in turn:
content = turn.get("value")
elif "role" in turn:
- content = (
- turn.get("content")
- if "content" in turn
- else turn.get("value")
- )
+ content = turn.get("content") if "content" in turn else turn.get("value")
elif "from" in turn:
content = turn.get("value")
else:
- content = (
- turn.get("content")
- if "content" in turn
- else turn.get("value")
- )
+ content = turn.get("content") if "content" in turn else turn.get("value")
if is_none_or_empty(content) and not (
role == "assistant" and turn.get("tool_calls")
):
@@ -827,12 +802,8 @@ examples:
python dataset_none_detect.py org/my-dataset --token hf_...
""",
)
- parser.add_argument(
- "dataset", help = "HuggingFace dataset repo id (e.g. org/my-dataset)"
- )
- parser.add_argument(
- "--split", default = "train", help = "Dataset split to load (default: train)"
- )
+ parser.add_argument("dataset", help = "HuggingFace dataset repo id (e.g. org/my-dataset)")
+ parser.add_argument("--split", default = "train", help = "Dataset split to load (default: train)")
parser.add_argument(
"--format",
default = "auto",
diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py
index 26378d64ee..792e0ea9bc 100644
--- a/studio/backend/utils/datasets/dataset_utils.py
+++ b/studio/backend/utils/datasets/dataset_utils.py
@@ -202,7 +202,11 @@ _CHATML_ROLE_ORDER = ("system", "user", "assistant")
_CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"}
-def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
+def _apply_user_mapping(
+ dataset,
+ mapping: dict,
+ batch_size: int = 1000,
+):
"""
Apply user-provided column mapping to convert dataset to conversations format.
@@ -279,7 +283,10 @@ def _extract_column_value(val, col: str, label_mapping: dict) -> str:
def _apply_template_mapping(
- dataset, column_roles: dict, meta: dict, batch_size: int = 1000
+ dataset,
+ column_roles: dict,
+ meta: dict,
+ batch_size: int = 1000,
):
"""
Apply advisor-driven mapping for non-conversational datasets.
@@ -324,9 +331,7 @@ def _apply_template_mapping(
user_parts = []
for col in role_groups["user"]:
if col in examples:
- user_parts.append(
- _extract_column_value(examples[col][i], col, label_mapping)
- )
+ user_parts.append(_extract_column_value(examples[col][i], col, label_mapping))
if user_parts:
convo.append({"role": "user", "content": "\n".join(user_parts)})
@@ -334,9 +339,7 @@ def _apply_template_mapping(
asst_parts = []
for col in role_groups["assistant"]:
if col in examples:
- asst_parts.append(
- _extract_column_value(examples[col][i], col, label_mapping)
- )
+ asst_parts.append(_extract_column_value(examples[col][i], col, label_mapping))
if asst_parts:
convo.append({"role": "assistant", "content": "\n".join(asst_parts)})
@@ -351,7 +354,11 @@ def _apply_template_mapping(
)
-def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
+def _apply_user_mapping_alpaca(
+ dataset,
+ mapping: dict,
+ batch_size: int = 1000,
+):
"""
Apply user-provided column mapping to convert dataset to Alpaca format.
@@ -382,11 +389,7 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
("output", outputs),
):
col = col_for[field]
- val = (
- str(examples[col][i])
- if col and col in examples and examples[col][i]
- else ""
- )
+ val = str(examples[col][i]) if col and col in examples and examples[col][i] else ""
dest.append(val)
return {"instruction": instructions, "input": inputs, "output": outputs}
@@ -464,9 +467,7 @@ def format_dataset(
else:
# auto / chatml / sharegpt / conversational — all produce chatml conversations
# (sharegpt is always standardized to role/content internally)
- mapped_dataset = _apply_user_mapping(
- dataset, custom_format_mapping, batch_size
- )
+ mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size)
final_format = "chatml_conversations"
chat_column = "conversations"
@@ -578,9 +579,7 @@ def format_dataset(
# Unknown - try standardization, if fails pass as is
else:
- warnings.append(
- f"Unknown format detected. Keys found: {detected['sample_keys']}"
- )
+ warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}")
# NEW: Try heuristic detection
if auto_detect_custom:
@@ -606,9 +605,7 @@ def format_dataset(
if role == target_role and col_name in examples:
content = examples[col_name][i]
if content and str(content).strip():
- convo.append(
- {"role": role, "content": str(content)}
- )
+ convo.append({"role": role, "content": str(content)})
conversations.append(convo)
return {"conversations": conversations, **preserved_columns}
@@ -656,9 +653,7 @@ def format_dataset(
"warnings": warnings,
}
except Exception as e:
- warnings.append(
- f"Could not standardize: {e}. Passing dataset as-is."
- )
+ warnings.append(f"Could not standardize: {e}. Passing dataset as-is.")
# Return as-is with warnings
return {
@@ -928,9 +923,7 @@ def format_and_template_dataset(
f"text='{user_vlm_text_column}') failed: {e} — "
f"falling back to auto-detection"
)
- logger.info(
- f"⚠️ User VLM mapping failed, falling back to auto-detection..."
- )
+ logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...")
custom_format_mapping = None # clear so auto-detection runs below
else:
errors.append(
@@ -984,9 +977,7 @@ def format_and_template_dataset(
dataset_name = dataset_name,
progress_callback = progress_callback,
)
- warnings.append(
- "Converted from ShareGPT+image format to standard VLM format"
- )
+ warnings.append("Converted from ShareGPT+image format to standard VLM format")
except Exception as e:
errors.append(f"Failed to convert ShareGPT+image format: {e}")
import traceback
@@ -1020,7 +1011,6 @@ def format_and_template_dataset(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
-
friendly = llm_generate_dataset_warning(
issues,
dataset_name = dataset_name,
@@ -1055,13 +1045,9 @@ def format_and_template_dataset(
)
if vlm_instruction:
- warnings.append(
- f"Using user-provided instruction: '{vlm_instruction}'"
- )
+ warnings.append(f"Using user-provided instruction: '{vlm_instruction}'")
else:
- warnings.append(
- "Auto-generated instruction based on dataset analysis"
- )
+ warnings.append("Auto-generated instruction based on dataset analysis")
except Exception as e:
errors.append(f"Failed to convert to VLM format: {e}")
@@ -1166,9 +1152,7 @@ def format_and_template_dataset(
summary = get_dataset_info_summary(dataset_info)
# Combine results
- all_warnings = dataset_info.get("warnings", []) + template_result.get(
- "warnings", []
- )
+ all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", [])
all_errors = template_result.get("errors", [])
# If format_dataset returned "unknown" but apply_chat_template rescued
diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py
index 289b30e55e..5433d3115c 100644
--- a/studio/backend/utils/datasets/format_conversion.py
+++ b/studio/backend/utils/datasets/format_conversion.py
@@ -140,7 +140,11 @@ def standardize_chat_format(
return dataset.map(_standardize_dataset, **dataset_map_kwargs)
-def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
+def convert_chatml_to_alpaca(
+ dataset,
+ batch_size = 1000,
+ num_proc = None,
+):
"""
Converts ChatML format (messages OR conversations) to Alpaca format.
Handles both standardized and ShareGPT formats.
@@ -151,7 +155,6 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
"""
try:
from torch.utils.data import IterableDataset
-
_is_torch_iterable = isinstance(dataset, IterableDataset)
except ImportError:
_is_torch_iterable = False
@@ -159,15 +162,11 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
def _convert(examples):
# Auto-detect which column name is used
chatml_data = (
- examples.get("messages")
- or examples.get("conversations")
- or examples.get("texts")
+ examples.get("messages") or examples.get("conversations") or examples.get("texts")
)
if chatml_data is None:
- raise ValueError(
- "No 'messages' or 'conversations' or 'texts' column found."
- )
+ raise ValueError("No 'messages' or 'conversations' or 'texts' column found.")
instructions = []
outputs = []
@@ -215,7 +214,11 @@ def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
return dataset.map(_convert, **dataset_map_kwargs)
-def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
+def convert_alpaca_to_chatml(
+ dataset,
+ batch_size = 1000,
+ num_proc = None,
+):
"""
Converts Alpaca format to ChatML format.
@@ -223,7 +226,6 @@ def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
"""
try:
from torch.utils.data import IterableDataset
-
_is_torch_iterable = isinstance(dataset, IterableDataset)
except ImportError:
_is_torch_iterable = False
@@ -328,16 +330,12 @@ def convert_to_vlm_format(
instruction_column = instruction_info.get("instruction_column")
uses_dynamic = instruction_info["uses_dynamic_instruction"]
- logger.info(
- f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}"
- )
+ logger.info(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}")
logger.info(f"📝 Confidence: {instruction_info['confidence']:.2f}")
if not uses_dynamic:
logger.info(f"📝 Using instruction: '{instruction}'")
else:
- logger.info(
- f"📝 Using dynamic instructions from column: '{instruction_column}'"
- )
+ logger.info(f"📝 Using dynamic instructions from column: '{instruction_column}'")
else:
instruction_column = None
uses_dynamic = False
@@ -351,13 +349,11 @@ def convert_to_vlm_format(
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
-
with fsspec.open(image_data, "rb", expand = True) as f:
image_data = Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
# Bare filename → resolve via HF repo lookup
from huggingface_hub import hf_hub_download
-
local_path = hf_hub_download(
dataset_name,
_image_lookup[image_data],
@@ -371,7 +367,6 @@ def convert_to_vlm_format(
text_data = sample[text_column]
if isinstance(text_data, list) and len(text_data) > 0:
import random
-
text_data = random.choice(text_data)
# Get instruction (static or dynamic)
@@ -397,9 +392,7 @@ def convert_to_vlm_format(
total = len(dataset)
first_image = next(iter(dataset))[image_column]
- has_urls = isinstance(first_image, str) and first_image.startswith(
- ("http://", "https://")
- )
+ has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://"))
# ── Bare-filename detection: images stored as filenames (e.g. "img_001.png")
# that don't exist locally. Build a basename→repo_path lookup so we can
@@ -449,9 +442,7 @@ def convert_to_vlm_format(
num_workers = safe_thread_num_proc()
_notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...")
- logger.info(
- f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..."
- )
+ logger.info(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...")
probe_samples = [dataset[i] for i in range(PROBE_SIZE)]
probe_ok = 0
@@ -459,9 +450,7 @@ def convert_to_vlm_format(
probe_start = time.time()
with ThreadPoolExecutor(max_workers = num_workers) as executor:
- futures = {
- executor.submit(_convert_single_sample, s): s for s in probe_samples
- }
+ futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples}
for future in as_completed(futures):
try:
future.result()
@@ -483,7 +472,6 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
-
friendly = llm_generate_dataset_warning(
issues,
dataset_name = dataset_name,
@@ -554,9 +542,7 @@ def convert_to_vlm_format(
except Exception as e:
failed_count += 1
if failed_count == 1:
- logger.info(
- f"First VLM conversion failure: {type(e).__name__}: {e}"
- )
+ logger.info(f"First VLM conversion failure: {type(e).__name__}: {e}")
converted_list.extend(r for r in batch_results if r is not None)
@@ -581,9 +567,7 @@ def convert_to_vlm_format(
failed_count += 1
if failed_count == 1:
# Log the first failure to aid debugging
- logger.info(
- f"First VLM conversion failure: {type(e).__name__}: {e}"
- )
+ logger.info(f"First VLM conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
@@ -601,7 +585,6 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
-
friendly = llm_generate_dataset_warning(
issues,
dataset_name = dataset_name,
@@ -627,7 +610,6 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
-
friendly = llm_generate_dataset_warning(
issues,
dataset_name = dataset_name,
@@ -739,12 +721,10 @@ def convert_sharegpt_with_images_to_vlm_format(
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
-
with fsspec.open(image_data, "rb", expand = True) as f:
return Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
from huggingface_hub import hf_hub_download
-
local_path = hf_hub_download(
dataset_name,
_image_lookup[image_data],
@@ -753,12 +733,9 @@ def convert_sharegpt_with_images_to_vlm_format(
return Image.open(local_path).convert("RGB")
else:
return Image.open(image_data).convert("RGB")
- if isinstance(image_data, dict) and (
- "bytes" in image_data or "path" in image_data
- ):
+ if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data):
if image_data.get("bytes"):
from io import BytesIO
-
return Image.open(BytesIO(image_data["bytes"])).convert("RGB")
if image_data.get("path"):
return Image.open(image_data["path"]).convert("RGB")
@@ -812,9 +789,7 @@ def convert_sharegpt_with_images_to_vlm_format(
pbar.close()
if failed_count > 0:
- logger.info(
- f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples"
- )
+ logger.info(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples")
if len(converted_list) == 0:
raise ValueError(
@@ -840,9 +815,7 @@ def convert_llava_to_vlm_format(dataset):
"""
from PIL import Image
- logger.info(
- f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format..."
- )
+ logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...")
def _convert_single_sample(sample):
"""Convert a single llava sample to standard VLM format."""
diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py
index 7b70ff3a76..829838064a 100644
--- a/studio/backend/utils/datasets/format_detection.py
+++ b/studio/backend/utils/datasets/format_detection.py
@@ -13,10 +13,7 @@ import re
def _keyword_in_column(keyword: str, col_name: str) -> bool:
"""Word-boundary keyword match to avoid false positives like 'pic' in 'topic'."""
- return (
- re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE)
- is not None
- )
+ return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None
def detect_dataset_format(dataset):
@@ -220,10 +217,7 @@ def detect_custom_format_heuristic(dataset):
return True
for pattern in metadata_prefix_patterns:
- if (
- col_lower.startswith(pattern.split("_")[0] + "_")
- and col_lower != pattern
- ):
+ if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern:
if "_" in col_lower:
prefix = col_lower.split("_")[0]
if prefix in ["generation", "pass", "inference"]:
@@ -267,9 +261,7 @@ def detect_custom_format_heuristic(dataset):
if role_type == "user":
col_lower = col_name.lower()
# If column is ONLY "task" (or task_xxx), give it lower priority for user role
- if "task" in col_lower and not any(
- kw in col_lower for kw in user_words_high_priority
- ):
+ if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
score -= 15 # Significant penalty so other user columns win
priority_bonus = get_priority_score(col_name)
@@ -301,17 +293,13 @@ def detect_custom_format_heuristic(dataset):
content_columns = [col for col in all_columns if not is_metadata(col)]
# Count candidates first
- assistant_potential = [
- col for col in content_columns if has_keyword(col, assistant_words)
- ]
+ assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
user_potential = [col for col in content_columns if has_keyword(col, user_words)]
# STEP 1: Find best ASSISTANT column
assistant_candidates = []
for col in assistant_potential:
- score = score_column(
- col, assistant_words, "assistant", len(assistant_potential)
- )
+ score = score_column(col, assistant_words, "assistant", len(assistant_potential))
if score > 0:
assistant_candidates.append((col, score))
@@ -518,7 +506,6 @@ def _is_image_value(value) -> bool:
# PIL Image instance
try:
from PIL.Image import Image as PILImage
-
if isinstance(value, PILImage):
return True
except ImportError:
@@ -647,9 +634,7 @@ def detect_vlm_dataset_structure(dataset):
if isinstance(content[0], dict) and "type" in content[0]:
# Check for llava format
has_index = any(
- "index" in item
- for item in content
- if isinstance(item, dict)
+ "index" in item for item in content if isinstance(item, dict)
)
has_images_column = "images" in column_names
@@ -664,9 +649,7 @@ def detect_vlm_dataset_structure(dataset):
# Standard VLM format
has_image = any(
- "image" in item
- for item in content
- if isinstance(item, dict)
+ "image" in item for item in content if isinstance(item, dict)
)
if has_image:
return {
@@ -777,9 +760,7 @@ def detect_vlm_dataset_structure(dataset):
return True
# Check prefixes
- if any(
- col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]
- ):
+ if any(col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]):
return True
return False
@@ -791,9 +772,7 @@ def detect_vlm_dataset_structure(dataset):
return 100
# Dict with image data (bytes/path from HF Image feature)
- if isinstance(sample_value, dict) and (
- "bytes" in sample_value or "path" in sample_value
- ):
+ if isinstance(sample_value, dict) and ("bytes" in sample_value or "path" in sample_value):
return 75
if isinstance(sample_value, str):
@@ -818,9 +797,7 @@ def detect_vlm_dataset_structure(dataset):
# Local file — check it exists
if not sample_value.startswith(("http://", "https://")):
- return os.path.exists(
- sample_value
- ) # bare filenames return False here, that's OK
+ return os.path.exists(sample_value) # bare filenames return False here, that's OK
# URL — quick HEAD request with short timeout
try:
diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py
index 758717f76b..10004ce3db 100644
--- a/studio/backend/utils/datasets/llm_assist.py
+++ b/studio/backend/utils/datasets/llm_assist.py
@@ -68,9 +68,7 @@ def precache_helper_gguf():
return
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
- variant = os.environ.get(
- "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
- )
+ variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
try:
from huggingface_hub import HfApi, hf_hub_download
@@ -86,9 +84,7 @@ def precache_helper_gguf():
# Find all GGUF files matching the variant (may be split into shards)
variant_lower = variant.lower().replace("-", "_")
- matching = sorted(
- f for f in gguf_files if variant_lower in f.lower().replace("-", "_")
- )
+ matching = sorted(f for f in gguf_files if variant_lower in f.lower().replace("-", "_"))
if matching:
logger.info(
@@ -119,9 +115,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
- variant = os.environ.get(
- "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
- )
+ variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
backend = None
try:
@@ -143,9 +137,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
return None
messages = [{"role": "user", "content": prompt}]
- logger.info(
- "Helper model request: enable_thinking=False (per-request override)"
- )
+ logger.info("Helper model request: enable_thinking=False (per-request override)")
cumulative = ""
for chunk in backend.generate_chat_completion(
messages = messages,
@@ -240,10 +232,7 @@ def llm_generate_vlm_instruction(
}
-def llm_classify_columns(
- column_names: list[str],
- samples: list[dict],
-) -> Optional[dict[str, str]]:
+def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Optional[dict[str, str]]:
"""
Ask a helper LLM to classify dataset columns into roles.
@@ -294,7 +283,6 @@ def llm_classify_columns(
except json.JSONDecodeError:
# Try to find JSON object in the response
import re
-
match = re.search(r"\{[^}]+\}", text)
if match:
try:
@@ -313,11 +301,7 @@ def llm_classify_columns(
valid_roles = {"user", "assistant", "system", "metadata"}
cleaned = {}
for col, role in mapping.items():
- if (
- col in column_names
- and isinstance(role, str)
- and role.lower() in valid_roles
- ):
+ if col in column_names and isinstance(role, str) and role.lower() in valid_roles:
cleaned[col] = role.lower()
if not cleaned:
@@ -420,7 +404,11 @@ def _parse_json_response(text: str) -> Optional[dict]:
return None
-def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) -> str:
+def _generate_with_backend(
+ backend,
+ messages: list[dict],
+ max_tokens: int = 512,
+) -> str:
"""Run one chat completion on an already-loaded backend. Returns raw text."""
logger.info("Advisor request: enable_thinking=False (per-request override)")
cumulative = ""
@@ -480,9 +468,7 @@ def fetch_hf_dataset_card(
if val is not None:
metadata[key] = val
- logger.info(
- f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields"
- )
+ logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields")
return readme, metadata
except Exception as e:
@@ -509,9 +495,7 @@ def _run_multi_pass_advisor(
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
- variant = os.environ.get(
- "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
- )
+ variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
backend = None
try:
@@ -541,9 +525,7 @@ def _run_multi_pass_advisor(
samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n"
metadata_str = (
- json.dumps(dataset_metadata, indent = 2, default = str)[:500]
- if dataset_metadata
- else "N/A"
+ json.dumps(dataset_metadata, indent = 2, default = str)[:500] if dataset_metadata else "N/A"
)
card_excerpt = (dataset_card or "")[:1200] or "N/A"
@@ -745,9 +727,7 @@ def _run_multi_pass_advisor(
# Validate: must have at least one user AND one assistant
roles_present = set(column_roles.values())
if "user" not in roles_present or "assistant" not in roles_present:
- logger.warning(
- f"Pass 2 sanity fail: missing user or assistant role: {column_roles}"
- )
+ logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}")
return None # triggers fallback to simple classification
# ── Pass 3: System prompt (non-conversational datasets only) ──
diff --git a/studio/backend/utils/datasets/raw_text.py b/studio/backend/utils/datasets/raw_text.py
index 353145fd5a..86b1963fc1 100644
--- a/studio/backend/utils/datasets/raw_text.py
+++ b/studio/backend/utils/datasets/raw_text.py
@@ -40,10 +40,7 @@ def _split_scope(split_name: str | None) -> str:
def _drop_invalid_text_rows(
- dataset: Dataset,
- *,
- mode_title: str,
- split_scope: str,
+ dataset: Dataset, *, mode_title: str, split_scope: str
) -> tuple[Dataset, list[RawTextNotice]]:
filtered_dataset = dataset.filter(lambda ex: isinstance(ex["text"], str))
dropped_rows = len(dataset) - len(filtered_dataset)
@@ -105,8 +102,7 @@ def prepare_raw_text_dataset(
notices.append(
RawTextNotice(
message = (
- f"{mode_title}: renaming column '{renamed_col}' -> 'text' "
- f"for {split_scope}"
+ f"{mode_title}: renaming column '{renamed_col}' -> 'text' " f"for {split_scope}"
),
level = "info",
)
diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py
index 7b63152ede..a0f1fd9f99 100644
--- a/studio/backend/utils/datasets/vlm_processing.py
+++ b/studio/backend/utils/datasets/vlm_processing.py
@@ -66,9 +66,7 @@ def generate_smart_vlm_instruction(
# OCR / Transcription
"ocr": {
"keywords": ["ocr", "transcribe", "transcript"],
- "content_hints": [
- r"[A-Za-z\u0600-\u06FF]{10,}"
- ], # Long text passages (Latin/Arabic)
+ "content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"], # Long text passages (Latin/Arabic)
"instruction": "Transcribe all the text shown in this image.",
"confidence": 0.9,
},
@@ -220,7 +218,6 @@ def generate_smart_vlm_instruction(
}
except Exception as e:
import logging
-
logging.getLogger(__name__).debug(f"LLM-assisted instruction skipped: {e}")
# ===== LEVEL 5: Generic Fallback =====
diff --git a/studio/backend/utils/downsample.py b/studio/backend/utils/downsample.py
index bccf6a23b7..2d340ca248 100644
--- a/studio/backend/utils/downsample.py
+++ b/studio/backend/utils/downsample.py
@@ -12,7 +12,5 @@ def downsample(values: list[float], target_count: int) -> list[float]:
return []
if target_count == 1:
return [values[-1]]
- indices = [
- round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count)
- ]
+ indices = [round(i * (len(values) - 1) / (target_count - 1)) for i in range(target_count)]
return [values[i] for i in indices]
diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py
index 48d5890399..04b9494a29 100644
--- a/studio/backend/utils/hardware/amd.py
+++ b/studio/backend/utils/hardware/amd.py
@@ -160,9 +160,7 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
# amd-smi metric output structure varies by version; try common paths
usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {}))
if isinstance(usage, dict):
- gpu_util = _parse_numeric(
- usage.get("gfx_activity", usage.get("gpu_use_percent"))
- )
+ gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent")))
else:
gpu_util = _parse_numeric(usage)
@@ -188,9 +186,7 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
power_data.get("average_socket_power", power_data.get("socket_power")),
)
)
- power_limit = _parse_numeric(
- power_data.get("power_cap", power_data.get("max_power_limit"))
- )
+ power_limit = _parse_numeric(power_data.get("power_cap", power_data.get("max_power_limit")))
else:
power_draw = None
power_limit = None
@@ -205,14 +201,10 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
)
if isinstance(vram_data, dict):
vram_used_mb = _parse_memory_mb(
- vram_data.get(
- "used_vram", vram_data.get("vram_used", vram_data.get("used"))
- )
+ vram_data.get("used_vram", vram_data.get("vram_used", vram_data.get("used")))
)
vram_total_mb = _parse_memory_mb(
- vram_data.get(
- "total_vram", vram_data.get("vram_total", vram_data.get("total"))
- )
+ vram_data.get("total_vram", vram_data.get("vram_total", vram_data.get("total")))
)
else:
vram_used_mb = None
@@ -220,9 +212,7 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
# Build the standardized dict (same shape as nvidia._build_gpu_metrics)
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
- vram_total_gb = (
- round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
- )
+ vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
vram_util = (
round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb is not None and vram_total_mb > 0
@@ -340,8 +330,7 @@ def get_primary_gpu_utilization() -> dict[str, Any]:
def get_visible_gpu_utilization(
- parent_visible_ids: Optional[list[int]],
- parent_cuda_visible_devices: Optional[str] = None,
+ parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None
) -> dict[str, Any]:
"""Return utilization metrics for visible AMD GPUs."""
if parent_visible_ids is None:
@@ -391,14 +380,11 @@ def get_visible_gpu_utilization(
# "unit": "none"}``, so route raw_id through ``_parse_numeric``
# which already handles bare ints, floats, strings, and that
# dict shape uniformly.
- raw_id = gpu_data.get(
- "gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx))
- )
+ raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx)))
parsed_id = _parse_numeric(raw_id)
if parsed_id is None:
logger.warning(
- "amd-smi GPU id %r could not be parsed; falling back to "
- "enumeration index %d",
+ "amd-smi GPU id %r could not be parsed; falling back to enumeration index %d",
raw_id,
fallback_idx,
)
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 180fde8f13..3f2823302a 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -51,9 +51,7 @@ class DeviceType(str, Enum):
DEVICE: Optional[DeviceType] = None
CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.)
-IS_ROCM: bool = (
- False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py
-)
+IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- routes GPU monitoring to amd.py
def _backend_label(device: DeviceType) -> str:
@@ -85,7 +83,6 @@ def _has_torch() -> bool:
"""Check if PyTorch is importable."""
try:
import torch
-
return True
except ImportError:
return False
@@ -95,7 +92,6 @@ def _has_mlx() -> bool:
"""Check if MLX is importable."""
try:
import mlx.core
-
return True
except ImportError:
return False
@@ -120,7 +116,6 @@ def detect_hardware() -> DeviceType:
# --- CUDA / ROCm: try PyTorch ---
if _has_torch():
import torch
-
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
CHAT_ONLY = False
@@ -142,7 +137,6 @@ def detect_hardware() -> DeviceType:
# --- XPU: Intel GPU ---
if _has_torch():
import torch
-
if hasattr(torch, "xpu") and torch.xpu.is_available():
DEVICE = DeviceType.XPU
CHAT_ONLY = False
@@ -198,7 +192,6 @@ def clear_gpu_cache():
torch.cuda.ipc_collect()
elif device == DeviceType.XPU:
import torch
-
torch.xpu.synchronize()
torch.xpu.empty_cache()
elif device == DeviceType.MLX:
@@ -379,7 +372,6 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# GPU runtime version bundled with torch
try:
import torch
-
versions["cuda"] = getattr(torch.version, "cuda", None)
versions["rocm"] = getattr(torch.version, "hip", None)
except Exception:
@@ -646,9 +638,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
result["backend"] = _backend_label(device)
if IS_ROCM:
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
- _reconcile_primary_rocm_unified_memory(
- result, _get_parent_visible_gpu_spec()
- )
+ _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec())
return result
# SMI tool unavailable or returned no usable data. On Windows, query
# the Performance Counter API (same source as Task Manager) for
@@ -709,9 +699,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
"temperature_c": None,
"vram_used_gb": _used,
"vram_total_gb": _total,
- "vram_utilization_pct": round((_used / _total) * 100, 1)
- if _total > 0
- else None,
+ "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
@@ -722,7 +710,6 @@ def get_gpu_utilization() -> Dict[str, Any]:
if device == DeviceType.MLX:
try:
import psutil
-
agx = _read_apple_gpu_stats()
total_bytes = psutil.virtual_memory().total
except Exception as e:
@@ -795,9 +782,7 @@ def _apply_unified_memory_correction(
device_metrics["vram_total_gb"] = torch_total_gb
device_metrics["vram_used_gb"] = torch_used_gb
device_metrics["vram_utilization_pct"] = (
- round((torch_used_gb / torch_total_gb) * 100, 1)
- if torch_total_gb > 0
- else None
+ round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None
)
logger.debug(
"ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with "
@@ -808,9 +793,7 @@ def _apply_unified_memory_correction(
)
-def _reconcile_rocm_unified_memory(
- utilization: Dict[str, Any], device_indices: list[int]
-) -> None:
+def _reconcile_rocm_unified_memory(utilization: Dict[str, Any], device_indices: list[int]) -> None:
"""Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo).
amd-smi reports only the dedicated slice (~512 MB); torch sees the full
@@ -969,9 +952,7 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]:
# stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES.
_is_rocm_spec = IS_ROCM or (
"CUDA_VISIBLE_DEVICES" not in os.environ
- and (
- "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
- )
+ and ("HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ)
)
if _is_rocm_spec:
hip_vis = os.environ.get("HIP_VISIBLE_DEVICES")
@@ -1064,9 +1045,7 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]:
max_parent_id = max(parent_visible_ids)
if physical_gpu_count > max_parent_id:
# Count is plausibly physical (not just visible), so enforce it
- out_of_range = [
- gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count
- ]
+ out_of_range = [gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count]
if out_of_range:
raise ValueError(
f"Invalid gpu_ids {requested_ids}: IDs must be physical GPU IDs "
@@ -1074,9 +1053,7 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]:
f"Rejected IDs: {out_of_range}. Parent-visible GPUs: {parent_visible_ids}"
)
- disallowed_ids = [
- gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids
- ]
+ disallowed_ids = [gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids]
if disallowed_ids:
raise ValueError(
f"Invalid gpu_ids {requested_ids}: requested GPUs {disallowed_ids} are "
@@ -1097,9 +1074,7 @@ def _resolve_model_identifier_for_gpu_estimate(
return config.base_model
return config.identifier if config else model_name
except Exception as e:
- logger.debug(
- "Could not resolve base model for GPU estimate '%s': %s", model_name, e
- )
+ logger.debug("Could not resolve base model for GPU estimate '%s': %s", model_name, e)
return model_name
@@ -1136,7 +1111,6 @@ def _get_hf_safetensors_total_params(
def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None):
try:
from transformers import AutoConfig
-
trust_remote_code = model_name.lower().startswith("unsloth/")
return AutoConfig.from_pretrained(
model_name,
@@ -1188,7 +1162,6 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str:
try:
import torch.distributed as _td
-
for _attr, _stub in (
("is_initialized", lambda: False),
("is_available", lambda: False),
@@ -1251,17 +1224,15 @@ def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]:
synthetic_total_bytes,
synthetic_total_bytes,
)
- _, _, _, memory_left_for_kv_cache_gb = (
- _vllm_utils.approximate_vllm_memory_usage(
- config,
- load_in_4bit = False,
- load_in_8bit = False,
- max_seq_length = 1,
- gpu_memory_utilization = 1.0,
- enable_lora = False,
- account_for_gradients = False,
- cuda_graph_overhead = False,
- )
+ _, _, _, memory_left_for_kv_cache_gb = _vllm_utils.approximate_vllm_memory_usage(
+ config,
+ load_in_4bit = False,
+ load_in_8bit = False,
+ max_seq_length = 1,
+ gpu_memory_utilization = 1.0,
+ enable_lora = False,
+ account_for_gradients = False,
+ cuda_graph_overhead = False,
)
finally:
_vllm_utils.get_mem_info = original_get_mem_info
@@ -1283,15 +1254,11 @@ def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]:
def estimate_fp16_model_size_bytes(
model_name: str, hf_token: Optional[str] = None
) -> tuple[Optional[int], str]:
- estimate_model = _resolve_model_identifier_for_gpu_estimate(
- model_name, hf_token = hf_token
- )
+ estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token)
total_params = None
if "/" in estimate_model and not Path(estimate_model).exists():
- total_params = _get_hf_safetensors_total_params(
- estimate_model, hf_token = hf_token
- )
+ total_params = _get_hf_safetensors_total_params(estimate_model, hf_token = hf_token)
if total_params:
return int(total_params * 2), "safetensors"
@@ -1346,9 +1313,7 @@ def estimate_required_model_memory_gb(
DEFAULT_TARGET_MODULES,
)
- model_size_bytes, source = estimate_fp16_model_size_bytes(
- model_name, hf_token = hf_token
- )
+ model_size_bytes, source = estimate_fp16_model_size_bytes(model_name, hf_token = hf_token)
metadata: Dict[str, Any] = {
"mode": "inference" if training_type is None else "training",
"model_size_source": source,
@@ -1371,9 +1336,7 @@ def estimate_required_model_memory_gb(
return required_gb, metadata
training_method = (
- "full"
- if training_type == "Full Finetuning"
- else ("qlora" if load_in_4bit else "lora")
+ "full" if training_type == "Full Finetuning" else ("qlora" if load_in_4bit else "lora")
)
vram_config = TrainingVramConfig(
training_method = training_method,
@@ -1386,14 +1349,12 @@ def estimate_required_model_memory_gb(
load_in_4bit = load_in_4bit,
)
- estimate_model = _resolve_model_identifier_for_gpu_estimate(
- model_name, hf_token = hf_token
- )
+ estimate_model = _resolve_model_identifier_for_gpu_estimate(model_name, hf_token = hf_token)
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
if config is not None:
try:
- vram_config.attention_implementation = (
- _determine_attention_impl_for_gpu_estimate(config)
+ vram_config.attention_implementation = _determine_attention_impl_for_gpu_estimate(
+ config
)
except Exception as e:
# Log at debug: on Windows ROCm the torch.distributed stub does
@@ -1587,9 +1548,7 @@ def auto_select_gpu_ids(
return selected, metadata
# Use only GPUs with verified VRAM data (from gpu_candidates, not raw devices)
- fallback_all = (
- [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids
- )
+ fallback_all = [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids
metadata["selection_mode"] = "fallback_all"
if ranked:
fallback_usable = ranked[0]["free_gb"] + sum(
@@ -1853,7 +1812,6 @@ def get_visible_gpu_count() -> int:
# No visibility env var set -- try torch, fall back to physical count
try:
import torch
-
if get_device() == DeviceType.XPU and hasattr(torch, "xpu"):
_visible_gpu_count = torch.xpu.device_count()
else:
@@ -1902,7 +1860,6 @@ def apply_gpu_ids(gpu_ids) -> None:
# Broad except: a probe failure must never crash a training worker.
try:
import torch as _torch
-
_is_rocm = (
getattr(_torch.version, "hip", None) is not None
or "rocm" in getattr(_torch, "__version__", "").lower()
@@ -1923,9 +1880,7 @@ def apply_gpu_ids(gpu_ids) -> None:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value)
-def get_device_map(
- gpu_ids: Optional[list[int]] = None,
-) -> str:
+def get_device_map(gpu_ids: Optional[list[int]] = None) -> str:
"""Return the Hugging Face ``device_map`` string for model loading.
Returns ``"balanced"`` (shard evenly across GPUs) when:
@@ -1949,10 +1904,7 @@ def get_device_map(
# UUID/MIG masks cannot be split into numeric IDs, so if multiple
# GPUs are visible we assume multi-GPU sharding is intended.
parent_visible_spec = _get_parent_visible_gpu_spec()
- if (
- parent_visible_spec["numeric_ids"] is None
- and get_visible_gpu_count() > 1
- ):
+ if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1:
multi_gpu = True
if multi_gpu:
@@ -1972,14 +1924,16 @@ def get_offloaded_device_map_entries(model) -> dict[str, str]:
}
-def raise_if_offloaded(model, device_map: str, context: str = "Loading") -> None:
+def raise_if_offloaded(
+ model,
+ device_map: str,
+ context: str = "Loading",
+) -> None:
"""Raise ``ValueError`` if *model* has modules offloaded to CPU or disk."""
offloaded = get_offloaded_device_map_entries(model)
if not offloaded:
return
- example = ", ".join(
- f"{name}={placement}" for name, placement in list(offloaded.items())[:5]
- )
+ example = ", ".join(f"{name}={placement}" for name, placement in list(offloaded.items())[:5])
raise ValueError(
f"{context} does not support models loaded with CPU or disk offload. "
f"device_map='{device_map}' produced offloaded modules: {example}"
diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py
index 099c5fa3a5..6cead61f08 100644
--- a/studio/backend/utils/hardware/nvidia.py
+++ b/studio/backend/utils/hardware/nvidia.py
@@ -25,20 +25,12 @@ def _parse_smi_value(raw: str):
def _build_gpu_metrics(
- vram_used_mb,
- vram_total_mb,
- power_draw,
- power_limit,
- **extra,
+ vram_used_mb, vram_total_mb, power_draw, power_limit, **extra
) -> dict[str, Any]:
return {
**extra,
- "vram_used_gb": round(vram_used_mb / 1024, 2)
- if vram_used_mb is not None
- else None,
- "vram_total_gb": round(vram_total_mb / 1024, 2)
- if vram_total_mb is not None
- else None,
+ "vram_used_gb": round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None,
+ "vram_total_gb": round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None,
"vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
else None,
@@ -50,9 +42,7 @@ def _build_gpu_metrics(
}
-def _visible_ordinal_map(
- parent_visible_ids: Optional[list[int]],
-) -> Optional[dict[int, int]]:
+def _visible_ordinal_map(parent_visible_ids: Optional[list[int]]) -> Optional[dict[int, int]]:
if parent_visible_ids is None:
return None
return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)}
@@ -118,8 +108,7 @@ def get_primary_gpu_utilization() -> dict[str, Any]:
def get_visible_gpu_utilization(
- parent_visible_ids: Optional[list[int]],
- parent_cuda_visible_devices: Optional[str] = None,
+ parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None
) -> dict[str, Any]:
# When parent_visible_ids is None (UUID/MIG mask), we cannot safely
# map nvidia-smi rows to the process's visible devices. Return empty
@@ -188,9 +177,7 @@ def get_visible_gpu_utilization(
index = idx,
index_kind = "physical",
visible_ordinal = (
- visible_ordinals[idx]
- if visible_ordinals is not None
- else len(devices)
+ visible_ordinals[idx] if visible_ordinals is not None else len(devices)
),
gpu_utilization_pct = _parse_smi_value(parts[1]),
temperature_c = _parse_smi_value(parts[2]),
@@ -207,8 +194,7 @@ def get_visible_gpu_utilization(
def get_backend_visible_gpu_info(
- parent_visible_ids: Optional[list[int]],
- backend_cuda_visible_devices: Optional[str],
+ parent_visible_ids: Optional[list[int]], backend_cuda_visible_devices: Optional[str]
) -> dict[str, Any]:
# When parent_visible_ids is None (UUID/MIG mask), we cannot safely
# map nvidia-smi rows to the process's visible devices.
@@ -274,9 +260,7 @@ def get_backend_visible_gpu_info(
"index": idx,
"index_kind": "physical",
"visible_ordinal": (
- visible_ordinals[idx]
- if visible_ordinals is not None
- else len(devices)
+ visible_ordinals[idx] if visible_ordinals is not None else len(devices)
),
"name": name,
"memory_total_gb": round(mem_total_mb / 1024, 2),
diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py
index ba1b1dfe61..ddc39733e3 100644
--- a/studio/backend/utils/hardware/vram_estimation.py
+++ b/studio/backend/utils/hardware/vram_estimation.py
@@ -16,9 +16,7 @@ from dataclasses import dataclass, field
from typing import Dict, Optional
QUANT_4BIT_FACTOR = 16 / 5
-DOUBLE_QUANT_4BIT_FACTOR = (
- 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1
-)
+DOUBLE_QUANT_4BIT_FACTOR = 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1
CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti
NON_FLASH_ATTENTION_FACTOR = (
12.0 # eager attention score+workspace overhead; see VRAM_ESTIMATION.md section 5
@@ -148,12 +146,7 @@ class VramBreakdown:
Weights/LoRA/optimizer/gradients shard across GPUs.
Activations do NOT shard (the GPU running a layer holds them).
"""
- shardable = (
- self.model_weights
- + self.lora_adapters
- + self.optimizer_states
- + self.gradients
- )
+ shardable = self.model_weights + self.lora_adapters + self.optimizer_states + self.gradients
per_gpu_fixed = self.activations + self.cuda_overhead
return shardable // max(n_gpus, 1) + per_gpu_fixed
@@ -194,9 +187,7 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple:
layer_types = getattr(text_config, "mlp_layer_types", None)
if layer_types:
return tuple(
- i
- for i, t in enumerate(layer_types[:total_layers])
- if str(t).lower() == "dense"
+ i for i, t in enumerate(layer_types[:total_layers]) if str(t).lower() == "dense"
)
# why: Llama4TextConfig.__init__ auto-populates self.moe_layers from
@@ -234,9 +225,7 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple:
if sparse_step is not None and sparse_step > 0:
mlp_only_set = {int(i) for i in mlp_only}
return tuple(
- i
- for i in range(total_layers)
- if i in mlp_only_set or (i + 1) % sparse_step != 0
+ i for i in range(total_layers) if i in mlp_only_set or (i + 1) % sparse_step != 0
)
return ()
@@ -264,8 +253,7 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
intermediate_size = hidden_size * 4
if not all(
- v is not None
- for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size)
+ v is not None for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size)
):
return None
if num_heads <= 0:
@@ -330,9 +318,7 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
# per MoE layer (modeling_llama4.py).
intermediate_size_mlp_raw = _first_scalar(_moe_attr("intermediate_size_mlp"))
dense_intermediate_size = (
- int(intermediate_size_mlp_raw)
- if intermediate_size_mlp_raw is not None
- else None
+ int(intermediate_size_mlp_raw) if intermediate_size_mlp_raw is not None else None
)
if (
intermediate_size_mlp_raw is not None
@@ -391,9 +377,7 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
None,
)
or 0,
- quantization_skip_modules = list(
- quantization_config.get("llm_int8_skip_modules", []) or []
- ),
+ quantization_skip_modules = list(quantization_config.get("llm_int8_skip_modules", []) or []),
quant_4bit_factor = quant_4bit_factor,
moe_has_dense_mlp = bool(getattr(text_config, "enable_moe_block", False)),
dense_layer_indices = dense_layer_indices,
@@ -475,11 +459,7 @@ def _per_layer_input_norm_elements(arch: ModelArchConfig) -> int:
return hd * n_layers + pli
-def _per_layer_input_lora_params(
- arch: ModelArchConfig,
- r: int,
- target_modules,
-) -> int:
+def _per_layer_input_lora_params(arch: ModelArchConfig, r: int, target_modules) -> int:
# why: Unsloth's get_peft_regex (unsloth_zoo/peft_utils.py) requires module
# names to contain a component tag (mlp/attn/...); PLE module names lack
# any tag, so all-linear training does NOT attach LoRA to them. Only count
@@ -487,11 +467,7 @@ def _per_layer_input_lora_params(
pli = arch.hidden_size_per_layer_input
if pli <= 0:
return 0
- targets = (
- {target_modules}
- if isinstance(target_modules, str)
- else set(target_modules or [])
- )
+ targets = {target_modules} if isinstance(target_modules, str) else set(target_modules or [])
n_layers = arch.num_hidden_layers
hd = arch.hidden_size
total = 0
@@ -508,11 +484,7 @@ def _layer_attention_dims(arch: ModelArchConfig, layer_idx: int) -> tuple:
layer_types = _layer_types(arch)
layer_type = layer_types[layer_idx]
is_sliding = layer_type == "sliding_attention"
- head_dim = (
- arch.global_head_dim
- if not is_sliding and arch.global_head_dim
- else _head_dim(arch)
- )
+ head_dim = arch.global_head_dim if not is_sliding and arch.global_head_dim else _head_dim(arch)
use_alt_attention = arch.attention_k_eq_v and not is_sliding
num_kv_heads = (
arch.num_global_key_value_heads
@@ -532,10 +504,7 @@ def _layer_mlp_size(arch: ModelArchConfig, layer_idx: int) -> int:
return _dense_mlp_size(arch)
-def _text_linear_dims(
- arch: ModelArchConfig,
- layer_idx: int,
-) -> Dict[str, tuple[int, int]]:
+def _text_linear_dims(arch: ModelArchConfig, layer_idx: int) -> Dict[str, tuple[int, int]]:
hd = arch.hidden_size
if _uses_structured_layer_shapes(arch):
q_size, kv_size, has_k, has_v = _layer_attention_dims(arch, layer_idx)
@@ -589,11 +558,7 @@ def _module_path_matches(skip_module: str, alias: str) -> bool:
return ".".join(prefix_parts) in _SKIP_MODULE_TEXT_PREFIXES
-def _add_module_aliases(
- aliases: Dict[str, str],
- canonical: str,
- suffix: str,
-) -> None:
+def _add_module_aliases(aliases: Dict[str, str], canonical: str, suffix: str) -> None:
for prefix in (
"",
"model",
@@ -607,9 +572,7 @@ def _add_module_aliases(
aliases[alias] = canonical
-def _build_text_module_elements(
- arch: ModelArchConfig,
-) -> tuple[Dict[str, int], Dict[str, str]]:
+def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int], Dict[str, str]]:
elements: Dict[str, int] = {}
aliases: Dict[str, str] = {}
@@ -620,12 +583,8 @@ def _build_text_module_elements(
for layer_idx in range(arch.num_hidden_layers):
layer_modules: Dict[str, int] = {}
dims = _text_linear_dims(arch, layer_idx)
- attn_dims = {
- name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES
- }
- mlp_dims = {
- name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES
- }
+ attn_dims = {name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES}
+ mlp_dims = {name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES}
if is_mla:
# why: _text_linear_dims uses (hd, hd) for q/o; MLA actually splits
@@ -677,10 +636,7 @@ def _build_text_module_elements(
)
else:
layer_modules.update(
- {
- f"mlp.{name}": in_dim * out_dim
- for name, (in_dim, out_dim) in mlp_dims.items()
- }
+ {f"mlp.{name}": in_dim * out_dim for name, (in_dim, out_dim) in mlp_dims.items()}
)
if pli > 0:
@@ -704,10 +660,7 @@ def _build_text_module_elements(
for name, value in layer_modules.items()
if (
name == "mlp"
- or (
- name.startswith("mlp.")
- and not (is_sibling_experts and name == "mlp.experts")
- )
+ or (name.startswith("mlp.") and not (is_sibling_experts and name == "mlp.experts"))
)
)
experts_total = layer_modules.get("mlp.experts", 0) if is_sibling_experts else 0
@@ -764,10 +717,7 @@ def _compute_skipped_quantizable_elements(arch: ModelArchConfig) -> int:
pruned = {
canonical
for canonical in matched
- if not any(
- canonical != parent and canonical.startswith(f"{parent}.")
- for parent in matched
- )
+ if not any(canonical != parent and canonical.startswith(f"{parent}.") for parent in matched)
}
return sum(module_elements[canonical] for canonical in pruned)
@@ -900,9 +850,7 @@ def _compute_layer_elements(arch: ModelArchConfig):
mlp_total = _compute_dense_mlp_elements(arch) * n_layers
layernorms = 2 * hd
- per_layer_embed = (
- arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers
- )
+ per_layer_embed = arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers
ple_text_linear = _per_layer_input_quantizable(arch)
ple_norms = _per_layer_input_norm_elements(arch)
embed_tokens = arch.vocab_size * hd + per_layer_embed + ple_norms
@@ -911,9 +859,7 @@ def _compute_layer_elements(arch: ModelArchConfig):
def compute_model_weights_bytes(
- arch: ModelArchConfig,
- training_method: str,
- load_in_4bit: bool,
+ arch: ModelArchConfig, training_method: str, load_in_4bit: bool
) -> int:
total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch)
n_layers = arch.num_hidden_layers
@@ -926,9 +872,7 @@ def compute_model_weights_bytes(
)
quantized = total_quantizable - skipped_quantizable
return int(
- quantized * 2 / arch.quant_4bit_factor
- + skipped_quantizable * 2
- + non_quantizable * 2
+ quantized * 2 / arch.quant_4bit_factor + skipped_quantizable * 2 + non_quantizable * 2
)
return int((total_quantizable + non_quantizable) * 2)
@@ -940,11 +884,7 @@ def compute_total_params(arch: ModelArchConfig) -> int:
return total_quantizable + layernorms * n_layers + embed_tokens + lm_head
-def _lora_attn_elements(
- arch: ModelArchConfig,
- r: int,
- target_modules: list,
-) -> int:
+def _lora_attn_elements(arch: ModelArchConfig, r: int, target_modules: list) -> int:
hd = arch.hidden_size
if arch.q_lora_rank is not None:
# MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o
@@ -974,11 +914,7 @@ def _lora_attn_elements(
def _lora_mlp_elements(
- hd: int,
- mlp_size: int,
- r: int,
- target_modules: list,
- expert_mult: int,
+ hd: int, mlp_size: int, r: int, target_modules: list, expert_mult: int
) -> int:
module_ab = {
"gate_proj": (hd * r, r * mlp_size),
@@ -992,11 +928,7 @@ def _lora_mlp_elements(
return total
-def compute_lora_params(
- arch: ModelArchConfig,
- lora_rank: int,
- target_modules: list,
-) -> int:
+def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: list) -> int:
all_linear = _targets_all_linear(target_modules)
selected_modules = list(DEFAULT_TARGET_MODULES) if all_linear else target_modules
hd = arch.hidden_size
@@ -1062,11 +994,7 @@ def compute_lora_params(
mlp_total = moe_mlp * n_moe + dense_only
else:
mlp_total = structured_dense_mlp
- return (
- attn_total
- + mlp_total
- + _per_layer_input_lora_params(arch, r, target_modules)
- )
+ return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules)
elif n_experts > 1:
attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers
n_dense = arch.num_dense_layers
@@ -1118,9 +1046,7 @@ def compute_lora_params(
* n_layers
)
- return (
- attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules)
- )
+ return attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules)
def compute_lora_adapter_bytes(lora_params: int) -> int:
@@ -1144,10 +1070,7 @@ def _is_linear_attention(attention_implementation: Optional[str]) -> bool:
def _compute_non_flash_attention_bytes(
- arch: ModelArchConfig,
- batch_size: int,
- seq_len: int,
- effective_layers: float,
+ arch: ModelArchConfig, batch_size: int, seq_len: int, effective_layers: float
) -> int:
score_elements = batch_size * arch.num_attention_heads * seq_len * seq_len
return int(score_elements * 2 * NON_FLASH_ATTENTION_FACTOR * effective_layers)
@@ -1190,10 +1113,7 @@ def _layer_qkv_mlp_sizes(arch: ModelArchConfig, layer_idx: int) -> tuple:
def _per_layer_activation_bytes(
- arch: ModelArchConfig,
- layer_idx: int,
- batch_size: int,
- seq_len: int,
+ arch: ModelArchConfig, layer_idx: int, batch_size: int, seq_len: int
) -> int:
qkv_size, mlp_size = _layer_qkv_mlp_sizes(arch, layer_idx)
activation_qkv = seq_len * batch_size * qkv_size
@@ -1204,9 +1124,7 @@ def _per_layer_activation_bytes(
# is set; see gemma4/modular_gemma4.py:1141-1145.
pli = arch.hidden_size_per_layer_input
activation_ple = seq_len * batch_size * (arch.hidden_size + pli) if pli > 0 else 0
- return int(
- (activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25
- )
+ return int((activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25)
def compute_activation_bytes(
@@ -1227,14 +1145,12 @@ def compute_activation_bytes(
if gc_multiplier is None:
effective_layers = n_layers
linear_bytes = sum(
- _per_layer_activation_bytes(arch, i, batch_size, seq_len)
- for i in range(n_layers)
+ _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers)
)
else:
effective_layers = gc_multiplier
max_layer_bytes = max(
- _per_layer_activation_bytes(arch, i, batch_size, seq_len)
- for i in range(n_layers)
+ _per_layer_activation_bytes(arch, i, batch_size, seq_len) for i in range(n_layers)
)
linear_bytes = int(max_layer_bytes * effective_layers)
@@ -1257,10 +1173,7 @@ def compute_activation_bytes(
)
-def estimate_training_vram(
- arch: ModelArchConfig,
- config: TrainingVramConfig,
-) -> VramBreakdown:
+def estimate_training_vram(arch: ModelArchConfig, config: TrainingVramConfig) -> VramBreakdown:
method = config.training_method.lower()
is_lora = method in ("qlora", "lora")
load_in_4bit = config.load_in_4bit or method == "qlora"
diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py
index 9efc281b0b..e07bc62cfd 100644
--- a/studio/backend/utils/inference/inference_config.py
+++ b/studio/backend/utils/inference/inference_config.py
@@ -34,10 +34,7 @@ def _load_family_defaults():
return
json_path = (
- Path(__file__).parent.parent.parent
- / "assets"
- / "configs"
- / "inference_defaults.json"
+ Path(__file__).parent.parent.parent / "assets" / "configs" / "inference_defaults.json"
)
try:
with open(json_path, "r", encoding = "utf-8") as f:
diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py
index 2c781f4a7b..f0e642a38e 100644
--- a/studio/backend/utils/llama_cpp_freshness.py
+++ b/studio/backend/utils/llama_cpp_freshness.py
@@ -38,7 +38,6 @@ def _cache_dir() -> Path:
"""Lazy import so tests can stub storage_roots."""
try:
from utils.paths.storage_roots import cache_root
-
return cache_root() / "llama_cpp_freshness"
except Exception:
return Path.home() / ".unsloth" / "studio" / "cache" / "llama_cpp_freshness"
@@ -136,9 +135,7 @@ def _fetch_latest_release_tag(repo: str, timeout: float = 5.0) -> Optional[str]:
return tag if isinstance(tag, str) and tag else None
-def latest_published_release(
- repo: str, *, force_refresh: bool = False
-) -> Optional[str]:
+def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optional[str]:
"""Latest release tag for `repo`. Memo + disk-cached (24h TTL).
None when offline and never previously cached."""
if not repo:
diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py
index b6b2e11c2e..99770ce7a6 100644
--- a/studio/backend/utils/models/checkpoints.py
+++ b/studio/backend/utils/models/checkpoints.py
@@ -99,9 +99,7 @@ def scan_checkpoints(
name_part = parts[0]
idx = name_part.find("_")
if idx > 0:
- metadata["base_model"] = (
- name_part[:idx] + "/" + name_part[idx + 1 :]
- )
+ metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1 :]
else:
metadata["base_model"] = name_part
@@ -131,9 +129,7 @@ def scan_checkpoints(
)
models.append((item.name, checkpoints, metadata))
- logger.debug(
- f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)"
- )
+ logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
# Sort by modification time (newest first)
models.sort(key = lambda x: Path(x[1][0][1]).stat().st_mtime, reverse = True)
diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index dc9c21226e..32618773b7 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -282,8 +282,7 @@ def is_mmproj_by_metadata(meta: Optional[Dict[str, str]]) -> Optional[bool]:
def pairing_score(
- weight_meta: Optional[Dict[str, str]],
- mmproj_meta: Optional[Dict[str, str]],
+ weight_meta: Optional[Dict[str, str]], mmproj_meta: Optional[Dict[str, str]]
) -> int:
"""Pairing confidence: 100 = base_model URL match, 80 = basename + org,
60 = basename, -1 = definitive mismatch, 0 = decide from filename."""
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index b488a19953..92960485a4 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -57,13 +57,9 @@ def _env_offline() -> bool:
# ── Model size extraction ────────────────────────────────────
import re as _re
-_MODEL_SIZE_RE = _re.compile(
- r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE
-)
+_MODEL_SIZE_RE = _re.compile(r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
# MoE active-parameter pattern: matches "A3B", "A3.5B", etc.
-_ACTIVE_SIZE_RE = _re.compile(
- r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE
-)
+_ACTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
def extract_model_size_b(model_id: str) -> float | None:
@@ -571,9 +567,7 @@ except Exception as exc:
"""
-def _is_vision_model_subprocess(
- model_name: str, hf_token: Optional[str] = None
-) -> Optional[bool]:
+def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
"""Run is_vision_model check in a subprocess with transformers 5.x.
Same pattern as training/inference workers: spawn a clean subprocess
@@ -715,9 +709,7 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
return False
-def _is_vision_model_uncached(
- model_name: str, hf_token: Optional[str] = None
-) -> Optional[bool]:
+def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
"""Uncached vision model detection -- called by is_vision_model().
Returns True/False for definitive results, or None when detection failed
@@ -775,9 +767,7 @@ def _is_vision_model_uncached(
# Check 5: Known VLM model_type values that may not match above checks
if hasattr(config, "model_type"):
if config.model_type in _VLM_MODEL_TYPES:
- logger.info(
- f"Model {model_name} detected as VLM: model_type={config.model_type}"
- )
+ logger.info(f"Model {model_name} detected as VLM: model_type={config.model_type}")
return True
return False
@@ -823,9 +813,7 @@ _AUDIO_TOKEN_PATTERNS = {
and "<|text_start|>" in tokens
and "<|text_end|>" in tokens
),
- "snac": lambda tokens: (
- sum(1 for t in tokens if t.startswith(" 10000
- ),
+ "snac": lambda tokens: (sum(1 for t in tokens if t.startswith(" 10000),
}
@@ -849,9 +837,7 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
return result
-def _detect_audio_from_tokenizer(
- model_name: str, hf_token: Optional[str] = None
-) -> Optional[str]:
+def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
"""Detect audio type from tokenizer special tokens (for LLM-based audio models).
First checks local HF cache, then fetches tokenizer_config.json from HuggingFace.
@@ -913,9 +899,7 @@ def _detect_audio_from_tokenizer(
return None
except Exception as e:
- logger.debug(
- f"Could not detect audio type from tokenizer for {model_name}: {e}"
- )
+ logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
return None
@@ -1346,9 +1330,7 @@ def _iter_hf_cache_snapshots(repo_id: str):
yield from snap_dirs
-def _list_gguf_variants_from_hf_cache(
- repo_id: str,
-) -> Optional[tuple[list[GgufVariantInfo], bool]]:
+def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]:
"""Variants from the local HF cache snapshot, or None if not cached."""
for snap in _iter_hf_cache_snapshots(repo_id):
variants, has_vision = list_local_gguf_variants(str(snap))
@@ -1358,8 +1340,7 @@ def _list_gguf_variants_from_hf_cache(
def list_gguf_variants(
- repo_id: str,
- hf_token: Optional[str] = None,
+ repo_id: str, hf_token: Optional[str] = None
) -> tuple[list[GgufVariantInfo], bool]:
"""
List all GGUF quantization variants in a HuggingFace repo.
@@ -1462,9 +1443,7 @@ def _resolve_gguf_dir(p: Path) -> Optional[Path]:
return None
-def list_local_gguf_variants(
- directory: str,
-) -> tuple[list[GgufVariantInfo], bool]:
+def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], bool]:
"""List GGUF quantization variants in a local directory.
Mirrors :func:`list_gguf_variants` but reads from the filesystem
@@ -1533,8 +1512,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
matches = sorted(
f
for f in _iter_gguf_files(p, recursive = True)
- if not _is_mmproj(f.name)
- and _extract_quant_label(f.relative_to(p).as_posix()) == variant
+ if not _is_mmproj(f.name) and _extract_quant_label(f.relative_to(p).as_posix()) == variant
)
if matches:
return str(matches[0].resolve())
@@ -1558,10 +1536,7 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
return None
-def detect_gguf_model_remote(
- repo_id: str,
- hf_token: Optional[str] = None,
-) -> Optional[str]:
+def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Optional[str]:
"""
Check if a HuggingFace repo contains GGUF files.
@@ -1617,9 +1592,7 @@ def detect_gguf_model_remote(
)
return cached
- logger.warning(
- f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}"
- )
+ logger.warning(f"Could not check GGUF files for '{repo_id}' after 3 attempts: {last_err}")
return None
@@ -1752,9 +1725,7 @@ def _looks_like_lora_adapter(model_dir: Path) -> bool:
)
-def scan_trained_models(
- outputs_dir: str = str(outputs_root()),
-) -> List[Tuple[str, str, str]]:
+def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]:
"""
Scan outputs folder for trained Studio models.
@@ -1823,9 +1794,7 @@ def scan_exported_models(
# Check for flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/)
# Filter out mmproj (vision projection) files — they aren't loadable as main models
- gguf_files = [
- f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name)
- ]
+ gguf_files = [f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name)]
if gguf_files:
base_model = None
export_meta = run_dir / "export_metadata.json"
@@ -1900,9 +1869,7 @@ def scan_exported_models(
# Fallback: read base model from the original training run's
# adapter_config.json in ./outputs/{run_name}/
if not base_model:
- outputs_adapter_cfg = (
- resolve_output_dir(run_dir.name) / "adapter_config.json"
- )
+ outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
@@ -1935,9 +1902,7 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
- logger.info(
- "Detected base model from adapter_config.json: %s", base_model
- )
+ logger.info("Detected base model from adapter_config.json: %s", base_model)
return base_model
config_path = checkpoint_path_obj / "config.json"
@@ -2010,9 +1975,7 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
- logger.info(
- f"Detected base model from adapter_config.json: {base_model}"
- )
+ logger.info(f"Detected base model from adapter_config.json: {base_model}")
return base_model
# Fallback: try training_args.bin (requires torch)
@@ -2084,9 +2047,7 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
if config_path.is_file():
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
- logger.info(
- f"Loaded model defaults from {config_path} (via mapping)"
- )
+ logger.info(f"Loaded model defaults from {config_path} (via mapping)")
return config
# If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
@@ -2156,14 +2117,10 @@ class ModelConfig:
is_lora: bool # Is this a lora adapter?
is_gguf: bool = False # Is this a GGUF model?
is_audio: bool = False # Is this a TTS audio model?
- audio_type: Optional[str] = (
- None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
- )
+ audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
- gguf_mmproj_file: Optional[str] = (
- None # Full path to the mmproj .gguf file (vision projection)
- )
+ gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
gguf_hf_repo: Optional[str] = (
None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
)
@@ -2172,7 +2129,9 @@ class ModelConfig:
@classmethod
def from_lora_path(
- cls, lora_path: str, hf_token: Optional[str] = None
+ cls,
+ lora_path: str,
+ hf_token: Optional[str] = None,
) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a local LoRA adapter path.
@@ -2321,9 +2280,7 @@ class ModelConfig:
gguf_is_vision = True
logger.info(f"Detected mmproj for vision: {mmproj_file}")
elif base_is_vision:
- logger.warning(
- f"Base model is vision but no mmproj file found in {gguf_dir}"
- )
+ logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
return cls(
identifier = identifier,
@@ -2385,15 +2342,11 @@ class ModelConfig:
# Auto-detect LoRA for local paths (check adapter_config.json on disk)
if not is_lora and is_local:
detected_base = (
- get_base_model_from_lora(path)
- if _looks_like_lora_adapter(Path(path))
- else None
+ get_base_model_from_lora(path) if _looks_like_lora_adapter(Path(path)) else None
)
if detected_base:
is_lora = True
- logger.info(
- f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})"
- )
+ logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})")
# Auto-detect LoRA for remote HF models. When offline, huggingface_hub
# raises OfflineModeIsEnabled in ~0ms; we fall through to the cache.
@@ -2407,18 +2360,14 @@ class ModelConfig:
is_lora = True
logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'")
except Exception as e:
- logger.debug(
- f"Could not check remote LoRA status for '{identifier}': {e}"
- )
+ logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
# API may have failed; adapter_config.json may still be cached.
if not is_lora:
for snap in _iter_hf_cache_snapshots(identifier):
if (snap / "adapter_config.json").is_file():
is_lora = True
- logger.info(
- f"Auto-detected cached LoRA adapter: '{identifier}'"
- )
+ logger.info(f"Auto-detected cached LoRA adapter: '{identifier}'")
break
# Handle LoRA adapters
@@ -2432,9 +2381,7 @@ class ModelConfig:
try:
from huggingface_hub import hf_hub_download
- config_path = hf_hub_download(
- identifier, "adapter_config.json", token = hf_token
- )
+ config_path = hf_hub_download(identifier, "adapter_config.json", token = hf_token)
with open(config_path, "r") as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
@@ -2498,9 +2445,7 @@ class ModelConfig:
# Use the correct 'local_models' parameter to resolve display names
if " (Active)" in selected or " (Ready)" in selected:
- clean_display_name = selected.replace(" (Active)", "").replace(
- " (Ready)", ""
- )
+ clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
if local_models:
for local_display, local_path in local_models:
if local_display == clean_display_name:
diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py
index a69dfab532..7d8514abc8 100644
--- a/studio/backend/utils/native_path_leases.py
+++ b/studio/backend/utils/native_path_leases.py
@@ -68,9 +68,7 @@ def native_path_leases_supported() -> bool:
return True
-def child_env_without_native_path_secret(
- env: Mapping[str, str] | None = None,
-) -> dict[str, str]:
+def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) -> dict[str, str]:
"""Return a child-process env with the native path lease secret removed."""
if env is None:
@@ -82,11 +80,7 @@ def child_env_without_native_path_secret(
return cleaned
-def run_without_native_path_secret(
- target: Callable[..., Any],
- *args: Any,
- **kwargs: Any,
-) -> Any:
+def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
"""Run a multiprocessing child target without the native path lease secret."""
global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET
@@ -153,9 +147,7 @@ def verify_native_path_lease(
raise NativePathLeaseError("Native path is no longer accessible.") from exc
_reject_network_or_device_path(resolved)
if not _same_native_path(resolved, path):
- raise NativePathLeaseError(
- "Native path grant no longer resolves to the selected path."
- )
+ raise NativePathLeaseError("Native path grant no longer resolves to the selected path.")
grant = NativePathGrant(
operation = str(payload["operation"]),
@@ -219,9 +211,7 @@ def _decode_secret() -> bytes:
if encoded is None and _SCRUB_SAVED_SECRET is not None:
encoded = _SCRUB_SAVED_SECRET
if not encoded:
- raise NativePathLeaseError(
- "Native path grants require the managed desktop backend."
- )
+ raise NativePathLeaseError("Native path grants require the managed desktop backend.")
try:
secret = _b64decode(encoded)
except Exception as exc:
@@ -272,9 +262,7 @@ def _validate_payload(
)
missing = [key for key in required if key not in payload]
if missing:
- raise NativePathLeaseError(
- "Native path grant payload is missing required fields."
- )
+ raise NativePathLeaseError("Native path grant payload is missing required fields.")
if _required_int(payload, "version") != 1:
raise NativePathLeaseError("Native path grant version is unsupported.")
if payload["operation"] != operation:
@@ -353,19 +341,13 @@ def _reject_network_or_device_path(path: Path) -> None:
rest = normalized[4:]
is_local_drive = len(rest) >= 3 and rest[0].isalpha() and rest[1:3] == ":\\"
if not is_local_drive:
- raise NativePathLeaseError(
- "Network paths are not supported for native grants."
- )
+ raise NativePathLeaseError("Network paths are not supported for native grants.")
elif normalized.startswith("\\\\"):
- raise NativePathLeaseError(
- "Network paths are not supported for native grants."
- )
+ raise NativePathLeaseError("Network paths are not supported for native grants.")
if os.name != "nt":
for root in ("/dev", "/proc", "/sys"):
if path.is_relative_to(root):
- raise NativePathLeaseError(
- "Device and virtual filesystem paths are not supported."
- )
+ raise NativePathLeaseError("Device and virtual filesystem paths are not supported.")
if "\x00" in text:
raise NativePathLeaseError("Native path contains invalid characters.")
@@ -397,9 +379,7 @@ def _optional_int(value: Any) -> int | None:
def _required_int(payload: dict[str, Any], key: str) -> int:
raw = payload.get(key)
if raw is None:
- raise NativePathLeaseError(
- "Native path grant payload is missing required fields."
- )
+ raise NativePathLeaseError("Native path grant payload is missing required fields.")
try:
return int(raw)
except (TypeError, ValueError) as exc:
diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py
index 9ef9a2dd92..22c6c46ee1 100644
--- a/studio/backend/utils/paths/path_utils.py
+++ b/studio/backend/utils/paths/path_utils.py
@@ -134,7 +134,6 @@ def _hf_hub_cache_dir() -> Path:
"""Return HF cache root honoring HF_HUB_CACHE when available."""
try:
from huggingface_hub.constants import HF_HUB_CACHE
-
return Path(HF_HUB_CACHE)
except Exception as exc:
logger.debug(
diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py
index 6319452ed2..b254c20f97 100644
--- a/studio/backend/utils/paths/storage_roots.py
+++ b/studio/backend/utils/paths/storage_roots.py
@@ -263,9 +263,7 @@ def _setup_cache_env() -> None:
Works on Linux, macOS, and Windows.
"""
root = cache_root()
- xdg_cache = Path(
- os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")
- ).expanduser()
+ xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser()
hf_default = xdg_cache / "huggingface"
defaults: dict[str, str] = {
"HF_HOME": str(hf_default),
@@ -298,9 +296,7 @@ def ensure_studio_directories() -> None:
_setup_cache_env()
-def _clean_relative_path(
- path_value: str, *, strip_prefixes: tuple[str, ...] = ()
-) -> Path:
+def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ()) -> Path:
path = Path(path_value).expanduser()
parts = [part for part in path.parts if part not in ("", ".")]
while parts and parts[0] in strip_prefixes:
@@ -319,8 +315,7 @@ def _assert_contained(resolved: Path, root: Path) -> None:
resolved_real.relative_to(root_real)
except ValueError as exc:
raise ValueError(
- f"path escapes root: {resolved!s} -> {resolved_real!s} "
- f"is not under {root_real!s}"
+ f"path escapes root: {resolved!s} -> {resolved_real!s} " f"is not under {root_real!s}"
) from exc
@@ -394,9 +389,7 @@ def resolve_dataset_path(path_value: str) -> Path:
return path
except ValueError:
continue
- raise ValueError(
- f"dataset path must be relative or under a dataset root: {raw!r}"
- )
+ raise ValueError(f"dataset path must be relative or under a dataset root: {raw!r}")
parts = [part for part in Path(path_value).parts if part not in ("", ".")]
if parts[:2] == ["assets", "datasets"]:
diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py
index 70059f8a3c..e59439a2cc 100644
--- a/studio/backend/utils/studio_version.py
+++ b/studio/backend/utils/studio_version.py
@@ -39,9 +39,7 @@ def _path_is_in_site_packages(path: Path) -> bool:
def _is_source_checkout(repo_root: Path) -> bool:
- return (repo_root / ".git").exists() and not _path_is_in_site_packages(
- Path(__file__).resolve()
- )
+ return (repo_root / ".git").exists() and not _path_is_in_site_packages(Path(__file__).resolve())
def _exact_git_studio_tag(repo_root: Path) -> str | None:
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index c23857e0a4..16f964d628 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -201,7 +201,6 @@ def _resolve_base_model(model_name: str) -> str:
if local_path.is_dir():
try:
from utils.models import get_base_model_from_lora
-
base = get_base_model_from_lora(model_name)
if base:
logger.info(
@@ -275,9 +274,7 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
_tokenizer_class_cache[model_name] = result
return result
except Exception as exc:
- logger.debug(
- "Could not fetch tokenizer_config.json for '%s': %s", model_name, exc
- )
+ logger.debug("Could not fetch tokenizer_config.json for '%s': %s", model_name, exc)
_tokenizer_class_cache[model_name] = False
return False
@@ -467,8 +464,7 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool:
pkg_name_norm = pkg_name.replace("-", "_")
# Check directory exists
if not any(
- (Path(venv_dir) / d).is_dir()
- for d in (pkg_name_norm, pkg_name_norm.replace("_", "-"))
+ (Path(venv_dir) / d).is_dir() for d in (pkg_name_norm, pkg_name_norm.replace("_", "-"))
):
return False
# For unpinned packages, existence is enough
@@ -563,9 +559,7 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo
if _venv_dir_is_valid(venv_dir, packages):
return True
- logger.warning(
- "%s not found or incomplete at %s -- installing at runtime", label, venv_dir
- )
+ logger.warning("%s not found or incomplete at %s -- installing at runtime", label, venv_dir)
shutil.rmtree(venv_dir, ignore_errors = True)
os.makedirs(venv_dir, exist_ok = True)
for pkg in packages:
@@ -577,16 +571,12 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo
def _ensure_venv_t5_530_exists() -> bool:
"""Ensure .venv_t5_530/ exists with transformers 5.3.0."""
- return _ensure_venv_dir(
- _VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0"
- )
+ return _ensure_venv_dir(_VENV_T5_530_DIR, _VENV_T5_530_PACKAGES, "transformers 5.3.0")
def _ensure_venv_t5_550_exists() -> bool:
"""Ensure .venv_t5_550/ exists with transformers 5.5.0."""
- return _ensure_venv_dir(
- _VENV_T5_550_DIR, _VENV_T5_550_PACKAGES, "transformers 5.5.0"
- )
+ return _ensure_venv_dir(_VENV_T5_550_DIR, _VENV_T5_550_PACKAGES, "transformers 5.5.0")
def _ensure_venv_t5_exists() -> bool:
@@ -693,15 +683,12 @@ def ensure_transformers_version(model_name: str) -> None:
_deactivate_5x()
if not ensure_fn():
raise RuntimeError(
- f"Cannot activate transformers {target_version}: "
- f"venv missing at {venv_dir}"
+ f"Cannot activate transformers {target_version}: " f"venv missing at {venv_dir}"
)
logger.info("Activating transformers %s…", target_version)
_activate_venv(venv_dir, f"transformers {target_version}")
else:
- logger.info(
- "Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION
- )
+ logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION)
_deactivate_5x()
final = _get_in_memory_version()
diff --git a/studio/backend/utils/update_status.py b/studio/backend/utils/update_status.py
index 9142203a69..9b71ff31a0 100644
--- a/studio/backend/utils/update_status.py
+++ b/studio/backend/utils/update_status.py
@@ -73,11 +73,7 @@ def detect_install_source() -> str:
try:
dist = distribution(PACKAGE_NAME)
except PackageNotFoundError:
- return (
- "local_repo"
- if _path_has_git_parent(_repo_root_from_this_file())
- else "unknown"
- )
+ return "local_repo" if _path_has_git_parent(_repo_root_from_this_file()) else "unknown"
try:
direct_url = dist.read_text("direct_url.json")
@@ -146,9 +142,7 @@ def get_studio_update_status(current_version: str) -> dict[str, Any]:
current_version = current_version,
latest_version = None,
install_source = install_source,
- reason = "invalid_current_version"
- if current_version != "dev"
- else "dev_build",
+ reason = "invalid_current_version" if current_version != "dev" else "dev_build",
)
latest_result = get_latest_pypi_version()
if latest_result.latest_version is None:
@@ -216,9 +210,7 @@ def get_latest_pypi_version() -> LatestVersionResult:
error = "Could not check PyPI update metadata.",
)
- ttl = (
- PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS
- )
+ ttl = PYPI_SUCCESS_TTL_SECONDS if result.latest_version else PYPI_FAILURE_TTL_SECONDS
with _cache_condition:
_latest_version_cache = _LatestVersionCacheEntry(
result = result,
@@ -262,9 +254,7 @@ def _fetch_latest_pypi_version() -> LatestVersionResult:
error = "Could not reach PyPI for update metadata.",
)
- latest = (
- payload.get("info", {}).get("version") if isinstance(payload, dict) else None
- )
+ latest = payload.get("info", {}).get("version") if isinstance(payload, dict) else None
if not isinstance(latest, str) or not latest.strip():
return LatestVersionResult(
latest_version = None,
@@ -366,9 +356,4 @@ def _parse_current_version(current_version: str) -> Version | None:
def _utc_now_iso() -> str:
- return (
- datetime.now(timezone.utc)
- .replace(microsecond = 0)
- .isoformat()
- .replace("+00:00", "Z")
- )
+ return datetime.now(timezone.utc).replace(microsecond = 0).isoformat().replace("+00:00", "Z")
diff --git a/studio/backend/utils/upload_limits.py b/studio/backend/utils/upload_limits.py
index b8a6a2474b..c21ea69af7 100644
--- a/studio/backend/utils/upload_limits.py
+++ b/studio/backend/utils/upload_limits.py
@@ -52,7 +52,6 @@ def validate_upload_limit_mb(value: Any) -> int:
def get_upload_limit_mb() -> int:
try:
from storage.studio_db import get_app_setting
-
stored = get_app_setting(UPLOAD_LIMIT_SETTING_KEY, None)
except Exception:
stored = None
diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py
index e95ef08a7d..7a7774be40 100644
--- a/studio/backend/utils/utils.py
+++ b/studio/backend/utils/utils.py
@@ -22,9 +22,7 @@ logger = get_logger(__name__)
# log the full exception server-side and return a generic message.
-def safe_error_detail(
- error: Exception, fallback: str = "An internal error occurred"
-) -> str:
+def safe_error_detail(error: Exception, fallback: str = "An internal error occurred") -> str:
"""Map a caught exception to a generic, client-safe message.
Never includes raw ``str(error)`` (which can leak internal paths or stack
@@ -44,9 +42,7 @@ def safe_error_detail(
return fallback
-def safe_curated_detail(
- error: Exception, fallback: str = "An internal error occurred"
-) -> str:
+def safe_curated_detail(error: Exception, fallback: str = "An internal error occurred") -> str:
"""Client-safe text for curated domain/validation exceptions meant for the user.
Keeps the message (paths stripped) instead of a generic fallback; use for known
diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py
index e0ce02261b..cca30bfd44 100644
--- a/studio/backend/utils/wheel_utils.py
+++ b/studio/backend/utils/wheel_utils.py
@@ -19,9 +19,7 @@ from utils.subprocess_compat import windows_hidden_subprocess_kwargs
_logger = logging.getLogger(__name__)
-FLASH_ATTN_RELEASE_BASE_URL = (
- "https://github.com/Dao-AILab/flash-attention/releases/download"
-)
+FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/releases/download"
@functools.lru_cache(maxsize = 1)
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index 62b7b2b290..09953843f0 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -70,7 +70,12 @@ def windows_hidden_subprocess_kwargs() -> dict[str, object]:
return kwargs
-def env_int(name: str, default: int, *, minimum: int | None = None) -> int:
+def env_int(
+ name: str,
+ default: int,
+ *,
+ minimum: int | None = None,
+) -> int:
raw = os.environ.get(name)
if raw is None:
value = default
@@ -104,9 +109,7 @@ UPSTREAM_REPO = "ggml-org/llama.cpp"
UPSTREAM_RELEASES_API = f"https://api.github.com/repos/{UPSTREAM_REPO}/releases/latest"
LEMONADE_ROCM_REPO = "lemonade-sdk/llamacpp-rocm"
-LEMONADE_ROCM_RELEASES_API = (
- f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/latest"
-)
+LEMONADE_ROCM_RELEASES_API = f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/latest"
def _lemonade_release_api_for(llama_tag: str) -> str:
@@ -135,9 +138,7 @@ def _lemonade_release_api_for(llama_tag: str) -> str:
)
-TEST_MODEL_URL = (
- "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
-)
+TEST_MODEL_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
TEST_MODEL_SHA256 = "270cba1bd5109f42d03350f60406024560464db173c0e387d91f0426d3bd256d"
VALIDATION_MODEL_CACHE_DIRNAME = ".cache"
VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf"
@@ -256,12 +257,8 @@ _BLACKWELL_MIN_SM = 120
# windows-cuda build at or above this already covers Blackwell and makes the
# older pinned 13.1 fallback unnecessary (cuda-12.4 is below it).
_BLACKWELL_MIN_TOOLKIT = (12, 8)
-_PINNED_BLACKWELL_LLAMA_SHA256 = (
- "31ddb8b42d7ab4a47cab8c48c397519f580ca502df7e73f3ab396eacc16c8e8d"
-)
-_PINNED_BLACKWELL_CUDART_SHA256 = (
- "f96935e7e385e3b2d0189239077c10fe8fd7e95690fea4afec455b1b6c7e3f18"
-)
+_PINNED_BLACKWELL_LLAMA_SHA256 = "31ddb8b42d7ab4a47cab8c48c397519f580ca502df7e73f3ab396eacc16c8e8d"
+_PINNED_BLACKWELL_CUDART_SHA256 = "f96935e7e385e3b2d0189239077c10fe8fd7e95690fea4afec455b1b6c7e3f18"
def _cuda_runtime_lines_for_major(major: int) -> list[str]:
@@ -279,9 +276,7 @@ def _resolve_linux_bundle_profile(bundle_profile: str) -> "dict[str, Any] | None
known = DIRECT_LINUX_BUNDLE_PROFILES.get(bundle_profile)
if known is not None:
return known
- m = re.fullmatch(
- r"cuda(?P\d+)-(?Polder|newer|portable)", bundle_profile
- )
+ m = re.fullmatch(r"cuda(?P\d+)-(?Polder|newer|portable)", bundle_profile)
if not m:
return None
base_key = max(
@@ -788,9 +783,9 @@ def refs_match(candidate_ref: str | None, requested_ref: str | None) -> bool:
candidate_commit = normalize_source_commit(candidate_ref)
requested_commit = normalize_source_commit(requested_ref)
if candidate_commit and requested_commit:
- return candidate_commit.startswith(
- requested_commit
- ) or requested_commit.startswith(candidate_commit)
+ return candidate_commit.startswith(requested_commit) or requested_commit.startswith(
+ candidate_commit
+ )
return False
@@ -820,9 +815,7 @@ def windows_cuda_upstream_asset_names(llama_tag: str, runtime: str) -> list[str]
def windows_cuda_asset_aliases(
- asset_name: str,
- *,
- compatibility_tag: str | None = None,
+ asset_name: str, *, compatibility_tag: str | None = None
) -> list[str]:
aliases: list[str] = []
legacy_match = re.fullmatch(
@@ -886,11 +879,7 @@ class DownloadProgress:
self.last_emit = 0.0
term_ok = os.environ.get("TERM", "").lower() != "dumb"
self.stream = (
- sys.stderr
- if sys.stderr.isatty()
- else sys.stdout
- if sys.stdout.isatty()
- else sys.stderr
+ sys.stderr if sys.stderr.isatty() else sys.stdout if sys.stdout.isatty() else sys.stderr
)
self.is_tty = term_ok and self.stream.isatty()
self.completed = False
@@ -898,7 +887,12 @@ class DownloadProgress:
self.last_milestone_bytes = 0
self.has_rendered_tty_progress = False
- def _render(self, downloaded_bytes: int, *, final: bool = False) -> str:
+ def _render(
+ self,
+ downloaded_bytes: int,
+ *,
+ final: bool = False,
+ ) -> str:
elapsed = max(time.monotonic() - self.start_time, 1e-6)
speed = downloaded_bytes / elapsed
speed_text = f"{format_byte_count(speed)}/s"
@@ -918,10 +912,7 @@ class DownloadProgress:
if self.is_tty:
elapsed = now - self.start_time
if not self.has_rendered_tty_progress:
- if (
- self.total_bytes is not None
- and downloaded_bytes >= self.total_bytes
- ):
+ if self.total_bytes is not None and downloaded_bytes >= self.total_bytes:
return
if elapsed < TTY_PROGRESS_START_DELAY_SECONDS:
return
@@ -943,10 +934,7 @@ class DownloadProgress:
if self.total_bytes is not None:
percent = int((downloaded_bytes * 100) / max(self.total_bytes, 1))
milestone_percent = min((percent // 25) * 25, 100)
- if (
- milestone_percent > self.last_milestone_percent
- and milestone_percent < 100
- ):
+ if milestone_percent > self.last_milestone_percent and milestone_percent < 100:
self.last_milestone_percent = milestone_percent
should_emit = True
else:
@@ -999,11 +987,7 @@ def download_bytes(
content_length = response.headers.get("Content-Length")
if content_length and content_length.isdigit():
total_bytes = int(content_length)
- progress = (
- DownloadProgress(progress_label, total_bytes)
- if progress_label
- else None
- )
+ progress = DownloadProgress(progress_label, total_bytes) if progress_label else None
data = bytearray()
while True:
chunk = response.read(1024 * 1024)
@@ -1033,17 +1017,13 @@ def fetch_json(url: str) -> Any:
data = download_bytes(
url,
timeout = 30,
- headers = github_api_headers(url)
- if is_github_api_url(url)
- else auth_headers(url),
+ headers = github_api_headers(url) if is_github_api_url(url) else auth_headers(url),
)
except urllib.error.HTTPError as exc:
if exc.code == 403 and is_github_api_url(url):
hint = ""
if not (os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")):
- hint = (
- "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits"
- )
+ hint = "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits"
raise RuntimeError(f"GitHub API returned 403 for {url}{hint}") from exc
raise
if not data:
@@ -1052,9 +1032,7 @@ def fetch_json(url: str) -> Any:
try:
payload = json.loads(data.decode("utf-8"))
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
- last_decode_exc = RuntimeError(
- f"downloaded invalid JSON from {url}: {exc}"
- )
+ last_decode_exc = RuntimeError(f"downloaded invalid JSON from {url}: {exc}")
else:
if not isinstance(payload, dict) and not isinstance(payload, list):
raise RuntimeError(
@@ -1088,9 +1066,7 @@ def download_file(url: str, destination: Path) -> None:
content_length = response.headers.get("Content-Length")
if content_length and content_length.isdigit():
total_bytes = int(content_length)
- progress = DownloadProgress(
- f"Downloading {destination.name}", total_bytes
- )
+ progress = DownloadProgress(f"Downloading {destination.name}", total_bytes)
downloaded_bytes = 0
while True:
chunk = response.read(1024 * 1024)
@@ -1115,27 +1091,19 @@ def download_file(url: str, destination: Path) -> None:
pass
if attempt >= HTTP_FETCH_ATTEMPTS or not is_retryable_url_error(exc):
raise
- log(
- f"download failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying"
- )
+ log(f"download failed ({attempt}/{HTTP_FETCH_ATTEMPTS}) for {url}: {exc}; retrying")
sleep_backoff(attempt, exc = exc)
assert last_exc is not None
raise last_exc
def download_file_verified(
- url: str,
- destination: Path,
- *,
- expected_sha256: str | None,
- label: str,
+ url: str, destination: Path, *, expected_sha256: str | None, label: str
) -> None:
normalized_expected = normalize_sha256_digest(expected_sha256)
if not normalized_expected:
download_file(url, destination)
- log(
- f"downloaded {label} without a published sha256; relying on install validation"
- )
+ log(f"downloaded {label} without a published sha256; relying on install validation")
return
for attempt in range(1, 3):
@@ -1219,9 +1187,7 @@ def latest_upstream_release_tag() -> str:
payload = fetch_json(UPSTREAM_RELEASES_API)
tag = payload.get("tag_name")
if not isinstance(tag, str) or not tag:
- raise RuntimeError(
- f"latest release tag was missing from {UPSTREAM_RELEASES_API}"
- )
+ raise RuntimeError(f"latest release tag was missing from {UPSTREAM_RELEASES_API}")
return tag
@@ -1256,19 +1222,13 @@ def iter_release_payloads_by_time(
yield github_release(repo, published_release_tag)
return
- if (
- requested_tag
- and requested_tag != "latest"
- and is_release_tag_like(requested_tag)
- ):
+ if requested_tag and requested_tag != "latest" and is_release_tag_like(requested_tag):
try:
yield github_release(repo, requested_tag)
return
except urllib.error.HTTPError as exc:
if exc.code == 404:
- log(
- f"release tag {requested_tag} not found in {repo}; scanning recent releases"
- )
+ log(f"release tag {requested_tag} not found in {repo}; scanning recent releases")
else:
raise
except Exception:
@@ -1276,21 +1236,15 @@ def iter_release_payloads_by_time(
releases = [
release
- for release in github_releases(
- repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES
- )
- if isinstance(release, dict)
- and not release.get("draft")
- and not release.get("prerelease")
+ for release in github_releases(repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES)
+ if isinstance(release, dict) and not release.get("draft") and not release.get("prerelease")
]
releases.sort(key = release_time_sort_key, reverse = True)
for release in releases:
yield release
-def direct_release_matches_request(
- *, release_tag: str, llama_tag: str, requested_tag: str
-) -> bool:
+def direct_release_matches_request(*, release_tag: str, llama_tag: str, requested_tag: str) -> bool:
if requested_tag == "latest":
return True
for candidate in (release_tag, llama_tag):
@@ -1392,10 +1346,7 @@ def parse_direct_linux_release_bundle(
def direct_linux_release_plan(
- release: dict[str, Any],
- host: HostInfo,
- repo: str,
- requested_tag: str,
+ release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
) -> InstallReleasePlan | None:
bundle = parse_direct_linux_release_bundle(repo, release)
if bundle is None:
@@ -1477,10 +1428,7 @@ def direct_linux_release_plan(
def direct_upstream_release_plan(
- release: dict[str, Any],
- host: HostInfo,
- repo: str,
- requested_tag: str,
+ release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
) -> InstallReleasePlan | None:
release_tag = release.get("tag_name")
if not isinstance(release_tag, str) or not release_tag:
@@ -1674,9 +1622,7 @@ def resolve_simple_install_release_plans(
f"{repo} ships only linux-x64 prebuilts; "
f"{host.machine or 'non-x64'} Linux falls back to source build"
)
- allow_older_release_fallback = (
- requested_tag == "latest" and not published_release_tag
- )
+ allow_older_release_fallback = requested_tag == "latest" and not published_release_tag
# macOS: pin the last upstream build that loads on a pre-26 host instead of
# fetching the latest (macOS 26 only) build and walking back release by
# release. No-op on macOS 26+, unknown version, non-macOS, and the fork.
@@ -1690,17 +1636,13 @@ def resolve_simple_install_release_plans(
last_error: PrebuiltFallback | None = None
try:
- releases = iter_release_payloads_by_time(
- repo, published_release_tag, requested_tag
- )
+ releases = iter_release_payloads_by_time(repo, published_release_tag, requested_tag)
for release in releases:
try:
if host.is_linux and repo == "unslothai/llama.cpp":
plan = direct_linux_release_plan(release, host, repo, requested_tag)
else:
- plan = direct_upstream_release_plan(
- release, host, repo, requested_tag
- )
+ plan = direct_upstream_release_plan(release, host, repo, requested_tag)
if plan is None:
continue
except PrebuiltFallback as exc:
@@ -1720,17 +1662,13 @@ def resolve_simple_install_release_plans(
except PrebuiltFallback:
raise
except Exception as exc:
- raise PrebuiltFallback(
- f"failed to inspect published releases in {repo}: {exc}"
- ) from exc
+ raise PrebuiltFallback(f"failed to inspect published releases in {repo}: {exc}") from exc
if plans:
return requested_tag, plans
if last_error is not None:
raise last_error
- raise PrebuiltFallback(
- f"no installable published llama.cpp releases were found in {repo}"
- )
+ raise PrebuiltFallback(f"no installable published llama.cpp releases were found in {repo}")
def normalized_requested_llama_tag(requested_tag: str | None) -> str:
@@ -1782,9 +1720,7 @@ def parse_cuda_visible_devices(value: str | None) -> list[str] | None:
return [token.strip() for token in raw.split(",") if token.strip()]
-def supports_explicit_visible_device_matching(
- visible_devices: list[str] | None,
-) -> bool:
+def supports_explicit_visible_device_matching(visible_devices: list[str] | None) -> bool:
if not visible_devices:
return False
for token in visible_devices:
@@ -1796,8 +1732,7 @@ def supports_explicit_visible_device_matching(
def select_visible_gpu_rows(
- gpu_rows: Iterable[tuple[str, str, str]],
- visible_devices: list[str] | None,
+ gpu_rows: Iterable[tuple[str, str, str]], visible_devices: list[str] | None
) -> list[tuple[str, str, str]]:
rows = list(gpu_rows)
if visible_devices is None:
@@ -1835,9 +1770,7 @@ def dir_provides_exact_library(directory: str | Path, library: str) -> bool:
return candidate.exists() and (candidate.is_file() or candidate.is_symlink())
-def linux_runtime_dirs_for_required_libraries(
- required_libraries: Iterable[str],
-) -> list[str]:
+def linux_runtime_dirs_for_required_libraries(required_libraries: Iterable[str]) -> list[str]:
required = [library for library in required_libraries if library]
candidates: list[str | Path] = []
@@ -1853,9 +1786,7 @@ def linux_runtime_dirs_for_required_libraries(
value = os.environ.get(name)
if value:
cuda_roots.append(Path(value))
- cuda_roots.extend(
- Path(path) for path in glob_paths("/usr/local/cuda", "/usr/local/cuda-*")
- )
+ cuda_roots.extend(Path(path) for path in glob_paths("/usr/local/cuda", "/usr/local/cuda-*"))
for root in cuda_roots:
candidates.extend(
@@ -1880,8 +1811,7 @@ def linux_runtime_dirs_for_required_libraries(
)
)
candidates.extend(
- Path(path)
- for path in glob_paths("/usr/local/lib/ollama/cuda_v*", "/usr/lib/wsl/lib")
+ Path(path) for path in glob_paths("/usr/local/lib/ollama/cuda_v*", "/usr/lib/wsl/lib")
)
candidates.extend(Path(path) for path in python_runtime_dirs())
candidates.extend(Path(path) for path in ldconfig_runtime_dirs(required))
@@ -1893,9 +1823,7 @@ def linux_runtime_dirs_for_required_libraries(
matched: list[tuple[int, str]] = []
for directory in resolved:
base = Path(directory)
- provided = sum(
- 1 for library in required if dir_provides_exact_library(directory, library)
- )
+ provided = sum(1 for library in required if dir_provides_exact_library(directory, library))
if provided:
matched.append((provided, directory))
@@ -1916,9 +1844,7 @@ def detected_linux_runtime_lines() -> tuple[list[str], dict[str, list[str]]]:
matching_dirs: list[str] = []
for library in required:
matched_dirs = [
- directory
- for directory in dirs
- if any(Path(directory).glob(f"{library}*"))
+ directory for directory in dirs if any(Path(directory).glob(f"{library}*"))
]
if not matched_dirs:
library_matches = {}
@@ -1955,17 +1881,13 @@ def parse_published_artifact(raw: Any) -> PublishedLlamaArtifact | None:
if not isinstance(asset_name, str) or not asset_name:
raise ValueError("artifact.asset_name was missing or not a string")
if not isinstance(install_kind, str) or not install_kind:
- raise ValueError(
- f"artifact {asset_name} install_kind was missing or not a string"
- )
+ raise ValueError(f"artifact {asset_name} install_kind was missing or not a string")
supported_sms_raw = raw.get("supported_sms", [])
if not isinstance(supported_sms_raw, (list, tuple)):
raise ValueError(f"artifact {asset_name} supported_sms must be a list or tuple")
if any(not isinstance(value, (int, str)) for value in supported_sms_raw):
- raise ValueError(
- f"artifact {asset_name} supported_sms entries must be ints or strings"
- )
+ raise ValueError(f"artifact {asset_name} supported_sms entries must be ints or strings")
supported_sms = normalize_compute_caps(supported_sms_raw)
min_sm_raw = raw.get("min_sm")
@@ -1974,9 +1896,7 @@ def parse_published_artifact(raw: Any) -> PublishedLlamaArtifact | None:
min_sm = int(min_sm_raw) if min_sm_raw is not None else None
max_sm = int(max_sm_raw) if max_sm_raw is not None else None
except (TypeError, ValueError) as exc:
- raise ValueError(
- f"artifact {asset_name} min_sm/max_sm were not integers"
- ) from exc
+ raise ValueError(f"artifact {asset_name} min_sm/max_sm were not integers") from exc
runtime_line = raw.get("runtime_line")
coverage_class = raw.get("coverage_class")
bundle_profile = raw.get("bundle_profile")
@@ -1994,9 +1914,7 @@ def parse_published_artifact(raw: Any) -> PublishedLlamaArtifact | None:
return PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = install_kind,
- runtime_line = runtime_line
- if isinstance(runtime_line, str) and runtime_line
- else None,
+ runtime_line = runtime_line if isinstance(runtime_line, str) and runtime_line else None,
coverage_class = coverage_class
if isinstance(coverage_class, str) and coverage_class
else None,
@@ -2071,9 +1989,7 @@ def parse_published_release_bundle(
try:
artifact = parse_published_artifact(raw_artifact)
except ValueError as exc:
- log(
- f"published artifact ignored for {repo}@{release_tag} artifact[{index}]: {exc}"
- )
+ log(f"published artifact ignored for {repo}@{release_tag} artifact[{index}]: {exc}")
continue
if artifact is not None:
artifacts.append(artifact)
@@ -2092,9 +2008,7 @@ def parse_published_release_bundle(
release_tag = release_tag,
upstream_tag = upstream_tag,
manifest_sha256 = manifest_sha256,
- source_repo = source_repo
- if isinstance(source_repo, str) and source_repo
- else None,
+ source_repo = source_repo if isinstance(source_repo, str) and source_repo else None,
source_repo_url = source_repo_url
if isinstance(source_repo_url, str) and source_repo_url
else None,
@@ -2117,9 +2031,7 @@ def parse_published_release_bundle(
def parse_approved_release_checksums(
- repo: str,
- release_tag: str,
- payload: Any,
+ repo: str, release_tag: str, payload: Any
) -> ApprovedReleaseChecksums:
if not isinstance(payload, dict):
raise RuntimeError(
@@ -2157,18 +2069,12 @@ def parse_approved_release_checksums(
artifacts: dict[str, ApprovedArtifactHash] = {}
for asset_name, raw_entry in artifacts_payload.items():
if not isinstance(asset_name, str) or not asset_name:
- raise RuntimeError(
- "published checksum asset used a non-string artifact key"
- )
+ raise RuntimeError("published checksum asset used a non-string artifact key")
if not isinstance(raw_entry, dict):
- raise RuntimeError(
- f"published checksum entry for {asset_name} was not an object"
- )
+ raise RuntimeError(f"published checksum entry for {asset_name} was not an object")
digest = normalize_sha256_digest(raw_entry.get("sha256"))
if not digest:
- raise RuntimeError(
- f"published checksum entry for {asset_name} omitted a valid sha256"
- )
+ raise RuntimeError(f"published checksum entry for {asset_name} omitted a valid sha256")
repo_value = raw_entry.get("repo")
kind_value = raw_entry.get("kind")
artifacts[asset_name] = ApprovedArtifactHash(
@@ -2189,9 +2095,7 @@ def parse_approved_release_checksums(
repo = repo,
release_tag = release_tag,
upstream_tag = upstream_tag,
- source_repo = source_repo
- if isinstance(source_repo, str) and source_repo
- else None,
+ source_repo = source_repo if isinstance(source_repo, str) and source_repo else None,
source_repo_url = source_repo_url
if isinstance(source_repo_url, str) and source_repo_url
else None,
@@ -2210,9 +2114,7 @@ def parse_approved_release_checksums(
)
-def load_approved_release_checksums(
- repo: str, release_tag: str
-) -> ApprovedReleaseChecksums:
+def load_approved_release_checksums(repo: str, release_tag: str) -> ApprovedReleaseChecksums:
try:
release = github_release(repo, release_tag)
except Exception as exc:
@@ -2246,9 +2148,7 @@ def iter_published_release_bundles(
else github_releases(repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES)
)
for release in releases:
- if not published_release_tag and (
- release.get("draft") or release.get("prerelease")
- ):
+ if not published_release_tag and (release.get("draft") or release.get("prerelease")):
continue
try:
bundle = parse_published_release_bundle(repo, release)
@@ -2301,13 +2201,9 @@ def linux_cuda_choice_from_release(
)
)
published_artifacts = [
- artifact
- for artifact in release.artifacts
- if artifact.install_kind == "linux-cuda"
+ artifact for artifact in release.artifacts if artifact.install_kind == "linux-cuda"
]
- published_asset_names = sorted(
- artifact.asset_name for artifact in published_artifacts
- )
+ published_asset_names = sorted(artifact.asset_name for artifact in published_artifacts)
selection_log.append(
"linux_cuda_selection: published_assets="
+ (",".join(published_asset_names) if published_asset_names else "none")
@@ -2343,9 +2239,7 @@ def linux_cuda_choice_from_release(
attempts: list[AssetChoice] = []
seen_attempts: set[str] = set()
- def add_attempt(
- artifact: PublishedLlamaArtifact, asset_url: str, reason: str
- ) -> None:
+ def add_attempt(artifact: PublishedLlamaArtifact, asset_url: str, reason: str) -> None:
asset_name = artifact.asset_name
if asset_name in seen_attempts:
return
@@ -2382,9 +2276,7 @@ def linux_cuda_choice_from_release(
asset_name = artifact.asset_name
asset_url = release.assets.get(asset_name)
if not asset_url:
- selection_log.append(
- f"linux_cuda_selection: reject {asset_name} missing asset"
- )
+ selection_log.append(f"linux_cuda_selection: reject {asset_name} missing asset")
continue
if not host_sms and artifact.coverage_class != "portable":
selection_log.append(
@@ -2412,9 +2304,7 @@ def linux_cuda_choice_from_release(
supported_sms = {str(value) for value in artifact.supported_sms}
missing_sms = [sm for sm in host_sms if sm not in supported_sms]
out_of_range_sms = [
- sm
- for sm in host_sms
- if not (artifact.min_sm <= int(sm) <= artifact.max_sm)
+ sm for sm in host_sms if not (artifact.min_sm <= int(sm) <= artifact.max_sm)
]
reasons: list[str] = []
if missing_sms:
@@ -2458,8 +2348,7 @@ def linux_cuda_choice_from_release(
return None
selection_log.append(
- "linux_cuda_selection: attempt_order="
- + ",".join(choice.name for choice in attempts)
+ "linux_cuda_selection: attempt_order=" + ",".join(choice.name for choice in attempts)
)
for attempt in attempts:
attempt.selection_log = list(selection_log) + [
@@ -2477,9 +2366,7 @@ def latest_published_linux_cuda_tag(host: HostInfo, published_repo: str) -> str
def iter_upstream_releases() -> Iterable[dict[str, Any]]:
- for release in github_releases(
- UPSTREAM_REPO, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES
- ):
+ for release in github_releases(UPSTREAM_REPO, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES):
if release.get("draft") or release.get("prerelease"):
continue
yield release
@@ -2514,9 +2401,7 @@ def validated_checksums_for_bundle(
return checksums
-def published_release_matches_request(
- bundle: PublishedReleaseBundle, requested_ref: str
-) -> bool:
+def published_release_matches_request(bundle: PublishedReleaseBundle, requested_ref: str) -> bool:
if requested_ref == "latest":
return True
for candidate in (
@@ -2571,9 +2456,7 @@ def resolve_published_release(
raise PrebuiltFallback(
f"no usable published llama.cpp releases were available in {repo}"
)
- raise PrebuiltFallback(
- f"no published llama.cpp releases were available in {repo}"
- )
+ raise PrebuiltFallback(f"no published llama.cpp releases were available in {repo}")
raise PrebuiltFallback(
f"no published prebuilt release in {repo} matched upstream tag {normalized_requested}"
@@ -2632,9 +2515,7 @@ def iter_resolved_published_releases(
return
if normalized_requested == "latest":
- raise PrebuiltFallback(
- f"no published llama.cpp releases were available in {repo}"
- )
+ raise PrebuiltFallback(f"no published llama.cpp releases were available in {repo}")
raise PrebuiltFallback(
f"no published prebuilt release in {repo} matched upstream tag {normalized_requested}"
@@ -2692,33 +2573,23 @@ def resolve_requested_install_tag(
).bundle.upstream_tag
-def exact_source_archive_hash(
- checksums: ApprovedReleaseChecksums,
-) -> ApprovedArtifactHash | None:
+def exact_source_archive_hash(checksums: ApprovedReleaseChecksums) -> ApprovedArtifactHash | None:
if not checksums.source_commit:
return None
- return checksums.artifacts.get(
- exact_source_archive_logical_name(checksums.source_commit)
- )
+ return checksums.artifacts.get(exact_source_archive_logical_name(checksums.source_commit))
def source_clone_url_from_checksums(checksums: ApprovedReleaseChecksums) -> str | None:
return source_repo_clone_url(checksums.source_repo, checksums.source_repo_url)
-def source_build_plan_for_release(
- release: ResolvedPublishedRelease,
-) -> SourceBuildPlan:
+def source_build_plan_for_release(release: ResolvedPublishedRelease) -> SourceBuildPlan:
checksums = release.checksums
exact_source = exact_source_archive_hash(checksums)
source_repo = checksums.source_repo or release.bundle.source_repo
source_repo_url = checksums.source_repo_url or release.bundle.source_repo_url
- requested_source_ref = (
- checksums.requested_source_ref or release.bundle.requested_source_ref
- )
- resolved_source_ref = (
- checksums.resolved_source_ref or release.bundle.resolved_source_ref
- )
+ requested_source_ref = checksums.requested_source_ref or release.bundle.requested_source_ref
+ resolved_source_ref = checksums.resolved_source_ref or release.bundle.resolved_source_ref
source_commit = checksums.source_commit or release.bundle.source_commit
source_ref_kind = checksums.source_ref_kind or release.bundle.source_ref_kind
source_url = source_repo_clone_url(source_repo, source_repo_url)
@@ -2734,14 +2605,8 @@ def source_build_plan_for_release(
resolved_source_ref = resolved_source_ref,
source_commit = source_commit,
)
- source_ref = checkout_friendly_ref(
- source_ref_kind, resolved_source_ref or requested_source_ref
- )
- if (
- source_url
- and source_ref
- and source_ref_kind in {"tag", "branch", "pull", "commit"}
- ):
+ source_ref = checkout_friendly_ref(source_ref_kind, resolved_source_ref or requested_source_ref)
+ if source_url and source_ref and source_ref_kind in {"tag", "branch", "pull", "commit"}:
return SourceBuildPlan(
source_url = source_url,
source_ref = source_ref,
@@ -2925,9 +2790,7 @@ def detect_host() -> HostInfo:
# ROCm host as NVIDIA and short-circuit the ROCm path.
try:
listing = run_capture([nvidia_smi, "-L"], timeout = 20)
- gpu_lines = [
- line for line in listing.stdout.splitlines() if line.startswith("GPU ")
- ]
+ gpu_lines = [line for line in listing.stdout.splitlines() if line.startswith("GPU ")]
if gpu_lines:
has_physical_nvidia = True
has_usable_nvidia = visible_device_tokens != []
@@ -3225,9 +3088,7 @@ def detect_torch_cuda_runtime_preference(host: HostInfo) -> CudaRuntimePreferenc
try:
cuda_available = bool(torch.cuda.is_available())
except Exception as exc:
- selection_log.append(
- f"torch_cuda_preference: torch.cuda.is_available() failed: {exc}"
- )
+ selection_log.append(f"torch_cuda_preference: torch.cuda.is_available() failed: {exc}")
return CudaRuntimePreference(runtime_line = None, selection_log = selection_log)
if not cuda_available:
@@ -3315,14 +3176,10 @@ def windows_cuda_attempts(
f"{preferred_runtime_line} unavailable_or_incompatible"
)
else:
- selection_log.append(
- "windows_cuda_selection: no Torch runtime preference available"
- )
+ selection_log.append("windows_cuda_selection: no Torch runtime preference available")
runtime_order.extend(
- runtime_line
- for runtime_line in normal_runtime_lines
- if runtime_line not in runtime_order
+ runtime_line for runtime_line in normal_runtime_lines if runtime_line not in runtime_order
)
# Keep every driver-compatible line reachable as a fallback, so a line gated
# out by the driver version still drops to an older major (cuda13 -> cuda12).
@@ -3346,9 +3203,7 @@ def windows_cuda_attempts(
# Track whatever minor llama.cpp actually ships for this major
# (cuda13 -> 13.1, 13.3, ...). Skip the line when the release has no
# matching asset instead of guessing a now-missing name.
- runtime = _published_windows_cuda_runtime(
- upstream_assets, major, host.driver_cuda_version
- )
+ runtime = _published_windows_cuda_runtime(upstream_assets, major, host.driver_cuda_version)
if runtime is None:
selection_log.append(
f"windows_cuda_selection: no driver-supported asset for {runtime_line}"
@@ -3414,9 +3269,7 @@ def _windows_cuda_attempt_covers_blackwell(attempt: AssetChoice) -> bool:
if attempt.install_kind != "windows-cuda":
return False
m = re.search(r"-bin-win-cuda-(\d+)\.(\d+)-x64\.zip$", attempt.name)
- return (
- m is not None and (int(m.group(1)), int(m.group(2))) >= _BLACKWELL_MIN_TOOLKIT
- )
+ return m is not None and (int(m.group(1)), int(m.group(2))) >= _BLACKWELL_MIN_TOOLKIT
def _pinned_windows_cuda_fallback(
@@ -3441,10 +3294,7 @@ def _pinned_windows_cuda_fallback(
caps = normalize_compute_caps(host.compute_caps)
if not caps or int(caps[-1]) < _BLACKWELL_MIN_SM:
return None
- if any(
- _windows_cuda_attempt_covers_blackwell(attempt)
- for attempt in existing_cuda_attempts
- ):
+ if any(_windows_cuda_attempt_covers_blackwell(attempt) for attempt in existing_cuda_attempts):
return None
tag = _PINNED_BLACKWELL_FALLBACK_TAG
runtime = _PINNED_BLACKWELL_FALLBACK_RUNTIME
@@ -3499,9 +3349,7 @@ def _augment_checksums_with_pin(
def _with_pinned_windows_cuda_fallback(
- host: HostInfo,
- attempts: list[AssetChoice],
- checksums: ApprovedReleaseChecksums,
+ host: HostInfo, attempts: list[AssetChoice], checksums: ApprovedReleaseChecksums
) -> tuple[list[AssetChoice], ApprovedReleaseChecksums]:
"""Insert the Blackwell pin ahead of the Windows CUDA attempts and keep it
through apply_approved_hashes, or return the inputs unchanged when dormant.
@@ -3544,9 +3392,7 @@ def published_windows_cuda_attempts(
selection_log,
)
published_artifacts = [
- artifact
- for artifact in release.artifacts
- if artifact.install_kind == "windows-cuda"
+ artifact for artifact in release.artifacts if artifact.install_kind == "windows-cuda"
]
artifacts_by_runtime: dict[str, list[PublishedLlamaArtifact]] = {}
for artifact in published_artifacts:
@@ -3642,15 +3488,10 @@ def resolve_linux_cuda_choice(
def published_asset_choice_for_kind(
- release: PublishedReleaseBundle,
- install_kind: str,
+ release: PublishedReleaseBundle, install_kind: str
) -> AssetChoice | None:
candidates = sorted(
- (
- artifact
- for artifact in release.artifacts
- if artifact.install_kind == install_kind
- ),
+ (artifact for artifact in release.artifacts if artifact.install_kind == install_kind),
key = lambda artifact: (artifact.rank, artifact.asset_name),
)
for artifact in candidates:
@@ -3666,9 +3507,7 @@ def published_asset_choice_for_kind(
install_kind = install_kind,
runtime_line = artifact.runtime_line,
selection_log = list(release.selection_log)
- + [
- f"published_selection: selected {artifact.asset_name} install_kind={install_kind}"
- ],
+ + [f"published_selection: selected {artifact.asset_name} install_kind={install_kind}"],
)
return None
@@ -3725,11 +3564,7 @@ def _detect_host_rocm_version() -> tuple[int, int] | None:
if result.returncode == 0:
raw = (result.stdout or "").strip().split("\n")[0]
parts = raw.split(".")
- if (
- len(parts) >= 2
- and parts[0].isdigit()
- and parts[1].split("-")[0].isdigit()
- ):
+ if len(parts) >= 2 and parts[0].isdigit() and parts[1].split("-")[0].isdigit():
return int(parts[0]), int(parts[1].split("-")[0])
except Exception:
pass
@@ -3888,9 +3723,7 @@ def resolve_lemonade_rocm_choice(
return None
release_tag = release.get("tag_name") if isinstance(release, dict) else None
if not isinstance(release_tag, str) or not release_tag:
- log(
- f"Unexpected {LEMONADE_ROCM_REPO} release payload; skipping lemonade prebuilt"
- )
+ log(f"Unexpected {LEMONADE_ROCM_REPO} release payload; skipping lemonade prebuilt")
return None
assets = release_asset_map(release)
asset_name = f"llama-{release_tag}-{os_prefix}-rocm-{gfx_family}-x64.zip"
@@ -3986,9 +3819,7 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
_compatible: list[tuple[tuple[int, ...], str]] = rocm_candidates
if _host_rocm_version is not None:
_compatible = [
- item
- for item in rocm_candidates
- if item[0][:2] <= _host_rocm_version
+ item for item in rocm_candidates if item[0][:2] <= _host_rocm_version
]
if rocm_candidates and not _compatible:
# Fall back to the newest candidate so a source build is
@@ -4052,9 +3883,7 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
hip_name = f"llama-{llama_tag}-bin-win-hip-radeon-x64.zip"
if hip_name in upstream_assets:
- log(
- f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}"
- )
+ log(f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}")
return AssetChoice(
repo = UPSTREAM_REPO,
tag = llama_tag,
@@ -4063,9 +3892,7 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
source_label = "upstream",
install_kind = "windows-hip",
)
- log(
- "AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU"
- )
+ log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU")
upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip"
if upstream_name not in upstream_assets:
@@ -4105,9 +3932,7 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
install_kind = "macos-x64",
)
- raise PrebuiltFallback(
- f"no prebuilt policy exists for {host.system} {host.machine}"
- )
+ raise PrebuiltFallback(f"no prebuilt policy exists for {host.system} {host.machine}")
def resolve_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice:
@@ -4185,18 +4010,14 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
normalized = member_name.replace("\\", "/")
member_path = Path(normalized)
if member_path.is_absolute():
- raise PrebuiltFallback(
- f"archive member used an absolute path: {member_name}"
- )
+ raise PrebuiltFallback(f"archive member used an absolute path: {member_name}")
target = (base / member_path).resolve()
base_resolved = base.resolve()
try:
target.relative_to(base_resolved)
except ValueError as exc:
- raise PrebuiltFallback(
- f"archive member escaped destination: {member_name}"
- ) from exc
+ raise PrebuiltFallback(f"archive member escaped destination: {member_name}") from exc
return target
def _try_repair_missing_slash(
@@ -4242,11 +4063,7 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
return candidates[0][len(prefix) :]
def safe_link_target(
- base: Path,
- member_name: str,
- link_name: str,
- target: Path,
- archive_names: set[str],
+ base: Path, member_name: str, link_name: str, target: Path, archive_names: set[str]
) -> tuple[str, Path]:
normalized = link_name.replace("\\", "/")
repaired = _try_repair_missing_slash(member_name, normalized, archive_names)
@@ -4306,9 +4123,7 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
target.parent.mkdir(parents = True, exist_ok = True)
extracted = archive.extractfile(member)
if extracted is None:
- raise PrebuiltFallback(
- f"tar archive entry could not be read: {member.name}"
- )
+ raise PrebuiltFallback(f"tar archive entry could not be read: {member.name}")
with extracted, target.open("wb") as dst:
shutil.copyfileobj(extracted, dst)
@@ -4342,9 +4157,7 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
details = ", ".join(
f"{member.name} -> {member.linkname}" for member, _ in next_round
)
- raise PrebuiltFallback(
- f"tar archive contained unresolved link entries: {details}"
- )
+ raise PrebuiltFallback(f"tar archive contained unresolved link entries: {details}")
unresolved = next_round
destination.mkdir(parents = True, exist_ok = True)
@@ -4358,7 +4171,11 @@ def extract_archive(archive_path: Path, destination: Path) -> None:
def copy_globs(
- source_dir: Path, destination: Path, patterns: list[str], *, required: bool = True
+ source_dir: Path,
+ destination: Path,
+ patterns: list[str],
+ *,
+ required: bool = True,
) -> None:
destination.mkdir(parents = True, exist_ok = True)
matched_sources: dict[str, Path] = {}
@@ -4457,9 +4274,7 @@ def hydrate_source_tree(
for index, source_url in enumerate(source_urls):
try:
if index > 0:
- log(
- f"retrying source tree download from fallback URL: {source_url}"
- )
+ log(f"retrying source tree download from fallback URL: {source_url}")
download_file_verified(
source_url,
archive_path,
@@ -4484,14 +4299,11 @@ def hydrate_source_tree(
source_root / "gguf-py",
]
missing = [
- str(path.relative_to(source_root))
- for path in required_paths
- if not path.exists()
+ str(path.relative_to(source_root)) for path in required_paths if not path.exists()
]
if missing:
raise PrebuiltFallback(
- "upstream source archive was missing required repo files: "
- + ", ".join(missing)
+ "upstream source archive was missing required repo files: " + ", ".join(missing)
)
copy_directory_contents(source_root, install_dir)
except PrebuiltFallback:
@@ -4518,9 +4330,7 @@ def discover_installed_executable(install_dir: Path, executable_name: str) -> Pa
direct = install_dir / executable_name
if direct.exists() and direct.is_file():
return direct
- candidate = next(
- (path for path in install_dir.rglob(executable_name) if path.is_file()), None
- )
+ candidate = next((path for path in install_dir.rglob(executable_name) if path.is_file()), None)
if candidate is None:
raise PrebuiltFallback(f"{executable_name} was not installed")
return candidate
@@ -4550,9 +4360,7 @@ def create_exec_entrypoint(entrypoint: Path, target: Path) -> None:
write_exec_wrapper(entrypoint, target)
-def overlay_directory_for_choice(
- install_dir: Path, choice: AssetChoice, host: HostInfo
-) -> Path:
+def overlay_directory_for_choice(install_dir: Path, choice: AssetChoice, host: HostInfo) -> Path:
if host.is_windows or choice.install_kind.startswith("windows"):
path = install_dir / "build" / "bin" / "Release"
else:
@@ -4590,9 +4398,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
"windows-arm64",
}:
return ["llama-server.exe", "llama-quantize.exe", "*.dll"]
- raise PrebuiltFallback(
- f"unsupported install kind for runtime overlay: {choice.install_kind}"
- )
+ raise PrebuiltFallback(f"unsupported install kind for runtime overlay: {choice.install_kind}")
def runtime_subdirs_for_choice(choice: AssetChoice) -> list[str]:
@@ -4801,9 +4607,7 @@ def confirm_install_tree(install_dir: Path, host: HostInfo) -> None:
expected.append(install_dir / "UNSLOTH_PREBUILT_INFO.json")
missing = [str(path) for path in expected if not path.exists()]
if missing:
- raise RuntimeError(
- "activated install was missing expected files: " + ", ".join(missing)
- )
+ raise RuntimeError("activated install was missing expected files: " + ", ".join(missing))
def activate_install_tree(staging_dir: Path, install_dir: Path, host: HostInfo) -> None:
@@ -4926,15 +4730,11 @@ def install_from_archives(
expected_sha256 = choice.runtime_sha256,
label = f"prebuilt runtime archive {choice.runtime_name}",
)
- runtime_extract_dir = Path(
- tempfile.mkdtemp(prefix = "extract-runtime-", dir = work_dir)
- )
+ runtime_extract_dir = Path(tempfile.mkdtemp(prefix = "extract-runtime-", dir = work_dir))
extract_archive(runtime_archive, runtime_extract_dir)
source_dir = extract_dir
overlay_dir = overlay_directory_for_choice(install_dir, choice, host)
- copy_globs(
- source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True
- )
+ copy_globs(source_dir, overlay_dir, runtime_patterns_for_choice(choice), required = True)
for _subdir in runtime_subdirs_for_choice(choice):
_src_subdir = source_dir / _subdir
if _src_subdir.is_dir():
@@ -4979,9 +4779,7 @@ def install_from_archives(
source_server = build_bin / "llama-server"
source_quantize = build_bin / "llama-quantize"
if not source_server.exists() or not source_quantize.exists():
- raise PrebuiltFallback(
- "unix executables were not installed correctly into build/bin"
- )
+ raise PrebuiltFallback("unix executables were not installed correctly into build/bin")
os.chmod(source_server, 0o755)
os.chmod(source_quantize, 0o755)
@@ -5007,13 +4805,9 @@ def ensure_repo_shape(install_dir: Path) -> None:
install_dir / "convert_hf_to_gguf.py",
install_dir / "gguf-py",
]
- missing = [
- str(path.relative_to(install_dir)) for path in required if not path.exists()
- ]
+ missing = [str(path.relative_to(install_dir)) for path in required if not path.exists()]
if missing:
- raise PrebuiltFallback(
- "hydrated llama.cpp source tree was missing: " + ", ".join(missing)
- )
+ raise PrebuiltFallback("hydrated llama.cpp source tree was missing: " + ", ".join(missing))
def validation_model_cache_path(install_dir: Path) -> Path:
@@ -5028,8 +4822,7 @@ def validated_validation_model_bytes(data: bytes) -> bytes:
digest = hashlib.sha256(data).hexdigest()
if digest != TEST_MODEL_SHA256:
raise RuntimeError(
- "validation model checksum mismatch: "
- f"expected={TEST_MODEL_SHA256} actual={digest}"
+ f"validation model checksum mismatch: expected={TEST_MODEL_SHA256} actual={digest}"
)
return data
@@ -5042,9 +4835,7 @@ def download_validation_model(path: Path, cache_path: Path | None = None) -> Non
data = validated_validation_model_bytes(cache_path.read_bytes())
log(f"using cached tiny GGUF validation model from {cache_path}")
except Exception as exc:
- log(
- f"cached tiny GGUF validation model was invalid; refreshing cache ({exc})"
- )
+ log(f"cached tiny GGUF validation model was invalid; refreshing cache ({exc})")
data = None
if data is None:
log("downloading tiny GGUF validation model")
@@ -5133,9 +4924,7 @@ def dedupe_existing_dirs(paths: Iterable[str | Path]) -> list[str]:
return unique
-def linux_missing_libraries(
- binary_path: Path, *, env: dict[str, str] | None = None
-) -> list[str]:
+def linux_missing_libraries(binary_path: Path, *, env: dict[str, str] | None = None) -> list[str]:
try:
result = run_capture(["ldd", str(binary_path)], timeout = 20, env = env)
except Exception:
@@ -5292,9 +5081,7 @@ def _macho_slice_minos(data: bytes, offset: int) -> tuple[int, int] | None:
return None
-def macho_minimum_macos(
- path: Path, host: HostInfo | None = None
-) -> tuple[int, int] | None:
+def macho_minimum_macos(path: Path, host: HostInfo | None = None) -> tuple[int, int] | None:
"""Minimum macOS (major, minor) a Mach-O binary or dylib requires.
Pure-Python so it works on consumer Macs without the Xcode command line
@@ -5333,9 +5120,7 @@ def macho_minimum_macos(
return None
if host is not None:
want = (
- _CPU_TYPE_ARM64
- if host.is_arm64
- else (_CPU_TYPE_X86_64 if host.is_x86_64 else None)
+ _CPU_TYPE_ARM64 if host.is_arm64 else (_CPU_TYPE_X86_64 if host.is_x86_64 else None)
)
for cputype, minos in slices:
if cputype == want:
@@ -5355,9 +5140,7 @@ def looks_like_macos_incompatibility(text: str) -> bool:
def macos_binary_minos_issues(
- binaries: Iterable[Path],
- install_dir: Path,
- host: HostInfo,
+ binaries: Iterable[Path], install_dir: Path, host: HostInfo
) -> list[str]:
"""Issue strings for every installed Mach-O whose minimum macOS exceeds the
host. Scans the given executables plus every bundled .dylib next to them --
@@ -5387,9 +5170,7 @@ def macos_binary_minos_issues(
def preflight_macos_installed_binaries(
- binaries: Iterable[Path],
- install_dir: Path,
- host: HostInfo,
+ binaries: Iterable[Path], install_dir: Path, host: HostInfo
) -> None:
"""Reject a macos prebuilt whose minimum-OS is newer than the host. The
upstream selector pins a loadable release up front, so here this is the
@@ -5400,15 +5181,12 @@ def preflight_macos_installed_binaries(
issues = macos_binary_minos_issues(binaries, install_dir, host)
if issues:
raise PrebuiltFallback(
- "macos prebuilt requires a newer macOS than this host:\n"
- + "\n".join(issues)
+ "macos prebuilt requires a newer macOS than this host:\n" + "\n".join(issues)
)
def preflight_linux_installed_binaries(
- binaries: Iterable[Path],
- install_dir: Path,
- host: HostInfo,
+ binaries: Iterable[Path], install_dir: Path, host: HostInfo
) -> None:
if not host.is_linux:
return
@@ -5419,18 +5197,14 @@ def preflight_linux_installed_binaries(
missing = linux_missing_libraries(binary_path, env = env)
if not missing:
continue
- runtime_dirs = [
- part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part
- ]
+ runtime_dirs = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part]
issues.append(
f"{binary_path.name}: missing={','.join(missing)} "
f"ld_library_path={','.join(runtime_dirs) if runtime_dirs else 'none'}"
)
if issues:
- raise PrebuiltFallback(
- "linux extracted binary preflight failed:\n" + "\n".join(issues)
- )
+ raise PrebuiltFallback("linux extracted binary preflight failed:\n" + "\n".join(issues))
def glob_paths(*patterns: str) -> list[str]:
@@ -5474,12 +5248,9 @@ def windows_runtime_dirs() -> list[str]:
def windows_runtime_dirs_for_patterns(
- required_patterns: Iterable[str],
- candidate_dirs: Iterable[str] | None = None,
+ required_patterns: Iterable[str], candidate_dirs: Iterable[str] | None = None
) -> list[str]:
- directories = (
- list(candidate_dirs) if candidate_dirs is not None else windows_runtime_dirs()
- )
+ directories = list(candidate_dirs) if candidate_dirs is not None else windows_runtime_dirs()
matching_dirs: list[str] = []
for pattern in required_patterns:
matched_dirs = [
@@ -5523,20 +5294,12 @@ def binary_env(
str(install_dir),
*linux_runtime_dirs(binary_path),
]
- existing = [
- part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part
- ]
- env["LD_LIBRARY_PATH"] = os.pathsep.join(
- dedupe_existing_dirs([*ld_dirs, *existing])
- )
+ existing = [part for part in env.get("LD_LIBRARY_PATH", "").split(os.pathsep) if part]
+ env["LD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*ld_dirs, *existing]))
elif host.is_macos:
dyld_dirs = [str(binary_path.parent), str(install_dir)]
- existing = [
- part for part in env.get("DYLD_LIBRARY_PATH", "").split(os.pathsep) if part
- ]
- env["DYLD_LIBRARY_PATH"] = os.pathsep.join(
- dedupe_existing_dirs([*dyld_dirs, *existing])
- )
+ existing = [part for part in env.get("DYLD_LIBRARY_PATH", "").split(os.pathsep) if part]
+ env["DYLD_LIBRARY_PATH"] = os.pathsep.join(dedupe_existing_dirs([*dyld_dirs, *existing]))
return env
@@ -5558,11 +5321,7 @@ def validate_quantize(
env = binary_env(quantize_path, install_dir, host, runtime_line = runtime_line),
**windows_hidden_subprocess_kwargs(),
)
- if (
- result.returncode != 0
- or not quantized_path.exists()
- or quantized_path.stat().st_size == 0
- ):
+ if result.returncode != 0 or not quantized_path.exists() or quantized_path.stat().st_size == 0:
combined = result.stdout + ("\n" + result.stderr if result.stderr else "")
# Backstop for prebuilts the static minos scan could not read: a dyld
# "built for macOS N" / missing Metal symbol failure means this binary
@@ -5572,9 +5331,7 @@ def validate_quantize(
if looks_like_macos_incompatibility(combined)
else ""
)
- raise PrebuiltFallback(
- prefix + "llama-quantize validation failed:\n" + combined
- )
+ raise PrebuiltFallback(prefix + "llama-quantize validation failed:\n" + combined)
def validate_server(
@@ -5630,9 +5387,7 @@ def validate_server(
# is exercised against the actual hardware rather than the
# CPU fallback. NVIDIA and macOS-arm64 are already covered.
_enable_gpu_layers = (
- host.has_usable_nvidia
- or host.has_rocm
- or (host.is_macos and host.is_arm64)
+ host.has_usable_nvidia or host.has_rocm or (host.is_macos and host.is_arm64)
)
if _enable_gpu_layers:
command.extend(["--n-gpu-layers", "1"])
@@ -5648,9 +5403,7 @@ def validate_server(
stdout = log_handle,
stderr = subprocess.STDOUT,
text = True,
- env = binary_env(
- server_path, install_dir, host, runtime_line = runtime_line
- ),
+ env = binary_env(server_path, install_dir, host, runtime_line = runtime_line),
**windows_hidden_subprocess_kwargs(),
)
deadline = time.time() + 60
@@ -5665,9 +5418,7 @@ def validate_server(
exited_quickly = (
time.time() - startup_started
) <= SERVER_BIND_RETRY_WINDOW_SECONDS
- failure = PrebuiltFallback(
- "llama-server exited during startup:\n" + output
- )
+ failure = PrebuiltFallback("llama-server exited during startup:\n" + output)
if (
port_attempt < SERVER_PORT_BIND_ATTEMPTS
and is_retryable_server_bind_error(
@@ -5684,9 +5435,7 @@ def validate_server(
break
raise failure
- payload = json.dumps({"prompt": "a", "n_predict": 1}).encode(
- "utf-8"
- )
+ payload = json.dumps({"prompt": "a", "n_predict": 1}).encode("utf-8")
request = urllib.request.Request(
f"http://127.0.0.1:{port}/completion",
data = payload,
@@ -5698,9 +5447,7 @@ def validate_server(
response_body = response.read().decode("utf-8", "replace")
if status_code == 200:
return
- last_error = RuntimeError(
- f"unexpected HTTP status {status_code}"
- )
+ last_error = RuntimeError(f"unexpected HTTP status {status_code}")
except urllib.error.HTTPError as exc:
response_body = exc.read().decode("utf-8", "replace")
last_error = exc
@@ -5734,9 +5481,7 @@ def validate_server(
raise PrebuiltFallback("llama-server validation failed unexpectedly")
-def collect_system_report(
- host: HostInfo, choice: AssetChoice | None, install_dir: Path
-) -> str:
+def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_dir: Path) -> str:
lines = [
f"platform={host.system} machine={host.machine}",
f"driver_cuda_version={host.driver_cuda_version}",
@@ -5750,8 +5495,7 @@ def collect_system_report(
if host.is_linux and host.has_physical_nvidia:
runtime_lines, runtime_dirs = detected_linux_runtime_lines()
lines.append(
- "linux_runtime_lines="
- + (",".join(runtime_lines) if runtime_lines else "none")
+ "linux_runtime_lines=" + (",".join(runtime_lines) if runtime_lines else "none")
)
for runtime_line in ("cuda13", "cuda12"):
lines.append(
@@ -5780,10 +5524,7 @@ def collect_system_report(
server_env = binary_env(server_binary, install_dir, host)
lines.append(
"linux_missing_libs="
- + (
- ",".join(linux_missing_libraries(server_binary, env = server_env))
- or "none"
- )
+ + (",".join(linux_missing_libraries(server_binary, env = server_env)) or "none")
)
lines.append(
"linux_runtime_dirs="
@@ -5791,9 +5532,7 @@ def collect_system_report(
",".join(
[
part
- for part in server_env.get("LD_LIBRARY_PATH", "").split(
- os.pathsep
- )
+ for part in server_env.get("LD_LIBRARY_PATH", "").split(os.pathsep)
if part
]
)
@@ -5801,21 +5540,16 @@ def collect_system_report(
)
)
try:
- ldd = run_capture(
- ["ldd", str(server_binary)], timeout = 20, env = server_env
- )
+ ldd = run_capture(["ldd", str(server_binary)], timeout = 20, env = server_env)
lines.append("ldd llama-server:")
lines.append((ldd.stdout + ldd.stderr).strip())
except Exception as exc:
lines.append(f"ldd error: {exc}")
elif host.is_windows:
- lines.append(
- "windows_runtime_dirs=" + (",".join(windows_runtime_dirs()) or "none")
- )
+ lines.append("windows_runtime_dirs=" + (",".join(windows_runtime_dirs()) or "none"))
runtime_lines, runtime_dirs = detected_windows_runtime_lines()
lines.append(
- "windows_runtime_lines="
- + (",".join(runtime_lines) if runtime_lines else "none")
+ "windows_runtime_lines=" + (",".join(runtime_lines) if runtime_lines else "none")
)
for runtime_line in ("cuda13", "cuda12"):
lines.append(
@@ -5840,8 +5574,7 @@ def collect_system_report(
def apply_approved_hashes(
- attempts: Iterable[AssetChoice],
- checksums: ApprovedReleaseChecksums,
+ attempts: Iterable[AssetChoice], checksums: ApprovedReleaseChecksums
) -> list[AssetChoice]:
def approved_hash_for_attempt(attempt: AssetChoice) -> ApprovedArtifactHash | None:
candidate_names = [attempt.name]
@@ -5949,8 +5682,7 @@ def preferred_source_archive(
def selected_source_archive_metadata(
- checksums: ApprovedReleaseChecksums,
- llama_tag: str,
+ checksums: ApprovedReleaseChecksums, llama_tag: str
) -> tuple[str, str | None]:
_source_repo, _source_ref, source_archive, _exact_source = preferred_source_archive(
checksums, llama_tag
@@ -5961,10 +5693,7 @@ def selected_source_archive_metadata(
def resolve_install_attempts(
- llama_tag: str,
- host: HostInfo,
- published_repo: str,
- published_release_tag: str,
+ llama_tag: str, host: HostInfo, published_repo: str, published_release_tag: str
) -> tuple[str, str, list[AssetChoice], ApprovedReleaseChecksums]:
requested_tag, plans = resolve_install_release_plans(
llama_tag,
@@ -5987,17 +5716,11 @@ def resolve_install_release_plans(
max_release_fallbacks: int = DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS,
) -> tuple[str, list[InstallReleasePlan]]:
requested_tag = normalized_requested_llama_tag(llama_tag)
- allow_older_release_fallback = (
- requested_tag == "latest" and not published_release_tag
- )
+ allow_older_release_fallback = requested_tag == "latest" and not published_release_tag
release_limit = max(1, max_release_fallbacks)
# macOS may need to walk past a run of too-new prebuilts. Only when the host
# version is known; otherwise keep the default (cannot tell up front).
- if (
- host.is_macos
- and allow_older_release_fallback
- and host.macos_version is not None
- ):
+ if host.is_macos and allow_older_release_fallback and host.macos_version is not None:
release_limit = max(release_limit, DEFAULT_MAX_MACOS_RELEASE_FALLBACKS)
plans: list[InstallReleasePlan] = []
last_error: PrebuiltFallback | None = None
@@ -6013,9 +5736,7 @@ def resolve_install_release_plans(
try:
if host.is_linux and host.is_x86_64 and host.has_usable_nvidia:
linux_cuda_selection = resolve_linux_cuda_choice(host, bundle)
- attempts = apply_approved_hashes(
- linux_cuda_selection.attempts, checksums
- )
+ attempts = apply_approved_hashes(linux_cuda_selection.attempts, checksums)
if not attempts:
raise PrebuiltFallback("no compatible Linux CUDA asset was found")
log_lines(linux_cuda_selection.selection_log)
@@ -6117,9 +5838,7 @@ def write_prebuilt_metadata(
"prebuilt_fallback_used": prebuilt_fallback_used,
"installed_at_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
}
- (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(
- json.dumps(metadata, indent = 2) + "\n"
- )
+ (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(metadata, indent = 2) + "\n")
def expected_install_fingerprint(
@@ -6236,9 +5955,7 @@ def install_runtime_dir(install_dir: Path, host: HostInfo) -> Path:
return install_dir / "build" / "bin"
-def runtime_payload_is_healthy(
- install_dir: Path, host: HostInfo, choice: AssetChoice
-) -> bool:
+def runtime_payload_is_healthy(install_dir: Path, host: HostInfo, choice: AssetChoice) -> bool:
runtime_dir = install_runtime_dir(install_dir, host)
if not runtime_dir.exists():
return False
@@ -6326,9 +6043,7 @@ def existing_install_matches_choice(
def existing_install_matches_plan(
- install_dir: Path,
- host: HostInfo,
- plan: InstallReleasePlan,
+ install_dir: Path, host: HostInfo, plan: InstallReleasePlan
) -> bool:
if not plan.attempts:
return False
@@ -6360,9 +6075,7 @@ def validate_prebuilt_choice(
approved_checksums, llama_tag
)
if exact_source:
- log(
- f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}"
- )
+ log(f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}")
else:
log(f"hydrating upstream llama.cpp source for {llama_tag} into {install_dir}")
hydrate_source_tree(
@@ -6379,9 +6092,7 @@ def validate_prebuilt_choice(
exact_source = exact_source,
)
log(f"overlaying prebuilt bundle {choice.name} into {install_dir}")
- server_path, quantize_path = install_from_archives(
- choice, host, install_dir, work_dir
- )
+ server_path, quantize_path = install_from_archives(choice, host, install_dir, work_dir)
preflight_linux_installed_binaries((server_path, quantize_path), install_dir, host)
preflight_macos_installed_binaries((server_path, quantize_path), install_dir, host)
ensure_repo_shape(install_dir)
@@ -6539,9 +6250,7 @@ def install_prebuilt(
published_repo,
published_release_tag,
)
- if release_plans and existing_install_matches_plan(
- install_dir, host, release_plans[0]
- ):
+ if release_plans and existing_install_matches_plan(install_dir, host, release_plans[0]):
current = release_plans[0]
log(
"existing llama.cpp install already matches selected release "
@@ -6551,9 +6260,7 @@ def install_prebuilt(
with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp:
work_dir = Path(tmp)
probe_path = work_dir / "stories260K.gguf"
- download_validation_model(
- probe_path, validation_model_cache_path(install_dir)
- )
+ download_validation_model(probe_path, validation_model_cache_path(install_dir))
release_count = len(release_plans)
for release_index, plan in enumerate(release_plans):
choice = plan.attempts[0]
@@ -6759,9 +6466,7 @@ def main() -> int:
)
emit_resolver_output(
{
- "requested_tag": normalized_requested_llama_tag(
- args.resolve_install_tag
- ),
+ "requested_tag": normalized_requested_llama_tag(args.resolve_install_tag),
"llama_tag": resolved,
},
output_format = args.output_format,
@@ -6776,9 +6481,7 @@ def main() -> int:
)
emit_resolver_output(
{
- "requested_tag": normalized_requested_llama_tag(
- args.resolve_source_build
- ),
+ "requested_tag": normalized_requested_llama_tag(args.resolve_source_build),
"source_url": plan.source_url,
"source_ref_kind": plan.source_ref_kind,
"source_ref": plan.source_ref,
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index da202f7ce6..5ff8e8572b 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -176,7 +176,6 @@ def _detect_rocm_version() -> tuple[int, int] | None:
)
if result.returncode == 0:
import re
-
m = re.search(r"ROCm version:\s*(\d+)\.(\d+)", result.stdout)
if m:
return int(m.group(1)), int(m.group(2))
@@ -196,11 +195,7 @@ def _detect_rocm_version() -> tuple[int, int] | None:
if result.returncode == 0:
raw = result.stdout.decode().strip().split("\n")[0]
parts = raw.split(".")
- if (
- len(parts) >= 2
- and parts[0].isdigit()
- and parts[1].split("-")[0].isdigit()
- ):
+ if len(parts) >= 2 and parts[0].isdigit() and parts[1].split("-")[0].isdigit():
return int(parts[0]), int(parts[1].split("-")[0])
except Exception:
pass
@@ -309,8 +304,7 @@ def _detect_windows_gfx_arch() -> str | None:
# findall picks every gcnArchName line so multi-GPU hosts
# are enumerable and HIP_VISIBLE_DEVICES selects correctly.
_tokens = [
- t.strip().lower()
- for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
+ t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
]
_pick = _dedup_pick(_tokens)
if _pick:
@@ -612,9 +606,7 @@ def _ensure_rocm_torch() -> None:
if not _torch_already_rocm:
index_url = _windows_rocm_index_url(gfx_arch)
if index_url is None:
- print(
- f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping"
- )
+ print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping")
return
print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
pip_install(
@@ -688,9 +680,7 @@ def _ensure_rocm_torch() -> None:
except (OSError, subprocess.TimeoutExpired):
probe = None
has_hip_torch = (
- probe is not None
- and probe.returncode == 0
- and probe.stdout.decode().strip() != ""
+ probe is not None and probe.returncode == 0 and probe.stdout.decode().strip() != ""
)
rocm_torch_ready = has_hip_torch
@@ -714,14 +704,11 @@ def _ensure_rocm_torch() -> None:
# specific index into gfx_codes, use that gfx; else default to the
# first listed GPU. Skip the override unless the resolved GPU is
# Strix.
- _runtime_gfx = (
- gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None
- )
+ _runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None
if _runtime_gfx in _strix_gfx:
_selected_gfx = _runtime_gfx
_amd_mirror = (
- os.environ.get("UNSLOTH_AMD_ROCM_MIRROR")
- or "https://repo.amd.com/rocm/whl"
+ os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") or "https://repo.amd.com/rocm/whl"
).rstrip("/")
_strix_override_url = f"{_amd_mirror}/{_selected_gfx}/"
_strix_override_pkgs = (
@@ -782,10 +769,7 @@ def _ensure_rocm_torch() -> None:
None,
)
if tag is None:
- print(
- f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- "
- f"skipping torch reinstall"
- )
+ print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall")
else:
index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
@@ -922,9 +906,7 @@ CONSTRAINTS = SINGLE_ENV / "constraints.txt"
LOCAL_DD_UNSTRUCTURED_PLUGIN = (
SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed"
)
-LOCAL_DD_GITHUB_PLUGIN = (
- SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
-)
+LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides file).
_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
@@ -1026,7 +1008,11 @@ def _title(msg: str) -> str:
_RULE = "\u2500" * 52
-def _step(label: str, value: str, color_fn = None) -> None:
+def _step(
+ label: str,
+ value: str,
+ color_fn = None,
+) -> None:
"""Print a single step line in the column format."""
if color_fn is None:
color_fn = _green
@@ -1046,16 +1032,17 @@ def _progress(label: str) -> None:
pad = " " * (_COL - len(_LABEL))
end = "\n" if _STEP >= _TOTAL else ""
try:
- sys.stdout.write(
- f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}"
- )
+ sys.stdout.write(f"\r {_dim(_LABEL)}{pad}[{bar}] {_STEP:2}/{_TOTAL} {label:<20}{end}")
sys.stdout.flush()
except OSError:
pass
def run(
- label: str, cmd: list[str], *, quiet: bool = True
+ label: str,
+ cmd: list[str],
+ *,
+ quiet: bool = True,
) -> subprocess.CompletedProcess[bytes]:
"""Run a command; on failure print output and exit."""
if VERBOSE:
@@ -1107,9 +1094,7 @@ def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None:
return flash_attn_wheel_url(env)
-def _print_optional_install_failure(
- label: str, result: subprocess.CompletedProcess[str]
-) -> None:
+def _print_optional_install_failure(label: str, result: subprocess.CompletedProcess[str]) -> None:
_step("warning", f"{label} failed (exit code {result.returncode})", _cyan)
if result.stdout:
print(result.stdout.strip())
@@ -1204,9 +1189,7 @@ def _filter_requirements(req: Path, skip: set[str]) -> Path:
"""Return a temp copy of a requirements file with certain packages removed."""
lines = req.read_text(encoding = "utf-8").splitlines(keepends = True)
filtered = [
- line
- for line in lines
- if not any(line.strip().lower().startswith(pkg) for pkg in skip)
+ line for line in lines if not any(line.strip().lower().startswith(pkg) for pkg in skip)
]
tmp = tempfile.NamedTemporaryFile(
mode = "w",
@@ -1416,9 +1399,7 @@ def install_python_stack() -> int:
if not IS_MACOS and not NO_TORCH:
base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms
if not IS_WINDOWS:
- base_total += (
- 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only
- )
+ base_total += 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only
_TOTAL = (base_total - 1) if skip_base else base_total
# 1. Try to use uv for faster installs (must happen before pip upgrade
diff --git a/tests/_zoo_aggressive_cuda_spoof.py b/tests/_zoo_aggressive_cuda_spoof.py
index eaafe445fb..9111aa9519 100644
--- a/tests/_zoo_aggressive_cuda_spoof.py
+++ b/tests/_zoo_aggressive_cuda_spoof.py
@@ -159,7 +159,11 @@ def apply() -> None:
if _orig is None:
continue
- def _wrap(*args: Any, _orig = _orig, **kwargs: Any):
+ def _wrap(
+ *args: Any,
+ _orig = _orig,
+ **kwargs: Any,
+ ):
kwargs.pop("pin_memory", None)
return _orig(*args, **kwargs)
diff --git a/tests/conftest.py b/tests/conftest.py
index 2d7038d5d4..ad58cb9706 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -105,13 +105,11 @@ def _patch_torch_cuda_for_import() -> None:
CPU like normal."""
try:
import torch.cuda.memory as _cuda_memory # type: ignore
-
_cuda_memory.mem_get_info = lambda *a, **k: (0, 80 * 1024**3)
except Exception:
pass
try:
import torch
-
torch.cuda.get_device_capability = lambda *a, **k: (8, 0)
torch.cuda.is_bf16_supported = lambda *a, **k: True
except Exception:
diff --git a/tests/python/conftest.py b/tests/python/conftest.py
index 9129e384e5..f7b125edf6 100644
--- a/tests/python/conftest.py
+++ b/tests/python/conftest.py
@@ -2,9 +2,5 @@
def pytest_configure(config):
- config.addinivalue_line(
- "markers", "server: heavyweight tests requiring studio venv"
- )
- config.addinivalue_line(
- "markers", "e2e: end-to-end tests requiring network and venv creation"
- )
+ config.addinivalue_line("markers", "server: heavyweight tests requiring studio venv")
+ config.addinivalue_line("markers", "e2e: end-to-end tests requiring network and venv creation")
diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py
index 7b9868c1f2..6c504579ce 100644
--- a/tests/python/test_cross_platform_parity.py
+++ b/tests/python/test_cross_platform_parity.py
@@ -28,17 +28,11 @@ class TestNoTorchBackendAutoInInstallSh:
for i, line in enumerate(lines):
if fallback_start is None and "GPU detection failed" in line:
fallback_start = i
- elif (
- fallback_start is not None
- and fallback_end is None
- and line.strip() == "fi"
- ):
+ elif fallback_start is not None and fallback_end is None and line.strip() == "fi":
fallback_end = i
break
fallback_range = (
- range(fallback_start or 0, (fallback_end or 0) + 1)
- if fallback_start
- else range(0)
+ range(fallback_start or 0, (fallback_end or 0) + 1) if fallback_start else range(0)
)
matches = [
diff --git a/tests/python/test_dpo_vision_processor_passthrough.py b/tests/python/test_dpo_vision_processor_passthrough.py
index a4f2e2e12a..a320cab935 100644
--- a/tests/python/test_dpo_vision_processor_passthrough.py
+++ b/tests/python/test_dpo_vision_processor_passthrough.py
@@ -33,7 +33,11 @@ class _Tok:
eos_token_id = 99
bos_token_id = None
- def __call__(self, t, add_special_tokens = False):
+ def __call__(
+ self,
+ t,
+ add_special_tokens = False,
+ ):
return {"input_ids": [10]}
@@ -46,7 +50,12 @@ class _Capture:
self.last_text = None
self.last_images = "__sentinel__"
- def __call__(self, images = None, text = None, add_special_tokens = False):
+ def __call__(
+ self,
+ images = None,
+ text = None,
+ add_special_tokens = False,
+ ):
self.last_text = text
self.last_images = images
out = {"input_ids": [[1, 2]]}
diff --git a/tests/python/test_e2e_no_torch_sandbox.py b/tests/python/test_e2e_no_torch_sandbox.py
index f36f69201d..382d0afe4e 100644
--- a/tests/python/test_e2e_no_torch_sandbox.py
+++ b/tests/python/test_e2e_no_torch_sandbox.py
@@ -247,12 +247,8 @@ class TestBeforeAfterImportChain:
exec(source)
""")
result = _run_in_sandbox(no_torch_venv, code)
- assert (
- result.returncode != 0
- ), "BEFORE chat_templates.py should crash without torch"
- assert (
- b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
- )
+ assert result.returncode != 0, "BEFORE chat_templates.py should crash without torch"
+ assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
def test_before_data_collators_crashes(self, no_torch_venv, sandbox_dir):
"""BEFORE: data_collators.py with top-level 'import torch' crashes."""
@@ -270,12 +266,8 @@ class TestBeforeAfterImportChain:
exec(open({str(before_file)!r}).read())
""")
result = _run_in_sandbox(no_torch_venv, code)
- assert (
- result.returncode != 0
- ), "BEFORE data_collators.py should crash without torch"
- assert (
- b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
- )
+ assert result.returncode != 0, "BEFORE data_collators.py should crash without torch"
+ assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
def test_before_full_import_chain_crashes(self, no_torch_venv, sandbox_dir):
"""BEFORE: full utils/datasets/ package with top-level torch imports crashes."""
@@ -320,12 +312,8 @@ class TestBeforeAfterImportChain:
from utils.datasets import detect_dataset_format
""")
result = _run_in_sandbox(no_torch_venv, code)
- assert (
- result.returncode != 0
- ), "BEFORE full import chain should crash without torch"
- assert (
- b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
- )
+ assert result.returncode != 0, "BEFORE full import chain should crash without torch"
+ assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
# -- AFTER: succeeds --
@@ -539,9 +527,7 @@ class TestEdgeCasesBrokenTorch:
print("OK: data_collators works despite broken torch on sys.path")
""")
result = _run_in_sandbox(no_torch_venv, code)
- assert (
- result.returncode == 0
- ), f"Should work with broken torch:\n{result.stderr.decode()}"
+ assert result.returncode == 0, f"Should work with broken torch:\n{result.stderr.decode()}"
assert b"OK:" in result.stdout
def test_torch_import_error_hardware_fallback(self, no_torch_venv, sandbox_dir):
@@ -604,14 +590,10 @@ class TestEdgeCasesBrokenTorch:
print("OK: detect_hardware returned CPU with fake torch (no CUDA)")
""")
result = _run_in_sandbox(no_torch_venv, code)
- assert (
- result.returncode == 0
- ), f"Should fall back to CPU:\n{result.stderr.decode()}"
+ assert result.returncode == 0, f"Should fall back to CPU:\n{result.stderr.decode()}"
assert b"OK:" in result.stdout
- def test_lazy_torch_fails_at_call_time_not_import_time(
- self, no_torch_venv, sandbox_dir
- ):
+ def test_lazy_torch_fails_at_call_time_not_import_time(self, no_torch_venv, sandbox_dir):
"""apply_chat_template_to_dataset is importable without torch.
Calling the alpaca branch triggers the lazy 'from torch.utils.data' inside
@@ -657,9 +639,7 @@ class TestEdgeCasesBrokenTorch:
print("OK: call succeeded (unexpected but not a crash)")
""")
result = _run_in_sandbox(no_torch_venv, code)
- assert (
- result.returncode == 0
- ), f"Should not crash at import time:\n{result.stderr.decode()}"
+ assert result.returncode == 0, f"Should not crash at import time:\n{result.stderr.decode()}"
assert b"OK: import succeeded" in result.stdout
@@ -1011,9 +991,7 @@ class TestInstallPythonStackFiltering:
source = Path(ips.__file__).read_text(encoding = "utf-8")
# NO_TORCH guard before overrides
- assert (
- "if NO_TORCH:" in source
- ), "NO_TORCH guard not found in install_python_stack.py"
+ assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py"
# macOS guard for triton
assert (
@@ -1037,7 +1015,6 @@ def _studio_venv_python() -> Path | None:
def _server_port() -> int:
"""Find an available port for the test server."""
import socket
-
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
return s.getsockname()[1]
@@ -1117,9 +1094,7 @@ class TestLiveServerStartup:
for _ in range(30):
time.sleep(1)
try:
- resp = urllib.request.urlopen(
- f"http://127.0.0.1:{port}/api/health", timeout = 2
- )
+ resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/api/health", timeout = 2)
if resp.status == 200:
ready = True
break
@@ -1143,12 +1118,8 @@ class TestLiveServerStartup:
capture_output = True,
timeout = 300,
)
- server_output = stdout.decode(errors = "replace") + stderr.decode(
- errors = "replace"
- )
- pytest.skip(
- f"Server failed to start within 30 seconds. Output:\n{server_output}"
- )
+ server_output = stdout.decode(errors = "replace") + stderr.decode(errors = "replace")
+ pytest.skip(f"Server failed to start within 30 seconds. Output:\n{server_output}")
yield proc, port
@@ -1192,9 +1163,7 @@ class TestLiveServerStartup:
import urllib.request
_, port = server_process
- resp = urllib.request.urlopen(
- f"http://127.0.0.1:{port}/openapi.json", timeout = 5
- )
+ resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/openapi.json", timeout = 5)
spec = json.loads(resp.read())
assert (
len(spec.get("paths", {})) >= 20
diff --git a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
index 31d86b09a4..3f5f235aae 100644
--- a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
+++ b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
@@ -62,7 +62,6 @@ class _RecordingTransformerOk:
def __init__(self, model_name, **kwargs):
from transformers import AutoModel, AutoProcessor, AutoTokenizer
-
type(self).last_calls = {
"model": AutoModel.from_pretrained(model_name),
"processor": AutoProcessor.from_pretrained(model_name),
@@ -73,7 +72,6 @@ class _RecordingTransformerOk:
class _RaisingTransformer:
def __init__(self, *a, **kw):
from transformers import AutoModel
-
AutoModel.from_pretrained(a[0] if a else kw.get("model_name_or_path"))
raise RuntimeError("simulated init failure")
@@ -129,18 +127,10 @@ def _build_driver(transformer_class):
return model if is_requested_model_name(a, kw) else original_model(*a, **kw)
def return_existing_tokenizer(*a, **kw):
- return (
- tokenizer
- if is_requested_model_name(a, kw)
- else original_tokenizer(*a, **kw)
- )
+ return tokenizer if is_requested_model_name(a, kw) else original_tokenizer(*a, **kw)
def return_existing_processor(*a, **kw):
- return (
- tokenizer
- if is_requested_model_name(a, kw)
- else original_processor(*a, **kw)
- )
+ return tokenizer if is_requested_model_name(a, kw) else original_processor(*a, **kw)
try:
AutoModel.from_pretrained = return_existing_model
@@ -190,7 +180,6 @@ def test_redirect_passes_through_for_other_model_names():
def __init__(self, model_name, **kw):
from transformers import AutoModel
-
type(self).captured = AutoModel.from_pretrained("some-other/aux-model")
driver, *_ = _build_driver(_OtherNameTransformer)
@@ -210,7 +199,6 @@ def test_is_requested_model_name_handles_pathlib_path(tmp_path):
def __init__(self, model_name, **kw):
from transformers import AutoModel
-
type(self).last_calls = AutoModel.from_pretrained(pathlib.Path(model_name))
driver, *_ = _build_driver(_PathTransformer)
@@ -228,7 +216,6 @@ def test_is_requested_model_name_trailing_slash_local_path(tmp_path):
def __init__(self, model_name, **kw):
from transformers import AutoModel
-
type(self).last_calls = AutoModel.from_pretrained(str(target) + "/")
driver, *_ = _build_driver(_SlashTransformer)
@@ -243,7 +230,6 @@ def test_is_requested_model_name_returns_false_when_no_identifier():
class _NoNameTransformer:
def __init__(self, model_name, **kw):
from transformers import AutoModel
-
captured["args"] = AutoModel.from_pretrained(some_other_kwarg = "x")
driver, *_ = _build_driver(_NoNameTransformer)
diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py
index 49f4350a7b..26ff03505a 100644
--- a/tests/python/test_flash_attn_install_python_stack.py
+++ b/tests/python/test_flash_attn_install_python_stack.py
@@ -33,64 +33,42 @@ class TestHasBlackwellGpu:
def test_returns_true_for_sm_100(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
- mock.patch.object(
- wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+ mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_true_for_sm_120(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
- mock.patch.object(
- wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+ mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_true_for_sm_121(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
- mock.patch.object(
- wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+ mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")),
):
assert wheel_utils.has_blackwell_gpu() is True
def test_returns_false_for_sm_90(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
- mock.patch.object(
- wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+ mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_returns_false_for_sm_89(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
- mock.patch.object(
- wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
+ mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_mixed_gpus_with_one_blackwell_returns_true(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
mock.patch.object(
wheel_utils.subprocess,
"run",
@@ -101,9 +79,7 @@ class TestHasBlackwellGpu:
def test_returns_false_when_nvidia_smi_fails(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
mock.patch.object(
wheel_utils.subprocess,
"run",
@@ -114,9 +90,7 @@ class TestHasBlackwellGpu:
def test_returns_false_on_subprocess_timeout(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
mock.patch.object(
wheel_utils.subprocess,
"run",
@@ -127,9 +101,7 @@ class TestHasBlackwellGpu:
def test_returns_false_on_malformed_output(self):
with (
- mock.patch.object(
- wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"
- ),
+ mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"),
mock.patch.object(
wheel_utils.subprocess,
"run",
@@ -161,10 +133,7 @@ class TestFlashAttnWheelSelection:
)
assert url is not None
assert "v2.8.1" in url
- assert (
- "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl"
- in url
- )
+ assert "flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
def test_missing_cuda_major_disables_wheel_lookup(self):
assert (
@@ -262,7 +231,11 @@ class TestEnsureFlashAttn:
step_messages: list[tuple[str, str]] = []
printed_failures: list[str] = []
- def fake_step(label: str, value: str, color_fn = None):
+ def fake_step(
+ label: str,
+ value: str,
+ color_fn = None,
+ ):
step_messages.append((label, value))
with (
@@ -313,7 +286,11 @@ class TestEnsureFlashAttn:
def test_wheel_missing_skips_install_at_setup_time(self):
step_messages: list[tuple[str, str]] = []
- def fake_step(label: str, value: str, color_fn = None):
+ def fake_step(
+ label: str,
+ value: str,
+ color_fn = None,
+ ):
step_messages.append((label, value))
with (
@@ -339,10 +316,7 @@ class TestEnsureFlashAttn:
ips._ensure_flash_attn()
mock_install_wheel.assert_not_called()
- assert (
- "warning",
- "No published flash-attn prebuilt wheel found",
- ) in step_messages
+ assert ("warning", "No published flash-attn prebuilt wheel found") in step_messages
def test_skip_env_disables_setup_install(self):
with (
@@ -362,7 +336,11 @@ class TestEnsureFlashAttn:
def test_blackwell_gpu_skips_install_with_warning(self):
step_messages: list[tuple[str, str]] = []
- def fake_step(label: str, value: str, color_fn = None):
+ def fake_step(
+ label: str,
+ value: str,
+ color_fn = None,
+ ):
step_messages.append((label, value))
with (
@@ -379,14 +357,16 @@ class TestEnsureFlashAttn:
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
- assert any(
- label == "warning" and "Blackwell" in msg for label, msg in step_messages
- )
+ assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
def test_blackwell_gpu_on_windows_emits_blackwell_warning(self):
step_messages: list[tuple[str, str]] = []
- def fake_step(label: str, value: str, color_fn = None):
+ def fake_step(
+ label: str,
+ value: str,
+ color_fn = None,
+ ):
step_messages.append((label, value))
with (
@@ -403,14 +383,16 @@ class TestEnsureFlashAttn:
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
- assert any(
- label == "warning" and "Blackwell" in msg for label, msg in step_messages
- )
+ assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages)
def test_non_blackwell_windows_does_not_emit_blackwell_warning(self):
step_messages: list[tuple[str, str]] = []
- def fake_step(label: str, value: str, color_fn = None):
+ def fake_step(
+ label: str,
+ value: str,
+ color_fn = None,
+ ):
step_messages.append((label, value))
with (
@@ -453,9 +435,7 @@ class TestInstallPythonStackFlashAttnIntegration:
mock.patch("subprocess.run", side_effect = fake_run),
mock.patch.object(ips, "_has_usable_nvidia_gpu", return_value = False),
mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
- mock.patch.object(
- ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")
- ),
+ mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")),
mock.patch("pathlib.Path.is_dir", return_value = True),
mock.patch("pathlib.Path.is_file", return_value = True),
mock.patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}, clear = False),
diff --git a/tests/python/test_gpu_init_ldconfig_guard.py b/tests/python/test_gpu_init_ldconfig_guard.py
index 081a6132b4..248bb84faa 100644
--- a/tests/python/test_gpu_init_ldconfig_guard.py
+++ b/tests/python/test_gpu_init_ldconfig_guard.py
@@ -19,9 +19,7 @@ def _find_geteuid_guard(tree: ast.AST):
def test_gpu_init_has_geteuid_guard():
tree = ast.parse(GPU_INIT.read_text())
guard = _find_geteuid_guard(tree)
- assert (
- guard is not None
- ), "_gpu_init.py must guard ldconfig recovery on os.geteuid()"
+ assert guard is not None, "_gpu_init.py must guard ldconfig recovery on os.geteuid()"
def test_ldconfig_calls_only_inside_geteuid_guard():
diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py
index 29cadf87ae..f7a90fc6a8 100644
--- a/tests/python/test_no_torch_filtering.py
+++ b/tests/python/test_no_torch_filtering.py
@@ -156,9 +156,7 @@ class TestFilterRequirements:
)
# First filter Windows packages, then NO_TORCH packages
intermediate = ips._filter_requirements(req, ips.WINDOWS_SKIP_PACKAGES)
- result = ips._filter_requirements(
- Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES
- )
+ result = ips._filter_requirements(Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES)
lines = Path(result).read_text(encoding = "utf-8").splitlines()
non_blank = [l.strip() for l in lines if l.strip()]
assert non_blank == [
@@ -177,9 +175,7 @@ class TestFilterRequirements:
result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
lines = Path(result).read_text(encoding = "utf-8").splitlines()
non_blank = [l.strip() for l in lines if l.strip()]
- assert non_blank == [
- "numpy"
- ], f"VCS URL line should be filtered, got: {non_blank}"
+ assert non_blank == ["numpy"], f"VCS URL line should be filtered, got: {non_blank}"
def test_env_marker_line_filtered(self, tmp_path):
"""Package lines with env markers are still filtered by prefix."""
@@ -193,9 +189,7 @@ class TestFilterRequirements:
result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
lines = Path(result).read_text(encoding = "utf-8").splitlines()
non_blank = [l.strip() for l in lines if l.strip()]
- assert non_blank == [
- "numpy"
- ], f"Env marker line should be filtered, got: {non_blank}"
+ assert non_blank == ["numpy"], f"Env marker line should be filtered, got: {non_blank}"
def test_git_plus_url_not_over_matched(self, tmp_path):
"""A git+ URL whose path contains a skip package name but does NOT start with it."""
@@ -247,9 +241,7 @@ class TestRealRequirementsFiltering:
expected = [
l
for l in original
- if not any(
- l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES
- )
+ if not any(l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES)
]
assert filtered == expected, (
f"Filtered extras.txt should match expected.\n"
@@ -259,9 +251,7 @@ class TestRealRequirementsFiltering:
def test_extras_no_deps_txt_torchcodec_and_dlpack_removed(self):
"""extras-no-deps.txt: torchcodec and torch-c-dlpack-ext must be removed."""
- result = ips._filter_requirements(
- EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES
- )
+ result = ips._filter_requirements(EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES)
filtered = self._non_blank_non_comment(Path(result))
original = self._non_blank_non_comment(EXTRAS_NO_DEPS_TXT)
@@ -273,9 +263,7 @@ class TestRealRequirementsFiltering:
expected = [
l
for l in original
- if not any(
- l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES
- )
+ if not any(l.strip().lower().startswith(p) for p in ips.NO_TORCH_SKIP_PACKAGES)
]
assert filtered == expected
@@ -291,9 +279,7 @@ class TestRealRequirementsFiltering:
def test_extras_no_deps_txt_trl_preserved(self):
"""trl should survive NO_TORCH filtering in extras-no-deps.txt."""
- result = ips._filter_requirements(
- EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES
- )
+ result = ips._filter_requirements(EXTRAS_NO_DEPS_TXT, ips.NO_TORCH_SKIP_PACKAGES)
filtered_text = Path(result).read_text(encoding = "utf-8").lower()
assert "trl" in filtered_text, "trl should survive NO_TORCH filtering"
@@ -370,7 +356,6 @@ class TestIsMacosConstant:
def test_is_macos_matches_platform(self):
import sys
-
expected = sys.platform == "darwin"
assert ips.IS_MACOS is expected
@@ -405,9 +390,7 @@ class TestInstallPythonStackSubprocessMock:
captured_cmds: list[list[str]] = []
def mock_run(cmd, **kw):
- captured_cmds.append(
- list(cmd) if isinstance(cmd, (list, tuple)) else [str(cmd)]
- )
+ captured_cmds.append(list(cmd) if isinstance(cmd, (list, tuple)) else [str(cmd)])
return subprocess.CompletedProcess(cmd, 0, b"", b"")
env = {"SKIP_STUDIO_BASE": "1"} if skip_base else {}
@@ -424,9 +407,7 @@ class TestInstallPythonStackSubprocessMock:
mock.patch.object(ips, "_has_rocm_gpu", return_value = False),
mock.patch("subprocess.run", side_effect = mock_run),
mock.patch.object(ips, "_bootstrap_uv", return_value = True),
- mock.patch.object(
- ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")
- ),
+ mock.patch.object(ips, "LOCAL_DD_UNSTRUCTURED_PLUGIN", Path("/fake/plugin")),
mock.patch("pathlib.Path.is_dir", return_value = True),
mock.patch("pathlib.Path.is_file", return_value = True),
):
@@ -469,9 +450,7 @@ class TestInstallPythonStackSubprocessMock:
has_extras_nd = self._cmds_contain_file(cmds, "extras-no-deps.txt") or any(
"-r" in cmd and "tmp" in cmd.lower() for cmd in cmds
)
- assert (
- has_extras_nd
- ), "extras-no-deps.txt (or its filtered temp) should be called"
+ assert has_extras_nd, "extras-no-deps.txt (or its filtered temp) should be called"
# -- IS_WINDOWS=True + NO_TORCH=True (stacked) --
@@ -570,17 +549,13 @@ class TestOverridesSkip:
def test_no_torch_guard_exists_in_source(self):
"""The install_python_stack source must contain a NO_TORCH guard around overrides."""
source = Path(ips.__file__).read_text(encoding = "utf-8")
- assert (
- "if NO_TORCH:" in source
- ), "NO_TORCH guard not found in install_python_stack.py"
+ assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py"
def test_overrides_skipped_when_no_torch(self):
"""With NO_TORCH=True on the module, pip_install should NOT be called for overrides."""
source = Path(ips.__file__).read_text(encoding = "utf-8")
overrides_match = re.search(r"if NO_TORCH:.*?overrides", source, re.DOTALL)
- assert (
- overrides_match is not None
- ), "Expected NO_TORCH conditional before overrides install"
+ assert overrides_match is not None, "Expected NO_TORCH conditional before overrides install"
# ── install.sh --no-torch flag tests ──────────────────────────────────
@@ -599,33 +574,21 @@ class TestInstallShNoTorchFlag:
def test_no_torch_flag_in_case_statement(self):
"""--no-torch must appear in the flag parser case statement."""
- assert (
- "--no-torch)" in self.source
- ), "--no-torch not found in install.sh flag parser"
+ assert "--no-torch)" in self.source, "--no-torch not found in install.sh flag parser"
def test_no_torch_flag_variable_initialized(self):
"""_NO_TORCH_FLAG must be initialized to false."""
- assert (
- "_NO_TORCH_FLAG=false" in self.source
- ), "_NO_TORCH_FLAG=false not found in install.sh"
+ assert "_NO_TORCH_FLAG=false" in self.source, "_NO_TORCH_FLAG=false not found in install.sh"
def test_skip_torch_variable_exists(self):
"""SKIP_TORCH variable must be defined."""
- assert (
- "SKIP_TORCH=false" in self.source
- ), "SKIP_TORCH=false not found in install.sh"
- assert (
- "SKIP_TORCH=true" in self.source
- ), "SKIP_TORCH=true not found in install.sh"
+ assert "SKIP_TORCH=false" in self.source, "SKIP_TORCH=false not found in install.sh"
+ assert "SKIP_TORCH=true" in self.source, "SKIP_TORCH=true not found in install.sh"
def test_skip_torch_driven_by_flag_and_mac_intel(self):
"""SKIP_TORCH must check both _NO_TORCH_FLAG and MAC_INTEL."""
- assert (
- "_NO_TORCH_FLAG" in self.source
- ), "_NO_TORCH_FLAG not referenced in SKIP_TORCH logic"
- assert (
- "MAC_INTEL" in self.source
- ), "MAC_INTEL not referenced in SKIP_TORCH logic"
+ assert "_NO_TORCH_FLAG" in self.source, "_NO_TORCH_FLAG not referenced in SKIP_TORCH logic"
+ assert "MAC_INTEL" in self.source, "MAC_INTEL not referenced in SKIP_TORCH logic"
def test_unsloth_no_torch_uses_skip_torch(self):
"""UNSLOTH_NO_TORCH must reference $SKIP_TORCH, not $MAC_INTEL."""
@@ -633,18 +596,12 @@ class TestInstallShNoTorchFlag:
matches = re.findall(r'UNSLOTH_NO_TORCH="\$(\w+)"', self.source)
for var in matches:
- assert (
- var == "SKIP_TORCH"
- ), f"UNSLOTH_NO_TORCH references ${var} instead of $SKIP_TORCH"
+ assert var == "SKIP_TORCH", f"UNSLOTH_NO_TORCH references ${var} instead of $SKIP_TORCH"
def test_cpu_hint_message_exists(self):
"""CPU hint message must exist in install.sh."""
- assert (
- "No GPU detected" in self.source
- ), "CPU hint message not found in install.sh"
- assert (
- "--no-torch" in self.source
- ), "--no-torch suggestion not found in CPU hint"
+ assert "No GPU detected" in self.source, "CPU hint message not found in install.sh"
+ assert "--no-torch" in self.source, "--no-torch suggestion not found in CPU hint"
def test_no_torch_flag_parsing_subprocess(self):
"""--no-torch flag sets _NO_TORCH_FLAG=true (subprocess test)."""
diff --git a/tests/python/test_orpo_processor_text_tokenizer.py b/tests/python/test_orpo_processor_text_tokenizer.py
index 44c4e26a86..9ae205c6b9 100644
--- a/tests/python/test_orpo_processor_text_tokenizer.py
+++ b/tests/python/test_orpo_processor_text_tokenizer.py
@@ -34,7 +34,12 @@ class _Tokenizer:
def __init__(self):
self.calls = []
- def __call__(self, text, add_special_tokens = False, **kwargs):
+ def __call__(
+ self,
+ text,
+ add_special_tokens = False,
+ **kwargs,
+ ):
self.calls.append((text, add_special_tokens, kwargs))
ids = [ord(c) % 31 + 3 for c in text]
return {"input_ids": ids, "attention_mask": [1] * len(ids)}
@@ -60,7 +65,11 @@ class _Trainer:
self.padding_value = 0
-def _exec_rewritten(function_name, source, extra_ns = None):
+def _exec_rewritten(
+ function_name,
+ source,
+ extra_ns = None,
+):
rewriter = _load_orpo_rewriter()
rewritten = rewriter(function_name, source)
ns = {} if extra_ns is None else dict(extra_ns)
diff --git a/tests/python/test_studio_import_no_torch.py b/tests/python/test_studio_import_no_torch.py
index 5592a282ff..c8b67023f5 100644
--- a/tests/python/test_studio_import_no_torch.py
+++ b/tests/python/test_studio_import_no_torch.py
@@ -23,15 +23,9 @@ from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
-DATA_COLLATORS = (
- REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "data_collators.py"
-)
-CHAT_TEMPLATES = (
- REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "chat_templates.py"
-)
-FORMAT_CONVERSION = (
- REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "format_conversion.py"
-)
+DATA_COLLATORS = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "data_collators.py"
+CHAT_TEMPLATES = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "chat_templates.py"
+FORMAT_CONVERSION = REPO_ROOT / "studio" / "backend" / "utils" / "datasets" / "format_conversion.py"
def _has_uv() -> bool:
@@ -72,9 +66,7 @@ def no_torch_venv(request, tmp_path_factory):
[str(venv_python), "-c", "import torch"],
capture_output = True,
)
- assert (
- check.returncode != 0
- ), f"torch should NOT be importable in fresh {py_version} venv"
+ assert check.returncode != 0, f"torch should NOT be importable in fresh {py_version} venv"
return str(venv_python)
@@ -223,9 +215,7 @@ class TestDataCollatorsNoTorchVenv:
capture_output = True,
timeout = 30,
)
- assert (
- result.returncode == 0
- ), f"DeepSeekOCRDataCollator failed:\n{result.stderr.decode()}"
+ assert result.returncode == 0, f"DeepSeekOCRDataCollator failed:\n{result.stderr.decode()}"
assert b"OK: DeepSeekOCRDataCollator instantiated" in result.stdout
def test_dataclass_vlm_collator_instantiable(self, no_torch_venv):
@@ -246,9 +236,7 @@ class TestDataCollatorsNoTorchVenv:
capture_output = True,
timeout = 30,
)
- assert (
- result.returncode == 0
- ), f"VLMDataCollator failed:\n{result.stderr.decode()}"
+ assert result.returncode == 0, f"VLMDataCollator failed:\n{result.stderr.decode()}"
assert b"OK: VLMDataCollator instantiated" in result.stdout
@@ -529,12 +517,9 @@ class TestNegativeControls:
capture_output = True,
timeout = 30,
)
+ assert result.returncode != 0, "Expected failure when 'import torch' is prepended"
assert (
- result.returncode != 0
- ), "Expected failure when 'import torch' is prepended"
- assert (
- b"ModuleNotFoundError" in result.stderr
- or b"ImportError" in result.stderr
+ b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
), f"Expected ImportError, got:\n{result.stderr.decode()}"
finally:
os.unlink(temp_file)
@@ -577,6 +562,4 @@ class TestNegativeControls:
timeout = 30,
)
assert result.returncode != 0, "import torch should fail in no-torch venv"
- assert (
- b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
- )
+ assert b"ModuleNotFoundError" in result.stderr or b"ImportError" in result.stderr
diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py
index 4be53d2d03..17f2d17f94 100644
--- a/tests/python/test_tokenizers_and_torch_constraint.py
+++ b/tests/python/test_tokenizers_and_torch_constraint.py
@@ -18,9 +18,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
_INSTALL_SH = _REPO_ROOT / "install.sh"
_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
-_NO_TORCH_RT = (
- _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
-)
+_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
def _read(path: pathlib.Path) -> str:
@@ -45,30 +43,23 @@ class TestStructuralTokenizers:
def test_tokenizers_present(self):
"""tokenizers must be a standalone package line."""
pkgs = _lines(_NO_TORCH_RT)
- bare_names = [
- p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
- ]
+ bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
assert "tokenizers" in bare_names
def test_tokenizers_before_transformers(self):
"""tokenizers should appear before transformers (install order intent)."""
pkgs = _lines(_NO_TORCH_RT)
- bare_names = [
- p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
- ]
+ bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
idx_tok = bare_names.index("tokenizers")
idx_tf = bare_names.index("transformers")
assert idx_tok < idx_tf, (
- f"tokenizers at index {idx_tok} should appear before "
- f"transformers at index {idx_tf}"
+ f"tokenizers at index {idx_tok} should appear before " f"transformers at index {idx_tf}"
)
def test_torch_not_in_no_torch_file(self):
"""torch itself must NOT be listed in the no-torch requirements."""
pkgs = _lines(_NO_TORCH_RT)
- bare_names = [
- p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
- ]
+ bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
assert "torch" not in bare_names
@@ -409,9 +400,7 @@ class TestE2ETokenizersFix:
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
- result = self._run_python(
- venv, "from transformers import AutoConfig; print('OK')"
- )
+ result = self._run_python(venv, "from transformers import AutoConfig; print('OK')")
assert (
result.returncode == 0
), f"AutoConfig import failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
@@ -441,22 +430,15 @@ class TestE2ETokenizersFix:
req_no_tokenizers = tmp_path / "no-tokenizers.txt"
req_no_tokenizers.write_text(
"\n".join(
- line
- for line in _read(_NO_TORCH_RT).splitlines()
- if line.strip() != "tokenizers"
+ line for line in _read(_NO_TORCH_RT).splitlines() if line.strip() != "tokenizers"
),
encoding = "utf-8",
)
r = self._pip_install(venv, "--no-deps", "-r", str(req_no_tokenizers))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "from transformers import AutoConfig")
- assert (
- result.returncode != 0
- ), "AutoConfig should fail without tokenizers installed"
- assert (
- "tokenizers" in result.stderr.lower()
- or "ModuleNotFoundError" in result.stderr
- )
+ assert result.returncode != 0, "AutoConfig should fail without tokenizers installed"
+ assert "tokenizers" in result.stderr.lower() or "ModuleNotFoundError" in result.stderr
# ======================================================================
@@ -535,9 +517,7 @@ class TestE2EFullNoTorchSandbox:
venv = self._create_venv(tmp_path, "full-no-torch")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
- result = self._run_python(
- venv, "from transformers import AutoConfig; print('OK')"
- )
+ result = self._run_python(venv, "from transformers import AutoConfig; print('OK')")
assert (
result.returncode == 0
), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py
index 6aff02494b..6e3e3a722d 100644
--- a/tests/python/test_unsloth_run_tool_policy_resolver.py
+++ b/tests/python/test_unsloth_run_tool_policy_resolver.py
@@ -141,15 +141,11 @@ class TestZeroHost:
class TestIsExternalHost:
- @pytest.mark.parametrize(
- "host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"]
- )
+ @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"])
def test_loopback_aliases_are_local(self, host):
assert is_external_host(host) is False
- @pytest.mark.parametrize(
- "host", ["0.0.0.0", "::", "192.168.1.5", "10.0.0.1", "example.com"]
- )
+ @pytest.mark.parametrize("host", ["0.0.0.0", "::", "192.168.1.5", "10.0.0.1", "example.com"])
def test_non_loopback_is_external(self, host):
assert is_external_host(host) is True
diff --git a/tests/qlora/test_hf_qlora_train_and_merge.py b/tests/qlora/test_hf_qlora_train_and_merge.py
index ae975b0266..0892627c46 100644
--- a/tests/qlora/test_hf_qlora_train_and_merge.py
+++ b/tests/qlora/test_hf_qlora_train_and_merge.py
@@ -91,9 +91,7 @@ if __name__ == "__main__":
print(training_args)
print(peft_config)
- trainer = setup_trainer(
- model, tokenizer, dataset, training_args, peft_config = peft_config
- )
+ trainer = setup_trainer(model, tokenizer, dataset, training_args, peft_config = peft_config)
with header_footer_context("Model"):
print(type(model.model))
diff --git a/tests/saving/gpt-oss-merge/test_merged_model.py b/tests/saving/gpt-oss-merge/test_merged_model.py
index 48f0ed2d3d..497c74debf 100644
--- a/tests/saving/gpt-oss-merge/test_merged_model.py
+++ b/tests/saving/gpt-oss-merge/test_merged_model.py
@@ -42,9 +42,7 @@ inputs = merged_tokenizer.apply_chat_template(
reasoning_effort = "low", # **NEW!** Set reasoning effort to low, medium or high
).to(merged_model.device)
-_ = merged_model.generate(
- **inputs, max_new_tokens = 512, streamer = TextStreamer(merged_tokenizer)
-)
+_ = merged_model.generate(**inputs, max_new_tokens = 512, streamer = TextStreamer(merged_tokenizer))
print("\n✅ Inference complete.")
# --- Final Cleanup ---
@@ -54,7 +52,5 @@ torch.cuda.empty_cache()
gc.collect()
safe_remove_directory("./gpt-oss-finetuned-merged")
-safe_remove_directory(
- "./unsloth_compiled_cache"
-) # Clean up cache created by this process
+safe_remove_directory("./unsloth_compiled_cache") # Clean up cache created by this process
print("✅ Final cleanup complete. Exiting inference script.")
diff --git a/tests/saving/gpt-oss-merge/train_and_merge.py b/tests/saving/gpt-oss-merge/train_and_merge.py
index 308d19bfb4..8c76ff9662 100644
--- a/tests/saving/gpt-oss-merge/train_and_merge.py
+++ b/tests/saving/gpt-oss-merge/train_and_merge.py
@@ -28,9 +28,7 @@ tokenizer = None
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -84,9 +82,7 @@ print("Fine-tuning complete.")
# --- Merge and Save ---
print("\n💾 Merging and saving the 16-bit model to './gpt-oss-finetuned-merged'...")
-model.save_pretrained_merged(
- save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer
-)
+model.save_pretrained_merged(save_directory = "./gpt-oss-finetuned-merged", tokenizer = tokenizer)
print("✅ Model merged and saved.")
# --- Cleanup ---
@@ -96,7 +92,5 @@ torch.cuda.empty_cache()
gc.collect()
safe_remove_directory("./outputs")
-safe_remove_directory(
- "./unsloth_compiled_cache"
-) # Clean up the cache created by this process
+safe_remove_directory("./unsloth_compiled_cache") # Clean up the cache created by this process
print("✅ Cleanup complete. Exiting training script.")
diff --git a/tests/saving/language_models/test_merge_4bit_validation.py b/tests/saving/language_models/test_merge_4bit_validation.py
index 343e737710..c889001706 100644
--- a/tests/saving/language_models/test_merge_4bit_validation.py
+++ b/tests/saving/language_models/test_merge_4bit_validation.py
@@ -16,9 +16,7 @@ from tests.utils.cleanup_utils import safe_remove_directory
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -51,9 +49,7 @@ tokenizer = get_chat_template(
)
# Load small dataset for quick training
-dataset_train = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "train[:100]"
-)
+dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train[:100]")
dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
print("✅ Base model loaded successfully!")
diff --git a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
index dd0e8c25c6..3f2b811b07 100644
--- a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
+++ b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
@@ -35,15 +35,17 @@ from tests.utils.perplexity_eval import (
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+ result_queue,
+ load_in_4bit = False,
+ load_in_8bit = False,
+):
"""Load model and compute perplexity in subprocess"""
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
@@ -63,17 +65,13 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
)
# Load dataset fresh in subprocess
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
# Format the dataset
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- merged_tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -130,12 +128,8 @@ if __name__ == "__main__":
from unsloth.chat_templates import standardize_sharegpt
- dataset_train = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "train"
- )
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
diff --git a/tests/saving/language_models/test_merge_model_perplexity_mistral.py b/tests/saving/language_models/test_merge_model_perplexity_mistral.py
index 14e657c68a..d17a20d755 100644
--- a/tests/saving/language_models/test_merge_model_perplexity_mistral.py
+++ b/tests/saving/language_models/test_merge_model_perplexity_mistral.py
@@ -30,7 +30,11 @@ from tests.utils.perplexity_eval import (
)
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+ result_queue,
+ load_in_4bit = False,
+ load_in_8bit = False,
+):
"""Load model and compute perplexity in subprocess"""
from unsloth import FastLanguageModel
from tests.utils.perplexity_eval import ppl_model
@@ -49,9 +53,7 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
# )
# Load dataset fresh in subprocess
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@@ -90,10 +92,7 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
outputs.append(assistant_message)
# Create formatted text
- text = (
- alpaca_prompt.format(instruction, user_message, assistant_message)
- + EOS_TOKEN
- )
+ text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
texts.append(text)
return {
@@ -186,10 +185,7 @@ if __name__ == "__main__":
outputs.append(assistant_message)
# Create formatted text
- text = (
- alpaca_prompt.format(instruction, user_message, assistant_message)
- + EOS_TOKEN
- )
+ text = alpaca_prompt.format(instruction, user_message, assistant_message) + EOS_TOKEN
texts.append(text)
return {
@@ -199,12 +195,8 @@ if __name__ == "__main__":
"text": texts,
}
- dataset_train = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "train"
- )
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
diff --git a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py
index bebea8168e..6dbdf36032 100644
--- a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py
+++ b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py
@@ -35,9 +35,7 @@ from tests.utils.perplexity_eval import (
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {
@@ -45,7 +43,11 @@ def formatting_prompts_func(examples):
}
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+ result_queue,
+ load_in_4bit = False,
+ load_in_8bit = False,
+):
"""Load model and compute perplexity in subprocess"""
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
@@ -65,17 +67,13 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
)
# Load dataset fresh in subprocess
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
# Format the dataset
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- merged_tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -130,12 +128,8 @@ if __name__ == "__main__":
chat_template = "phi-4",
)
- dataset_train = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "train"
- )
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
diff --git a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
index c6da9e2ca6..a0624f0c2c 100644
--- a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
+++ b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
@@ -34,15 +34,17 @@ from tests.utils.perplexity_eval import (
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+ result_queue,
+ load_in_4bit = False,
+ load_in_8bit = False,
+):
"""Load model and compute perplexity in subprocess"""
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
@@ -62,17 +64,13 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
)
# Load dataset fresh in subprocess
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
# Format the dataset
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- merged_tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ merged_tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -129,12 +127,8 @@ if __name__ == "__main__":
from unsloth.chat_templates import standardize_sharegpt
- dataset_train = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "train"
- )
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
diff --git a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
index d63bb9fe09..0b377eca81 100644
--- a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
+++ b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
@@ -78,7 +78,11 @@ def formatting_prompts_func(examples):
}
-def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit = False):
+def load_and_compute_8bit_ppl(
+ result_queue,
+ load_in_4bit = False,
+ load_in_8bit = False,
+):
"""Load model and compute perplexity in subprocess"""
from unsloth import FastLanguageModel
from tests.utils.perplexity_eval import ppl_model
@@ -97,9 +101,7 @@ def load_and_compute_8bit_ppl(result_queue, load_in_4bit = False, load_in_8bit =
# )
# Load dataset fresh in subprocess
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@@ -191,12 +193,8 @@ if __name__ == "__main__":
attn_implementation = attn_implementation,
)
- dataset_train = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "train"
- )
- dataset_ppl = load_dataset(
- "allenai/openassistant-guanaco-reformatted", split = "eval"
- )
+ dataset_train = load_dataset("allenai/openassistant-guanaco-reformatted", split = "train")
+ dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
dataset_train = dataset_train.map(formatting_prompts_func, batched = True)
dataset_ppl = dataset_ppl.map(formatting_prompts_func, batched = True)
diff --git a/tests/saving/language_models/test_push_to_hub_merged.py b/tests/saving/language_models/test_push_to_hub_merged.py
index 58d589305a..aa79394556 100644
--- a/tests/saving/language_models/test_push_to_hub_merged.py
+++ b/tests/saving/language_models/test_push_to_hub_merged.py
@@ -36,9 +36,7 @@ from tests.utils.perplexity_eval import (
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -176,9 +174,7 @@ try:
print("=== TESTING MODEL DOWNLOAD ===".center(80))
print("=" * 80 + "\n")
# Force download even if cached
- model, tokenizer = FastLanguageModel.from_pretrained(
- f"{hf_username}/merged_llama_text_model"
- )
+ model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
success["download"] = True
print("✅ Model downloaded successfully!")
except Exception as e:
diff --git a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
index 038565d170..38b82c5469 100644
--- a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
+++ b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
@@ -36,9 +36,7 @@ from tests.utils.perplexity_eval import (
def formatting_prompts_func(examples):
convos = examples["messages"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {"text": texts}
@@ -195,9 +193,7 @@ try:
print("=== TESTING MODEL DOWNLOAD ===".center(80))
print("=" * 80 + "\n")
# Force download even if cached
- model, tokenizer = FastLanguageModel.from_pretrained(
- f"{hf_username}/merged_llama_text_model"
- )
+ model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
success["download"] = True
print("✅ Model downloaded successfully!")
except Exception as e:
diff --git a/tests/saving/language_models/test_save_merged_grpo_model.py b/tests/saving/language_models/test_save_merged_grpo_model.py
index 67b649305a..b5d8025fcb 100644
--- a/tests/saving/language_models/test_save_merged_grpo_model.py
+++ b/tests/saving/language_models/test_save_merged_grpo_model.py
@@ -24,7 +24,11 @@ max_seq_length = 2048 # Can increase for longer reasoning traces
lora_rank = 64 # Larger rank = smarter, but slower
-def evaluate_merged_model(result_queue, load_in_4bit = False, load_in_8bit = False):
+def evaluate_merged_model(
+ result_queue,
+ load_in_4bit = False,
+ load_in_8bit = False,
+):
from unsloth import FastLanguageModel
from tests.utils.aime_eval import evaluate_model_aime
@@ -176,12 +180,14 @@ def training_run(result_queue):
avg_length = sum(lengths) / len(lengths)
min_length = min(lengths)
- print(
- f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}"
- )
+ print(f"Prompt lengths - Min: {min_length}, Max: {max_length}, Avg: {avg_length:.1f}")
return max_length, avg_length
- def extract_unsloth_answer(text, start_tag = "", end_tag = ""):
+ def extract_unsloth_answer(
+ text,
+ start_tag = "",
+ end_tag = "",
+ ):
"""Extract answer from Unsloth SOLUTION tags"""
pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag)
matches = re.findall(pattern, text, re.DOTALL)
@@ -265,9 +271,7 @@ def training_run(result_queue):
ground_truth_num = float(norm_ground_truth)
if ground_truth_num != 0:
- relative_error = abs(extracted_num - ground_truth_num) / abs(
- ground_truth_num
- )
+ relative_error = abs(extracted_num - ground_truth_num) / abs(ground_truth_num)
if relative_error < 0.01:
return True, True, 0.9
@@ -302,10 +306,7 @@ def training_run(result_queue):
)
responses = [completion[0]["content"] for completion in completions]
- rewards = [
- 3.0 if re.match(pattern, response, re.DOTALL) else 0.0
- for response in responses
- ]
+ rewards = [3.0 if re.match(pattern, response, re.DOTALL) else 0.0 for response in responses]
return rewards
def match_format_approximately(completions, **kwargs):
@@ -405,9 +406,7 @@ def training_run(result_queue):
format_improvement = (
result["correct_format_pct"] - base_result["correct_format_pct"]
)
- exact_improvement = (
- result["exact_match_pct"] - base_result["exact_match_pct"]
- )
+ exact_improvement = result["exact_match_pct"] - base_result["exact_match_pct"]
plausible_improvement = (
result["plausible_match_pct"] - base_result["plausible_match_pct"]
)
@@ -440,9 +439,7 @@ def training_run(result_queue):
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3
- print(
- f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB"
- )
+ print(f"GPU memory - Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB")
"""#### Data Loading and Preparation"""
@@ -486,9 +483,7 @@ def training_run(result_queue):
def formatting_prompts_func(examples):
convos = examples["prompt"]
texts = [
- tokenizer.apply_chat_template(
- convo, tokenize = False, add_generation_prompt = False
- )
+ tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False)
for convo in convos
]
return {
@@ -715,9 +710,7 @@ def training_run(result_queue):
# Save as merged model
try:
- model.save_pretrained_merged(
- "final_merged_model", tokenizer, save_method = "merged_16bit"
- )
+ model.save_pretrained_merged("final_merged_model", tokenizer, save_method = "merged_16bit")
print("✅ Merged model saved to: final_merged_model/")
except Exception as e:
print(f"⚠️ Could not save merged model: {e}")
diff --git a/tests/saving/test_fix_sentencepiece_gguf_robustness.py b/tests/saving/test_fix_sentencepiece_gguf_robustness.py
index 49bd70fa2f..9c61ca4067 100644
--- a/tests/saving/test_fix_sentencepiece_gguf_robustness.py
+++ b/tests/saving/test_fix_sentencepiece_gguf_robustness.py
@@ -44,9 +44,7 @@ def test_user_defined_special_piece_is_not_retyped(tmp_path):
]
(tmp_path / "tokenizer.model").write_bytes(_build(pieces))
(tmp_path / "tokenizer.json").write_text(
- json.dumps(
- {"added_tokens": [{"id": 2, "content": "", "special": True}]}
- )
+ json.dumps({"added_tokens": [{"id": 2, "content": "", "special": True}]})
)
fix_sentencepiece_gguf(str(tmp_path))
got = dict(_read(str(tmp_path / "tokenizer.model")))
@@ -87,10 +85,7 @@ def test_save_py_except_clause_is_broad_exception():
with open(_SAVE_PY) as f:
tree = ast.parse(f.read())
for node in ast.walk(tree):
- if (
- isinstance(node, ast.FunctionDef)
- and node.name == "unsloth_save_pretrained_gguf"
- ):
+ if isinstance(node, ast.FunctionDef) and node.name == "unsloth_save_pretrained_gguf":
for subnode in ast.walk(node):
if isinstance(subnode, ast.Try):
body_src = "\n".join(ast.unparse(s) for s in subnode.body)
diff --git a/tests/saving/test_preserve_tokenizer_eos_token.py b/tests/saving/test_preserve_tokenizer_eos_token.py
index 6e2f8c7f9d..2ea40ab778 100644
--- a/tests/saving/test_preserve_tokenizer_eos_token.py
+++ b/tests/saving/test_preserve_tokenizer_eos_token.py
@@ -16,8 +16,7 @@ def _load_preserve_helper():
helper = next(
node
for node in tree.body
- if isinstance(node, ast.FunctionDef)
- and node.name == "_preserve_tokenizer_eos_token"
+ if isinstance(node, ast.FunctionDef) and node.name == "_preserve_tokenizer_eos_token"
)
module = ast.Module(body = [helper], type_ignores = [])
ast.fix_missing_locations(module)
@@ -46,9 +45,7 @@ def test_preserve_tokenizer_eos_token_supports_processor_tokenizer(tmp_path):
preserve = _load_preserve_helper()
tokenizer_config = tmp_path / "tokenizer_config.json"
tokenizer_config.write_text(json.dumps({"eos_token": ""}), encoding = "utf-8")
- processor = types.SimpleNamespace(
- tokenizer = types.SimpleNamespace(eos_token = "")
- )
+ processor = types.SimpleNamespace(tokenizer = types.SimpleNamespace(eos_token = ""))
preserve(processor, tmp_path)
diff --git a/tests/saving/test_save_shell_injection.py b/tests/saving/test_save_shell_injection.py
index c6c2c8fe15..b02748c250 100644
--- a/tests/saving/test_save_shell_injection.py
+++ b/tests/saving/test_save_shell_injection.py
@@ -19,10 +19,7 @@ def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None:
popen_calls = []
for call in calls:
if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen":
- if (
- isinstance(call.func.value, ast.Name)
- and call.func.value.id == "subprocess"
- ):
+ if isinstance(call.func.value, ast.Name) and call.func.value.id == "subprocess":
popen_calls.append(call)
assert popen_calls, "Expected at least one subprocess.Popen call"
@@ -54,9 +51,7 @@ def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None:
assert call.args, "subprocess.Popen must receive argv as a positional argument"
argv = call.args[0]
- assert isinstance(
- argv, ast.List
- ), "subprocess.Popen must be called with an argv list"
+ assert isinstance(argv, ast.List), "subprocess.Popen must be called with an argv list"
assert len(argv.elts) == 5, "GGML conversion argv should have five elements"
second_arg = argv.elts[1]
diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py
index 35fdad6ba0..c7dd712734 100644
--- a/tests/saving/test_unsloth_save.py
+++ b/tests/saving/test_unsloth_save.py
@@ -132,20 +132,14 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
model.config._name_or_path.replace("/", "_"),
)
- model.save_pretrained_merged(
- save_path, tokenizer = tokenizer, save_method = "merged_16bit"
- )
+ model.save_pretrained_merged(save_path, tokenizer = tokenizer, save_method = "merged_16bit")
# Check model files
assert os.path.isdir(save_path), f"Directory {save_path} does not exist."
- assert os.path.isfile(
- os.path.join(save_path, "config.json")
- ), "config.json not found."
+ assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found."
weight_files = [
- f
- for f in os.listdir(save_path)
- if f.endswith(".bin") or f.endswith(".safetensors")
+ f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")
]
assert len(weight_files) > 0, "No weight files found in the save directory."
@@ -160,9 +154,7 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
with open(config_path, "r") as f:
config = json.load(f)
- assert (
- "quantization_config" not in config
- ), "Quantization config not found in the model config."
+ assert "quantization_config" not in config, "Quantization config not found in the model config."
# Store the size of the model files
total_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files)
@@ -185,20 +177,14 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
model.config._name_or_path.replace("/", "_"),
)
- model.save_pretrained_merged(
- save_path, tokenizer = tokenizer, save_method = "merged_4bit_forced"
- )
+ model.save_pretrained_merged(save_path, tokenizer = tokenizer, save_method = "merged_4bit_forced")
# Check model files
assert os.path.isdir(save_path), f"Directory {save_path} does not exist."
- assert os.path.isfile(
- os.path.join(save_path, "config.json")
- ), "config.json not found."
+ assert os.path.isfile(os.path.join(save_path, "config.json")), "config.json not found."
weight_files = [
- f
- for f in os.listdir(save_path)
- if f.endswith(".bin") or f.endswith(".safetensors")
+ f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")
]
assert len(weight_files) > 0, "No weight files found in the save directory."
@@ -223,9 +209,7 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
with open(config_path, "r") as f:
config = json.load(f)
- assert (
- "quantization_config" in config
- ), "Quantization config not found in the model config."
+ assert "quantization_config" in config, "Quantization config not found in the model config."
# Test loading the model from the saved path
loaded_model, loaded_tokenizer = FastModel.from_pretrained(
@@ -257,29 +241,19 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
)
weight_files_16bit = [
- f
- for f in os.listdir(save_path)
- if f.endswith(".bin") or f.endswith(".safetensors")
+ f for f in os.listdir(save_path) if f.endswith(".bin") or f.endswith(".safetensors")
]
- total_16bit_size = sum(
- os.path.getsize(os.path.join(save_path, f)) for f in weight_files_16bit
- )
+ total_16bit_size = sum(os.path.getsize(os.path.join(save_path, f)) for f in weight_files_16bit)
save_file_sizes["merged_16bit"][model.config._name_or_path] = total_16bit_size
torchao_save_path = save_path + "-torchao"
# Check model files
- assert os.path.isdir(
- torchao_save_path
- ), f"Directory {torchao_save_path} does not exist."
- assert os.path.isfile(
- os.path.join(torchao_save_path, "config.json")
- ), "config.json not found."
+ assert os.path.isdir(torchao_save_path), f"Directory {torchao_save_path} does not exist."
+ assert os.path.isfile(os.path.join(torchao_save_path, "config.json")), "config.json not found."
weight_files = [
- f
- for f in os.listdir(torchao_save_path)
- if f.endswith(".bin") or f.endswith(".safetensors")
+ f for f in os.listdir(torchao_save_path) if f.endswith(".bin") or f.endswith(".safetensors")
]
assert len(weight_files) > 0, "No weight files found in the save directory."
@@ -290,9 +264,7 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
), f"{file} not found in the save directory."
# Store the size of the model files
- total_size = sum(
- os.path.getsize(os.path.join(torchao_save_path, f)) for f in weight_files
- )
+ total_size = sum(os.path.getsize(os.path.join(torchao_save_path, f)) for f in weight_files)
save_file_sizes["torchao"][model.config._name_or_path] = total_size
assert (
@@ -304,9 +276,7 @@ def test_save_torchao(fp16_model_tokenizer, temp_save_dir: str):
with open(config_path, "r") as f:
config = json.load(f)
- assert (
- "quantization_config" in config
- ), "Quantization config not found in the model config."
+ assert "quantization_config" in config, "Quantization config not found in the model config."
# Test loading the model from the saved path
# can't set `load_in_4bit` to True because the model is torchao quantized
@@ -332,9 +302,7 @@ def test_save_and_inference_torchao(fp16_model_tokenizer, temp_save_dir: str):
print(f"Testing TorchAO save and inference for: {model_name}")
- save_path = os.path.join(
- temp_save_dir, "torchao_models", model_name.replace("/", "_")
- )
+ save_path = os.path.join(temp_save_dir, "torchao_models", model_name.replace("/", "_"))
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
diff --git a/tests/saving/text_to_speech_models/test_csm.py b/tests/saving/text_to_speech_models/test_csm.py
index c1a892a8d3..dd2287d1d6 100644
--- a/tests/saving/text_to_speech_models/test_csm.py
+++ b/tests/saving/text_to_speech_models/test_csm.py
@@ -134,9 +134,7 @@ import torch
output_audio_path = "csm_audio.wav"
try:
- text = (
- "We just finished fine tuning a text to speech model... and it's pretty good!"
- )
+ text = "We just finished fine tuning a text to speech model... and it's pretty good!"
speaker_id = 0
inputs = processor(f"[{speaker_id}]{text}", add_special_tokens = True).to("cuda")
audio_values = model.generate(
diff --git a/tests/saving/text_to_speech_models/test_lasa.py b/tests/saving/text_to_speech_models/test_lasa.py
index 804ff512f9..c0c4f80e0e 100644
--- a/tests/saving/text_to_speech_models/test_lasa.py
+++ b/tests/saving/text_to_speech_models/test_lasa.py
@@ -167,9 +167,7 @@ def extract_speech_ids(speech_tokens_str):
# TTS start!
with torch.inference_mode():
with torch.amp.autocast("cuda", dtype = model.dtype):
- formatted_text = (
- f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
- )
+ formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
# Tokenize the text
chat = [
diff --git a/tests/saving/text_to_speech_models/test_orpheus.py b/tests/saving/text_to_speech_models/test_orpheus.py
index bd8bf14979..2915749d99 100644
--- a/tests/saving/text_to_speech_models/test_orpheus.py
+++ b/tests/saving/text_to_speech_models/test_orpheus.py
@@ -152,9 +152,7 @@ for prompt in prompts_:
all_input_ids.append(input_ids)
start_token = torch.tensor([[128259]], dtype = torch.int64) # Start of human
-end_tokens = torch.tensor(
- [[128009, 128260]], dtype = torch.int64
-) # End of text, End of human
+end_tokens = torch.tensor([[128009, 128260]], dtype = torch.int64) # End of text, End of human
all_modified_input_ids = []
for input_ids in all_input_ids:
@@ -165,9 +163,7 @@ for input_ids in all_input_ids:
all_padded_tensors = []
all_attention_masks = []
-max_length = max(
- [modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids]
-)
+max_length = max([modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids])
for modified_input_ids in all_modified_input_ids:
padding = max_length - modified_input_ids.shape[1]
padded_tensor = torch.cat(
diff --git a/tests/saving/text_to_speech_models/test_whisper.py b/tests/saving/text_to_speech_models/test_whisper.py
index 55f6d98ca0..d0eeb49d17 100644
--- a/tests/saving/text_to_speech_models/test_whisper.py
+++ b/tests/saving/text_to_speech_models/test_whisper.py
@@ -181,13 +181,9 @@ expected_phrases = [
]
transcribed_lower = transcribed_text["text"].lower()
-all_phrases_found = all(
- phrase.lower() in transcribed_lower for phrase in expected_phrases
-)
+all_phrases_found = all(phrase.lower() in transcribed_lower for phrase in expected_phrases)
-assert (
- all_phrases_found
-), f"Expected phrases not found in transcription: {transcribed_text['text']}"
+assert all_phrases_found, f"Expected phrases not found in transcription: {transcribed_text['text']}"
print("✅ Transcription contains all expected phrases!")
diff --git a/tests/saving/vision_models/test_index_file_sharded_model.py b/tests/saving/vision_models/test_index_file_sharded_model.py
index 8d107463e0..79a25ec666 100644
--- a/tests/saving/vision_models/test_index_file_sharded_model.py
+++ b/tests/saving/vision_models/test_index_file_sharded_model.py
@@ -138,9 +138,7 @@ try:
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
gradient_checkpointing = True,
- gradient_checkpointing_kwargs = {
- "use_reentrant": False
- }, # use reentrant checkpointing
+ gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
max_grad_norm = 0.3, # max gradient norm based on QLoRA paper
warmup_ratio = 0.03,
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
diff --git a/tests/saving/vision_models/test_push_to_hub_merged.py b/tests/saving/vision_models/test_push_to_hub_merged.py
index fb2af4b4fe..86c7b56cf2 100644
--- a/tests/saving/vision_models/test_push_to_hub_merged.py
+++ b/tests/saving/vision_models/test_push_to_hub_merged.py
@@ -139,9 +139,7 @@ try:
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
gradient_checkpointing = True,
- gradient_checkpointing_kwargs = {
- "use_reentrant": False
- }, # use reentrant checkpointing
+ gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
max_grad_norm = 0.3, # max gradient norm based on QLoRA paper
warmup_ratio = 0.03,
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
diff --git a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
index 2b24bc4a32..391d7bacca 100644
--- a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
+++ b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
@@ -134,9 +134,7 @@ trainer = SFTTrainer(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
gradient_checkpointing = True,
- gradient_checkpointing_kwargs = {
- "use_reentrant": False
- }, # use reentrant checkpointing
+ gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
max_grad_norm = 0.3, # max gradient norm based on QLoRA paper
warmup_ratio = 0.03,
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
diff --git a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
index 16914707c2..b4812cdfa8 100644
--- a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
+++ b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
@@ -134,9 +134,7 @@ trainer = SFTTrainer(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
gradient_checkpointing = True,
- gradient_checkpointing_kwargs = {
- "use_reentrant": False
- }, # use reentrant checkpointing
+ gradient_checkpointing_kwargs = {"use_reentrant": False}, # use reentrant checkpointing
max_grad_norm = 0.3, # max gradient norm based on QLoRA paper
warmup_ratio = 0.03,
# num_train_epochs = 2, # Set this instead of max_steps for full training runs
diff --git a/tests/security/test_lockfile_supply_chain_audit.py b/tests/security/test_lockfile_supply_chain_audit.py
index 483bb9e763..cec07aea28 100644
--- a/tests/security/test_lockfile_supply_chain_audit.py
+++ b/tests/security/test_lockfile_supply_chain_audit.py
@@ -143,7 +143,6 @@ def test_lockfile_auditor_blocked_versions_match_scanner():
comment until the next PR factors them into a shared module).
"""
from scripts import scan_npm_packages as snp
-
assert (
lsa.BLOCKED_NPM_VERSIONS == snp.BLOCKED_NPM_VERSIONS
), "auditor and scanner BLOCKED_NPM_VERSIONS tables drifted"
@@ -275,9 +274,7 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
npm_lockfiles = [FIXTURES / "clean_lockfile.json"],
cargo_lockfiles = [lockfile],
)
- warning_lines = [
- line for line in proc.stderr.splitlines() if line.startswith("::warning::")
- ]
+ warning_lines = [line for line in proc.stderr.splitlines() if line.startswith("::warning::")]
assert warning_lines, (
"expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}"
)
diff --git a/tests/security/test_new_install_scripts.py b/tests/security/test_new_install_scripts.py
index 32340d2536..9a73f4b0d9 100644
--- a/tests/security/test_new_install_scripts.py
+++ b/tests/security/test_new_install_scripts.py
@@ -18,7 +18,12 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
SCRIPT = REPO_ROOT / "scripts" / "check_new_install_scripts.py"
-def _run(base: Path, head: Path, *, timeout: int = 30) -> subprocess.CompletedProcess:
+def _run(
+ base: Path,
+ head: Path,
+ *,
+ timeout: int = 30,
+) -> subprocess.CompletedProcess:
return subprocess.run(
[
sys.executable,
@@ -103,9 +108,7 @@ def test_new_dep_with_postinstall_exits_1(tmp_path: Path):
head_pkgs = dict(base_pkgs)
head_pkgs["node_modules/evil-postinstall"] = {
"version": "1.0.0",
- "resolved": (
- "https://registry.npmjs.org/evil-postinstall/-/evil-postinstall-1.0.0.tgz"
- ),
+ "resolved": ("https://registry.npmjs.org/evil-postinstall/-/evil-postinstall-1.0.0.tgz"),
"integrity": "sha512-fake",
"hasInstallScript": True,
}
@@ -166,8 +169,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
"node_modules/v2-postinstall-dep": {
"version": "2.0.0",
"resolved": (
- "https://registry.npmjs.org/v2-postinstall-dep/-/"
- "v2-postinstall-dep-2.0.0.tgz"
+ "https://registry.npmjs.org/v2-postinstall-dep/-/v2-postinstall-dep-2.0.0.tgz"
),
"integrity": "sha512-fake",
"hasInstallScript": True,
@@ -177,8 +179,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
"v2-postinstall-dep": {
"version": "2.0.0",
"resolved": (
- "https://registry.npmjs.org/v2-postinstall-dep/-/"
- "v2-postinstall-dep-2.0.0.tgz"
+ "https://registry.npmjs.org/v2-postinstall-dep/-/v2-postinstall-dep-2.0.0.tgz"
),
"integrity": "sha512-fake",
},
@@ -187,8 +188,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
head = _write(tmp_path / "head.json", _v2_lockfile(head_pkgs, head_deps))
result = _run(base, head)
assert result.returncode == 1, (
- f"expected exit 1 for v2 lockfile, got {result.returncode}; "
- f"stderr:\n{result.stderr}"
+ f"expected exit 1 for v2 lockfile, got {result.returncode}; " f"stderr:\n{result.stderr}"
)
assert "v2-postinstall-dep" in result.stderr
diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py
index c632c618b0..fb575b730d 100644
--- a/tests/security/test_scan_npm_packages.py
+++ b/tests/security/test_scan_npm_packages.py
@@ -106,16 +106,10 @@ def test_blocked_npm_versions_complete():
table = snp.BLOCKED_NPM_VERSIONS
tanstack_keys = [k for k in table if k.startswith("@tanstack/")]
assert len(tanstack_keys) == 42, (
- f"expected 42 @tanstack/* entries, got {len(tanstack_keys)}: "
- f"{sorted(tanstack_keys)}"
+ f"expected 42 @tanstack/* entries, got {len(tanstack_keys)}: " f"{sorted(tanstack_keys)}"
)
assert "@opensearch-project/opensearch" in table
- assert table["@opensearch-project/opensearch"] == {
- "3.5.3",
- "3.6.2",
- "3.7.0",
- "3.8.0",
- }
+ assert table["@opensearch-project/opensearch"] == {"3.5.3", "3.6.2", "3.7.0", "3.8.0"}
squawk = [k for k in table if k.startswith("@squawk/")]
assert len(squawk) >= 22, (
f"expected at least 22 @squawk/* entries (full safedep.io enumeration), "
diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py
index 6ef10f12eb..b35d89ce48 100644
--- a/tests/security/test_scan_packages.py
+++ b/tests/security/test_scan_packages.py
@@ -118,9 +118,7 @@ def test_clean_wheel_no_findings():
str(FIXTURES / "clean_wheel.whl"),
"clean_fixture",
)
- assert (
- findings == []
- ), f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
+ assert findings == [], f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
# ---------------------------------------------------------------------------
@@ -245,8 +243,7 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
assert findings, "scan_archive returned 0 findings on corrupt wheel"
corrupted = [f for f in findings if f.check == "archive_corrupted"]
assert corrupted, (
- "no archive_corrupted finding; got "
- f"{[(f.severity, f.check) for f in findings]}"
+ "no archive_corrupted finding; got " f"{[(f.severity, f.check) for f in findings]}"
)
assert all(f.severity == sp.CRITICAL for f in corrupted)
diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py
index b190b2b3e1..831153bf96 100644
--- a/tests/studio/_playwright_robust.py
+++ b/tests/studio/_playwright_robust.py
@@ -182,9 +182,7 @@ def wait_for_health(
# but accept any 200 -- different Studio builds report differently.
if status == 200:
if info is not None:
- info(
- f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}"
- )
+ info(f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}")
return True
time.sleep(0.5)
if info is not None:
@@ -230,9 +228,7 @@ def recover_or_replace_page(
info(f"recovery: page.is_closed() check failed: {exc!r}")
if goto_url is not None:
try:
- page.goto(
- goto_url, wait_until = "domcontentloaded", timeout = default_timeout_ms
- )
+ page.goto(goto_url, wait_until = "domcontentloaded", timeout = default_timeout_ms)
if settle_networkidle:
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
diff --git a/tests/studio/install/smoke_test_llama_prebuilt.py b/tests/studio/install/smoke_test_llama_prebuilt.py
index d87537dc94..f7fd58aaa4 100644
--- a/tests/studio/install/smoke_test_llama_prebuilt.py
+++ b/tests/studio/install/smoke_test_llama_prebuilt.py
@@ -15,9 +15,7 @@ INSTALLER_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
def load_installer_module():
- spec = importlib.util.spec_from_file_location(
- "studio_install_llama_prebuilt", INSTALLER_PATH
- )
+ spec = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", INSTALLER_PATH)
if spec is None or spec.loader is None:
raise RuntimeError(f"unable to load installer module from {INSTALLER_PATH}")
module = importlib.util.module_from_spec(spec)
@@ -112,17 +110,13 @@ def main() -> int:
published_release_tag = args.published_release_tag,
)
print(f"[smoke] PASS install_dir={install_dir}")
- print(
- "[smoke] note=This was a real prebuilt install into an isolated temp directory."
- )
+ print("[smoke] note=This was a real prebuilt install into an isolated temp directory.")
return installer.EXIT_SUCCESS
except SystemExit as exc:
code = int(exc.code) if isinstance(exc.code, int) else installer.EXIT_ERROR
if code == installer.EXIT_FALLBACK:
print(f"[smoke] FALLBACK install_dir={install_dir}")
- print(
- "[smoke] note=Prebuilt path failed and would fall back to source build in setup."
- )
+ print("[smoke] note=Prebuilt path failed and would fall back to source build in setup.")
print(installer.collect_system_report(host, choice, install_dir))
else:
print(f"[smoke] ERROR exit_code={code} install_dir={install_dir}")
diff --git a/tests/studio/install/smoke_test_parallel_studio_home.py b/tests/studio/install/smoke_test_parallel_studio_home.py
index 133591fb33..c4fc250655 100644
--- a/tests/studio/install/smoke_test_parallel_studio_home.py
+++ b/tests/studio/install/smoke_test_parallel_studio_home.py
@@ -86,12 +86,7 @@ def _free_port() -> int:
def _run_one_install(
- label: str,
- repo: Path,
- studio_home: Path,
- fake_home: Path,
- uv_cache: Path,
- log_path: Path,
+ label: str, repo: Path, studio_home: Path, fake_home: Path, uv_cache: Path, log_path: Path
) -> tuple[str, int]:
studio_home.mkdir(parents = True, exist_ok = True)
fake_home.mkdir(parents = True, exist_ok = True)
@@ -159,12 +154,14 @@ def _wait_for_health(port: int, timeout: float) -> dict:
except (urllib.error.URLError, ConnectionError, OSError) as e:
last_err = e
time.sleep(HEALTH_POLL_INTERVAL_S)
- raise TestFailure(
- f"port {port}: /api/health never returned 200 (last_err={last_err})"
- )
+ raise TestFailure(f"port {port}: /api/health never returned 200 (last_err={last_err})")
-def _http_status(port: int, path: str, timeout: float = 5.0) -> int:
+def _http_status(
+ port: int,
+ path: str,
+ timeout: float = 5.0,
+) -> int:
url = f"http://127.0.0.1:{port}{path}"
try:
with urllib.request.urlopen(url, timeout = timeout) as r:
@@ -211,9 +208,7 @@ def _check_install_layout(label: str, studio_home: Path) -> dict:
raise TestFailure(f"[{label}] launch-studio.sh kept @@DATA_DIR@@ placeholder")
expected_data_dir_line = f"DATA_DIR='{studio_home}/share'"
if expected_data_dir_line not in launcher:
- raise TestFailure(
- f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}"
- )
+ raise TestFailure(f"[{label}] launch-studio.sh missing {expected_data_dir_line!r}")
return {"label": label, "studio_home": str(studio_home), "install_id": install_id}
@@ -230,9 +225,7 @@ def _check_fake_home_clean(fake_home: Path) -> None:
]
leaked = [str(p) for p in forbidden if (fake_home / p).exists()]
if leaked:
- raise TestFailure(
- f"redirected HOME picked up persistent install pollution: {leaked}"
- )
+ raise TestFailure(f"redirected HOME picked up persistent install pollution: {leaked}")
def _backend_pid_python(pid: int) -> Path | None:
@@ -256,9 +249,7 @@ def run(n_installs: int, keep: bool) -> int:
repo = PACKAGE_ROOT
if not (repo / "install.sh").is_file():
- raise TestFailure(
- f"install.sh not found at {repo}; " "run from a clone of unslothai/unsloth"
- )
+ raise TestFailure(f"install.sh not found at {repo}; run from a clone of unslothai/unsloth")
test_root = Path(tempfile.mkdtemp(prefix = "unsloth_studio_clash_"))
_log(f"test root: {test_root}")
@@ -346,8 +337,7 @@ def run(n_installs: int, keep: bool) -> int:
raise TestFailure(f"[{label}] chat_only is not true under --no-torch")
if health["studio_root_id"] in seen_root_ids:
raise TestFailure(
- f"[{label}] studio_root_id collision at runtime: "
- f"{health['studio_root_id']}"
+ f"[{label}] studio_root_id collision at runtime: " f"{health['studio_root_id']}"
)
seen_root_ids.add(health["studio_root_id"])
@@ -358,9 +348,7 @@ def run(n_installs: int, keep: bool) -> int:
exe = _backend_pid_python(proc.pid)
if exe is not None:
- expected_python = (
- studio_home / "unsloth_studio" / "bin" / "python"
- ).resolve()
+ expected_python = (studio_home / "unsloth_studio" / "bin" / "python").resolve()
if exe != expected_python:
raise TestFailure(
f"[{label}] PID {proc.pid} exe={exe}, expected {expected_python}"
@@ -370,10 +358,7 @@ def run(n_installs: int, keep: bool) -> int:
if len(versions) != 1:
raise TestFailure(f"version mismatch across installs: {versions}")
- _log(
- f"PASS: all install + runtime invariants hold "
- f"(version={next(iter(versions))})"
- )
+ _log(f"PASS: all install + runtime invariants hold " f"(version={next(iter(versions))})")
return 0
except TestFailure as e:
diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py
index 4a427d6f54..de3808469c 100644
--- a/tests/studio/install/test_install_llama_prebuilt_logic.py
+++ b/tests/studio/install/test_install_llama_prebuilt_logic.py
@@ -12,9 +12,7 @@ import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
-SPEC = importlib.util.spec_from_file_location(
- "studio_install_llama_prebuilt", MODULE_PATH
-)
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
assert SPEC is not None and SPEC.loader is not None
INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
@@ -274,12 +272,8 @@ def test_validate_prebuilt_choice_creates_repo_shaped_linux_install(
"preflight_linux_installed_binaries",
lambda *args, **kwargs: None,
)
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None
- )
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None
- )
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None)
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None)
host = HostInfo(
system = "Linux",
@@ -381,9 +375,7 @@ def test_simple_linux_direct_release_uses_published_source_checksums_for_branch(
INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
source_commit
): ApprovedArtifactHash(
- asset_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
- source_commit
- ),
+ asset_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(source_commit),
sha256 = "b" * 64,
repo = "ggml-org/llama.cpp",
kind = "exact-source",
@@ -429,9 +421,7 @@ def test_simple_linux_direct_release_uses_published_source_checksums_for_branch(
assert plan.approved_checksums.source_commit == source_commit
assert plan.attempts[0].expected_sha256 == "a" * 64
source_repo, source_ref, _source_archive, exact_source = (
- INSTALL_LLAMA_PREBUILT.preferred_source_archive(
- plan.approved_checksums, plan.llama_tag
- )
+ INSTALL_LLAMA_PREBUILT.preferred_source_archive(plan.approved_checksums, plan.llama_tag)
)
assert source_repo == "ggml-org/llama.cpp"
assert source_ref == source_commit
@@ -468,9 +458,7 @@ def test_simple_linux_direct_release_honors_torch_cudart_preference(
["cuda13", "cuda12"],
{
"cuda13": ["/usr/local/lib/python3.13/site-packages/nvidia/cu13/lib"],
- "cuda12": [
- "/venv/lib/python3.13/site-packages/nvidia/cuda_runtime/lib"
- ],
+ "cuda12": ["/venv/lib/python3.13/site-packages/nvidia/cuda_runtime/lib"],
},
),
)
@@ -517,24 +505,20 @@ def test_simple_linux_direct_release_honors_torch_cudart_preference(
[
# Missing source_commit.
(
- lambda c: setattr(c, "source_commit", None)
- or setattr(c, "source_commit_short", None),
+ lambda c: setattr(c, "source_commit", None) or setattr(c, "source_commit_short", None),
"exact source provenance",
),
# source_commit present, but no exact-source archive hash.
(
lambda c: c.artifacts.pop(
- INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
- c.source_commit
- ),
+ INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(c.source_commit),
None,
),
"exact source provenance",
),
# source_commit + exact-source archive present, but no source_repo.
(
- lambda c: setattr(c, "source_repo", None)
- or setattr(c, "source_repo_url", None),
+ lambda c: setattr(c, "source_repo", None) or setattr(c, "source_repo_url", None),
"exact source provenance",
),
],
@@ -545,9 +529,7 @@ def test_simple_linux_direct_release_honors_torch_cudart_preference(
],
)
def test_simple_linux_direct_release_rejects_branch_without_exact_source_metadata(
- monkeypatch: pytest.MonkeyPatch,
- mutate,
- expected_match,
+ monkeypatch: pytest.MonkeyPatch, mutate, expected_match
):
source_commit = "25b1bc9c2f9aa0a390b968ee1ffd9ff01340a3fe"
release = {
@@ -583,9 +565,7 @@ def test_simple_linux_direct_release_rejects_branch_without_exact_source_metadat
INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
source_commit
): ApprovedArtifactHash(
- asset_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(
- source_commit
- ),
+ asset_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name(source_commit),
sha256 = "b" * 64,
repo = "ggml-org/llama.cpp",
kind = "exact-source",
@@ -646,9 +626,7 @@ def test_simple_linux_direct_release_keeps_legacy_b_tag_path_without_checksums(
}
def unexpected_checksum_load(repo: str, release_tag: str):
- raise AssertionError(
- "legacy b-tag direct releases should not require checksum metadata"
- )
+ raise AssertionError("legacy b-tag direct releases should not require checksum metadata")
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
@@ -741,12 +719,8 @@ def test_validate_prebuilt_choice_creates_repo_shaped_windows_install(
"preflight_linux_installed_binaries",
lambda *args, **kwargs: None,
)
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None
- )
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None
- )
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_quantize", lambda *args, **kwargs: None)
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_server", lambda *args, **kwargs: None)
host = HostInfo(
system = "Windows",
@@ -809,9 +783,7 @@ def test_validate_prebuilt_choice_creates_repo_shaped_windows_install(
def test_activate_install_tree_restores_existing_install_after_activation_failure(
- tmp_path: Path,
- monkeypatch: pytest.MonkeyPatch,
- capsys: pytest.CaptureFixture[str],
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
@@ -839,9 +811,7 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"confirm_install_tree",
- lambda *_args, **_kwargs: (_ for _ in ()).throw(
- RuntimeError("activation confirm failed")
- ),
+ lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("activation confirm failed")),
)
with pytest.raises(
@@ -862,9 +832,7 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur
def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
- tmp_path: Path,
- monkeypatch: pytest.MonkeyPatch,
- capsys: pytest.CaptureFixture[str],
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
@@ -892,9 +860,7 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"confirm_install_tree",
- lambda *_args, **_kwargs: (_ for _ in ()).throw(
- RuntimeError("activation confirm failed")
- ),
+ lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("activation confirm failed")),
)
original_replace = INSTALL_LLAMA_PREBUILT.os.replace
@@ -921,10 +887,7 @@ def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails(
captured = capsys.readouterr()
output = captured.out + captured.err
assert "rollback after failed activation also failed: restore failed" in output
- assert (
- "cleaning staging, install, and rollback paths before source build fallback"
- in output
- )
+ assert "cleaning staging, install, and rollback paths before source build fallback" in output
assert "removing failed install path" in output
assert "removing rollback path" in output
@@ -1108,9 +1071,7 @@ def write_linux_install_shape(install_dir: Path) -> None:
(runtime_dir / "libggml-base.so.0").write_bytes(b"DLL")
(runtime_dir / "libggml-cpu-x64.so.0").write_bytes(b"DLL")
(runtime_dir / "libmtmd.so.0").write_bytes(b"DLL")
- (install_dir / "convert_hf_to_gguf.py").write_text(
- "#!/usr/bin/env python3\n", encoding = "utf-8"
- )
+ (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8")
(install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True)
@@ -1134,9 +1095,7 @@ def write_windows_install_shape(
(runtime_dir / "cudart64_12.dll").write_bytes(b"DLL")
(runtime_dir / "cublas64_12.dll").write_bytes(b"DLL")
(runtime_dir / "cublasLt64_12.dll").write_bytes(b"DLL")
- (install_dir / "convert_hf_to_gguf.py").write_text(
- "#!/usr/bin/env python3\n", encoding = "utf-8"
- )
+ (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8")
(install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True)
@@ -1159,9 +1118,7 @@ def write_macos_install_shape(
(runtime_dir / "libggml.0.dylib").write_bytes(b"DLL")
if include_libmtmd:
(runtime_dir / "libmtmd.0.dylib").write_bytes(b"DLL")
- (install_dir / "convert_hf_to_gguf.py").write_text(
- "#!/usr/bin/env python3\n", encoding = "utf-8"
- )
+ (install_dir / "convert_hf_to_gguf.py").write_text("#!/usr/bin/env python3\n", encoding = "utf-8")
(install_dir / "gguf-py" / "gguf").mkdir(parents = True, exist_ok = True)
@@ -1240,8 +1197,7 @@ def test_existing_install_matches_plan_false_without_fingerprint(tmp_path: Path)
install_dir.mkdir()
write_linux_install_shape(install_dir)
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(
- json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"})
- + "\n",
+ json.dumps({"tag": "b9001", "asset": "llama-b9001-bin-ubuntu-x64.tar.gz"}) + "\n",
encoding = "utf-8",
)
@@ -1304,9 +1260,7 @@ def test_existing_install_matches_plan_false_with_malformed_metadata(tmp_path: P
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
write_linux_install_shape(install_dir)
- (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(
- "{not-json\n", encoding = "utf-8"
- )
+ (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text("{not-json\n", encoding = "utf-8")
host = HostInfo(
system = "Linux",
@@ -1437,9 +1391,7 @@ def test_existing_install_matches_plan_windows_cpu_requires_llama_dll(tmp_path:
def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path: Path):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
- write_windows_install_shape(
- install_dir, include_llama_dll = True, include_cuda_dll = True
- )
+ write_windows_install_shape(install_dir, include_llama_dll = True, include_cuda_dll = True)
host = HostInfo(
system = "Windows",
@@ -1508,9 +1460,7 @@ def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path:
assert existing_install_matches_plan(install_dir, host, plan) is False
-def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(
- tmp_path: Path,
-):
+def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_path: Path):
"""When the choice ships a paired cudart bundle (#5106), the install
is considered stale unless cudart64_*.dll and cublas64_*.dll are
actually on disk. Otherwise existing broken installs would keep
@@ -1626,9 +1576,7 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(
assert existing_install_matches_plan(install_dir, host, plan) is False
-def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(
- tmp_path: Path,
-):
+def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(tmp_path: Path):
"""If the choice has no paired runtime archive (manifest dropped it,
or upstream did not ship cudart), legacy installs without cudart on
disk must still pass the health check -- otherwise the installer
@@ -1708,9 +1656,7 @@ def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(
assert existing_install_matches_plan(install_dir, host, plan) is True
-def test_existing_install_fingerprint_changes_when_cudart_pair_added(
- tmp_path: Path,
-):
+def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: Path):
"""Existing pre-#5322 Windows CUDA installs (no paired cudart) must
be treated as stale once the choice gains a runtime archive,
otherwise the fingerprint match would keep skipping the reinstall
@@ -1985,9 +1931,7 @@ def test_install_prebuilt_skips_download_when_existing_install_matches(
INSTALL_LLAMA_PREBUILT,
"download_validation_model",
lambda *args, **kwargs: (_ for _ in ()).throw(
- AssertionError(
- "matching install should skip before validation model download"
- )
+ AssertionError("matching install should skip before validation model download")
),
)
@@ -2503,9 +2447,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
(staging_dir / "marker.txt").write_text("ready\n")
return attempts[0], staging_dir, initial_fallback_used
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate
- )
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "validate_prebuilt_attempts", fake_validate)
activated = {}
monkeypatch.setattr(
@@ -2523,10 +2465,7 @@ def test_install_prebuilt_same_tag_upstream_failure_uses_older_unsloth_release_p
install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
- assert attempted == [
- ("b9002", "release-2", "upstream"),
- ("b9001", "release-1", "upstream"),
- ]
+ assert attempted == [("b9002", "release-2", "upstream"), ("b9001", "release-1", "upstream")]
assert activated["install_dir"] == install_dir
@@ -2535,7 +2474,11 @@ def io_bytes(data: bytes):
def add_bytes_to_tar(
- archive: tarfile.TarFile, name: str, data: bytes, *, mode: int = 0o644
+ archive: tarfile.TarFile,
+ name: str,
+ data: bytes,
+ *,
+ mode: int = 0o644,
) -> None:
info = tarfile.TarInfo(name)
info.size = len(data)
@@ -2550,9 +2493,7 @@ def add_symlink_to_tar(archive: tarfile.TarFile, name: str, target: str) -> None
archive.addfile(info)
-def test_existing_install_matches_choice_fails_when_install_tree_incomplete(
- tmp_path: Path,
-):
+def test_existing_install_matches_choice_fails_when_install_tree_incomplete(tmp_path: Path):
"""confirm_install_tree guard rejects installs missing critical files."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
@@ -2641,9 +2582,7 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete(
)
-def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos(
- tmp_path: Path,
-):
+def test_existing_install_matches_choice_fails_when_install_tree_incomplete_macos(tmp_path: Path):
"""confirm_install_tree guard rejects macOS arm64 installs missing critical files."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
@@ -2780,9 +2719,7 @@ def test_paired_runtime_dll_patterns_excludes_executables() -> None:
assert paired_runtime_dll_patterns(non_windows) == []
-def test_runtime_overlay_cannot_overwrite_main_archive_payload(
- tmp_path: Path,
-) -> None:
+def test_runtime_overlay_cannot_overwrite_main_archive_payload(tmp_path: Path) -> None:
"""End-to-end: a malformed runtime archive containing
``llama-server.exe`` alongside the real cudart DLLs must NOT
replace the main archive's ``llama-server.exe``.
@@ -2846,15 +2783,20 @@ def test_runtime_overlay_cannot_overwrite_main_archive_payload(
orig_download = INSTALL_LLAMA_PREBUILT.download_file_verified
- def fake_download(url, target_path, *, expected_sha256 = None, label = None, **kw):
+ def fake_download(
+ url,
+ target_path,
+ *,
+ expected_sha256 = None,
+ label = None,
+ **kw,
+ ):
src = main_zip if "cudart" not in url else runtime_zip
_shutil.copy2(src, target_path)
if expected_sha256:
actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest()
if actual != expected_sha256:
- raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(
- f"sha256 mismatch on {label}"
- )
+ raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(f"sha256 mismatch on {label}")
INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download
try:
@@ -2866,16 +2808,13 @@ def test_runtime_overlay_cannot_overwrite_main_archive_payload(
server = release_dir / "llama-server.exe"
assert server.exists()
assert server.read_bytes() == b"MAIN-SERVER", (
- "runtime archive overwrote main llama-server.exe; "
- f"got {server.read_bytes()!r}"
+ "runtime archive overwrote main llama-server.exe; " f"got {server.read_bytes()!r}"
)
for name in ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"):
assert (release_dir / name).exists(), f"missing {name}"
-def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(
- tmp_path: Path,
-) -> None:
+def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(tmp_path: Path) -> None:
install_from_archives = INSTALL_LLAMA_PREBUILT.install_from_archives
work = tmp_path / "work"
@@ -2939,14 +2878,19 @@ def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(
orig_download = INSTALL_LLAMA_PREBUILT.download_file_verified
- def fake_download(url, target_path, *, expected_sha256 = None, label = None, **kw):
+ def fake_download(
+ url,
+ target_path,
+ *,
+ expected_sha256 = None,
+ label = None,
+ **kw,
+ ):
_shutil.copy2(bundle, target_path)
if expected_sha256:
actual = hashlib.sha256(Path(target_path).read_bytes()).hexdigest()
if actual != expected_sha256:
- raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(
- f"sha256 mismatch on {label}"
- )
+ raise INSTALL_LLAMA_PREBUILT.PrebuiltFallback(f"sha256 mismatch on {label}")
INSTALL_LLAMA_PREBUILT.download_file_verified = fake_download
try:
@@ -2964,9 +2908,7 @@ def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(
assert not (runtime_dir / "llama-cli").exists()
-def test_python_runtime_dirs_covers_cu13_and_library_bin(
- monkeypatch, tmp_path: Path
-) -> None:
+def test_python_runtime_dirs_covers_cu13_and_library_bin(monkeypatch, tmp_path: Path) -> None:
"""Installer-side runtime DLL discovery must scan the same path
set as the backend ``_windows_pip_nvidia_dll_dirs``: legacy
``nvidia//bin``, current ``nvidia//bin/x86_64``
diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py
index 2d7c038861..17c58cba51 100644
--- a/tests/studio/install/test_llama_pr_force_and_source.py
+++ b/tests/studio/install/test_llama_pr_force_and_source.py
@@ -35,7 +35,10 @@ requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not availa
# Helpers
# ---------------------------------------------------------------------------
def run_bash(
- script: str, *, timeout: int = 60, env: dict | None = None
+ script: str,
+ *,
+ timeout: int = 60,
+ env: dict | None = None,
) -> subprocess.CompletedProcess:
"""Run a bash script fragment and return the CompletedProcess.
60s default tolerates slow shell startup on heavily-loaded CI
@@ -53,7 +56,10 @@ def run_bash(
def run_pwsh(
- script: str, *, timeout: int = 60, env: dict | None = None
+ script: str,
+ *,
+ timeout: int = 60,
+ env: dict | None = None,
) -> subprocess.CompletedProcess:
"""Run a PowerShell script fragment and return the CompletedProcess.
60s default tolerates slow pwsh startup on heavily-loaded CI
@@ -383,10 +389,7 @@ class TestSourcePatternsSh:
assert '_DEFAULT_LLAMA_PR_FORCE=""' in self.content
def test_has_default_source(self):
- assert (
- '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"'
- in self.content
- )
+ assert '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"' in self.content
def test_has_pr_force_env_read(self):
assert "UNSLOTH_LLAMA_PR_FORCE" in self.content
@@ -416,8 +419,7 @@ class TestSourcePatternsSh:
def test_clone_urls_parameterized_pr_path(self):
"""PR clone path uses ${_LLAMA_SOURCE}.git, not hardcoded URL."""
pr_clone_idx = self.content.index(
- 'if [ -n "$_LLAMA_PR" ]; then\n'
- ' run_quiet_no_exit "clone llama.cpp"'
+ 'if [ -n "$_LLAMA_PR" ]; then\n run_quiet_no_exit "clone llama.cpp"'
)
else_idx = self.content.index("else\n", pr_clone_idx)
pr_block = self.content[pr_clone_idx:else_idx]
@@ -437,9 +439,7 @@ class TestSourcePatternsSh:
lines = self.content.splitlines()
for i, line in enumerate(lines, 1):
if "git clone" in line and "ggml-org/llama.cpp.git" in line:
- pytest.fail(
- f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}"
- )
+ pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
# =========================================================================
@@ -456,10 +456,7 @@ class TestSourcePatternsPs1:
assert '$DefaultLlamaPrForce = ""' in self.content
def test_has_default_source(self):
- assert (
- '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"'
- in self.content
- )
+ assert '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"' in self.content
def test_has_pr_force_env_read(self):
assert "$env:UNSLOTH_LLAMA_PR_FORCE" in self.content
@@ -469,10 +466,7 @@ class TestSourcePatternsPs1:
assert "$LlamaSource = $DefaultLlamaSource" in self.content
def test_release_repo_override_removed(self):
- assert (
- "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)"
- not in self.content
- )
+ assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content
assert '$HelperReleaseRepo = "ggml-org/llama.cpp"' in self.content
def test_force_compile_skips_prebuilt_resolution_early(self):
@@ -491,9 +485,7 @@ class TestSourcePatternsPs1:
def test_clone_urls_parameterized_pr_path(self):
"""PR clone path uses $LlamaSource.git, not hardcoded URL."""
- pr_idx = self.content.index(
- "if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp")
- )
+ pr_idx = self.content.index("if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp"))
else_idx = self.content.index("} else {", pr_idx)
pr_block = self.content[pr_idx:else_idx]
assert '"$LlamaSource.git"' in pr_block
@@ -511,9 +503,7 @@ class TestSourcePatternsPs1:
lines = self.content.splitlines()
for i, line in enumerate(lines, 1):
if "git clone" in line and "ggml-org/llama.cpp.git" in line:
- pytest.fail(
- f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}"
- )
+ pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
# =========================================================================
diff --git a/tests/studio/install/test_macos_version_compat.py b/tests/studio/install/test_macos_version_compat.py
index 1b87e5af65..7f93b295eb 100644
--- a/tests/studio/install/test_macos_version_compat.py
+++ b/tests/studio/install/test_macos_version_compat.py
@@ -20,9 +20,7 @@ import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
-SPEC = importlib.util.spec_from_file_location(
- "studio_install_llama_prebuilt_macos", MODULE_PATH
-)
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt_macos", MODULE_PATH)
assert SPEC is not None and SPEC.loader is not None
ILP = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = ILP
@@ -54,7 +52,12 @@ def make_macos_host(macos_version, *, arm64 = True):
)
-def thin_macho(minos = (14, 0), *, cputype = _CPU_TYPE_ARM64, build_version = True):
+def thin_macho(
+ minos = (14, 0),
+ *,
+ cputype = _CPU_TYPE_ARM64,
+ build_version = True,
+):
"""Synthesize a minimal little-endian 64-bit Mach-O carrying a macOS
minimum-version load command."""
encoded = (minos[0] << 16) | (minos[1] << 8)
@@ -139,10 +142,7 @@ class TestMachoMinimumMacos:
)
)
assert ILP.macho_minimum_macos(path, make_macos_host((14, 0))) == (14, 0)
- assert ILP.macho_minimum_macos(path, make_macos_host((26, 0), arm64 = False)) == (
- 26,
- 0,
- )
+ assert ILP.macho_minimum_macos(path, make_macos_host((26, 0), arm64 = False)) == (26, 0)
def test_non_macho_returns_none(self, tmp_path):
path = tmp_path / "script.sh"
@@ -185,23 +185,17 @@ class TestPreflightMacosInstalledBinaries:
def test_rejects_too_new_dylib(self, tmp_path):
install_dir, binaries = self._install_dir(tmp_path, (26, 0))
with pytest.raises(PrebuiltFallback, match = "newer macOS"):
- ILP.preflight_macos_installed_binaries(
- binaries, install_dir, make_macos_host((14, 0))
- )
+ ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host((14, 0)))
def test_accepts_compatible_prebuilt(self, tmp_path):
install_dir, binaries = self._install_dir(tmp_path, (14, 0))
# Must not raise on a macOS 15 host.
- ILP.preflight_macos_installed_binaries(
- binaries, install_dir, make_macos_host((15, 5))
- )
+ ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host((15, 5)))
def test_skips_when_host_version_unknown(self, tmp_path):
install_dir, binaries = self._install_dir(tmp_path, (26, 0))
# Unknown host version -> defer to runtime validation, do not raise.
- ILP.preflight_macos_installed_binaries(
- binaries, install_dir, make_macos_host(None)
- )
+ ILP.preflight_macos_installed_binaries(binaries, install_dir, make_macos_host(None))
def test_noop_on_non_macos_host(self, tmp_path):
install_dir, binaries = self._install_dir(tmp_path, (26, 0))
diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py
index 34b144a905..5b1c4a44e5 100644
--- a/tests/studio/install/test_pr4562_bugfixes.py
+++ b/tests/studio/install/test_pr4562_bugfixes.py
@@ -29,9 +29,7 @@ import pytest
# ---------------------------------------------------------------------------
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
-SPEC = importlib.util.spec_from_file_location(
- "studio_install_llama_prebuilt", MODULE_PATH
-)
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
assert SPEC is not None and SPEC.loader is not None
MOD = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = MOD
@@ -74,7 +72,12 @@ def make_host(*, system: str) -> HostInfo:
BASH = "/bin/bash"
-def run_bash(script: str, *, timeout: int = 10, env: dict | None = None) -> str:
+def run_bash(
+ script: str,
+ *,
+ timeout: int = 10,
+ env: dict | None = None,
+) -> str:
"""Run a bash script fragment and return its stdout."""
run_env = os.environ.copy()
if env:
@@ -113,9 +116,7 @@ class TestBinaryEnvCrossPlatform:
env = binary_env(binary_path, install_dir, host)
ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
assert str(bin_dir) in ld_dirs, f"build/bin not in LD_LIBRARY_PATH: {ld_dirs}"
- assert (
- str(install_dir) in ld_dirs
- ), f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}"
+ assert str(install_dir) in ld_dirs, f"install_dir not in LD_LIBRARY_PATH: {ld_dirs}"
def test_linux_binary_parent_comes_before_install_dir(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@@ -134,9 +135,7 @@ class TestBinaryEnvCrossPlatform:
ld_dirs = env["LD_LIBRARY_PATH"].split(os.pathsep)
bin_idx = ld_dirs.index(str(bin_dir))
install_idx = ld_dirs.index(str(install_dir))
- assert (
- bin_idx < install_idx
- ), "binary_path.parent should come before install_dir"
+ assert bin_idx < install_idx, "binary_path.parent should come before install_dir"
def test_linux_deduplicates_when_binary_parent_equals_install_dir(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@@ -195,17 +194,13 @@ class TestBinaryEnvCrossPlatform:
binary_path.write_bytes(b"MZ")
host = make_host(system = "Windows")
- monkeypatch.setattr(
- MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: []
- )
+ monkeypatch.setattr(MOD, "windows_runtime_dirs_for_runtime_line", lambda _rt: [])
env = binary_env(binary_path, install_dir, host)
path_dirs = env["PATH"].split(os.pathsep)
assert str(bin_dir) in path_dirs, f"build/bin/Release not in PATH: {path_dirs}"
- def test_macos_sets_dyld_library_path(
- self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
- ):
+ def test_macos_sets_dyld_library_path(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir(parents = True)
bin_dir = install_dir / "build" / "bin"
@@ -218,12 +213,8 @@ class TestBinaryEnvCrossPlatform:
env = binary_env(binary_path, install_dir, host)
dyld_parts = [p for p in env["DYLD_LIBRARY_PATH"].split(os.pathsep) if p]
- assert (
- str(bin_dir) in dyld_parts
- ), f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}"
- assert (
- str(install_dir) in dyld_parts
- ), f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}"
+ assert str(bin_dir) in dyld_parts, f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}"
+ assert str(install_dir) in dyld_parts, f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}"
# binary_path.parent (build/bin) should come before install_dir
assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir))
@@ -303,7 +294,11 @@ class TestResolveRequestedLlamaTag:
):
captured = {}
- def fake_resolve(requested_tag, published_repo, published_release_tag = ""):
+ def fake_resolve(
+ requested_tag,
+ published_repo,
+ published_release_tag = "",
+ ):
captured["requested_tag"] = requested_tag
captured["published_repo"] = published_repo
captured["published_release_tag"] = published_release_tag
@@ -350,9 +345,7 @@ class TestResolveRequestedLlamaTag:
class TestFetchJsonRetries:
- def test_fetch_json_retries_invalid_github_api_json(
- self, monkeypatch: pytest.MonkeyPatch
- ):
+ def test_fetch_json_retries_invalid_github_api_json(self, monkeypatch: pytest.MonkeyPatch):
calls = {"count": 0}
def fake_download_bytes(url, **kwargs):
@@ -593,11 +586,7 @@ class TestLatestTagResolution:
""")
def _run_resolve(
- self,
- tmp_path: Path,
- requested_tag: str,
- resolved_tag: str,
- resolve_status: int,
+ self, tmp_path: Path, requested_tag: str, resolved_tag: str, resolve_status: int
) -> str:
script = self.RESOLVE_TEMPLATE.format(
requested_tag = requested_tag,
@@ -691,10 +680,7 @@ class TestSourceCodePatterns:
content = SETUP_SH.read_text()
assert "--resolve-source-build" not in content
assert "--resolve-install-tag" not in content
- assert (
- '--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"'
- in content
- )
+ assert '--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"' in content
assert "--output-format json" in content
assert "_RESOLVED_SOURCE_URL" in content
assert "_RESOLVED_SOURCE_REF_KIND" in content
@@ -758,9 +744,7 @@ class TestSourceCodePatterns:
# Delivered via NVCC_PREPEND_FLAGS (covers the configure-time compiler
# probe too), not embedded in the word-split CMAKE_ARGS string.
assert "export NVCC_PREPEND_FLAGS=" in content
- cmake_args_lines = [
- line for line in content.splitlines() if "CMAKE_ARGS=" in line
- ]
+ cmake_args_lines = [line for line in content.splitlines() if "CMAKE_ARGS=" in line]
assert all(
"-allow-unsupported-compiler" not in line for line in cmake_args_lines
), "flag must stay out of CMAKE_ARGS (bash word-splitting safety)"
@@ -776,9 +760,7 @@ class TestSourceCodePatterns:
# Delivered via the process environment, not the $CmakeArgs array, so it
# reaches both the configure-time compiler probe and `cmake --build`.
assert "$env:NVCC_PREPEND_FLAGS" in content
- cmake_args_lines = [
- line for line in content.splitlines() if "$CmakeArgs +=" in line
- ]
+ cmake_args_lines = [line for line in content.splitlines() if "$CmakeArgs +=" in line]
assert all(
"-allow-unsupported-compiler" not in line for line in cmake_args_lines
), "flag must not be pushed into the $CmakeArgs array"
@@ -794,15 +776,10 @@ class TestSourceCodePatterns:
def test_macos_arm64_cpu_fallback_args_exclude_rpath(self):
"""CPU fallback args must NOT contain Metal-only RPATH flags at runtime."""
- script = (
- '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n'
- + _GPU_BACKEND_FRAGMENT
- )
+ script = '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT
output = run_bash(script)
fallback_line = next(
- line
- for line in output.splitlines()
- if line.startswith("CPU_FALLBACK_CMAKE_ARGS=")
+ line for line in output.splitlines() if line.startswith("CPU_FALLBACK_CMAKE_ARGS=")
)
assert "-DGGML_METAL=OFF" in fallback_line
assert (
@@ -823,8 +800,7 @@ class TestSourceCodePatterns:
assert (
"x86_64"
not in content[
- content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON")
- + 200
+ content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON") + 200
]
)
@@ -854,9 +830,7 @@ class TestSourceCodePatterns:
# Allow git pull in other contexts
context = "\n".join(lines[max(0, i - 5) : i + 5])
if "LlamaCppDir" in context:
- pytest.fail(
- f"Found 'git pull' in llama.cpp build section at line {i+1}"
- )
+ pytest.fail(f"Found 'git pull' in llama.cpp build section at line {i+1}")
def test_setup_ps1_prebuilt_install_uses_simple_policy_only(self):
"""PS1 prebuilt path should use the simplified helper install entrypoint."""
@@ -883,8 +857,7 @@ class TestSourceCodePatterns:
assert "--resolve-source-build" not in content
assert "--resolve-install-tag" not in content
assert (
- '"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"'
- in content
+ '"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"' in content
)
assert '--output-format", "json"' in content
assert "$ResolvedSourceUrl" in content
@@ -898,10 +871,7 @@ class TestSourceCodePatterns:
block = content[max(0, install_idx - 800) : install_idx + 800]
assert "$PSNativeCommandUseErrorActionPreference = $false" in block
assert "$restoreNativeErrorPreference = $true" in block
- assert (
- "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference"
- in block
- )
+ assert "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference" in block
def test_setup_ps1_helper_disables_error_action_abort(self):
"""Helper resolution should suppress terminating NativeCommandError on PS 5.1."""
@@ -922,9 +892,7 @@ class TestSourceCodePatterns:
"""The unconstrained nvcc fallback should not sort toolkit dirs lexicographically."""
content = SETUP_PS1.read_text()
assert "Sort-Object Name | Select-Object -Last 1" not in content
- assert (
- "Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content
- )
+ assert "Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content
def test_binary_env_linux_has_binary_parent(self):
"""The Linux branch of binary_env should include binary_path.parent."""
@@ -985,10 +953,7 @@ class TestMacOSMetalBuildLogic:
def test_macos_arm64_cmake_args_contain_metal_flags(self):
"""macOS arm64 should enable Metal, not CUDA."""
- script = (
- '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n'
- + _GPU_BACKEND_FRAGMENT
- )
+ script = '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT
output = run_bash(script)
assert "-DGGML_METAL=ON" in output
assert "-DGGML_CUDA=ON" not in output
@@ -996,10 +961,7 @@ class TestMacOSMetalBuildLogic:
def test_intel_macos_no_metal_flags(self):
"""Intel macOS (not arm64) should not get Metal flags."""
- script = (
- '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n'
- + _GPU_BACKEND_FRAGMENT
- )
+ script = '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n' + _GPU_BACKEND_FRAGMENT
output = run_bash(script)
assert "-DGGML_METAL=ON" not in output
assert "BUILD_DESC=building (CPU)" in output
@@ -1085,18 +1047,14 @@ class TestMacOSMetalBuildLogic:
# Verify cmake args: first call has Metal ON, second has Metal OFF
calls = calls_file.read_text().splitlines()
assert len(calls) >= 2, f"Expected >= 2 cmake calls, got {len(calls)}"
- assert (
- "-DGGML_METAL=ON" in calls[0]
- ), f"First cmake call should have Metal ON: {calls[0]}"
+ assert "-DGGML_METAL=ON" in calls[0], f"First cmake call should have Metal ON: {calls[0]}"
assert (
"-DGGML_METAL=OFF" in calls[1]
), f"Second cmake call should have Metal OFF: {calls[1]}"
assert (
"-DGGML_METAL=ON" not in calls[1]
), f"Second cmake call should NOT have Metal ON: {calls[1]}"
- assert (
- "@loader_path" not in calls[1]
- ), f"CPU fallback should not have RPATH: {calls[1]}"
+ assert "@loader_path" not in calls[1], f"CPU fallback should not have RPATH: {calls[1]}"
assert (
"-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[1]
), f"CPU fallback should not have RPATH build flag: {calls[1]}"
@@ -1204,9 +1162,7 @@ class TestMacOSMetalBuildLogic:
# Third call: re-configure with Metal OFF and no RPATH flags
assert "-DGGML_METAL=OFF" in calls[2]
assert "-DGGML_METAL=ON" not in calls[2]
- assert (
- "@loader_path" not in calls[2]
- ), f"CPU fallback should not have RPATH: {calls[2]}"
+ assert "@loader_path" not in calls[2], f"CPU fallback should not have RPATH: {calls[2]}"
assert (
"-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[2]
), f"CPU fallback should not have RPATH build flag: {calls[2]}"
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index 4aa7fb3a34..bcfc42c69f 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -40,9 +40,7 @@ _normalize_forwarded_gfx = prebuilt_mod._normalize_forwarded_gfx
# install_python_stack.py
_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
-_STACK_SPEC = importlib.util.spec_from_file_location(
- "studio_install_python_stack", _STACK_PATH
-)
+_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH)
assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
sys.modules[_STACK_SPEC.name] = stack_mod
@@ -304,9 +302,7 @@ class TestResolveUpstreamAssetChoice:
def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets):
"""AMD ROCm host should fall back to source build when no ROCm prebuilt exists."""
# Remove the ROCm asset from available assets
- assets_without_rocm = {
- k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k
- }
+ assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k}
mock_assets.return_value = assets_without_rocm
host = rocm_host()
with pytest.raises(PrebuiltFallback, match = "ROCm detected"):
@@ -573,9 +569,7 @@ class TestEnsureRocmTorch:
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
- def test_torch_already_has_cuda_skips(
- self, mock_ver, mock_gpu, mock_nvidia, mock_pip
- ):
+ def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
"""If torch already has CUDA, should skip ROCm reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@@ -589,9 +583,7 @@ class TestEnsureRocmTorch:
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
- def test_torch_already_has_hip_skips(
- self, mock_ver, mock_gpu, mock_nvidia, mock_pip
- ):
+ def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
"""If torch already has HIP, should skip ROCm reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@@ -629,9 +621,7 @@ class TestEnsureRocmTorch:
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3))
- def test_rocm_63_selects_correct_tag(
- self, mock_ver, mock_gpu, mock_nvidia, mock_pip
- ):
+ def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
"""ROCm 6.3 should select rocm6.3 tag."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@@ -698,9 +688,7 @@ class TestEnsureRocmTorch:
):
"""Probe subprocess timeout should not crash; should proceed to reinstall."""
with patch("os.path.isdir", return_value = True):
- with patch(
- "subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)
- ):
+ with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)):
_ensure_rocm_torch()
# If probe times out, the function should treat torch as unusable and reinstall
# both torch (via pip_install) and bitsandbytes (via pip_install_try).
@@ -788,25 +776,19 @@ class TestHardwareRocmFlag:
def test_hardware_py_has_is_rocm(self):
"""hardware.py should define IS_ROCM."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
assert "IS_ROCM: bool" in source and "False" in source
def test_hardware_py_sets_is_rocm_on_hip(self):
"""detect_hardware() should set IS_ROCM when torch.version.hip is set."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
assert 'torch.version, "hip"' in source or "torch.version.hip" in source
def test_hardware_py_still_returns_cuda_for_rocm(self):
"""DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
# Ensure ROCM is NOT a DeviceType member
enum_section = source.split("class DeviceType")[1].split("\n\n")[0]
@@ -814,17 +796,13 @@ class TestHardwareRocmFlag:
def test_hardware_py_has_rocm_in_package_versions(self):
"""get_package_versions() should include 'rocm' key."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
assert '"rocm"' in source
def test_hardware_py_device_type_cuda_references_intact(self):
"""All existing DeviceType.CUDA references should still be present."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
# Key functions that must still reference DeviceType.CUDA
assert "DeviceType.CUDA" in source
@@ -832,26 +810,20 @@ class TestHardwareRocmFlag:
def test_is_rocm_exported_from_init(self):
"""IS_ROCM should be exported from hardware __init__.py."""
- init_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
- )
+ init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
source = init_path.read_text(encoding = "utf-8")
assert "IS_ROCM" in source
def test_is_rocm_in_all_list(self):
"""IS_ROCM should be in __all__ list in __init__.py."""
- init_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
- )
+ init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
source = init_path.read_text(encoding = "utf-8")
# Extract __all__ section
assert '"IS_ROCM"' in source
def test_get_package_versions_returns_rocm_key(self):
"""get_package_versions() source should return both 'cuda' and 'rocm' keys."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
# Find the get_package_versions function body
func_start = source.find("def get_package_versions")
@@ -866,22 +838,16 @@ class TestHardwareRocmFlag:
Windows ROCm where torch.distributed ships without that helper, causing
a warning: 'module torch.distributed has no attribute is_torchelastic_launched'.
"""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
assert "is_torchelastic_launched" in source
def test_distributed_stubs_cover_core_helpers(self):
"""_determine_attention_impl_for_gpu_estimate must stub the four core distributed helpers."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
for attr in ("is_initialized", "is_available", "get_rank", "get_world_size"):
- assert (
- attr in source
- ), f"distributed stub for '{attr}' missing from hardware.py"
+ assert attr in source, f"distributed stub for '{attr}' missing from hardware.py"
# =============================================================================
@@ -947,12 +913,8 @@ class TestInstallShStructure:
nvidia_call = body.find("_has_usable_nvidia_gpu")
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
rocm_call = body.find("_has_amd_rocm_gpu")
- assert (
- nvidia_call >= 0
- ), "get_torch_index_url should call _has_usable_nvidia_gpu"
- assert (
- no_nvidia_branch >= 0
- ), "get_torch_index_url should gate ROCm on no-nvidia-smi"
+ assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
+ assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia-smi"
assert (
rocm_call > no_nvidia_branch
), "ROCm detection should sit inside the 'no nvidia-smi' branch"
@@ -1013,9 +975,7 @@ class TestInstallShStructure:
continue
# Remove POSIX character classes [[:foo:]] before checking for [[ ]]
cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line)
- assert (
- "[[" not in cleaned
- ), f"get_torch_index_url line {i} uses non-POSIX [["
+ assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [["
def test_no_arithmetic_expansion_in_rocm_block(self):
"""ROCm detection block should not use (( )) (bash-only)."""
@@ -1063,8 +1023,7 @@ class TestLiveRegression:
[
"bash",
"-c",
- "nvidia-smi -L 2>/dev/null | "
- "awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'",
+ "nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'",
],
capture_output = True,
)
@@ -1098,9 +1057,7 @@ class TestLiveRegression:
# Load worker.py module
_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py"
-_EXPORT_WORKER_PATH = (
- PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py"
-)
+_EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py"
# The torchao Windows-ROCm stub was de-duplicated out of the export/training
# workers into a shared module; both workers now call into it.
_TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py"
@@ -1126,9 +1083,7 @@ class TestWorkerRocmMambaSsm:
def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch):
"""_direct_wheel_url should return None when cuda_major is empty (ROCm)."""
# Load module for function access
- _worker_spec = importlib.util.spec_from_file_location(
- "test_worker", _WORKER_PATH
- )
+ _worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH)
assert _worker_spec is not None and _worker_spec.loader is not None
worker_mod = importlib.util.module_from_spec(_worker_spec)
@@ -1347,9 +1302,7 @@ class TestHardwareAmdBranching:
def test_hardware_imports_amd_module(self):
"""hardware.py should import from amd module when IS_ROCM."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
assert "from . import amd" in source
@@ -1357,17 +1310,13 @@ class TestHardwareAmdBranching:
"""get_gpu_utilization should dispatch to amd.py via _smi_query
when IS_ROCM, and the dispatcher itself must check IS_ROCM and
import the amd backend."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def get_gpu_utilization")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert '_smi_query("get_primary_gpu_utilization"' in func_body
smi = source[
- source.find("def _smi_query") : source.find(
- "\ndef ", source.find("def _smi_query") + 1
- )
+ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1)
]
assert "IS_ROCM" in smi
assert "from . import amd" in smi
@@ -1375,9 +1324,7 @@ class TestHardwareAmdBranching:
def test_hardware_branches_on_is_rocm_for_visible(self):
"""get_visible_gpu_utilization should dispatch to amd.py via
_smi_query when IS_ROCM."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def get_visible_gpu_utilization")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1387,18 +1334,14 @@ class TestHardwareAmdBranching:
assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body)
smi = source[
- source.find("def _smi_query") : source.find(
- "\ndef ", source.find("def _smi_query") + 1
- )
+ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1)
]
assert "IS_ROCM" in smi
assert "from . import amd" in smi
def test_hardware_branches_on_is_rocm_for_physical_count(self):
"""get_physical_gpu_count should try amd.py when IS_ROCM."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def get_physical_gpu_count")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1417,9 +1360,7 @@ class TestApplyGpuIdsRocmFallback:
def test_apply_gpu_ids_falls_back_to_torch_version_hip(self):
"""apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1427,9 +1368,7 @@ class TestApplyGpuIdsRocmFallback:
def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self):
"""apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1438,9 +1377,7 @@ class TestApplyGpuIdsRocmFallback:
def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self):
"""torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes."""
- hw_path = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
- )
+ hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
@@ -1591,9 +1528,7 @@ class TestWindowsRocmIndexUrl:
assert "repo.amd.com" in url
def test_mirror_env_var_overrides_base(self, monkeypatch):
- monkeypatch.setenv(
- "UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl"
- )
+ monkeypatch.setenv("UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl")
# Reload module-level constant by calling helper directly
url = stack_mod._windows_rocm_index_url("gfx1200")
# The env var is read at module load time for _ROCM_WINDOWS_INDEX_BASE,
@@ -1722,9 +1657,7 @@ class TestInstallBnbWindowsRocm:
with patch.dict(os.environ, {}, clear = False):
os.environ.pop("BNB_ROCM_VERSION", None)
with patch.object(stack_mod, "pip_install_try", return_value = True):
- with patch.object(
- stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"
- ):
+ with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"):
stack_mod._install_bnb_windows_rocm()
assert os.environ.get("BNB_ROCM_VERSION") == "72"
@@ -1733,9 +1666,7 @@ class TestInstallBnbWindowsRocm:
with patch.dict(os.environ, {}, clear = False):
os.environ.pop("BNB_ROCM_VERSION", None)
with patch.object(stack_mod, "pip_install_try", return_value = True):
- with patch.object(
- stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"
- ):
+ with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"):
stack_mod._install_bnb_windows_rocm()
assert os.environ.get("BNB_ROCM_VERSION") == "713"
@@ -1744,9 +1675,7 @@ class TestInstallBnbWindowsRocm:
with patch.dict(os.environ, {}, clear = False):
os.environ.pop("BNB_ROCM_VERSION", None)
with patch.object(stack_mod, "pip_install_try", return_value = True):
- with patch.object(
- stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None
- ):
+ with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None):
stack_mod._install_bnb_windows_rocm()
assert os.environ.get("BNB_ROCM_VERSION") == "72"
@@ -1764,7 +1693,6 @@ class TestDetectBnbRocmDllVer:
def test_returns_none_when_bnb_not_installed(self):
"""Returns None if bitsandbytes is not importable."""
import importlib.util
-
with patch.object(importlib.util, "find_spec", return_value = None):
assert stack_mod._detect_bnb_rocm_dll_ver() is None
@@ -1973,9 +1901,7 @@ class TestWorkerWindowsRocmPatches:
# entry-point function (not the trainer helper which has its own "# ── 2.").
idx_sec2 = source.find("# ── 2. Now import ML libraries")
assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py"
- assert (
- idx_sec2 != -1
- ), "'# ── 2. Now import ML libraries' marker not found in worker.py"
+ assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py"
assert idx_bnb < idx_sec2, (
"BNB_ROCM_VERSION must be set before section 2 ML imports "
f"(found at {idx_bnb}, section 2 at {idx_sec2})"
@@ -2269,16 +2195,12 @@ class TestHipSdkEnvPathResolution:
"""setup.ps1 must tell the user how to add the HIP bin dir to PATH."""
source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
# Should mention adding to PATH or SetEnvironmentVariable
- assert "PATH" in source and (
- "SetEnvironmentVariable" in source or "Add" in source
- )
+ assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source)
def test_install_provides_path_fix_hint(self):
"""install.ps1 must tell the user how to add the HIP bin dir to PATH."""
source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
- assert "PATH" in source and (
- "SetEnvironmentVariable" in source or "Add" in source
- )
+ assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source)
# =============================================================================
@@ -2448,9 +2370,7 @@ class TestSetupShGccInstallDir:
# =============================================================================
_MAIN_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py"
-_HARDWARE_PY_PATH = (
- PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
-)
+_HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
class TestServerStartupRocmFixes:
@@ -2664,9 +2584,7 @@ class TestApplyHostOverrides:
assert out.rocm_gfx_target is None
def test_malformed_forwarded_gfx_falls_back_to_has_rocm(self):
- out = _apply_host_overrides(
- cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk"
- )
+ out = _apply_host_overrides(cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk")
assert out.has_rocm is True
assert out.rocm_gfx_target is None
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index f1c3130a77..1ada35888c 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -22,9 +22,7 @@ import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
RUN_MODULE_PATH = PACKAGE_ROOT / "studio" / "backend" / "run.py"
-SPEC = importlib.util.spec_from_file_location(
- "studio_install_llama_prebuilt", MODULE_PATH
-)
+SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
assert SPEC is not None and SPEC.loader is not None
INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
@@ -49,15 +47,11 @@ supports_explicit_visible_device_matching = (
select_visible_gpu_rows = INSTALL_LLAMA_PREBUILT.select_visible_gpu_rows
compatible_linux_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_linux_runtime_lines
pick_windows_cuda_runtime = INSTALL_LLAMA_PREBUILT.pick_windows_cuda_runtime
-compatible_windows_runtime_lines = (
- INSTALL_LLAMA_PREBUILT.compatible_windows_runtime_lines
-)
+compatible_windows_runtime_lines = INSTALL_LLAMA_PREBUILT.compatible_windows_runtime_lines
runtime_line_from_cuda_version = INSTALL_LLAMA_PREBUILT.runtime_line_from_cuda_version
apply_approved_hashes = INSTALL_LLAMA_PREBUILT.apply_approved_hashes
linux_cuda_choice_from_release = INSTALL_LLAMA_PREBUILT.linux_cuda_choice_from_release
-parse_direct_linux_release_bundle = (
- INSTALL_LLAMA_PREBUILT.parse_direct_linux_release_bundle
-)
+parse_direct_linux_release_bundle = INSTALL_LLAMA_PREBUILT.parse_direct_linux_release_bundle
windows_cuda_attempts = INSTALL_LLAMA_PREBUILT.windows_cuda_attempts
resolve_upstream_asset_choice = INSTALL_LLAMA_PREBUILT.resolve_upstream_asset_choice
resolve_requested_install_tag = INSTALL_LLAMA_PREBUILT.resolve_requested_install_tag
@@ -66,19 +60,11 @@ resolve_install_release_plans = INSTALL_LLAMA_PREBUILT.resolve_install_release_p
resolve_published_release = INSTALL_LLAMA_PREBUILT.resolve_published_release
resolve_source_build_plan = INSTALL_LLAMA_PREBUILT.resolve_source_build_plan
validated_checksums_for_bundle = INSTALL_LLAMA_PREBUILT.validated_checksums_for_bundle
-parse_approved_release_checksums = (
- INSTALL_LLAMA_PREBUILT.parse_approved_release_checksums
-)
-published_release_matches_request = (
- INSTALL_LLAMA_PREBUILT.published_release_matches_request
-)
-exact_source_archive_logical_name = (
- INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name
-)
+parse_approved_release_checksums = INSTALL_LLAMA_PREBUILT.parse_approved_release_checksums
+published_release_matches_request = INSTALL_LLAMA_PREBUILT.published_release_matches_request
+exact_source_archive_logical_name = INSTALL_LLAMA_PREBUILT.exact_source_archive_logical_name
source_archive_logical_name = INSTALL_LLAMA_PREBUILT.source_archive_logical_name
-windows_cuda_upstream_asset_names = (
- INSTALL_LLAMA_PREBUILT.windows_cuda_upstream_asset_names
-)
+windows_cuda_upstream_asset_names = INSTALL_LLAMA_PREBUILT.windows_cuda_upstream_asset_names
env_int = INSTALL_LLAMA_PREBUILT.env_int
direct_upstream_release_plan = INSTALL_LLAMA_PREBUILT.direct_upstream_release_plan
_pinned_windows_cuda_fallback = INSTALL_LLAMA_PREBUILT._pinned_windows_cuda_fallback
@@ -89,9 +75,7 @@ _windows_cuda_attempt_covers_blackwell = (
)
resolve_release_asset_choice = INSTALL_LLAMA_PREBUILT.resolve_release_asset_choice
pinned_macos_release_tag = INSTALL_LLAMA_PREBUILT.pinned_macos_release_tag
-resolve_simple_install_release_plans = (
- INSTALL_LLAMA_PREBUILT.resolve_simple_install_release_plans
-)
+resolve_simple_install_release_plans = INSTALL_LLAMA_PREBUILT.resolve_simple_install_release_plans
def load_studio_run_module(monkeypatch):
@@ -244,9 +228,7 @@ def make_checksums_with_source(
kind = "upstream-source",
),
}
- normalized_source_commit = (
- source_commit.lower() if isinstance(source_commit, str) else None
- )
+ normalized_source_commit = source_commit.lower() if isinstance(source_commit, str) else None
if normalized_source_commit:
artifacts[exact_source_archive_logical_name(normalized_source_commit)] = (
ApprovedArtifactHash(
@@ -266,9 +248,7 @@ def make_checksums_with_source(
requested_source_ref = requested_source_ref,
resolved_source_ref = resolved_source_ref,
source_commit = normalized_source_commit,
- source_commit_short = normalized_source_commit[:7]
- if normalized_source_commit
- else None,
+ source_commit_short = normalized_source_commit[:7] if normalized_source_commit else None,
artifacts = artifacts,
)
@@ -363,10 +343,7 @@ class TestStudioLocalhostIpv6Warning:
lambda host, port, timeout = 1.0: True,
)
- assert (
- run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888)
- == "http://127.0.0.1:8888"
- )
+ assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) == "http://127.0.0.1:8888"
@pytest.mark.parametrize("host", ["0.0.0.0", "::"])
def test_network_bind_suppresses_warning(self, monkeypatch, host):
@@ -431,9 +408,7 @@ class TestStudioLocalhostIpv6Warning:
monkeypatch.setattr(
run_module,
"_verify_global_reachability",
- lambda display_host, port: calls["reachability"].append(
- (display_host, port)
- ),
+ lambda display_host, port: calls["reachability"].append((display_host, port)),
)
return calls
@@ -456,9 +431,7 @@ class TestStudioLocalhostIpv6Warning:
def test_emit_startup_output_plain_localhost(self, monkeypatch):
run_module = load_studio_run_module(monkeypatch)
calls = self._wire_recorders(run_module, monkeypatch)
- monkeypatch.setattr(
- run_module, "_localhost_ipv6_mismatch_url", lambda host, port: None
- )
+ monkeypatch.setattr(run_module, "_localhost_ipv6_mismatch_url", lambda host, port: None)
run_module._emit_startup_output("127.0.0.1", 8888, "127.0.0.1")
@@ -471,9 +444,7 @@ class TestStudioLocalhostIpv6Warning:
def test_emit_startup_output_wildcard_runs_reachability(self, monkeypatch, host):
run_module = load_studio_run_module(monkeypatch)
calls = self._wire_recorders(run_module, monkeypatch)
- monkeypatch.setattr(
- run_module, "_localhost_ipv6_mismatch_url", lambda h, port: None
- )
+ monkeypatch.setattr(run_module, "_localhost_ipv6_mismatch_url", lambda h, port: None)
run_module._emit_startup_output(host, 8888, "203.0.113.5")
@@ -651,9 +622,7 @@ class TestParseDirectLinuxReleaseBundle:
names = [f"app-bTEST-linux-x64-{t}.tar.gz" for t in targets]
return {
"tag_name": "bTEST",
- "assets": [
- {"name": n, "browser_download_url": "https://x/" + n} for n in names
- ],
+ "assets": [{"name": n, "browser_download_url": "https://x/" + n} for n in names],
}
def _cuda_artifact(self, bundle):
@@ -885,9 +854,7 @@ class TestPublishedReleaseResolution:
def fake_load(repo, release_tag):
if release_tag == "v2.0":
raise PrebuiltFallback("checksum asset missing")
- return make_checksums_with_source(
- [], release_tag = "v1.0", upstream_tag = "b8999"
- )
+ return make_checksums_with_source([], release_tag = "v1.0", upstream_tag = "b8999")
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
@@ -917,9 +884,7 @@ class TestPublishedReleaseResolution:
),
)
- assert (
- resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") == "b8508"
- )
+ assert resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp") == "b8508"
def test_concrete_tag_without_matching_release_raises(self, monkeypatch):
release = make_release([], release_tag = "release-b9000", upstream_tag = "b9000")
@@ -933,9 +898,7 @@ class TestPublishedReleaseResolution:
resolve_requested_install_tag("b8508", "", "unslothai/llama.cpp")
def test_pinned_release_must_match_requested_upstream_tag(self, monkeypatch):
- bundle = make_release(
- [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000"
- )
+ bundle = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000")
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"pinned_published_release_bundle",
@@ -1110,15 +1073,13 @@ class TestSourceBuildPlanResolution:
assert plan.source_ref == "main"
assert plan.compatibility_upstream_tag == "b9000"
- def test_direct_main_request_without_published_release_uses_branch_kind(
- self, monkeypatch
- ):
+ def test_direct_main_request_without_published_release_uses_branch_kind(self, monkeypatch):
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"resolve_published_release",
- lambda requested_tag, published_repo, published_release_tag = "": (
- _ for _ in ()
- ).throw(PrebuiltFallback("missing")),
+ lambda requested_tag, published_repo, published_release_tag = "": (_ for _ in ()).throw(
+ PrebuiltFallback("missing")
+ ),
)
plan = resolve_source_build_plan("main", "unslothai/llama.cpp")
@@ -1193,9 +1154,7 @@ class TestValidatedChecksumsForBundle:
def test_rejects_manifest_checksum_mismatch(self, monkeypatch):
bundle = make_release([], release_tag = "r1", upstream_tag = "b8508")
bundle.manifest_sha256 = "a" * 64
- checksums = make_checksums_with_source(
- [], release_tag = "r1", upstream_tag = "b8508"
- )
+ checksums = make_checksums_with_source([], release_tag = "r1", upstream_tag = "b8508")
checksums.artifacts[bundle.manifest_asset_name] = ApprovedArtifactHash(
asset_name = bundle.manifest_asset_name,
sha256 = "b" * 64,
@@ -1249,9 +1208,7 @@ class TestLinuxCudaChoiceFromRelease:
art12 = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12")
art13 = make_artifact("bundle-cuda13.tar.gz", runtime_line = "cuda13")
release = make_release([art12, art13])
- result = linux_cuda_choice_from_release(
- host, release, preferred_runtime_line = "cuda12"
- )
+ result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
assert result is not None
assert result.primary.runtime_line == "cuda12"
@@ -1260,9 +1217,7 @@ class TestLinuxCudaChoiceFromRelease:
host = make_host(driver_cuda_version = (12, 8))
art = make_artifact("bundle-cuda12.tar.gz", runtime_line = "cuda12")
release = make_release([art])
- result = linux_cuda_choice_from_release(
- host, release, preferred_runtime_line = "cuda13"
- )
+ result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda13")
assert result is not None
assert result.primary.runtime_line == "cuda12"
log_entries = result.selection_log
@@ -1273,9 +1228,7 @@ class TestLinuxCudaChoiceFromRelease:
def test_exact_sm_match(self, monkeypatch):
mock_linux_runtime(monkeypatch, ["cuda12"])
host = make_host(compute_caps = ["86"])
- art = make_artifact(
- "bundle.tar.gz", supported_sms = ["75", "86", "89"], min_sm = 75, max_sm = 89
- )
+ art = make_artifact("bundle.tar.gz", supported_sms = ["75", "86", "89"], min_sm = 75, max_sm = 89)
release = make_release([art])
result = linux_cuda_choice_from_release(host, release)
assert result is not None
@@ -1284,9 +1237,7 @@ class TestLinuxCudaChoiceFromRelease:
def test_sm_not_in_supported_sms(self, monkeypatch):
mock_linux_runtime(monkeypatch, ["cuda12"])
host = make_host(compute_caps = ["86"])
- art = make_artifact(
- "bundle.tar.gz", supported_sms = ["75", "80", "89"], min_sm = 75, max_sm = 89
- )
+ art = make_artifact("bundle.tar.gz", supported_sms = ["75", "80", "89"], min_sm = 75, max_sm = 89)
release = make_release([art])
result = linux_cuda_choice_from_release(host, release)
assert result is None
@@ -1294,9 +1245,7 @@ class TestLinuxCudaChoiceFromRelease:
def test_sm_outside_min_range(self, monkeypatch):
mock_linux_runtime(monkeypatch, ["cuda12"])
host = make_host(compute_caps = ["50"])
- art = make_artifact(
- "bundle.tar.gz", supported_sms = ["50", "75", "86"], min_sm = 75, max_sm = 90
- )
+ art = make_artifact("bundle.tar.gz", supported_sms = ["50", "75", "86"], min_sm = 75, max_sm = 90)
release = make_release([art])
result = linux_cuda_choice_from_release(host, release)
assert result is None
@@ -1365,9 +1314,7 @@ class TestLinuxCudaChoiceFromRelease:
def test_multi_gpu_not_all_covered(self, monkeypatch):
mock_linux_runtime(monkeypatch, ["cuda12"])
host = make_host(compute_caps = ["50", "89"])
- art = make_artifact(
- "bundle.tar.gz", supported_sms = ["75", "89"], min_sm = 75, max_sm = 89
- )
+ art = make_artifact("bundle.tar.gz", supported_sms = ["75", "89"], min_sm = 75, max_sm = 89)
release = make_release([art])
result = linux_cuda_choice_from_release(host, release)
assert result is None
@@ -1559,9 +1506,7 @@ class TestBlackwellUltraSm103Coverage:
class TestResolveInstallAttempts:
- def test_windows_cuda_prefers_published_asset_from_selected_release(
- self, monkeypatch
- ):
+ def test_windows_cuda_prefers_published_asset_from_selected_release(self, monkeypatch):
host = make_host(system = "Windows", machine = "AMD64")
host.driver_cuda_version = (12, 4)
mock_windows_runtime(monkeypatch, ["cuda12"])
@@ -1605,9 +1550,7 @@ class TestResolveInstallAttempts:
INSTALL_LLAMA_PREBUILT,
"github_release_assets",
lambda repo, tag: (_ for _ in ()).throw(
- AssertionError(
- "published Windows CUDA choice should not query upstream"
- )
+ AssertionError("published Windows CUDA choice should not query upstream")
),
)
@@ -1629,9 +1572,7 @@ class TestResolveInstallAttempts:
host = make_host(system = "Windows", machine = "AMD64")
host.driver_cuda_version = (12, 4)
mock_windows_runtime(monkeypatch, ["cuda12"])
- release = make_release(
- [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000"
- )
+ release = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000")
checksums = make_checksums_with_source(
["llama-b9000-bin-win-cuda-12.4-x64.zip"],
release_tag = release.release_tag,
@@ -1692,9 +1633,7 @@ class TestResolveInstallAttempts:
has_physical_nvidia = False,
nvidia_smi = None,
)
- release = make_release(
- [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000"
- )
+ release = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000")
checksums = make_checksums_with_source(
["llama-b9000-bin-ubuntu-x64.tar.gz"],
release_tag = release.release_tag,
@@ -1735,9 +1674,7 @@ class TestResolveInstallAttempts:
def test_linux_cuda_does_not_fall_back_to_upstream_cpu(self, monkeypatch):
host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"])
- release = make_release(
- [], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000"
- )
+ release = make_release([], release_tag = "llama-prebuilt-latest", upstream_tag = "b9000")
checksums = make_checksums_with_source(
[],
release_tag = release.release_tag,
@@ -1758,9 +1695,7 @@ class TestResolveInstallAttempts:
)
mock_linux_runtime(monkeypatch, ["cuda12"])
- with pytest.raises(
- PrebuiltFallback, match = "no compatible published Linux CUDA bundle"
- ):
+ with pytest.raises(PrebuiltFallback, match = "no compatible published Linux CUDA bundle"):
resolve_install_attempts("latest", host, "unslothai/llama.cpp", "")
def test_windows_cpu_prefers_published_asset(self, monkeypatch):
@@ -1956,9 +1891,7 @@ class TestResolveInstallAttempts:
class TestResolveInstallReleasePlans:
- def test_latest_collects_multiple_older_release_plans_up_to_limit(
- self, monkeypatch
- ):
+ def test_latest_collects_multiple_older_release_plans_up_to_limit(self, monkeypatch):
host = make_host(
has_usable_nvidia = False,
has_physical_nvidia = False,
@@ -1994,9 +1927,7 @@ class TestResolveInstallReleasePlans:
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"iter_resolved_published_releases",
- lambda requested_tag, published_repo, published_release_tag = "": iter(
- releases
- ),
+ lambda requested_tag, published_repo, published_release_tag = "": iter(releases),
)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
@@ -2018,9 +1949,7 @@ class TestResolveInstallReleasePlans:
assert [plan.release_tag for plan in plans] == ["r3", "r2"]
assert [plan.llama_tag for plan in plans] == ["b9003", "b9002"]
- def test_latest_skips_non_installable_release_and_keeps_searching(
- self, monkeypatch
- ):
+ def test_latest_skips_non_installable_release_and_keeps_searching(self, monkeypatch):
host = make_host(
has_usable_nvidia = False,
has_physical_nvidia = False,
@@ -2048,9 +1977,7 @@ class TestResolveInstallReleasePlans:
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"iter_resolved_published_releases",
- lambda requested_tag, published_repo, published_release_tag = "": iter(
- releases
- ),
+ lambda requested_tag, published_repo, published_release_tag = "": iter(releases),
)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
@@ -2078,13 +2005,9 @@ class TestResolveInstallReleasePlans:
def test_malformed_release_fallback_env_uses_default(self, monkeypatch):
monkeypatch.setenv("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", "not-an-int")
- assert (
- env_int("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", 3, minimum = 1) == 3
- )
+ assert env_int("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", 3, minimum = 1) == 3
- def test_import_with_malformed_release_fallback_env_does_not_crash(
- self, monkeypatch
- ):
+ def test_import_with_malformed_release_fallback_env_does_not_crash(self, monkeypatch):
monkeypatch.setenv("UNSLOTH_LLAMA_MAX_PREBUILT_RELEASE_FALLBACKS", "bad-value")
spec = importlib.util.spec_from_file_location(
"studio_install_llama_prebuilt_env_reload",
@@ -2108,7 +2031,11 @@ class TestResolveInstallReleasePlans:
class TestWindowsCudaAttempts:
TAG = "b8508"
- def _upstream(self, *runtime_versions, current_names: bool = False):
+ def _upstream(
+ self,
+ *runtime_versions,
+ current_names: bool = False,
+ ):
assets = {}
for rv in runtime_versions:
if current_names:
@@ -2350,23 +2277,14 @@ class TestPinnedBlackwellCudaFallback:
assert pin.runtime_sha256 and len(pin.runtime_sha256) == 64
def test_pin_offered_for_driver_13_2(self):
- assert (
- _pinned_windows_cuda_fallback(self._win_host((13, 2), ["120"]), [])
- is not None
- )
+ assert _pinned_windows_cuda_fallback(self._win_host((13, 2), ["120"]), []) is not None
def test_pin_offered_for_sm121_variant(self):
# sm_121 is Blackwell-family and also needs toolkit >= 12.8.
- assert (
- _pinned_windows_cuda_fallback(self._win_host((13, 1), ["121"]), [])
- is not None
- )
+ assert _pinned_windows_cuda_fallback(self._win_host((13, 1), ["121"]), []) is not None
def test_pin_uses_max_of_multi_gpu_caps(self):
- assert (
- _pinned_windows_cuda_fallback(self._win_host((13, 1), ["86", "120"]), [])
- is not None
- )
+ assert _pinned_windows_cuda_fallback(self._win_host((13, 1), ["86", "120"]), []) is not None
@pytest.mark.parametrize("sm", ["89", "90", "100"])
def test_pin_not_offered_to_non_blackwell(self, sm):
@@ -2377,16 +2295,11 @@ class TestPinnedBlackwellCudaFallback:
# b9360 is native sm_120a SASS (no JIT) and ships a cuda-13.1 cudart,
# both of which run on a 13.0 r580+ driver via CUDA minor-version
# compatibility. 13.0 is the mainstream Blackwell branch, so it must fire.
- assert (
- _pinned_windows_cuda_fallback(self._win_host((13, 0), ["120"]), [])
- is not None
- )
+ assert _pinned_windows_cuda_fallback(self._win_host((13, 0), ["120"]), []) is not None
def test_pin_not_offered_below_floor(self):
# 12.x predates Blackwell entirely; the pin stays dormant below 13.0.
- assert (
- _pinned_windows_cuda_fallback(self._win_host((12, 9), ["120"]), []) is None
- )
+ assert _pinned_windows_cuda_fallback(self._win_host((12, 9), ["120"]), []) is None
def test_pin_not_offered_without_driver(self):
assert _pinned_windows_cuda_fallback(self._win_host(None, ["120"]), []) is None
@@ -2460,10 +2373,7 @@ class TestPinnedBlackwellCudaFallback:
],
)
def test_attempt_covers_blackwell(self, minor, covers):
- assert (
- _windows_cuda_attempt_covers_blackwell(self._win_cuda_attempt(minor))
- is covers
- )
+ assert _windows_cuda_attempt_covers_blackwell(self._win_cuda_attempt(minor)) is covers
def test_attempt_covers_blackwell_ignores_non_cuda_kind(self):
cpu = AssetChoice(
@@ -2500,8 +2410,7 @@ class TestDirectUpstreamBlackwellPin:
return {
"tag_name": self.TAG,
"assets": [
- {"name": n, "browser_download_url": f"https://example.com/{n}"}
- for n in names
+ {"name": n, "browser_download_url": f"https://example.com/{n}"} for n in names
],
}
@@ -2521,15 +2430,9 @@ class TestDirectUpstreamBlackwellPin:
driver_cuda_version = (13, 1),
compute_caps = ["120"],
)
- plan = direct_upstream_release_plan(
- self._release(), host, UPSTREAM_REPO, "latest"
- )
+ plan = direct_upstream_release_plan(self._release(), host, UPSTREAM_REPO, "latest")
order = [(a.tag, a.runtime_line or a.install_kind) for a in plan.attempts]
- assert order == [
- ("b9360", "cuda13"),
- (self.TAG, "cuda12"),
- (self.TAG, "windows-cpu"),
- ]
+ assert order == [("b9360", "cuda13"), (self.TAG, "cuda12"), (self.TAG, "windows-cpu")]
assert plan.attempts[0].name == "llama-b9360-bin-win-cuda-13.1-x64.zip"
# Direct/upstream path stays unverified-by-manifest (no approved hashes).
assert plan.approved_checksums.artifacts == {}
@@ -2543,9 +2446,7 @@ class TestDirectUpstreamBlackwellPin:
driver_cuda_version = (13, 3),
compute_caps = ["120"],
)
- plan = direct_upstream_release_plan(
- self._release(), host, UPSTREAM_REPO, "latest"
- )
+ plan = direct_upstream_release_plan(self._release(), host, UPSTREAM_REPO, "latest")
assert "b9360" not in [a.tag for a in plan.attempts]
assert plan.attempts[0].tag == self.TAG
assert plan.attempts[0].runtime_line == "cuda13"
@@ -2581,9 +2482,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
# the old hardcoded cuda12/cuda13 seed would never order it (the cuda14
# line would be skipped for want of a 14.x asset in the seed).
mock_windows_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"])
- release = self._release(
- [("14.0", "cuda14"), ("13.3", "cuda13"), ("12.4", "cuda12")]
- )
+ release = self._release([("14.0", "cuda14"), ("13.3", "cuda13"), ("12.4", "cuda12")])
host = make_host(
system = "Windows",
machine = "AMD64",
@@ -2809,18 +2708,14 @@ class TestResolveUpstreamAssetChoice:
def test_linux_x86_64_cpu(self, monkeypatch):
name = f"llama-{self.TAG}-bin-ubuntu-x64.tar.gz"
self._mock_github_assets(monkeypatch, {name: f"https://x/{name}"})
- host = make_host(
- has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False
- )
+ host = make_host(has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False)
result = resolve_upstream_asset_choice(host, self.TAG)
assert result.install_kind == "linux-cpu"
assert result.name == name
def test_linux_cpu_missing(self, monkeypatch):
self._mock_github_assets(monkeypatch, {})
- host = make_host(
- has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False
- )
+ host = make_host(has_usable_nvidia = False, nvidia_smi = None, has_physical_nvidia = False)
with pytest.raises(PrebuiltFallback, match = "Linux CPU"):
resolve_upstream_asset_choice(host, self.TAG)
@@ -2907,9 +2802,7 @@ class TestResolveUpstreamAssetChoice:
has_physical_nvidia = False,
has_usable_nvidia = False,
)
- with pytest.raises(
- PrebuiltFallback, match = "no prebuilt policy exists for Linux aarch64"
- ):
+ with pytest.raises(PrebuiltFallback, match = "no prebuilt policy exists for Linux aarch64"):
resolve_upstream_asset_choice(host, self.TAG)
def test_windows_usable_nvidia_delegates(self, monkeypatch):
@@ -3022,7 +2915,11 @@ class TestResolveSimpleMacosPin:
],
}
- def fake_iter(repo, published_release_tag = "", requested_tag = ""):
+ def fake_iter(
+ repo,
+ published_release_tag = "",
+ requested_tag = "",
+ ):
calls.append((repo, published_release_tag, requested_tag))
# Emulate the real iterator: a specific tag yields only that release.
if requested_tag and requested_tag != "latest":
@@ -3031,9 +2928,7 @@ class TestResolveSimpleMacosPin:
for tag in self.TAGS:
yield _release(tag)
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", fake_iter
- )
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", fake_iter)
return calls
def test_pre26_host_pins_b9415_without_walkback(self, monkeypatch):
@@ -3081,14 +2976,10 @@ class TestLinuxArm64ForkFallsBackToSource:
def _boom(*_a, **_k):
raise AssertionError("iterator must not run for arm64 fork hosts")
- monkeypatch.setattr(
- INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", _boom
- )
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", _boom)
host = make_host(system = "Linux", machine = "aarch64")
with pytest.raises(PrebuiltFallback, match = "linux-x64 prebuilts"):
- resolve_simple_install_release_plans(
- "latest", host, "unslothai/llama.cpp", ""
- )
+ resolve_simple_install_release_plans("latest", host, "unslothai/llama.cpp", "")
def test_x86_64_fork_is_not_blocked_by_the_arch_guard(self, monkeypatch):
# x64 host must pass the guard and reach the iterator (here empty, so it
@@ -3100,9 +2991,7 @@ class TestLinuxArm64ForkFallsBackToSource:
)
host = make_host(system = "Linux", machine = "x86_64")
with pytest.raises(PrebuiltFallback) as exc:
- resolve_simple_install_release_plans(
- "latest", host, "unslothai/llama.cpp", ""
- )
+ resolve_simple_install_release_plans("latest", host, "unslothai/llama.cpp", "")
assert "linux-x64 prebuilts" not in str(exc.value)
def test_arm64_cpu_on_ggml_org_is_not_blocked(self, monkeypatch):
@@ -3123,9 +3012,7 @@ class TestLinuxArm64ForkFallsBackToSource:
has_usable_nvidia = False,
)
with pytest.raises(PrebuiltFallback) as exc:
- resolve_simple_install_release_plans(
- "latest", host, "ggml-org/llama.cpp", ""
- )
+ resolve_simple_install_release_plans("latest", host, "ggml-org/llama.cpp", "")
assert "linux-x64 prebuilts" not in str(exc.value)
@@ -3212,9 +3099,7 @@ class TestCpuFallback:
has_physical_nvidia = False,
has_usable_nvidia = False,
)
- plan = direct_upstream_release_plan(
- release, cpu_host, "ggml-org/llama.cpp", "latest"
- )
+ plan = direct_upstream_release_plan(release, cpu_host, "ggml-org/llama.cpp", "latest")
assert plan.attempts[0].install_kind == "linux-arm64"
assert plan.attempts[0].name == f"llama-{tag}-bin-ubuntu-arm64.tar.gz"
diff --git a/tests/studio/load_freeze/llama_server_shim.py b/tests/studio/load_freeze/llama_server_shim.py
index 1166c7521d..bb9e119820 100644
--- a/tests/studio/load_freeze/llama_server_shim.py
+++ b/tests/studio/load_freeze/llama_server_shim.py
@@ -41,7 +41,11 @@ class _Handler(BaseHTTPRequestHandler):
self.wfile.write(payload)
def _send_raw(
- self, status: int, body: bytes, *, content_type: str = "application/json"
+ self,
+ status: int,
+ body: bytes,
+ *,
+ content_type: str = "application/json",
) -> None:
self.send_response(status)
self.send_header("Content-Type", content_type)
@@ -119,9 +123,7 @@ class _Handler(BaseHTTPRequestHandler):
self._send_raw(srv.config.detok_status, srv.config.detok_body)
return
tids = body.get("tokens") or []
- content = "".join(
- srv.config.detok_map.get(int(t), f"") for t in tids
- )
+ content = "".join(srv.config.detok_map.get(int(t), f"") for t in tids)
self._send_json(srv.config.detok_status, {"content": content})
return
if path == "/completion":
@@ -224,9 +226,7 @@ class FakeLlamaServer:
def start(self) -> "FakeLlamaServer":
# port=0 lets ThreadingHTTPServer pick a free port atomically
# (avoids find-port-then-bind race); read back via server_address[1].
- self._server = FakeLlamaServer._Server(
- (self.host, self._requested_port), _Handler
- )
+ self._server = FakeLlamaServer._Server((self.host, self._requested_port), _Handler)
self._server.config = self.config
bound_port = self._server.server_address[1]
self._thread = threading.Thread(
diff --git a/tests/studio/load_freeze/test_load_orchestrator.py b/tests/studio/load_freeze/test_load_orchestrator.py
index 91b70f0b03..b76c4c361b 100644
--- a/tests/studio/load_freeze/test_load_orchestrator.py
+++ b/tests/studio/load_freeze/test_load_orchestrator.py
@@ -103,14 +103,18 @@ def _free_port() -> int:
class _UvicornServerThread:
- def __init__(self, app, *, host: str = "127.0.0.1", port: int) -> None:
+ def __init__(
+ self,
+ app,
+ *,
+ host: str = "127.0.0.1",
+ port: int,
+ ) -> None:
import uvicorn
self.host = host
self.port = port
- cfg = uvicorn.Config(
- app, host = host, port = port, log_level = "warning", access_log = False
- )
+ cfg = uvicorn.Config(app, host = host, port = port, log_level = "warning", access_log = False)
self._server = uvicorn.Server(cfg)
self._server.install_signal_handlers = lambda: None # type: ignore[assignment]
self._thread: threading.Thread | None = None
@@ -169,7 +173,12 @@ def _build_app(backend, *, wrap_in_thread: bool):
return app
-def _drive_concurrent_probe_and_health(base_url, *, n_health = 12, gap = 0.05):
+def _drive_concurrent_probe_and_health(
+ base_url,
+ *,
+ n_health = 12,
+ gap = 0.05,
+):
elapsed = -1.0
latencies: list[float] = []
@@ -211,9 +220,7 @@ def test_buggy_route_blocks_event_loop():
app = _build_app(backend, wrap_in_thread = False)
port = _free_port()
with _UvicornServerThread(app, port = port) as uv:
- max_lat, probe_t, _ = _drive_concurrent_probe_and_health(
- f"http://127.0.0.1:{uv.port}"
- )
+ max_lat, probe_t, _ = _drive_concurrent_probe_and_health(f"http://127.0.0.1:{uv.port}")
assert probe_t >= 0.5
assert max_lat >= 0.4, f"expected >=0.4s stall, got {max_lat:.3f}s"
@@ -442,9 +449,7 @@ def test_50_concurrent_probes_complete_without_deadlock():
with ThreadPoolExecutor(max_workers = 50) as pool:
futs = [
pool.submit(
- lambda: httpx.get(
- f"http://127.0.0.1:{uv.port}/probe", timeout = 30.0
- )
+ lambda: httpx.get(f"http://127.0.0.1:{uv.port}/probe", timeout = 30.0)
)
for _ in range(50)
]
@@ -660,7 +665,6 @@ def test_response_shape_matches_pre_fix_for_no_match():
bodies for the no-match scenario (the dominant code path in
practice for non-audio models)."""
import json as _json
-
with FakeLlamaServer(
detok_map = {128258: "abc", 128259: "def"},
tok_response_map = {
diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py
index efcd048b44..828187d462 100644
--- a/tests/studio/playwright_chat_ime_i18n.py
+++ b/tests/studio/playwright_chat_ime_i18n.py
@@ -246,8 +246,7 @@ with sync_playwright() as p:
dir_attr = composer.evaluate("(el) => el.getAttribute('dir')")
if dir_attr != "auto":
soft_fail(
- f'composer is missing dir="auto" (got {dir_attr!r}); RTL '
- "languages will render LTR."
+ f'composer is missing dir="auto" (got {dir_attr!r}); RTL ' "languages will render LTR."
)
else:
info('composer dir="auto" present')
@@ -258,9 +257,7 @@ with sync_playwright() as p:
_thread_src = (
_repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx"
).read_text()
- _shared_src = (
- _repo_root / "studio/frontend/src/features/chat/shared-composer.tsx"
- ).read_text()
+ _shared_src = (_repo_root / "studio/frontend/src/features/chat/shared-composer.tsx").read_text()
_edit_idx = _thread_src.find("aui-edit-composer-input")
if _edit_idx == -1 or 'dir="auto"' not in _thread_src[_edit_idx : _edit_idx + 600]:
soft_fail('edit composer source is missing dir="auto"')
@@ -269,8 +266,7 @@ with sync_playwright() as p:
_compare_idx = _shared_src.find("Send to both models")
if (
_compare_idx == -1
- or 'dir="auto"'
- not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400]
+ or 'dir="auto"' not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400]
):
soft_fail('compare composer source is missing dir="auto"')
else:
@@ -484,9 +480,7 @@ with sync_playwright() as p:
# handleSubmit / blockSend guards keep refusing. The Send button stays
# visually enabled (watchdog has already cleared the React state); the
# refusal happens at form.requestSubmit() time, not at the button.
- step(
- "BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)"
- )
+ step("BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)")
clear()
composer.click()
composer.evaluate(
@@ -533,9 +527,7 @@ with sync_playwright() as p:
"Form submitted after an IME keydown -- preedit text leaked "
"through the watchdog gap (#5546 follow-up regression)."
)
- info(
- f"Form submit refused after IME keydown; textarea retained {submit_probe.get('after')!r}"
- )
+ info(f"Form submit refused after IME keydown; textarea retained {submit_probe.get('after')!r}")
shoot("06c-keydown-repin")
info("keydown re-pin gate PASS")
clear()
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index b1279e64b9..dc62194be8 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -141,10 +141,7 @@ def expected_default_model():
for node in tree.body:
if not isinstance(node, ast.Assign):
continue
- if not any(
- isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF"
- for t in node.targets
- ):
+ if not any(isinstance(t, ast.Name) and t.id == "DEFAULT_MODELS_GGUF" for t in node.targets):
continue
try:
models = ast.literal_eval(node.value)
@@ -321,9 +318,7 @@ with sync_playwright() as p:
form_err: Exception | None = None
for _form_attempt in range(3):
try:
- page.goto(
- f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000
- )
+ page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000)
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
except Exception:
@@ -382,9 +377,7 @@ with sync_playwright() as p:
flush = True,
)
if page_errors:
- print(
- f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True
- )
+ print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True)
try:
shoot(f"01-change-password-attempt-{_form_attempt + 1}-fail")
except Exception:
@@ -458,9 +451,7 @@ with sync_playwright() as p:
flush = True,
)
if page_errors:
- print(
- f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True
- )
+ print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True)
try:
shoot(f"03-composer-wait-attempt-{_attempt + 1}-fail")
except Exception:
@@ -557,9 +548,7 @@ with sync_playwright() as p:
try:
sel_text = (selector_btn.text_content(timeout = 2_000) or "").strip()
except Exception as _sel_err:
- info(
- f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}"
- )
+ info(f"WARN: model-selector probe skipped: {type(_sel_err).__name__}: {_sel_err}")
if sel_text:
info(f"model selector button text: {sel_text!r}")
shoot("03b-default-model-button")
@@ -595,10 +584,7 @@ with sync_playwright() as p:
if load_resp.get("error"):
fail(f"/api/inference/load wedged: {load_resp['error']!r}")
if load_resp["status"] != 200:
- fail(
- f"/api/inference/load returned {load_resp['status']}: "
- f"{load_resp.get('body')!r}"
- )
+ fail(f"/api/inference/load returned {load_resp['status']}: " f"{load_resp.get('body')!r}")
info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}")
# Studio caches the per-context model state in zustand; reload
@@ -845,8 +831,7 @@ with sync_playwright() as p:
# Look for either "Disable X" or "Enable X" -- whichever
# is currently rendered.
toggle = page.locator(
- f'button[aria-label="Disable {feature}"], '
- f'button[aria-label="Enable {feature}"]'
+ f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
).first
if toggle.count() == 0:
info(f"toggle '{feature}' not present on this layout")
@@ -862,8 +847,7 @@ with sync_playwright() as p:
page.wait_for_timeout(200)
after = (
page.locator(
- f'button[aria-label="Disable {feature}"], '
- f'button[aria-label="Enable {feature}"]'
+ f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
).first.get_attribute("aria-label")
or ""
)
@@ -874,8 +858,7 @@ with sync_playwright() as p:
# Flip back so test state is unchanged.
try:
page.locator(
- f'button[aria-label="Disable {feature}"], '
- f'button[aria-label="Enable {feature}"]'
+ f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
).first.click()
except Exception:
pass
@@ -968,8 +951,7 @@ with sync_playwright() as p:
except Exception as exc:
if attempt == 1:
soft_fail(
- f"theme cycle {cycle + 1}: account-menu click failed "
- f"({exc!r})"
+ f"theme cycle {cycle + 1}: account-menu click failed " f"({exc!r})"
)
continue
try:
@@ -1020,8 +1002,7 @@ with sync_playwright() as p:
if click_err is not None:
page.keyboard.press("Escape")
soft_fail(
- f"theme cycle {cycle + 1}: theme menuitem click failed "
- f"({click_err!r})"
+ f"theme cycle {cycle + 1}: theme menuitem click failed " f"({click_err!r})"
)
break
# Settle. The ".dark" class on is the ground
@@ -1078,9 +1059,7 @@ with sync_playwright() as p:
# progressively more permissive locators so the test stays
# green on both platforms.
candidates = [
- page.get_by_role(
- "button", name = re.compile(rf"^\s*{label}\s*$", re.I)
- ).first,
+ page.get_by_role("button", name = re.compile(rf"^\s*{label}\s*$", re.I)).first,
page.locator(f'button:has-text("{label}")').first,
page.locator(f'a:has-text("{label}")').first,
page.locator(f'[data-sidebar="menu-button"]:has-text("{label}")').first,
@@ -1116,15 +1095,11 @@ with sync_playwright() as p:
click_nav("New Chat", r"/chat")
shoot("11-new-chat")
# Compare moved into the composer + menu (Tools and attachments).
- plus_btn = page.get_by_role(
- "button", name = re.compile(r"Tools and attachments", re.I)
- ).first
+ plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first
if plus_btn.count() > 0:
plus_btn.click(force = True)
page.wait_for_timeout(400)
- compare_item = page.get_by_role(
- "menuitem", name = re.compile(r"Compare chat", re.I)
- ).first
+ compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
if compare_item.count() > 0:
compare_item.click(force = True)
page.wait_for_timeout(800)
@@ -1159,9 +1134,7 @@ with sync_playwright() as p:
step("Developer (API) tab via account menu")
acct.click()
page.wait_for_timeout(400)
- dev = page.get_by_role(
- "menuitem", name = re.compile(r"developer|api", re.I)
- ).first
+ dev = page.get_by_role("menuitem", name = re.compile(r"developer|api", re.I)).first
if dev.count() > 0:
dev.click()
page.wait_for_timeout(800)
@@ -1178,9 +1151,7 @@ with sync_playwright() as p:
re.compile(r"api keys|developer", re.I),
).first
if keys_section.count() > 0:
- info(
- f"OK API tab text: {(keys_section.text_content() or '').strip()[:80]!r}"
- )
+ info(f"OK API tab text: {(keys_section.text_content() or '').strip()[:80]!r}")
# Close dialog with Escape.
page.keyboard.press("Escape")
page.wait_for_timeout(300)
@@ -1198,9 +1169,7 @@ with sync_playwright() as p:
page.wait_for_timeout(1500)
# Recipe cards are rendered as or button elements; count
# all clickable headings under main + screenshot.
- headings = page.locator(
- "main h2, main h3, [data-recipe], a[href*='/data-recipes/']"
- )
+ headings = page.locator("main h2, main h3, [data-recipe], a[href*='/data-recipes/']")
n_cards = headings.count()
info(f"Recipes route headings/cards: {n_cards}")
shoot("15b-recipes-cards")
@@ -1289,10 +1258,7 @@ with sync_playwright() as p:
info(f"recent-thread click {i} failed: {_click_err!s}")
continue
if not clicked_recent:
- soft_fail(
- f"no Recents entry was clickable within 30s deadline "
- f"(n_threads={n_threads})"
- )
+ soft_fail(f"no Recents entry was clickable within 30s deadline " f"(n_threads={n_threads})")
# Back to chat.
page.goto(f"{BASE}/chat")
composer = page.locator('textarea[aria-label="Message input"]')
diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py
index 0ac9f67a7e..20c7bda87c 100644
--- a/tests/studio/playwright_extra_ui.py
+++ b/tests/studio/playwright_extra_ui.py
@@ -170,9 +170,7 @@ with sync_playwright() as p:
form_err: Exception | None = None
for _form_attempt in range(3):
try:
- page.goto(
- f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000
- )
+ page.goto(f"{BASE}/change-password", wait_until = "domcontentloaded", timeout = 60_000)
try:
page.wait_for_load_state("networkidle", timeout = 30_000)
except Exception:
@@ -329,15 +327,11 @@ with sync_playwright() as p:
step("Compare tab: send to two panes")
# Compare moved into the composer + menu (Tools and attachments).
compare_opened = False
- plus_btn = page.get_by_role(
- "button", name = re.compile(r"Tools and attachments", re.I)
- ).first
+ plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first
if plus_btn.count() > 0:
plus_btn.click(force = True)
page.wait_for_timeout(400)
- compare_item = page.get_by_role(
- "menuitem", name = re.compile(r"Compare chat", re.I)
- ).first
+ compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
if compare_item.count() > 0:
compare_item.click(force = True)
compare_opened = True
@@ -416,9 +410,7 @@ with sync_playwright() as p:
arg = ok_count_before + 4,
timeout = 60_000,
)
- info(
- "OK Compare: 4 total new assistant bubbles after second prompt"
- )
+ info("OK Compare: 4 total new assistant bubbles after second prompt")
except Exception as exc:
runtime_warn(
f"Compare: 4 bubbles didn't appear (panes likely "
@@ -439,9 +431,7 @@ with sync_playwright() as p:
page.wait_for_timeout(1500)
shoot("05-recipes-list")
# Template cards render as