Compare commits
9 commits
main
...
harden-llm
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2aad8cbe5e | ||
|
|
4dfcdeb156 | ||
|
|
1f4bed5b35 | ||
|
|
413e5f844d | ||
|
|
b86d9ff821 | ||
|
|
fe5eee6a36 | ||
|
|
7d91e51ac9 | ||
|
|
0a28144f96 | ||
|
|
9ce0b6f39e |
2 changed files with 159 additions and 8 deletions
112
tests/saving/test_llm_compressor_install_pin.py
Normal file
112
tests/saving/test_llm_compressor_install_pin.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
"""Static guards (no import/network/GPU, like test_save_shell_injection.py) that
|
||||
install_llm_compressor()'s first-use auto-install of llm-compressor stays version-pinned to a vetted
|
||||
range and keeps its opt-out env gate, so a compromised/inflated release can't be auto-pulled."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
|
||||
|
||||
_ENV_FLAG = "UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL"
|
||||
|
||||
|
||||
def _module() -> ast.Module:
|
||||
return ast.parse(SAVE_PY.read_text(encoding = "utf-8"), filename = str(SAVE_PY))
|
||||
|
||||
|
||||
def _get_function(name: str) -> ast.FunctionDef:
|
||||
for node in ast.walk(_module()):
|
||||
if isinstance(node, ast.FunctionDef) and node.name == name:
|
||||
return node
|
||||
raise AssertionError(f"Function {name} not found in save.py")
|
||||
|
||||
|
||||
def _spec_value():
|
||||
for node in ast.walk(_module()):
|
||||
if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant):
|
||||
if any(
|
||||
isinstance(t, ast.Name) and t.id == "_LLM_COMPRESSOR_SPEC" for t in node.targets
|
||||
):
|
||||
return node.value.value
|
||||
return None
|
||||
|
||||
|
||||
def _first_lineno(fn: ast.AST, predicate) -> int | None:
|
||||
lines = [n.lineno for n in ast.walk(fn) if predicate(n) and hasattr(n, "lineno")]
|
||||
return min(lines) if lines else None
|
||||
|
||||
|
||||
def test_spec_is_a_bounded_pin() -> None:
|
||||
spec = _spec_value()
|
||||
assert spec is not None, "_LLM_COMPRESSOR_SPEC must be defined at module scope"
|
||||
assert "llmcompressor" in spec, f"spec must name llmcompressor, got {spec!r}"
|
||||
# A lower and an upper bound: pip cannot jump to an arbitrary (e.g. inflated) future release.
|
||||
assert ">=" in spec and "<" in spec, f"spec must have lower and upper bounds, got {spec!r}"
|
||||
|
||||
|
||||
def test_ceiling_blocks_inflated_versions() -> None:
|
||||
"""Cap to the exact vetted patch: block an inflated 0.x, a new major, and any higher in-range patch."""
|
||||
from packaging.requirements import Requirement
|
||||
|
||||
spec = Requirement(_spec_value()).specifier
|
||||
assert spec.contains("0.12.0"), "the current vetted release must resolve"
|
||||
assert not spec.contains("0.999.0"), "an inflated 0.x must be blocked"
|
||||
assert not spec.contains("1.0.0"), "a new major must not be auto-installed"
|
||||
assert not spec.contains(
|
||||
"0.12.1"
|
||||
), "a higher in-range patch must be blocked (cap to the vetted patch)"
|
||||
assert not spec.contains(
|
||||
"0.12.999"
|
||||
), "a crafted higher in-range patch (e.g. on a mirror) must be blocked"
|
||||
|
||||
|
||||
def test_floor_stays_compatible_with_supported_torch() -> None:
|
||||
"""Floor must stay <=0.6.0: 0.7+ need torch>=2.7, but the pinned torch can be as old as 2.4."""
|
||||
from packaging.requirements import Requirement
|
||||
from packaging.version import Version
|
||||
|
||||
req = Requirement(_spec_value())
|
||||
lowers = [Version(s.version) for s in req.specifier if s.operator in (">=", "==", "~=")]
|
||||
assert lowers, "spec must declare a lower bound"
|
||||
assert max(lowers) <= Version("0.6.0"), (
|
||||
f"floor {max(lowers)} requires a torch newer than Unsloth's minimum (2.4); "
|
||||
"llm-compressor >0.6.0 needs torch>=2.7. Keep the floor <= 0.6.0."
|
||||
)
|
||||
|
||||
|
||||
def test_install_command_uses_pinned_spec_not_bare_name() -> None:
|
||||
fn = _get_function("install_llm_compressor")
|
||||
# No argv list may pass the bare, unpinned package literal "llmcompressor".
|
||||
for node in ast.walk(fn):
|
||||
if isinstance(node, ast.List):
|
||||
for elt in node.elts:
|
||||
if isinstance(elt, ast.Constant) and elt.value == "llmcompressor":
|
||||
raise AssertionError(
|
||||
"install command must not pass an unpinned 'llmcompressor' literal; "
|
||||
"use the bounded _LLM_COMPRESSOR_SPEC"
|
||||
)
|
||||
names = {n.id for n in ast.walk(fn) if isinstance(n, ast.Name)}
|
||||
assert "_LLM_COMPRESSOR_SPEC" in names, "install command must reference _LLM_COMPRESSOR_SPEC"
|
||||
|
||||
|
||||
def test_optout_env_gate_precedes_subprocess_install() -> None:
|
||||
fn = _get_function("install_llm_compressor")
|
||||
env_line = _first_lineno(fn, lambda n: isinstance(n, ast.Constant) and n.value == _ENV_FLAG)
|
||||
assert env_line is not None, f"{_ENV_FLAG} opt-out must be checked in install_llm_compressor"
|
||||
|
||||
def _is_check_call(n: ast.AST) -> bool:
|
||||
return (
|
||||
isinstance(n, ast.Call)
|
||||
and isinstance(n.func, ast.Attribute)
|
||||
and n.func.attr == "check_call"
|
||||
and isinstance(n.func.value, ast.Name)
|
||||
and n.func.value.id == "subprocess"
|
||||
)
|
||||
|
||||
install_line = _first_lineno(fn, _is_check_call)
|
||||
assert install_line is not None, "expected a subprocess.check_call install in the function"
|
||||
assert (
|
||||
env_line < install_line
|
||||
), "the auto-install opt-out must be evaluated before any package install runs"
|
||||
|
|
@ -1363,11 +1363,18 @@ def install_python_non_blocking(packages = []):
|
|||
return run_installer
|
||||
|
||||
|
||||
# Bound the first-use auto-install so no unvetted release is pulled: not an inflated "0.999.0", nor
|
||||
# a crafted higher in-range patch like "0.12.999" from a mirror. Cap to the exact vetted patch and
|
||||
# bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below).
|
||||
_LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0"
|
||||
|
||||
|
||||
def install_llm_compressor():
|
||||
"""Import llm-compressor, installing it on first use for FP8/FP4 export.
|
||||
|
||||
Pins the current torch + transformers so pip does not upgrade them (a plain install pulls
|
||||
transformers>=5 and breaks Unsloth). Returns (oneshot, QuantizationModifier).
|
||||
Installs a version-pinned llm-compressor, pinning the current torch + transformers so pip does
|
||||
not upgrade them. Set UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL=1 to forbid the auto-install.
|
||||
Returns (oneshot, QuantizationModifier).
|
||||
"""
|
||||
try:
|
||||
from llmcompressor import oneshot
|
||||
|
|
@ -1376,9 +1383,41 @@ def install_llm_compressor():
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
# Already installed but not importable in THIS process? Do not reinstall. The in-process import
|
||||
# can fail under Unsloth's transformers patches (the compressed export quantizes in an isolated
|
||||
# subprocess that re-imports cleanly), and reinstalling the version-capped, torch/transformers-
|
||||
# pinned spec makes pip re-resolve and backtrack destructively (it drags in numpy<2 built from
|
||||
# source and the export fails). Only fall through to a pip install when it is genuinely absent.
|
||||
try:
|
||||
from importlib.metadata import version as _iv, PackageNotFoundError as _PNF
|
||||
try:
|
||||
_iv("llmcompressor")
|
||||
# Present; the compressed-export subprocess performs the real import. The caller only
|
||||
# uses this to trigger the install and fail fast, so returning None here is safe.
|
||||
return None, None
|
||||
except _PNF:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Opt-out for locked-down / air-gapped setups: forbid the auto-install, require a manual one.
|
||||
if os.environ.get("UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL", "0").lower() not in (
|
||||
"0",
|
||||
"",
|
||||
"false",
|
||||
"no",
|
||||
):
|
||||
raise RuntimeError(
|
||||
"Unsloth: llm-compressor is required for FP8/FP4 compressed export but is not "
|
||||
"installed, and automatic installation is disabled via "
|
||||
"UNSLOTH_DISABLE_LLM_COMPRESSOR_AUTOINSTALL. Install it manually with:\n"
|
||||
f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n"
|
||||
"(pin torch and transformers to your current versions to avoid upgrading them)."
|
||||
)
|
||||
|
||||
print(
|
||||
"Unsloth: Installing llm-compressor for FP8/FP4 export "
|
||||
"(pinning your torch + transformers so they are not upgraded). "
|
||||
f"({_LLM_COMPRESSOR_SPEC}; pinning your torch + transformers so they are not upgraded). "
|
||||
"This can take a few minutes..."
|
||||
)
|
||||
import importlib
|
||||
|
|
@ -1401,13 +1440,13 @@ def install_llm_compressor():
|
|||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("pip") is not None:
|
||||
cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"]
|
||||
cmd = [sys.executable, "-m", "pip", "install", _LLM_COMPRESSOR_SPEC]
|
||||
elif shutil.which("uv") is not None:
|
||||
cmd = ["uv", "pip", "install", "--python", sys.executable, "llmcompressor"]
|
||||
cmd = ["uv", "pip", "install", "--python", sys.executable, _LLM_COMPRESSOR_SPEC]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Unsloth: cannot install llm-compressor because this environment has neither pip nor "
|
||||
f"uv. Install it manually with:\n uv pip install --python {sys.executable} llmcompressor\n"
|
||||
f"uv. Install it manually with:\n uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n"
|
||||
"(pin torch and transformers to your current versions to avoid upgrading them)."
|
||||
)
|
||||
cpath = None
|
||||
|
|
@ -1421,8 +1460,8 @@ def install_llm_compressor():
|
|||
except subprocess.CalledProcessError as e:
|
||||
raise RuntimeError(
|
||||
"Unsloth: Failed to install llm-compressor. Install it manually with:\n"
|
||||
f" uv pip install --python {sys.executable} llmcompressor\n"
|
||||
f"or, if pip is available:\n {sys.executable} -m pip install llmcompressor\n"
|
||||
f" uv pip install --python {sys.executable} '{_LLM_COMPRESSOR_SPEC}'\n"
|
||||
f"or, if pip is available:\n {sys.executable} -m pip install '{_LLM_COMPRESSOR_SPEC}'\n"
|
||||
"(pin torch and transformers to your current versions to avoid upgrading them).\n"
|
||||
f"Underlying error: {e}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue