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

Merge remote-tracking branch 'origin/main' into r5748

Browse source
This commit is contained in:
Daniel Han 2026-07-16 04:54:39 +00:00
commit 1c39a283f6
426 changed files with 66427 additions and 6621 deletions
Download patch file Download diff file Expand all files Collapse all files

122
tests/python/test_fast_sentence_transformer_embedding_parity.py Normal file
View file

@ -0,0 +1,122 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text
like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a
"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building
via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings;
`_create_transformer_module` uses `Transformer.load(...)` instead.
Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST
is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity,
opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected).
"""
from __future__ import annotations
import inspect
import os
import pytest
def test_transformer_load_signature_supports_unsloth_kwargs():
"""Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs
the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back
to Transformer(...) there, so mirror that gate and skip."""
models = pytest.importorskip("sentence_transformers.models")
load = getattr(models.Transformer, "load", None)
assert callable(load), (
"sentence_transformers Transformer.load is missing; the #6881 fix in "
"unsloth.models.sentence_transformer._create_transformer_module depends on it."
)
params = inspect.signature(load).parameters
accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
# Mirror _create_transformer_module's hub_capable gate.
hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
if not hub_capable:
pytest.skip(
"legacy Transformer.load(input_path); production path falls back to Transformer(...)"
)
unsupported = [
k
for k in ("token", "cache_folder", "revision", "trust_remote_code")
if not (accepts_var_kw or k in params)
]
assert not unsupported, (
f"installed sentence_transformers Transformer.load no longer accepts {unsupported} "
f"and has no **kwargs; update _create_transformer_module (#6881) before it silently "
f"falls back to Transformer(...)."
)
def _probe_texts():
return [
"roasted chickpeas in 20 kg bags",
"The capital of France is Paris.",
"A fast brown fox jumps over the lazy dog.",
"recette de tarte aux pommes traditionnelle",
]
def test_fast_sentence_transformer_matches_stock_st():
"""End-to-end: FastSentenceTransformer embeddings and tokenization must match a
stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and
GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners."""
model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL")
if not model_id:
pytest.skip(
"set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model "
"(HF id or local path) to run the #6881 parity test"
)
torch = pytest.importorskip("torch")
if not torch.cuda.is_available():
pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
np = pytest.importorskip("numpy")
pytest.importorskip("sentence_transformers")
from sentence_transformers import SentenceTransformer
device = "cuda"
# Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native
# embedders such as EmbeddingGemma (Gemma3), which would mask real parity.
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
texts = _probe_texts()
max_seq_length = 256
# Control FIRST, before importing unsloth, so its global import patches never
# touch the stock reference (mirrors the issue's "restart runtime" repro).
ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
ctrl.max_seq_length = max_seq_length
ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
ctrl_emb = np.asarray(
ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
)
import unsloth # noqa: F401
from unsloth import FastSentenceTransformer
fast = FastSentenceTransformer.from_pretrained(
model_id,
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = False,
load_in_16bit = True,
)
fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist()
fast_emb = np.asarray(
fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
)
# Identical tokenization = no chat-template wrapping slipped in (the #6881 defect).
assert fast_ids == ctrl_ids, (
f"tokenization diverged (chat-template wrapping regressed?):\n"
f" stock: {ctrl_ids}\n fast: {fast_ids}"
)
cos = (ctrl_emb * fast_emb).sum(1) / (
np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1)
)
assert float(cos.min()) > 0.99, (
f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 "
f"(per-text {[round(float(c), 5) for c in cos]})"
)

203
tests/python/test_flash_attn_install_python_stack.py
View file

@ -13,102 +13,35 @@ sys.path.insert(0, str(STUDIO_DIR))
sys.path.insert(0, str(STUDIO_DIR / "backend"))
import install_python_stack as ips
from backend.utils import wheel_utils
from utils import wheel_utils
def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess:
return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "")
class TestPrebuiltWheelTorchMapping:
def test_torch_211_maps_to_torch210(self):
assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10"
def test_other_versions_pass_through(self):
for torch_mm in ("2.9", "2.10", "2.12"):
assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm
class TestHasBlackwellGpu:
def setup_method(self):
wheel_utils.has_blackwell_gpu.cache_clear()
def teardown_method(self):
wheel_utils.has_blackwell_gpu.cache_clear()
def test_returns_false_when_nvidia_smi_missing(self):
with mock.patch.object(wheel_utils.shutil, "which", return_value = None):
assert wheel_utils.has_blackwell_gpu() is False
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")),
):
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")),
):
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")),
):
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")),
):
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")),
):
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.subprocess,
"run",
return_value = _smi_result("8.0\n10.0\n"),
),
):
assert wheel_utils.has_blackwell_gpu() is True
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.subprocess,
"run",
return_value = _smi_result("", returncode = 1),
),
):
assert wheel_utils.has_blackwell_gpu() is False
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.subprocess,
"run",
side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10),
),
):
assert wheel_utils.has_blackwell_gpu() is False
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.subprocess,
"run",
return_value = _smi_result("not-a-number\n\n"),
),
):
assert wheel_utils.has_blackwell_gpu() is False
def test_direct_wheel_url_reuses_torch210_on_211(self):
# causal-conv1d / mamba go through direct_wheel_url; torch 2.11 reuses the
# torch2.10 wheel filename just like flash-attn does.
url = wheel_utils.direct_wheel_url(
filename_prefix = "causal_conv1d",
package_version = "1.6.1",
release_tag = "v1.6.1.post4",
release_base_url = "https://example.test/download",
env = {
"python_tag": "cp313",
"torch_mm": "2.11",
"cuda_major": "13",
"cxx11abi": "TRUE",
"platform_tag": "linux_x86_64",
},
)
assert url is not None
assert "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
class TestFlashAttnWheelSelection:
@ -118,9 +51,24 @@ class TestFlashAttnWheelSelection:
def test_torch_29_maps_to_v283(self):
assert ips._select_flash_attn_version("2.9") == "2.8.3"
def test_unsupported_torch_has_no_wheel_mapping(self):
def test_torch_211_has_no_native_version_entry(self):
# The raw version table has no torch2.11-tagged wheel; the URL builder
# reuses the torch2.10 wheel instead (see test_torch_211_reuses_torch210_wheel).
assert ips._select_flash_attn_version("2.11") is None
def test_torch_211_reuses_torch210_wheel(self):
url = ips._build_flash_attn_wheel_url(
{
"python_tag": "cp313",
"torch_mm": "2.11",
"cuda_major": "13",
"cxx11abi": "TRUE",
"platform_tag": "linux_x86_64",
}
)
assert url is not None
assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url
def test_exact_wheel_url_uses_full_env_tuple(self):
url = ips._build_flash_attn_wheel_url(
{
@ -333,83 +281,22 @@ class TestEnsureFlashAttn:
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
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,
):
step_messages.append((label, value))
with (
mock.patch.object(ips, "NO_TORCH", False),
mock.patch.object(ips, "IS_WINDOWS", False),
mock.patch.object(ips, "IS_MACOS", False),
mock.patch.object(ips, "has_blackwell_gpu", return_value = True),
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
mock.patch.object(ips, "_step", side_effect = fake_step),
mock.patch("subprocess.run", return_value = self._import_check()),
):
ips._ensure_flash_attn()
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)
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,
):
step_messages.append((label, value))
def test_windows_skips_install_without_probing(self):
# flash-attn is Linux-only: on Windows the installer returns before
# probing the torch env or resolving a wheel (no Windows wheels are
# published upstream).
with (
mock.patch.object(ips, "NO_TORCH", False),
mock.patch.object(ips, "IS_WINDOWS", True),
mock.patch.object(ips, "IS_MACOS", False),
mock.patch.object(ips, "has_blackwell_gpu", return_value = True),
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
mock.patch.object(ips, "_step", side_effect = fake_step),
mock.patch("subprocess.run", return_value = self._import_check()),
):
ips._ensure_flash_attn()
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)
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,
):
step_messages.append((label, value))
with (
mock.patch.object(ips, "NO_TORCH", False),
mock.patch.object(ips, "IS_WINDOWS", True),
mock.patch.object(ips, "IS_MACOS", False),
mock.patch.object(ips, "has_blackwell_gpu", return_value = False),
mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe,
mock.patch.object(ips, "install_wheel") as mock_install_wheel,
mock.patch.object(ips, "_step", side_effect = fake_step),
mock.patch("subprocess.run", return_value = self._import_check()),
):
ips._ensure_flash_attn()
mock_probe.assert_not_called()
mock_install_wheel.assert_not_called()
assert not any("Blackwell" in msg for _, msg in step_messages)
class TestInstallPythonStackFlashAttnIntegration:

58
tests/python/test_mlx_public_trainer_api.py
View file

@ -115,6 +115,9 @@ def test_mlx_training_arguments_accept_trl_style_kwargs():
def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras():
"""Implemented and falsey inert compatibility kwargs should stay quiet."""
unsloth = _import_mlx_unsloth()
supported_eval_kwargs = {}
if "eval_strategy" in unsloth._MLX_TRAINING_CONFIG_FIELDS:
supported_eval_kwargs = {"eval_strategy": "no", "eval_delay": 1}
with warnings.catch_warnings(record = True) as caught:
warnings.simplefilter("always")
@ -125,12 +128,16 @@ def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras():
remove_unused_columns = False,
assistant_only_loss = False,
completion_only_loss = False,
**supported_eval_kwargs,
)
assert args.warmup_steps == 2
assert args.padding_free is False
assert args.remove_unused_columns is False
assert args.completion_only_loss is False
if supported_eval_kwargs:
assert args.eval_strategy == "no"
assert args.eval_delay == 1
assert caught == []
@ -705,17 +712,10 @@ def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs():
)
def test_mlx_trainer_rejects_metrics_and_callbacks():
"""Trainer hooks should fail because MLXTrainer cannot honor them yet."""
def test_mlx_trainer_rejects_compute_metrics():
"""compute_metrics is still unsupported by MLXTrainer."""
unsloth = _import_mlx_unsloth()
with pytest.raises(NotImplementedError, match = "callbacks"):
unsloth.UnslothTrainer(
model = _DummyModel(),
tokenizer = None,
train_dataset = [],
callbacks = [object()],
)
with pytest.raises(NotImplementedError, match = "compute_metrics"):
unsloth.UnslothTrainer(
model = _DummyModel(),
@ -725,6 +725,46 @@ def test_mlx_trainer_rejects_metrics_and_callbacks():
)
def test_mlx_trainer_accepts_callbacks():
"""Callbacks are routed to MLXTrainer when the zoo backend supports them."""
unsloth = _import_mlx_unsloth()
from transformers import TrainerCallback
if not unsloth._mlx_trainer_supports_kwarg("callbacks"):
pytest.skip("requires unsloth-zoo MLXTrainer callback support")
class Callback(TrainerCallback):
pass
trainer = unsloth.UnslothTrainer(
model = _DummyModel(),
tokenizer = None,
train_dataset = [],
callbacks = [Callback()],
)
assert any(isinstance(cb, Callback) for cb in trainer.callback_handler.callbacks)
def test_mlx_trainer_rejects_callbacks_with_old_zoo(monkeypatch):
"""Older unsloth-zoo builds should fail clearly instead of TypeError."""
unsloth = _import_mlx_unsloth()
from transformers import TrainerCallback
monkeypatch.setattr(
unsloth,
"_mlx_trainer_supports_kwarg",
lambda name: name != "callbacks",
)
with pytest.raises(NotImplementedError, match = "callbacks require"):
unsloth.UnslothTrainer(
model = _DummyModel(),
tokenizer = None,
train_dataset = [],
callbacks = [TrainerCallback()],
)
def test_mlx_trainer_rejects_custom_data_collator():
"""MLXTrainer owns batching; custom SFT data collators must not be ignored."""
unsloth = _import_mlx_unsloth()

46
tests/python/test_remove_special_tokens_no_bos.py Normal file
View file

@ -0,0 +1,46 @@
import ast
from pathlib import Path
def _load_remove_special_tokens():
# Extract remove_special_tokens without importing unsloth (importing unsloth
# needs unsloth_zoo / a GPU). The function is pure Python and uses no imports,
# so it execs cleanly in an empty namespace.
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
funcs = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "remove_special_tokens"
]
namespace = {}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace["remove_special_tokens"]
class _StubTokenizer:
def __init__(self, bos_token):
self.bos_token = bos_token
def test_no_bos_tokenizer_does_not_crash():
# Tokenizers such as Qwen2 / Qwen2.5, GPT-2, Falcon and GPT-NeoX have no BOS
# token, so tokenizer.bos_token is None. remove_special_tokens must leave the
# prompt untouched instead of raising
# "TypeError: startswith first arg must be str or a tuple of str, not NoneType".
remove_special_tokens = _load_remove_special_tokens()
assert remove_special_tokens(_StubTokenizer(None), "Hello world") == "Hello world"
def test_double_bos_is_stripped():
# A tokenizer with a BOS token still has a single leading BOS removed.
remove_special_tokens = _load_remove_special_tokens()
assert remove_special_tokens(_StubTokenizer("<s>"), "<s>Hello world") == "Hello world"
def test_prompt_without_leading_bos_unchanged():
# A BOS-bearing tokenizer leaves a prompt that does not start with BOS alone.
remove_special_tokens = _load_remove_special_tokens()
assert remove_special_tokens(_StubTokenizer("<s>"), "Hello world") == "Hello world"

97
tests/python/test_to_sharegpt_optional_none.py Normal file
View file

@ -0,0 +1,97 @@
import ast
import re
from pathlib import Path
def _load_formatter_builders():
# Extract _parse_combined_prompt and _create_formatter without importing
# unsloth (importing unsloth needs unsloth_zoo / a GPU). Both are pure
# Python and only use the `re` module.
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
wanted = {"_parse_combined_prompt", "_create_formatter"}
funcs = [
node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted
]
namespace = {"re": re}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace["_parse_combined_prompt"], namespace["_create_formatter"]
class _StubDataset:
def __init__(self, column_names):
self.column_names = column_names
def _render(merged_prompt, columns, batch):
parse, create = _load_formatter_builders()
possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns))
processor = create(possible_columns, final_optional_prompts, "text")
return processor(batch)["text"]
def test_optional_block_missing_second_column_does_not_render_none():
# A [[...]] block may reference several columns; only the first gates the
# block. A later column that is None must not render as the literal "None".
merged_prompt = "Location: [[{city}, {country}]] end"
out = _render(
merged_prompt,
["city", "country"],
{"city": ["Paris"], "country": [None]},
)
assert out[0] == "Location: Paris, end"
assert "None" not in out[0]
def test_optional_block_all_columns_present_unchanged():
merged_prompt = "Location: [[{city}, {country}]] end"
out = _render(
merged_prompt,
["city", "country"],
{"city": ["Paris"], "country": ["France"]},
)
assert out[0] == "Location: Paris, France end"
def test_optional_block_gating_column_empty_is_dropped():
# When the gating (first) column is empty the whole block is omitted; this
# behaviour is unchanged by the None coercion.
merged_prompt = "Location: [[{city}, {country}]] end"
out = _render(
merged_prompt,
["city", "country"],
{"city": [""], "country": ["France"]},
)
assert out[0] == "Location: end"
def test_single_column_optional_block_gated_out_on_none():
# Single-column blocks were already gated correctly (the sole column is the
# gate); confirm they stay unaffected.
merged_prompt = "Name: [[{name}]]!"
out = _render(merged_prompt, ["name"], {"name": [None, "Bob"]})
assert out == ["Name: !", "Name: Bob!"]
def test_required_column_none_does_not_render_none():
# A required (non-[[...]]) column that is None must not render as the
# literal "None" either; coercion happens at the row source, so both the
# required and optional branches are covered.
merged_prompt = "Location: {city}, {country} end"
out = _render(
merged_prompt,
["city", "country"],
{"city": ["Paris"], "country": [None]},
)
assert out[0] == "Location: Paris, end"
assert "None" not in out[0]
def test_optional_block_falsy_but_present_gating_value_still_renders():
# The gate keeps a block whenever the first column is not "". A falsy but
# real value (0) must not be treated as absent, so the block still renders.
merged_prompt = "Count: [[{n}]]!"
out = _render(merged_prompt, ["n"], {"n": [0]})
assert out[0] == "Count: 0!"

51
tests/python/test_tokenizers_and_torch_constraint.py
View file

@ -14,6 +14,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"
_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged:
assert '"torch>=2.4,<2.11.0"' in self._ps1
class TestInstallPs1UvDefaultIndex:
"""Installer-managed torch indexes must override inherited uv defaults."""
_ps1 = _read(_INSTALL_PS1)
def test_torch_installs_use_default_index(self):
assert "--default-index $TorchIndexUrl" in self._ps1
assert "--default-index $ROCmIndexUrl" in self._ps1
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert "--index-url $TorchIndexUrl" not in self._ps1
assert "--index-url $ROCmIndexUrl" not in self._ps1
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# Extra-index vars outrank --default-index, so pinned installs must clear them.
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
assert 'Remove-Item "Env:$n"' in self._ps1
class TestSetupPs1FastInstallIndex:
"""setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
_ps1 = _read(_SETUP_PS1)
def test_fast_install_clears_all_uv_index_env_vars(self):
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
# Must truly remove the vars (child sees no value), not set them empty.
assert 'Remove-Item "Env:$n"' in self._ps1
class TestInstallShUvDefaultIndex:
"""Linux/Mac installer torch indexes must override inherited uv defaults."""
_sh = _read(_INSTALL_SH)
def test_torch_installs_use_default_index(self):
assert '--default-index "$TORCH_INDEX_URL"' in self._sh
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# --default-index installs run with all uv index env vars unset via `env -u`.
assert (
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
)
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
class TestTorchConstraintShell:
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""