a00fe86c13* Studio: read model text as utf-8 so umlauts survive on Windows Chat rejects or mangles non-ASCII on Windows: "ä ö ü" in a prompt, a chat template, or a model path comes back as mojibake, or the load dies with UnicodeDecodeError. open() and Path.read_text() fall back to locale.getencoding() when no encoding is passed. On Windows that is the ANSI codepage (cp1252, cp932, cp1251, ... by system locale), never UTF-8. Hugging Face writes these files as raw UTF-8, so every read of one decodes with the wrong codec: - tokenizer_config.json, which holds the chat template. Templates routinely carry -> arrows, smart quotes and CJK, so this is the common path into chat - config.json and adapter_config.json - modules.json, Ollama manifests, and the .py sources the remote-code scanner reads before a model is allowed to load The llama-server and embedding-server stdout readers have the same problem via subprocess(text = True); they now decode utf-8 with errors = "replace" so a stray byte cannot kill a log reader. Encoding arguments only, no logic changes. tests/test_chat_text_encoding.py covers a config.json and a chat template holding umlauts, arrows and CJK, plus the remote-code scanner reading a source file with umlauts. Those pass anywhere the locale is already UTF-8, so a fourth test re-runs the readers under -X warn_default_encoding and fails on any platform if an encoding argument goes missing again. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: name utf-8 explicitly on the remaining text I/O, with an AST guard (#7465) * Studio: name utf-8 explicitly on the remaining text I/O Follow-up to the model-text reads in #7467, covering the rest of the backend: system probes (nvidia-smi, amd-smi, powershell, git, node), package installers, /proc and /sys readers, and internal marker files (pid, install id, bootstrap password, Colab credentials). Same reason as #7467. open(), Path.read_text()/write_text() and subprocess(text = True) fall back to locale.getencoding(), which on Windows is the ANSI codepage rather than UTF-8. These paths are mostly ASCII today, so this is hardening, not a live bug. Encoding arguments only, no logic changes. Adds tests/test_text_io_encoding.py: an AST guard walking every backend source and asserting text I/O names its encoding, so the class of bug cannot creep back in one call at a time. 275 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Catch aliased subprocess and positional Path.open, migrate legacy JSONL The guard only matched a receiver literally named subprocess, so worker.py's `import subprocess as _sp` hid three text = True installs that decode pip output with the ANSI codepage. It also skipped any .open() with more than one positional argument, though Path.open takes buffering/encoding/errors/newline positionally. Resuming a scrape written by an older release is the other half: those JSONL lines are in the locale codepage, so the UTF-8 preload raised, the dedup keys were silently forgotten and duplicates were appended to a now mixed-encoding file. Decode with the locale codepage as fallback and rewrite as UTF-8 before the append handle opens, since Windows cannot replace a file it holds open. * Stream the JSONL preload and keep a torn line from relabelling the shard Reading the whole shard to migrate it was wrong twice over. These files reach gigabytes on a large scrape, so the preload now streams line by line and the rewrite streams through a temp file. Worse, one interrupted append used to condemn the file: the whole-file UTF-8 decode failed, every byte was retried as cp1252, and the rewrite persisted mojibake over records that were fine. A line now counts as legacy only if the locale codepage both decodes it and yields valid JSON, which a torn UTF-8 line does not. Damaged lines are skipped and copied through byte for byte. When the rewrite cannot be written at all, the append handle opens with the legacy encoding rather than mixing UTF-8 into the file. install_wheel takes run = subprocess.run as a parameter, so the guard cannot see it. Both wheel installs there now name their encoding. * Decide the shard's encoding from the file, not one line at a time Some byte strings parse both ways. cp1251 `Р°` is D0 B0, which is also valid UTF-8 for `а`, so a UTF-8-first parse quietly showed the wrong text instead of migrating it. A line now yields both readings, and the file decides. Any line that parses under the codepage but not as UTF-8 is unambiguous evidence, and ambiguous lines then follow that verdict, which is enough for any real shard: ordinary Cyrillic or Japanese prose is invalid UTF-8 several times per line. Keys for ambiguous lines are re-derived from the legacy reading during the rewrite. A shard is undecidable only if every line is ambiguous, and nothing can tell those apart. latin-1 is also tried after the locale codepage, so a scrape carried from Windows to a UTF-8 machine still has a reading rather than none. Requiring valid JSON, not just a decode, keeps that from claiming torn lines. * Weigh the whole shard, and never lose a record on the fallback path One structurally valid JSON line carrying a stray 0x96 parses as cp1252, so a single-line verdict let it relabel a healthy shard and mojibake every good record in it. Each line with non-ASCII bytes now votes: parsing only under the codepage is evidence for legacy, parsing as UTF-8 is evidence against, since codepage text rarely forms valid multibyte UTF-8. Ties leave the file alone. When the migration cannot be written the append handle uses the legacy codepage, and errors = "replace" quietly turned characters it cannot hold into question marks while write() still reported success. That path now escapes to \uXXXX instead, which is ASCII, so every codepage holds it and json.loads returns the exact characters. Nothing needs replacing, so errors = "strict" is safe. stream_installer runs sys.executable, so its output is now decoded as UTF-8 by utf8_child_env rather than read as the ANSI codepage. * Only rewrite a shard we can attribute, and append ASCII when we cannot latin-1 was doing too much work. It reads any byte, so it gave a moved shard a reading, but it is the right text only for cp1252: cp1251 Привет came back as Ïðèâåò and the rewrite made that permanent. The codepage is now trusted only when it is the locale's, and an untrusted reading is never written back. That leaves three cases where the file holds bytes UTF-8 cannot read and we are not converting it: no codepage to attribute it to, ambiguous lines outvoting the unambiguous ones, and a preload that could not read the file at all. All three used to append UTF-8 into it. They now append pure ASCII, which every ASCII-compatible codepage stores identically, so the file keeps decoding exactly as it did and no record is lost. Keys from the two readings are also kept apart. A damaged line in a healthy shard was marked seen through its codepage reading, so the retry that would have replaced the unreadable record was refused as a duplicate. * Let the flash-attn install stub take the kwargs the installer now passes _run_kwargs gained encoding and errors, so the one stub in this file that spelled its signature out rejected the call. The other four here already take **kwargs; this one now matches. * Do not let a stuck temp file mask the migration failure unlink() on the failure path could raise in its own right, on a stale .utf8.tmp directory or a temp another process holds. That escaped the constructor instead of returning False, so the caller never reached the ASCII append fallback that keeps the shard single-encoding. The pip fallback in install_wheel also spawns a Python child, so it gets utf8_child_env like the probe above it already had. The uv and nvidia-smi children are native binaries, where PYTHONIOENCODING would do nothing. * Stop converting legacy shards; the encoding that wrote them is unknowable trusted only ever meant that the bytes parse under this machine's codepage, which for a single-byte codepage is nearly always true. A cp1251 shard opened on a cp1252 Windows box decodes cleanly and would have been rewritten with Привет as Ïðèâåò. That is the fourth way this rewrite could corrupt a shard, and the common cause is that a file's encoding cannot be recovered from its bytes. So the rewrite is gone. The shard is left exactly as found, and appends are pure ASCII whenever it holds bytes UTF-8 cannot read, which is what actually delivered the no-mixed-encoding guarantee the rewrite was added for. Dedup keys still come from whichever reading parses, since ids are ASCII either way. This also removes the temp file, so there is no longer any file mode or ACL to carry across. * Scan the sandbox shim; it is shipped code, not a build artifact sandbox_site is on the sandboxed child's PYTHONPATH for every Python run (tools.py:332, 2660), so excluding it let two unannotated text calls through in code we ship. Both read and write the remap sidecar, which holds file paths. The exclusion list is meant for build output only, so the directory comes off it and the two calls name their encoding. * Force the worker's pip children to UTF-8, and read DBCS keys with a DBCS codec The three installer calls run sys.executable -m pip with an inherited environment, so the parent decoded UTF-8 while the child emitted the ANSI codepage. They now go through utf8_child_env like the other Python children. Two tests asserted no env kwarg was passed as a stand-in for no HIP flag being injected. They now assert the flag itself, which is the guarantee they were written for and does not depend on how the env is delivered. Separately, latin-1 cannot stand in for a double-byte codepage while recovering dedup keys: cp932 表 is 95 5C, and the trail byte reads as a JSON backslash, so the record failed to parse and its id was forgotten, appending a duplicate on resume. cp932, cp936, cp949 and cp950 are tried too. The reading is still only ever used for keys, which are ASCII and identical whichever codec parses. * Require more than one legacy line before trusting its dedup keys A shard whose valid records are all ASCII casts no UTF-8 votes, so a single damaged line won the vote by itself, its key was remembered, and the retry that would have replaced the unreadable record was refused. One such line is genuinely undecidable: a legacy record with one accented character and an ASCII record with one stray byte are the same shape. Reading it as damage costs a duplicate; reading it as legacy loses the record for good. Only one of those is recoverable, so it is now read as damage. A real legacy shard has a legacy line for every record carrying an umlaut, so its dedup is unaffected. * Append ASCII whenever the shard already holds non-ASCII bytes The gate asked whether any line was undecodable as UTF-8, which misses a shard where every legacy line happens to be valid UTF-8 too. A cp1251 shard of Р° records is bytes D0 B0 throughout, so appending 世界 as UTF-8 left a file where cp1251 reads the old records correctly and the new one as mojibake, and UTF-8 does the reverse. No single decoding recovered the whole scrape. The gate is now simply whether the shard holds any non-ASCII byte at all, which covers both cases and is easier to reason about: if what is already there reads differently under different encodings, do not add more bytes that do. Appending ASCII costs only \uXXXX escapes, which json.loads turns back into the exact characters, and it leaves the new record correct under either reading. * Skip the two Linux-gated flash-attn tests off Linux _should_try_runtime_flash_attn_install ends in sys.platform.startswith( "linux"), and the threshold test one line above already asserts exactly that, so the two tests that drive _ensure_flash_attn_for_long_context past the gate cannot pass anywhere else: the call returns before it reports a status. They were written on Linux and only surface once the suite actually runs on Windows or macOS, where both fail on an empty status list. This PR is about making the backend behave on Windows, so its own suite should be runnable there. * Fail closed when a KFD topology node does not decode This PR pins that read to utf-8, which turns an undecodable byte into UnicodeDecodeError. That is a ValueError, not an OSError, so it slips past the handler one line below and escapes a helper whose docstring promises to fail closed on any unreadable node. The caller would then lose the whole HIP-order map on a machine that has AMD GPUs, and the reason the helper fails closed is that dropping a node shifts every later ordinal and lets a similar-capacity GPU pass the total-size guard while showing another card's usage. Widening the handler is the same one-line change main already made in #7487, so the two agree and the eventual merge is clean. * Tighten the comments added in this branch * Treat an undecodable marker and undecodable metadata as malformed, not fatal Two more places where pinning the decode changed the failure mode. A UnicodeDecodeError is a ValueError, so neither `except OSError` nor `except (JSONDecodeError, OSError)` catches it, and both sites had a documented fallback that stopped being reached. An undecodable .transport marker used to read as an unknown value, and the caller then safely purged and restarted the partial download. It now aborts prepare_cache_for_transport instead, so the transfer fails rather than retrying. Undecodable .meta.json used to fall back to the file's own name, the same way invalid JSON does. It now aborts URI construction for the entire unstructured seed, so one corrupt byte in original_filename takes out the whole dataset. Both handlers are widened, matching the KFD fix earlier on this branch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Widen two more decode guards, and pin the kernel installer's pipe Same shape as the ones already fixed here: the read was pinned to UTF-8 while the handler around it still only catches OSError, and UnicodeDecodeError is a ValueError. hf_cache_snapshot_dir answers whether a model is already on disk, and the offline embedding checks turn a raise into a 500. A torn refs/main used to decode into a nonsense commit and miss the snapshot dir; it now skips that cache root and keeps looking. _remove_pid_file runs first in _graceful_shutdown, so a corrupt studio.pid raising there abandoned the inference, export, training and tunnel children the rest of that function exists to kill. ssm_runtime's source-build path builds its subprocess kwargs in a dict and splats them through _run_with_heartbeat, so neither the encoding guard nor the earlier sweep saw the text = True in it: pip's output was still decoded with the Windows ANSI codepage, where a non-ASCII path or a compiler diagnostic mojibakes or raises over an install that was going fine. It now pins the same utf-8/replace pair install_wheel uses, and the HIP branch extends that env rather than replacing it. The guard learned the dict-literal shape and reddens on the old code (ssm_runtime.py:253). * Tighten the comments around the UTF-8 text I/O pins Collapse the multi-line rationales added with the encoding pins down to a line or two each, drop what the code already says, and use one wording for the repeated child-env note. * Do not let an unreadable bootstrap password stop startup, and narrow the kwargs guard ensure_default_admin calls _load_bootstrap_password for every existing admin and the lifespan calls that with no handler, so pinning the decode turned a damaged or pre-pin .bootstrap_password file into a backend that will not start. We write that file ourselves in UTF-8, so a byte that will not decode belongs to a file whose plaintext is worthless anyway; it now reads as no bootstrap password, the same answer as an absent file. A readable one still loads. The new kwargs check also judged every dict literal in the tree, so an unrelated payload carrying "text": True would have been reported as subprocess configuration with a misleading message, and a dict that fills in its encoding on a later line would have been reported too. It now only judges a dict that actually reaches a call, either splatted through a name or written at the call site, and treats a later kw["encoding"] assignment as satisfying it. The ssm_runtime shape it was written for is still caught, and a test pins both directions. * Stop reading a UTF-8 record a second time _read_line always parsed the line under the codepage as well, even when it had already read as UTF-8. Both callers take the UTF-8 reading when there is one and never look at the other, so on a healthy shard the second parse is pure waste, and this file reads all of one on every resume of a scrape it expects to reach gigabytes. Measured on 200,000 records, 76 MB: 1.96s before, 0.81s after, so the double reading was costing 2.8x. The early return is limited to a record, since the key lookup deliberately falls through to the codepage reading when UTF-8 yields something that is not one. A line UTF-8 cannot read still tries the codepage, latin-1 and the double-byte encodings as before, which is what the second reading is for. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pin the scanned source fixture's line endings test_remote_code_scan_reads_non_ascii_sources compared a file's contents against the string it wrote, but wrote it in text mode, so Windows translated the line ends on the way out and the read back differed by a carriage return. That is the writer's doing, not the encoding the test is about, and it was the one failure on the Windows runner that belonged to this branch. The fixture now writes with newline = "" so the bytes on disk are the string on every platform. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim the newer comments to their point Shorten the widened-guard and state store notes added since the last pass, and collapse the line-ending note on the scanned source fixture. * Read the scraper checkpoint as UTF-8 only, never as a codepage A checkpoint holds nothing but base64 cursors and booleans, so one written by an older locale-encoded release is byte-identical to a UTF-8 one and already reads back. The codepage fallback can therefore only ever contribute non-ASCII: if a single-byte reading of the file were all ASCII, the UTF-8 read would have succeeded first. So the only file it changes the answer for is a damaged one, and there it turns a safe reset into a resume on a mojibaked cursor. GitHub answers that with INVALID_CURSOR_ARGUMENTS at HTTP 200, gh_client returns the partial document, and the scraper reads zero nodes and an empty pageInfo, which marks the stream done. Every later resume then skips it entirely. Reading UTF-8 only restores the earlier behaviour of dropping a checkpoint that will not decode, which re-scrapes from the first page while the writers dedup the replay. The shard scan below keeps its codepage reading; those records do carry non-ASCII. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate the remaining tilelang install tests to Linux _tilelang_platform_supported() returns False off Linux, so _ensure_tilelang_backend returns before the install and the subprocess mock these six assert on is never called. They fail on macOS runners for that reason alone. The rest of the file already carries this marker; these were missed. * Gate the Windows-incompatible worker and ROCm tests Two different gates, because the production code has two. The causal-conv1d and flash-linear-attention installers bail out on sys.platform == 'win32' alone and run everywhere else including macOS, so those cases get not_on_windows; marking them linux_only would skip tests that legitimately pass off Linux. The DRM and KFD readers return early unless platform.system() is Linux, and their fixtures build a fake sysfs tree needing PCI addresses like 0000:00:02.0 as directory names, which Windows cannot represent, so those get linux_only. The two visible-utilization cases failed for a different reason: on Windows get_visible_gpu_utilization takes the AMD adapter branch ahead of the torch fallback under test, and probing it imports torch, which the runner lacks. Stubbing that branch empty leaves every other platform unchanged. * Treat unparseable JSON nesting as a parse failure, and guard os.fdopen json.loads answers nesting it cannot descend with RecursionError, a RuntimeError, so _parse let it escape where the catch-all it replaced discarded the record. Both callers run _parse outside any further handler, so one damaged checkpoint or shard line aborted the scraper at startup. The encoding guard also missed os.fdopen, which is open() on a descriptor and takes the same locale default in text mode. It flags exactly the two text-mode calls that were left unencoded; the swap lock file's reader was already pinned to UTF-8 while its writer still used the codepage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Write the non-ASCII source fixture without a 3.10-only argument Path.write_text() only grew newline in 3.10, and pyproject declares requires-python >=3.9, so this raised TypeError there. open() takes the same argument on every supported version and pins the bytes on disk the same way. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten encoding comments * Follow subprocess calls through callable aliases in the encoding guard --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com> --------- Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothshared@gmail.com>
1685 lines
65 KiB
Python
1685 lines
65 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import builtins
|
|
import subprocess
|
|
import sys
|
|
from typing import Any
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from core.training import worker
|
|
|
|
# The runtime install is Linux-only, so elsewhere these return before any status.
|
|
linux_only = pytest.mark.skipif(
|
|
not sys.platform.startswith("linux"),
|
|
reason = "the runtime flash-attn install is gated to Linux",
|
|
)
|
|
|
|
# causal-conv1d and flash-linear-attention are NOT Linux-gated: both installers bail out
|
|
# on `sys.platform == "win32"` alone (no prebuilt wheel for Windows) and run everywhere
|
|
# else, macOS included. linux_only here would skip cases that legitimately pass off Linux.
|
|
not_on_windows = pytest.mark.skipif(
|
|
sys.platform == "win32",
|
|
reason = (
|
|
"mirrors the sys.platform == 'win32' bail-out in "
|
|
"_ensure_flash_linear_attention_unconditional and "
|
|
"_ensure_causal_conv1d_fast_path"
|
|
),
|
|
)
|
|
|
|
|
|
def _missing_flash_attn_import():
|
|
real_import = builtins.__import__
|
|
|
|
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)
|
|
|
|
return fake_import
|
|
|
|
|
|
def _missing_module_import(missing: str):
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(
|
|
name,
|
|
globals = None,
|
|
locals = None,
|
|
fromlist = (),
|
|
level = 0,
|
|
):
|
|
if name == missing:
|
|
raise ImportError
|
|
return real_import(name, globals, locals, fromlist, level)
|
|
|
|
return fake_import
|
|
|
|
|
|
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")
|
|
|
|
monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1")
|
|
assert worker._should_try_runtime_flash_attn_install(32768) is False
|
|
|
|
|
|
@linux_only
|
|
def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch):
|
|
statuses: list[str] = []
|
|
|
|
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"flash_attn_wheel_url",
|
|
lambda env: "https://example.com/fa.whl",
|
|
)
|
|
monkeypatch.setattr(worker, "url_exists", lambda url: True)
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_send_status",
|
|
lambda queue, message: statuses.append(message),
|
|
)
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"install_wheel",
|
|
lambda *args, **kwargs: [("pip", subprocess.CompletedProcess(["pip"], 0, ""))],
|
|
)
|
|
|
|
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768)
|
|
|
|
assert statuses == ["Installing flash-attn for faster training..."]
|
|
|
|
|
|
@linux_only
|
|
def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch):
|
|
calls: list[list[str]] = []
|
|
statuses: list[str] = []
|
|
|
|
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import())
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"probe_torch_wheel_env",
|
|
lambda timeout = 30: {
|
|
"python_tag": "cp313",
|
|
"torch_mm": "2.10",
|
|
"cuda_major": "13",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"flash_attn_wheel_url",
|
|
lambda env: "https://example.com/fa.whl",
|
|
)
|
|
monkeypatch.setattr(worker, "url_exists", lambda url: False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: None)
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_send_status",
|
|
lambda queue, message: statuses.append(message),
|
|
)
|
|
monkeypatch.setattr(worker, "install_wheel", mock.Mock())
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
calls.append(list(cmd))
|
|
return subprocess.CompletedProcess(cmd, 0, "")
|
|
|
|
monkeypatch.setattr(worker._sp, "run", fake_run)
|
|
|
|
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768)
|
|
|
|
assert statuses == ["Installing flash-attn from PyPI for long-context training..."]
|
|
assert calls == [[sys.executable, "-m", "pip", "install", "flash-attn"]]
|
|
|
|
|
|
def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch):
|
|
monkeypatch.setenv(worker._FLASH_ATTN_SKIP_ENV, "1")
|
|
monkeypatch.setattr(worker._sp, "run", mock.Mock())
|
|
|
|
worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 32768)
|
|
|
|
worker._sp.run.assert_not_called()
|
|
|
|
|
|
@not_on_windows
|
|
def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch):
|
|
install_mock = mock.Mock(return_value = True)
|
|
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
|
|
|
worker._ensure_causal_conv1d_fast_path(
|
|
event_queue = [],
|
|
model_name = "tiiuae/Falcon-H1-0.5B-Instruct",
|
|
)
|
|
|
|
install_mock.assert_called_once_with(
|
|
event_queue = [],
|
|
import_name = "causal_conv1d",
|
|
display_name = "causal-conv1d",
|
|
pypi_name = "causal-conv1d",
|
|
pypi_version = worker._CAUSAL_CONV1D_PACKAGE_VERSION,
|
|
filename_prefix = "causal_conv1d",
|
|
release_tag = worker._CAUSAL_CONV1D_RELEASE_TAG,
|
|
release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download",
|
|
)
|
|
|
|
|
|
@not_on_windows
|
|
def test_causal_conv1d_fast_path_includes_qwen3_6_variants(monkeypatch):
|
|
install_mock = mock.Mock(return_value = True)
|
|
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
|
|
|
worker._ensure_causal_conv1d_fast_path(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.6-4B",
|
|
)
|
|
worker._ensure_causal_conv1d_fast_path(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3_6-4B",
|
|
)
|
|
|
|
assert install_mock.call_count == 2
|
|
|
|
|
|
def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch):
|
|
install_mock = mock.Mock(return_value = True)
|
|
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
|
|
|
worker._ensure_mamba_ssm(
|
|
event_queue = [],
|
|
model_name = "tiiuae/Falcon-H1-0.5B-Instruct",
|
|
)
|
|
|
|
install_mock.assert_called_once_with(
|
|
event_queue = [],
|
|
import_name = "mamba_ssm",
|
|
display_name = "mamba-ssm",
|
|
pypi_name = "mamba-ssm",
|
|
pypi_version = worker._MAMBA_SSM_PACKAGE_VERSION,
|
|
filename_prefix = "mamba_ssm",
|
|
release_tag = worker._MAMBA_SSM_RELEASE_TAG,
|
|
release_base_url = "https://github.com/state-spaces/mamba/releases/download",
|
|
)
|
|
|
|
|
|
def _force_missing_fla_imports(monkeypatch):
|
|
"""Force fla.modules / fla.ops imports to raise ImportError."""
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *a, **kw):
|
|
if name.startswith("fla.modules") or name.startswith("fla.ops"):
|
|
raise ImportError
|
|
return real_import(name, *a, **kw)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
|
|
def _pin_fla_model_types(monkeypatch):
|
|
"""Pin the auto-discovered FLA allowlist to the Qwen GDN families.
|
|
|
|
`_discover_fla_model_types` scans the *installed* transformers, and
|
|
`models/qwen3_5/` only exists from 5.x. The backend supports
|
|
`transformers>=4.51`, so on a 4.x install the gate returns False and every
|
|
Qwen3.5 assertion below silently no-ops. Pinning keeps these tests hermetic
|
|
across the supported range.
|
|
"""
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_discover_fla_model_types",
|
|
lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_6", "qwen3_next"}),
|
|
)
|
|
|
|
|
|
@not_on_windows
|
|
def test_flash_linear_attention_installs_pinned_pair_for_qwen3_5(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
_force_missing_fla_imports(monkeypatch)
|
|
statuses: list[str] = []
|
|
monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_called_once()
|
|
args = run_mock.call_args[0][0]
|
|
assert f"flash-linear-attention=={worker._FLA_PACKAGE_VERSION}" in args
|
|
assert f"fla-core=={worker._FLA_CORE_PACKAGE_VERSION}" in args
|
|
assert "--no-deps" in args
|
|
assert run_mock.call_args.kwargs["timeout"] == worker._TILELANG_INSTALL_TIMEOUT_S
|
|
assert any("flash-linear-attention" in s for s in statuses)
|
|
|
|
|
|
def test_flash_linear_attention_skips_for_unrelated_models(monkeypatch):
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "meta-llama/Llama-3.2-1B-Instruct",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
def test_flash_linear_attention_skips_for_ssm_only_models(monkeypatch):
|
|
# Nemotron-H / Falcon-H1 / Granite-H / LFM2 take the mamba_ssm path,
|
|
# never FLA's gated_delta_rule kernels.
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
for name in (
|
|
"tiiuae/Falcon-H1-0.5B-Instruct",
|
|
"nvidia/Nemotron-H-8B-Base",
|
|
"ibm-granite/granite-4.0-h-tiny",
|
|
"LiquidAI/LFM2-1.2B-Instruct",
|
|
):
|
|
worker._ensure_flash_linear_attention(event_queue = [], model_name = name)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
@not_on_windows
|
|
def test_flash_linear_attention_matches_full_qwen3_family(monkeypatch):
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
_force_missing_fla_imports(monkeypatch)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
# Hermetic discovery: pretend transformers ships all Qwen GDN families.
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_discover_fla_model_types",
|
|
lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_6", "qwen3_next"}),
|
|
)
|
|
|
|
for name in (
|
|
"unsloth/Qwen3.5-2B",
|
|
"unsloth/Qwen3_5-MoE-A22B",
|
|
"unsloth/Qwen3.6-4B",
|
|
"unsloth/Qwen3_6-4B",
|
|
"unsloth/Qwen3-Next-80B-A3B",
|
|
"unsloth/Qwen3_Next-80B-A3B",
|
|
):
|
|
worker._ensure_flash_linear_attention(event_queue = [], model_name = name)
|
|
|
|
assert run_mock.call_count == 6
|
|
|
|
|
|
def test_flash_linear_attention_skipped_below_python_3_10(monkeypatch):
|
|
# sys.version_info is a structseq, not constructible; substitute a
|
|
# plain tuple so the `< _FLA_MIN_PYTHON` comparison still works.
|
|
monkeypatch.setattr(worker.sys, "version_info", (3, 9, 0, "final", 0))
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
def test_flash_linear_attention_skipped_via_env(monkeypatch):
|
|
monkeypatch.setenv(worker._FLA_SKIP_ENV, "1")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
@not_on_windows
|
|
def test_flash_linear_attention_skipped_below_torch_2_7(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 5))
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
statuses: list[str] = []
|
|
monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
assert any("torch>=" in s for s in statuses)
|
|
|
|
|
|
@not_on_windows
|
|
def test_flash_linear_attention_install_includes_einops(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9))
|
|
monkeypatch.setattr(worker, "_flash_linear_attention_importable", lambda: False)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
args = run_mock.call_args[0][0]
|
|
assert "--no-deps" in args
|
|
# packaging and triton are added because fla/utils.py imports them at load
|
|
# but neither is in fla-core's METADATA (an upstream FLA gap).
|
|
assert "einops" in args
|
|
assert "packaging" in args
|
|
assert "triton" in args
|
|
assert f"flash-linear-attention=={worker._FLA_PACKAGE_VERSION}" in args
|
|
assert f"fla-core=={worker._FLA_CORE_PACKAGE_VERSION}" in args
|
|
|
|
|
|
@not_on_windows
|
|
def test_flash_linear_attention_logs_post_install_import_failure(monkeypatch):
|
|
"""pip exits 0 but `import fla.modules` still fails (missing transitive)."""
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9))
|
|
import_calls = {"count": 0}
|
|
|
|
def fake_importable():
|
|
import_calls["count"] += 1
|
|
# Pre-install probe -> False (attempt install); post-install
|
|
# verify -> still False.
|
|
return False
|
|
|
|
monkeypatch.setattr(worker, "_flash_linear_attention_importable", fake_importable)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
statuses: list[str] = []
|
|
monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
|
|
|
|
worker._ensure_flash_linear_attention(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
assert import_calls["count"] == 2
|
|
assert any("not importable" in s for s in statuses)
|
|
|
|
|
|
def test_tilelang_backend_skipped_on_unsupported_linux_arch(monkeypatch):
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.sys, "platform", "linux")
|
|
import platform as _platform
|
|
|
|
monkeypatch.setattr(_platform, "machine", lambda: "ppc64le")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
@linux_only
|
|
def test_tilelang_backend_pins_only_binary(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
|
monkeypatch.setattr(worker, "_tilelang_importable", lambda: False)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
# Bypass the post-install probe too.
|
|
probe_calls = {"count": 0}
|
|
|
|
def fake_probe():
|
|
probe_calls["count"] += 1
|
|
# Pre-install probe: False (install runs); post-install: True
|
|
# (success branch taken).
|
|
return probe_calls["count"] > 1
|
|
|
|
monkeypatch.setattr(worker, "_tilelang_importable", fake_probe)
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
args = run_mock.call_args[0][0]
|
|
assert "--only-binary=:all:" in args
|
|
assert "--no-deps" not in args
|
|
|
|
|
|
def _force_missing_tilelang_imports(monkeypatch):
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *a, **kw):
|
|
if name in ("tilelang", "tvm_ffi"):
|
|
raise ImportError
|
|
return real_import(name, *a, **kw)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
|
|
@linux_only
|
|
def test_tilelang_backend_installs_pinned_pair_for_qwen3_5(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
_force_missing_tilelang_imports(monkeypatch)
|
|
statuses: list[str] = []
|
|
monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_called_once()
|
|
args = run_mock.call_args[0][0]
|
|
assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in args
|
|
assert f"tilelang=={worker._TILELANG_PACKAGE_VERSION}" in args
|
|
assert run_mock.call_args.kwargs["timeout"] == worker._TILELANG_INSTALL_TIMEOUT_S
|
|
assert any("Installing TileLang" in s for s in statuses)
|
|
|
|
|
|
@linux_only
|
|
def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch):
|
|
"""Repair path issues TWO pip calls:
|
|
|
|
1 (repair): --force-reinstall --no-deps apache-tvm-ffi -- downgrades only
|
|
the broken package; --no-deps stops the cascade through its deps to torch.
|
|
2 (install): plain apache-tvm-ffi + tilelang -- resolves missing transitive
|
|
deps without --force-reinstall, so it never replaces correct packages.
|
|
"""
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
assert run_mock.call_count == 2
|
|
repair_args, install_args = (call[0][0] for call in run_mock.call_args_list)
|
|
|
|
# Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY.
|
|
assert "--force-reinstall" in repair_args
|
|
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"
|
|
|
|
# Install: regular dep-resolving install, no --force-reinstall.
|
|
assert "--force-reinstall" not in install_args
|
|
assert "--no-deps" not in install_args
|
|
assert "--only-binary=:all:" in install_args
|
|
assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in install_args
|
|
assert f"tilelang=={worker._TILELANG_PACKAGE_VERSION}" in install_args
|
|
|
|
|
|
def test_tilelang_backend_skipped_below_python_3_10(monkeypatch):
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
# sys.version_info is a structseq, not constructible; substitute a
|
|
# plain tuple so the `< _FLA_MIN_PYTHON` comparison still works.
|
|
monkeypatch.setattr(worker.sys, "version_info", (3, 9, 0, "final", 0))
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
def test_tilelang_backend_skipped_on_windows(monkeypatch):
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.sys, "platform", "win32")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
@linux_only
|
|
def test_tilelang_backend_swallows_install_timeout(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
|
_force_missing_tilelang_imports(monkeypatch)
|
|
|
|
def raise_timeout(*a, **kw):
|
|
raise subprocess.TimeoutExpired(cmd = "pip", timeout = 1)
|
|
|
|
monkeypatch.setattr(worker._sp, "run", raise_timeout)
|
|
statuses: list[str] = []
|
|
monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
|
|
|
|
# Must not raise.
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
assert any("timed out" in s.lower() for s in statuses)
|
|
|
|
|
|
def test_tilelang_backend_skipped_for_ssm_models(monkeypatch):
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
# Nemotron-H / Falcon-H1 / Granite-H take the mamba_ssm path, not FLA's
|
|
# gated_delta_rule -> tilelang doesn't affect them.
|
|
for name in (
|
|
"tiiuae/Falcon-H1-0.5B-Instruct",
|
|
"nvidia/Nemotron-H-8B-Base",
|
|
"ibm-granite/granite-4.0-h-tiny",
|
|
"meta-llama/Llama-3.2-1B-Instruct",
|
|
):
|
|
worker._ensure_tilelang_backend(event_queue = [], model_name = name)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
def test_tilelang_backend_skipped_via_env(monkeypatch):
|
|
monkeypatch.setenv(worker._TILELANG_SKIP_ENV, "1")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
@linux_only
|
|
def test_tilelang_backend_swallows_install_failure(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: None)
|
|
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: None)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 1, stdout = "boom"))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
_force_missing_tilelang_imports(monkeypatch)
|
|
statuses: list[str] = []
|
|
monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
|
|
|
|
# Should not raise even when pip exits non-zero.
|
|
worker._ensure_tilelang_backend(
|
|
event_queue = [],
|
|
model_name = "unsloth/Qwen3.5-2B",
|
|
)
|
|
|
|
run_mock.assert_called_once()
|
|
assert any("failed" in s.lower() for s in statuses)
|
|
|
|
|
|
# Runtime hook on is_flash_linear_attention_available /
|
|
# is_causal_conv1d_available -- the primary gate in normal operation. The
|
|
# substring tests above cover the SKIP_FAST_PATH_HOOKS=1 fallback.
|
|
|
|
|
|
class _FakeQueue(list):
|
|
"""List with `.put` so worker._send_status can send into it in tests."""
|
|
|
|
def put(self, item):
|
|
self.append(item)
|
|
|
|
|
|
def _make_fake_gate(initial_return: bool):
|
|
"""Callable mimicking transformers' lru_cache-decorated gates.
|
|
|
|
Tracks call count and exposes `cache_clear`. Flip `.next_return` to
|
|
mimic install-then-True behaviour.
|
|
"""
|
|
|
|
class Gate:
|
|
def __init__(self, initial: bool) -> None:
|
|
self.next_return = initial
|
|
self.call_count = 0
|
|
self.cache_clear_count = 0
|
|
|
|
def __call__(self) -> bool:
|
|
self.call_count += 1
|
|
return self.next_return
|
|
|
|
def cache_clear(self) -> None:
|
|
self.cache_clear_count += 1
|
|
|
|
return Gate(initial_return)
|
|
|
|
|
|
def _patch_iu_gates(monkeypatch, fla_gate, conv_gate):
|
|
"""Drop fake gates onto transformers.utils.import_utils for the test."""
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
monkeypatch.setattr(_iu, "is_flash_linear_attention_available", fla_gate)
|
|
monkeypatch.setattr(_iu, "is_causal_conv1d_available", conv_gate)
|
|
|
|
|
|
@not_on_windows
|
|
def test_hook_installs_when_gate_returns_false(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = False)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
def _fla_install_side_effect(eq):
|
|
fla_gate.next_return = True
|
|
return True
|
|
|
|
fla_install = mock.Mock(side_effect = _fla_install_side_effect)
|
|
tile_install = mock.Mock(side_effect = lambda eq: None)
|
|
|
|
def _conv_install_side_effect(**kw):
|
|
conv_gate.next_return = True
|
|
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_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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
# Both gates wrapped; calling them should drive the install.
|
|
assert _iu.is_flash_linear_attention_available() is True
|
|
fla_install.assert_called_once()
|
|
tile_install.assert_called_once()
|
|
assert _iu.is_causal_conv1d_available() is True
|
|
conv_install.assert_called_once()
|
|
|
|
|
|
def test_hook_skips_install_when_gate_already_true(monkeypatch):
|
|
"""Both gates already True AND tilelang healthy -> zero install work.
|
|
(Tilelang repair on the already-True path is covered by
|
|
test_hook_runs_tilelang_repair_when_fla_already_true.)
|
|
"""
|
|
fla_gate = _make_fake_gate(initial_return = True)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
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_tilelang_backend_unconditional", tile_install)
|
|
monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install)
|
|
# Tilelang healthy -> post_available path is a no-op (otherwise it
|
|
# would call tile_install, correct but out of scope here).
|
|
monkeypatch.setattr(worker, "_tilelang_importable", lambda: True)
|
|
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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
assert _iu.is_flash_linear_attention_available() is True
|
|
assert _iu.is_causal_conv1d_available() is True
|
|
fla_install.assert_not_called()
|
|
tile_install.assert_not_called()
|
|
conv_install.assert_not_called()
|
|
|
|
|
|
def test_hook_idempotent_on_repeat_call(monkeypatch):
|
|
_pin_fla_model_types(monkeypatch)
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = False)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
def _fla_install_side_effect(eq):
|
|
fla_gate.next_return = True
|
|
return True
|
|
|
|
fla_install = mock.Mock(side_effect = _fla_install_side_effect)
|
|
tile_install = mock.Mock()
|
|
|
|
def _conv_install_side_effect(**kw):
|
|
conv_gate.next_return = True
|
|
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_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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
# First call: hook fires.
|
|
_iu.is_flash_linear_attention_available()
|
|
# Later calls: must not re-trigger the installer.
|
|
_iu.is_flash_linear_attention_available()
|
|
_iu.is_flash_linear_attention_available()
|
|
assert fla_install.call_count == 1
|
|
assert tile_install.call_count == 1
|
|
|
|
|
|
def test_hook_handles_install_failure_gracefully(monkeypatch):
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = True) # bypass to focus on FLA
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
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, "_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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
# Must not raise; returns False so transformers uses the torch loop.
|
|
assert _iu.is_flash_linear_attention_available() is False
|
|
|
|
|
|
def test_hook_can_be_disabled_via_env(monkeypatch):
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = False)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
fla_install = mock.Mock()
|
|
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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
# Hook not installed; gates remain the fakes.
|
|
assert _iu.is_flash_linear_attention_available is fla_gate
|
|
assert _iu.is_causal_conv1d_available is conv_gate
|
|
fla_install.assert_not_called()
|
|
|
|
|
|
def test_hook_clears_lru_cache_before_first_check(monkeypatch):
|
|
fla_gate = _make_fake_gate(initial_return = True)
|
|
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, "_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")
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
_iu.is_flash_linear_attention_available()
|
|
# Wrapper called cache_clear at least once before delegating.
|
|
assert fla_gate.cache_clear_count >= 1
|
|
|
|
|
|
def test_hook_rewrites_previously_imported_module_bindings(monkeypatch):
|
|
"""Modeling files bind is_flash_linear_attention_available locally via
|
|
`from ... import is_X`. Reassigning the attribute on import_utils alone
|
|
misses those; the hook installer sweeps sys.modules and rebinds them.
|
|
"""
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
# Fake modeling module that did `from ... import is_flash_linear_attention_available`.
|
|
fake_mod = sys.modules.setdefault(
|
|
"_test_fake_modeling_qwen35", type(sys)("_test_fake_modeling_qwen35")
|
|
)
|
|
fake_mod.is_flash_linear_attention_available = fla_gate
|
|
|
|
def fake_install(eq):
|
|
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, "_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")
|
|
|
|
# The fake module's local binding is rewritten to the wrapper.
|
|
assert fake_mod.is_flash_linear_attention_available is not fla_gate
|
|
# Calling through the fake module's reference triggers install.
|
|
assert fake_mod.is_flash_linear_attention_available() is True
|
|
|
|
del sys.modules["_test_fake_modeling_qwen35"]
|
|
|
|
|
|
def test_hook_skips_when_import_utils_unavailable(monkeypatch):
|
|
"""If transformers.utils.import_utils can't be imported, the hook
|
|
installer must log and return cleanly rather than crash the worker."""
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *a, **kw):
|
|
if name == "transformers.utils" or name == "transformers.utils.import_utils":
|
|
raise ImportError("transformers missing in worker venv")
|
|
return real_import(name, *a, **kw)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
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")
|
|
|
|
|
|
def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch):
|
|
"""Hook disabled -> legacy gate falls back to auto-discovered 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.setenv(worker._FAST_PATH_HOOKS_SKIP_ENV, "1")
|
|
|
|
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")
|
|
assert install_mock.call_count == 1
|
|
|
|
|
|
# Regression tests for the reviewer findings:
|
|
# 1. tilelang Qwen-guard on hook path (non-Qwen FLA models)
|
|
# 2. tilelang repair must not replace torch / CUDA stack
|
|
# 3. hook must trust installer's bool, not transformers metadata
|
|
# 4. causal-conv1d must stay eager for SSM models that bypass the gate
|
|
# 5. rebind sweep must not invoke lazy module __getattr__
|
|
# 6. tilelang skipped when FLA was skipped / failed
|
|
# 7. tilelang repair runs when FLA is already True
|
|
# 8. older FLA detected as stale and reinstalled
|
|
|
|
|
|
def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch):
|
|
"""A model not in the auto-discovered FLA allowlist calls
|
|
is_flash_linear_attention_available but must NOT get tilelang."""
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
def _fla_install(eq):
|
|
fla_gate.next_return = True
|
|
return True
|
|
|
|
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_tilelang_backend_unconditional", tile_install)
|
|
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
|
|
# 5.4.0). Test semantic: "outside-allowlist -> no tilelang".
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_discover_fla_model_types",
|
|
lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_next"}),
|
|
)
|
|
|
|
worker._install_fast_path_hooks(
|
|
event_queue = _FakeQueue(),
|
|
model_name = "fake-org/Fictional-FLA-Only-Model-7B",
|
|
)
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
assert _iu.is_flash_linear_attention_available() is True
|
|
fla_install.assert_called_once()
|
|
tile_install.assert_not_called()
|
|
|
|
|
|
def test_hook_does_install_tilelang_for_qwen35(monkeypatch):
|
|
"""Positive control for finding #1: Qwen3.5 still gets tilelang."""
|
|
_pin_fla_model_types(monkeypatch)
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
def _fla_install(eq):
|
|
fla_gate.next_return = True
|
|
return True
|
|
|
|
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_tilelang_backend_unconditional", tile_install)
|
|
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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
_iu.is_flash_linear_attention_available()
|
|
fla_install.assert_called_once()
|
|
tile_install.assert_called_once()
|
|
|
|
|
|
@linux_only
|
|
def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch):
|
|
"""Finding #2: the broken-tvm-ffi repair must use --no-deps on the
|
|
forced step so --force-reinstall doesn't cascade through
|
|
apache-tvm-ffi's dep graph and pull a different torch wheel.
|
|
"""
|
|
_pin_fla_model_types(monkeypatch)
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.10")
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
|
|
worker._ensure_tilelang_backend(event_queue = [], model_name = "unsloth/Qwen3.5-2B")
|
|
|
|
assert run_mock.call_count == 2
|
|
repair_args = run_mock.call_args_list[0][0][0]
|
|
# Forced step MUST be --no-deps so torch / CUDA stack is untouched.
|
|
assert "--force-reinstall" in repair_args and "--no-deps" in repair_args
|
|
# Touches ONLY apache-tvm-ffi, not tilelang / torch.
|
|
assert all("tilelang" not in a for a in repair_args)
|
|
assert all("torch" not in a for a in repair_args)
|
|
|
|
|
|
def test_hook_trusts_installer_bool_not_metadata(monkeypatch):
|
|
"""Finding #3: if pip exits 0 but deep imports fail, the installer returns
|
|
False; the hook must propagate that False even if the metadata-only gate
|
|
returns True after pip succeeds, so transformers takes the torch fallback.
|
|
"""
|
|
# Gate flips True after install (simulating "metadata sees fla").
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
# Installer "succeeds" at pip and flips the gate to True (metadata
|
|
# sees fla post-install), but returns False (deep import broken).
|
|
def _bad_install(eq):
|
|
fla_gate.next_return = True # metadata says yes after pip
|
|
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_tilelang_backend_unconditional", 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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
# Hook MUST return False (installer's verdict), not True (metadata lies).
|
|
assert _iu.is_flash_linear_attention_available() is False
|
|
fake_fla_install.assert_called_once()
|
|
|
|
|
|
def test_rebind_does_not_trigger_module_getattr(monkeypatch):
|
|
"""Finding #5: the rebind sweep must use __dict__, not getattr(), to
|
|
avoid invoking transformers' lazy module __getattr__ which spits out
|
|
hundreds of "Accessing X from .models..." warnings.
|
|
"""
|
|
original = object()
|
|
replacement = object()
|
|
|
|
class _GetattrTripwire(type(sys)):
|
|
getattr_called = False
|
|
|
|
def __getattr__(self, name):
|
|
type(self).getattr_called = True
|
|
raise AttributeError(name)
|
|
|
|
lazy = _GetattrTripwire("_lazy_test_module")
|
|
sys.modules["_lazy_test_module"] = lazy
|
|
try:
|
|
# No `is_flash_linear_attention_available` in __dict__, so the
|
|
# sweep must NOT trip the tripwire.
|
|
worker._rebind_in_already_imported_modules(
|
|
attr_name = "is_flash_linear_attention_available",
|
|
old_obj = original,
|
|
new_obj = replacement,
|
|
)
|
|
assert (
|
|
not _GetattrTripwire.getattr_called
|
|
), "Rebind sweep invoked __getattr__ — should use __dict__ probe"
|
|
finally:
|
|
sys.modules.pop("_lazy_test_module", None)
|
|
|
|
|
|
def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch):
|
|
"""Finding #6: env-skipped FLA returns False from
|
|
_ensure_flash_linear_attention_unconditional; tilelang must NOT
|
|
install then.
|
|
"""
|
|
fla_gate = _make_fake_gate(initial_return = False)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
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.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
|
|
|
|
worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
# FLA gate stays False (env-skipped, install never ran).
|
|
assert _iu.is_flash_linear_attention_available() is False
|
|
tile_install.assert_not_called()
|
|
|
|
|
|
def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch):
|
|
"""Finding #7: when FLA is already importable (gate True at first
|
|
probe) but tilelang is missing or apache-tvm-ffi is on the broken
|
|
list, the post-available action must still run tilelang.
|
|
"""
|
|
_pin_fla_model_types(monkeypatch)
|
|
fla_gate = _make_fake_gate(initial_return = True)
|
|
conv_gate = _make_fake_gate(initial_return = True)
|
|
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
|
|
|
|
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_tilelang_backend_unconditional", tile_install)
|
|
monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
|
|
# tilelang missing AND tvm-ffi 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")
|
|
|
|
from transformers.utils import import_utils as _iu
|
|
|
|
_iu.is_flash_linear_attention_available()
|
|
# FLA install NOT needed; tilelang repair still triggered.
|
|
fla_install.assert_not_called()
|
|
tile_install.assert_called_once()
|
|
|
|
|
|
@not_on_windows
|
|
def test_fla_installer_force_reinstalls_when_older_version_present(monkeypatch):
|
|
"""Finding #8: an older `flash-linear-attention` that is importable
|
|
but below the pin must force a reinstall (not no-op).
|
|
"""
|
|
monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
|
|
monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9))
|
|
# Importable but stale (current()=False though importable()=True).
|
|
monkeypatch.setattr(worker, "_flash_linear_attention_importable", lambda: True)
|
|
monkeypatch.setattr(worker, "_flash_linear_attention_current", lambda **kw: False)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
|
|
worker._ensure_flash_linear_attention_unconditional(event_queue = [])
|
|
|
|
run_mock.assert_called_once()
|
|
args = run_mock.call_args[0][0]
|
|
assert (
|
|
"--force-reinstall" in args
|
|
), "Stale FLA must trigger --force-reinstall, otherwise pip is a no-op"
|
|
# --no-deps still applies so torch stays untouched.
|
|
assert "--no-deps" in args
|
|
|
|
|
|
def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode():
|
|
"""Finding #4: SSM modeling files use lazy_load_kernel and never call
|
|
is_causal_conv1d_available(), so the hook won't fire; the orchestrator must
|
|
always run the eager installer regardless of hook mode. Reads the worker
|
|
source and asserts the eager install is OUTSIDE the if/else hook branch.
|
|
"""
|
|
import inspect
|
|
|
|
src = inspect.getsource(worker.run_training_process)
|
|
# Orchestration block.
|
|
assert "_ensure_causal_conv1d_fast_path(event_queue, model_name)" in src
|
|
assert "_install_fast_path_hooks(event_queue, model_name)" in src
|
|
# Eager causal_conv1d call must come BEFORE the hook-mode if/else, not
|
|
# nested inside the `if _FAST_PATH_HOOKS_SKIP_ENV` branch.
|
|
eager_pos = src.find("_ensure_causal_conv1d_fast_path(event_queue, model_name)")
|
|
skip_check_pos = src.find('os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1"')
|
|
assert eager_pos < skip_check_pos, (
|
|
"_ensure_causal_conv1d_fast_path must be called BEFORE the hook-mode "
|
|
"branch, so SSM models that bypass is_causal_conv1d_available() still "
|
|
"get the eager install"
|
|
)
|
|
|
|
|
|
# HIP / ROCm regression coverage (Strix Halo report).
|
|
# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch crashes
|
|
# mid-backward on AMD ("Unsupported target for gemm: hip"). Fix: skip install on
|
|
# HIP torch AND setdefault FLA_TILELANG=0 so an existing tilelang isn't used.
|
|
|
|
|
|
def test_tilelang_platform_unsupported_on_hip_torch(monkeypatch):
|
|
"""Strix Halo / MI300 with ROCm torch: linux + x86_64 looks identical
|
|
to a CUDA box at the OS level, so the platform check must consult
|
|
torch.version.hip explicitly.
|
|
"""
|
|
monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
|
|
assert worker._tilelang_platform_supported() is False
|
|
|
|
|
|
def test_tilelang_install_skipped_on_hip_torch(monkeypatch):
|
|
"""End-to-end: the unconditional installer must not call pip on HIP torch."""
|
|
monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
|
|
monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
|
|
run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
|
|
monkeypatch.setattr(worker._sp, "run", run_mock)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
|
|
result = worker._ensure_tilelang_backend_unconditional(event_queue = [])
|
|
|
|
assert result is False
|
|
run_mock.assert_not_called()
|
|
|
|
|
|
def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch):
|
|
"""On HIP torch, the hook installer must setdefault FLA_TILELANG=0
|
|
(respecting user override) so a PRE-EXISTING tilelang install isn't
|
|
used by FLA's dispatcher.
|
|
"""
|
|
import os as _os
|
|
|
|
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, "_install_package_wheel_first", lambda **kw: True)
|
|
|
|
worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
|
|
|
|
assert _os.environ.get("FLA_TILELANG") == "0"
|
|
|
|
|
|
def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch):
|
|
"""If the user set FLA_TILELANG (even on HIP), don't overwrite — they
|
|
may have a HIP-aware tilelang fork.
|
|
"""
|
|
import os as _os
|
|
|
|
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, "_install_package_wheel_first", lambda **kw: True)
|
|
|
|
worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
|
|
|
|
assert _os.environ["FLA_TILELANG"] == "1"
|
|
|
|
|
|
def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch):
|
|
"""CUDA path must NOT set FLA_TILELANG (tilelang is wanted there)."""
|
|
import os as _os
|
|
|
|
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, "_install_package_wheel_first", lambda **kw: True)
|
|
|
|
worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
|
|
|
|
assert _os.environ.get("FLA_TILELANG") is None
|
|
|
|
|
|
# ───────────────────────────────────────────────────────────────────
|
|
# Auto-discovery of FLA model_types from the installed transformers
|
|
# ───────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def _make_fake_transformers_tree(tmp_path, fla_types: list[str], non_fla_types: list[str]):
|
|
"""Lay out tmp dir as `transformers/models/{type}/modeling_{type}.py`."""
|
|
pkg = tmp_path / "transformers"
|
|
models = pkg / "models"
|
|
models.mkdir(parents = True)
|
|
(pkg / "__init__.py").write_text("")
|
|
for t in fla_types:
|
|
d = models / t
|
|
d.mkdir()
|
|
(d / f"modeling_{t}.py").write_text(
|
|
"from ...utils.import_utils import is_flash_linear_attention_available\n"
|
|
"if is_flash_linear_attention_available():\n"
|
|
" from fla.modules import FusedRMSNormGated\n"
|
|
" from fla.ops.gated_delta_rule import chunk_gated_delta_rule\n"
|
|
)
|
|
for t in non_fla_types:
|
|
d = models / t
|
|
d.mkdir()
|
|
(d / f"modeling_{t}.py").write_text("class Foo: pass\n")
|
|
return pkg
|
|
|
|
|
|
def _reset_fla_cache(monkeypatch):
|
|
monkeypatch.setattr(worker, "_TRANSFORMERS_FLA_MODEL_TYPES_CACHE", None)
|
|
|
|
|
|
def test_discover_fla_model_types_returns_only_fla_users(tmp_path, monkeypatch):
|
|
pkg = _make_fake_transformers_tree(
|
|
tmp_path,
|
|
fla_types = ["qwen3_5", "qwen3_5_moe", "qwen3_next"],
|
|
non_fla_types = ["llama", "gpt2", "mistral"],
|
|
)
|
|
fake = mock.MagicMock(__file__ = str(pkg / "__init__.py"))
|
|
monkeypatch.setitem(sys.modules, "transformers", fake)
|
|
_reset_fla_cache(monkeypatch)
|
|
|
|
result = worker._discover_fla_model_types()
|
|
assert result == frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_next"})
|
|
assert "llama" not in result
|
|
assert "gpt2" not in result
|
|
|
|
|
|
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 = [])
|
|
fake = mock.MagicMock(__file__ = str(pkg / "__init__.py"))
|
|
monkeypatch.setitem(sys.modules, "transformers", fake)
|
|
_reset_fla_cache(monkeypatch)
|
|
|
|
from pathlib import Path as _Path
|
|
|
|
read_calls = [0]
|
|
real_read = _Path.read_text
|
|
|
|
def counting_read(self, *a, **kw):
|
|
read_calls[0] += 1
|
|
return real_read(self, *a, **kw)
|
|
|
|
monkeypatch.setattr(_Path, "read_text", counting_read)
|
|
|
|
first = worker._discover_fla_model_types()
|
|
after_first = read_calls[0]
|
|
second = worker._discover_fla_model_types()
|
|
|
|
assert first == second
|
|
assert read_calls[0] == after_first # cache hit: no extra reads
|
|
|
|
|
|
def test_discover_fla_model_types_handles_missing_transformers(monkeypatch):
|
|
_reset_fla_cache(monkeypatch)
|
|
|
|
real_import = builtins.__import__
|
|
|
|
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)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
result = worker._discover_fla_model_types()
|
|
assert result == frozenset()
|
|
|
|
|
|
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 = [])
|
|
fake = mock.MagicMock(__file__ = str(pkg / "__init__.py"))
|
|
monkeypatch.setitem(sys.modules, "transformers", fake)
|
|
_reset_fla_cache(monkeypatch)
|
|
|
|
from pathlib import Path as _Path
|
|
|
|
real_read = _Path.read_text
|
|
|
|
def boom_read(self, *a, **kw):
|
|
if "modeling_qwen3_5.py" in str(self):
|
|
raise OSError("permission denied")
|
|
return real_read(self, *a, **kw)
|
|
|
|
monkeypatch.setattr(_Path, "read_text", boom_read)
|
|
result = worker._discover_fla_model_types()
|
|
assert result == frozenset() # unreadable file doesn't contribute
|
|
|
|
|
|
def test_model_wants_tilelang_handles_real_repo_names(monkeypatch):
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_discover_fla_model_types",
|
|
lambda: frozenset({"qwen3_5", "qwen3_5_moe", "qwen3_next"}),
|
|
)
|
|
cases = [
|
|
("unsloth/Qwen3.5-2B", True),
|
|
("Qwen/Qwen3.5-MoE-A3B", True),
|
|
("mlx-community/qwen3-next-80b", True),
|
|
("unsloth/qwen3_5_moe_a3b_lora", True),
|
|
("meta-llama/Llama-3.1-8B", False),
|
|
("nvidia/Nemotron-H-4B", False),
|
|
("mistralai/Mistral-7B-v0.3", False),
|
|
("", False),
|
|
]
|
|
for name, expected in cases:
|
|
assert worker._model_wants_tilelang(name) is expected, name
|
|
|
|
|
|
def test_model_wants_tilelang_empty_when_transformers_has_no_fla(monkeypatch):
|
|
monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset())
|
|
assert worker._model_wants_tilelang("unsloth/Qwen3.5-2B") is False
|
|
assert worker._model_wants_tilelang("meta-llama/Llama-3.1-8B") is False
|
|
|
|
|
|
def test_model_wants_tilelang_normalizes_separators(monkeypatch):
|
|
monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_next"}))
|
|
for variant in (
|
|
"qwen3-next",
|
|
"Qwen3.Next",
|
|
"Qwen/Qwen3 Next",
|
|
"anyone/qwen3_next",
|
|
"qwen3.next-80b",
|
|
):
|
|
assert worker._model_wants_tilelang(variant) is True, variant
|
|
|
|
|
|
# HIP source-build gcc-install-dir coverage (Strix Halo).
|
|
# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14, so ROCm
|
|
# clang-20 picks it and fails ('cstdlib' not found) building causal-conv1d.
|
|
# _hipcc_gcc_install_dir() finds a gcc dir with both halves; the HIP branch of
|
|
# _install_package_wheel_first passes it via HIPCC_COMPILE_FLAGS_APPEND.
|
|
# Parallels bbf004c's setup.sh fix for the llama.cpp HIP build (PR #5301).
|
|
|
|
|
|
def _isdir_for_layout(*existing: str):
|
|
"""os.path.isdir replacement treating only the given absolute paths as
|
|
directories, to simulate which gcc runtime / C++ header dirs exist."""
|
|
valid = set(existing)
|
|
|
|
def fake_isdir(path: str) -> bool:
|
|
return path in valid
|
|
|
|
return fake_isdir
|
|
|
|
|
|
def test_hipcc_gcc_install_dir_picks_highest_with_headers(monkeypatch):
|
|
"""gcc-14 has runtime but no /usr/include/c++/14; loop falls through
|
|
to gcc-13 which has both. The exact Ubuntu 24.04 layout."""
|
|
monkeypatch.setattr(sys, "platform", "linux")
|
|
import platform as _platform
|
|
|
|
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
|
|
monkeypatch.setattr(
|
|
worker.os.path,
|
|
"isdir",
|
|
_isdir_for_layout(
|
|
"/usr/lib/gcc/x86_64-linux-gnu/14/include", # runtime present
|
|
# but no /usr/include/c++/14 — typical Ubuntu 24.04 default
|
|
"/usr/lib/gcc/x86_64-linux-gnu/13/include",
|
|
"/usr/include/c++/13", # libstdc++-13-dev installed
|
|
),
|
|
)
|
|
assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/13"
|
|
|
|
|
|
def test_hipcc_gcc_install_dir_picks_14_when_headers_exist(monkeypatch):
|
|
"""If the user has libstdc++-14-dev installed, prefer gcc-14."""
|
|
monkeypatch.setattr(sys, "platform", "linux")
|
|
import platform as _platform
|
|
|
|
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
|
|
monkeypatch.setattr(
|
|
worker.os.path,
|
|
"isdir",
|
|
_isdir_for_layout(
|
|
"/usr/lib/gcc/x86_64-linux-gnu/14/include",
|
|
"/usr/include/c++/14",
|
|
),
|
|
)
|
|
assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/14"
|
|
|
|
|
|
def test_hipcc_gcc_install_dir_returns_none_when_no_match(monkeypatch):
|
|
"""No gcc dir has both halves → return None and skip env injection
|
|
rather than guessing wrong and causing a confusing build failure."""
|
|
monkeypatch.setattr(sys, "platform", "linux")
|
|
import platform as _platform
|
|
|
|
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
|
|
monkeypatch.setattr(worker.os.path, "isdir", lambda path: False)
|
|
assert worker._hipcc_gcc_install_dir() is None
|
|
|
|
|
|
def test_hipcc_gcc_install_dir_returns_none_on_non_linux(monkeypatch):
|
|
"""Don't probe gcc layout on macOS / Windows — early-return."""
|
|
monkeypatch.setattr(sys, "platform", "darwin")
|
|
|
|
def _isdir_should_not_be_called(_path):
|
|
raise AssertionError("isdir should not be called on non-Linux")
|
|
|
|
monkeypatch.setattr(worker.os.path, "isdir", _isdir_should_not_be_called)
|
|
assert worker._hipcc_gcc_install_dir() is None
|
|
|
|
|
|
def test_hipcc_gcc_install_dir_returns_none_on_non_x86_64(monkeypatch):
|
|
"""ROCm clang-20 on aarch64 has a different libstdc++ layout."""
|
|
monkeypatch.setattr(sys, "platform", "linux")
|
|
import platform as _platform
|
|
|
|
monkeypatch.setattr(_platform, "machine", lambda: "aarch64")
|
|
assert worker._hipcc_gcc_install_dir() is None
|
|
|
|
|
|
def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None):
|
|
"""Scaffolding for end-to-end tests of the HIP source-build branch of
|
|
_install_package_wheel_first: package not installed, no prebuilt
|
|
wheel, hipcc on PATH, fake env reports HIP torch."""
|
|
monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d"))
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"probe_torch_wheel_env",
|
|
lambda timeout = 30: {
|
|
"hip_version": "7.13.26176",
|
|
"python_tag": "cp312",
|
|
"torch_mm": "2.11",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None)
|
|
monkeypatch.setattr(
|
|
worker.shutil,
|
|
"which",
|
|
lambda name: "/opt/rocm/bin/hipcc" if name == "hipcc" else None,
|
|
)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
monkeypatch.setattr(worker, "_hipcc_gcc_install_dir", lambda: gcc_dir)
|
|
|
|
|
|
def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch):
|
|
"""HIP source-build with no user-set HIPCC_COMPILE_FLAGS_APPEND →
|
|
subprocess env carries --gcc-install-dir=<detected path>."""
|
|
monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False)
|
|
_make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
|
|
|
|
captured: dict[str, str] = {}
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
captured.update(kwargs.get("env") or {})
|
|
return subprocess.CompletedProcess(cmd, 0, "")
|
|
|
|
monkeypatch.setattr(worker._sp, "run", fake_run)
|
|
|
|
worker._install_package_wheel_first(
|
|
event_queue = [],
|
|
import_name = "causal_conv1d",
|
|
display_name = "causal-conv1d",
|
|
pypi_name = "causal-conv1d",
|
|
pypi_version = "1.6.2.post1",
|
|
filename_prefix = "causal_conv1d",
|
|
release_tag = "v1.6.2.post1",
|
|
release_base_url = "https://example.com",
|
|
)
|
|
|
|
assert (
|
|
captured.get("HIPCC_COMPILE_FLAGS_APPEND")
|
|
== "--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13"
|
|
)
|
|
|
|
|
|
def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch):
|
|
"""User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' → final value keeps
|
|
the user's flags AND appends --gcc-install-dir."""
|
|
monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO")
|
|
_make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
|
|
|
|
captured: dict[str, str] = {}
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
captured.update(kwargs.get("env") or {})
|
|
return subprocess.CompletedProcess(cmd, 0, "")
|
|
|
|
monkeypatch.setattr(worker._sp, "run", fake_run)
|
|
|
|
worker._install_package_wheel_first(
|
|
event_queue = [],
|
|
import_name = "causal_conv1d",
|
|
display_name = "causal-conv1d",
|
|
pypi_name = "causal-conv1d",
|
|
pypi_version = "1.6.2.post1",
|
|
filename_prefix = "causal_conv1d",
|
|
release_tag = "v1.6.2.post1",
|
|
release_base_url = "https://example.com",
|
|
)
|
|
|
|
assert captured.get("HIPCC_COMPILE_FLAGS_APPEND") == (
|
|
"-O3 -DFOO --gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13"
|
|
)
|
|
|
|
|
|
def test_install_respects_user_gcc_install_dir(monkeypatch):
|
|
"""User explicitly set --gcc-install-dir=… already → don't touch it.
|
|
Avoids two competing --gcc-install-dir flags on the clang command line."""
|
|
monkeypatch.setenv(
|
|
"HIPCC_COMPILE_FLAGS_APPEND",
|
|
"--gcc-install-dir=/opt/custom/gcc-13",
|
|
)
|
|
_make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
|
|
|
|
captured: dict[str, str] = {}
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
captured.update(kwargs.get("env") or {})
|
|
return subprocess.CompletedProcess(cmd, 0, "")
|
|
|
|
monkeypatch.setattr(worker._sp, "run", fake_run)
|
|
|
|
worker._install_package_wheel_first(
|
|
event_queue = [],
|
|
import_name = "causal_conv1d",
|
|
display_name = "causal-conv1d",
|
|
pypi_name = "causal-conv1d",
|
|
pypi_version = "1.6.2.post1",
|
|
filename_prefix = "causal_conv1d",
|
|
release_tag = "v1.6.2.post1",
|
|
release_base_url = "https://example.com",
|
|
)
|
|
|
|
assert captured["HIPCC_COMPILE_FLAGS_APPEND"] == "--gcc-install-dir=/opt/custom/gcc-13"
|
|
|
|
|
|
def test_install_does_not_inject_env_on_cuda(monkeypatch):
|
|
"""CUDA path (no hip_version in env) → no HIP flag injected."""
|
|
monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False)
|
|
monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d"))
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"probe_torch_wheel_env",
|
|
lambda timeout = 30: {
|
|
"python_tag": "cp312",
|
|
"torch_mm": "2.11",
|
|
"cuda_major": "12",
|
|
"cxx11abi": "TRUE",
|
|
"platform_tag": "linux_x86_64",
|
|
},
|
|
)
|
|
monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None)
|
|
monkeypatch.setattr(worker.shutil, "which", lambda name: None)
|
|
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
|
|
# _hipcc_gcc_install_dir must not be called on CUDA.
|
|
monkeypatch.setattr(
|
|
worker,
|
|
"_hipcc_gcc_install_dir",
|
|
lambda: (_ for _ in ()).throw(AssertionError("must not run on CUDA")),
|
|
)
|
|
|
|
captured: dict[str, Any] = {}
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
captured.update(kwargs.get("env") or {})
|
|
return subprocess.CompletedProcess(cmd, 0, "")
|
|
|
|
monkeypatch.setattr(worker._sp, "run", fake_run)
|
|
|
|
worker._install_package_wheel_first(
|
|
event_queue = [],
|
|
import_name = "causal_conv1d",
|
|
display_name = "causal-conv1d",
|
|
pypi_name = "causal-conv1d",
|
|
pypi_version = "1.6.2.post1",
|
|
filename_prefix = "causal_conv1d",
|
|
release_tag = "v1.6.2.post1",
|
|
release_base_url = "https://example.com",
|
|
)
|
|
|
|
# env is always passed (to force UTF-8), but never the HIP flag.
|
|
assert "HIPCC_COMPILE_FLAGS_APPEND" not in captured
|