diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5c063cd1c07..31eb89c3123 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -33,6 +33,7 @@ Changelog - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. - Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. +- Add a ``--deepgemm-url`` option to ``examples/vllm_serve/vllm_serve_fakequant.py`` that serves with a prebuilt DeepGEMM wheel, such as a DeepGEMM fork's, from a package index, without changing the Python environment; see `examples/vllm_serve/README.md `_. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. - vLLM fake-quant serving now runs on pre-quantized checkpoints such as FP8 when the recipe leaves those layers unquantized (for example a KV-cache-only recipe), and on MLA models with an FP8 KV cache; both previously failed during quantization. - ``KV_QUANT_CFG`` presets in vLLM fake-quant serving now quantize the MLA KV cache; on vLLM 0.16 and later they silently quantized nothing for MLA models. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 8eaea64700d..299ad9950d0 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -49,6 +49,8 @@ corresponding environment variable when omitted: | `--modelopt-quant-dataset` | `QUANT_DATASET` | Calibration dataset | | `--modelopt-quant-calib-size` | `QUANT_CALIB_SIZE` | Calibration sample count | | `--modelopt-calib-batch-size` | `CALIB_BATCH_SIZE` | Calibration batch size | +| `--deepgemm-url` | `DEEPGEMM_URL` | Prebuilt DeepGEMM to serve with: a package index with its wheels, or a wheel (see [Serve with a prebuilt DeepGEMM](#serve-with-a-prebuilt-deepgemm)) | +| `--deepgemm-version` | `DEEPGEMM_VERSION` | Version to install from the `DEEPGEMM_URL` package index instead of the newest one for the installed torch and CUDA | `QUANT_CFG` and `KV_QUANT_CFG` (and their CLI flags) will be deprecated in a future release. They still work today. For new runs, use a PTQ recipe through `RECIPE_PATH` or @@ -223,6 +225,51 @@ Tracking needs the client: `pip install nvidia-modelopt[mlflow]` (already in thi `--distributed-executor-backend ray` those are forwarded to the workers along with the tracking settings, since a Ray worker starts with a clean environment. +## Serve with a prebuilt DeepGEMM + +vLLM serves with its own copy of DeepGEMM unless another `deep_gemm` comes first on the Python +path. To serve with a DeepGEMM build instead, such as a DeepGEMM fork that publishes prebuilt wheels +to a package index, pass the index as `--deepgemm-url` or `DEEPGEMM_URL`: + +```bash +DEEPGEMM_URL="https:///simple" \ + python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 +``` + +Before vLLM starts, pip selects the newest `deep_gemm` version on the index that has a wheel for +this Python, CPU architecture and glibc (its tag, such as `manylinux_2_38_x86_64`, names the oldest +glibc it runs on), or the version that `--deepgemm-version` or `DEEPGEMM_VERSION` pins, e.g. +`2.8.0.post12` or exactly `2.8.0.post12+g1a2b3c4.torch2.13.cu130`. The wheel must be built for the +installed torch and CUDA: if the index has builds for others, pin one. The launcher installs the +wheel with `pip install --no-deps --target` into `/wheels/---`, +once per version and source and one launch at a time, checks in a subprocess that `import deep_gemm` +loads it with the installed torch, and puts that folder first on `PYTHONPATH`, so that vLLM's +workers import it instead of vLLM's own copy. Nothing is built, and the Python environment does not +change. `` is `MODELOPT_DEEPGEMM_CACHE_DIR`, else `$VLLM_CACHE_ROOT/modelopt/deepgemm`, else +`~/.cache/modelopt/deepgemm`. + +`DEEPGEMM_URL` also takes a single wheel (a path, or an `https://` or `file://` URL), installed the +same way; a local wheel once per content, so a rebuilt wheel of the same version is installed +again. + +To serve a DeepGEMM source tree, such as a local branch, install it into the serving environment +instead; vLLM then uses it, without `DEEPGEMM_URL`, for every serve from that environment until it +is uninstalled: + +```bash +pip install --no-deps --no-build-isolation +python vllm_serve_fakequant.py ... +``` + +Notes: + +- Every installed version adds a folder of tens of MB to `/wheels`, which can be deleted + when no server uses it. +- pip runs with your pip settings, e.g. `PIP_TRUSTED_HOST` for an `http://` index on another host. +- Ray workers that vLLM does not start from the launcher's environment, such as those of a running + Ray cluster, need the installed folder (printed at launch) on their `PYTHONPATH`, on a filesystem + that all nodes share. + ## Load QAT/PTQ model and serve in vLLM (WIP) Step 1: export the model with bf16 weights and quantizer state. To export the model: diff --git a/examples/vllm_serve/vllm_deepgemm_utils.py b/examples/vllm_serve/vllm_deepgemm_utils.py new file mode 100644 index 00000000000..5937eb50b6d --- /dev/null +++ b/examples/vllm_serve/vllm_deepgemm_utils.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Serve with a custom DeepGEMM: a prebuilt wheel, e.g. of a fork with the indexer scorer numerics. + +vLLM imports a ``deep_gemm`` on the Python path before its vendored copy. ``--deepgemm-url`` (or +``$DEEPGEMM_URL``) names a package index with prebuilt DeepGEMM wheels, or one wheel. From an index, +pip selects the newest version that has a wheel for this Python and machine, or the version that +``--deepgemm-version`` (``$DEEPGEMM_VERSION``) pins. Before vLLM starts its workers, the launcher +installs the wheel with ``pip install --target`` into a folder of its cache, once per version and +source and under a lock, checks in a subprocess that ``import deep_gemm`` loads it with the installed +torch, and puts the folder first on the Python path of the server's processes. Nothing is built, +and the Python environment itself does not change. + +Nothing here imports vLLM, so this can be tested without it. +""" + +import argparse +import contextlib +import fcntl +import hashlib +import importlib +import json +import os +import platform +import shutil +import subprocess +import sys +import tempfile +import time +from collections.abc import Iterator +from pathlib import Path +from urllib.parse import unquote, urlsplit + +import torch +from packaging import tags +from packaging.utils import InvalidWheelFilename, canonicalize_name, parse_wheel_filename +from packaging.version import InvalidVersion, Version + +# TODO: make this generic when another prebuilt kernel package is needed. Only the names below, the +# flags and the import check are DeepGEMM's. It only helps packages that vLLM imports by name from +# the Python path; vLLM builds some, such as FlashMLA, into its own wheel. +DEEPGEMM_URL_ENV = "DEEPGEMM_URL" +DEEPGEMM_VERSION_ENV = "DEEPGEMM_VERSION" +# The root of the cache of installed DeepGEMM wheels (see deepgemm_cache_dir). +DEEPGEMM_CACHE_ENV = "MODELOPT_DEEPGEMM_CACHE_DIR" + +_PROJECT = "deep-gemm" # the normalized name of deep_gemm on an index +# Wheels only: nothing is built, e.g. from the unrelated deep_gemm source distribution on PyPI. An +# installed deep_gemm does not count, or pip would select nothing that it satisfies. +_PIP_OPTIONS = ( + "--no-deps", + "--no-input", + "--disable-pip-version-check", + "--ignore-installed", + "--only-binary", + ":all:", +) +# A pip run reads an index and downloads a wheel of tens of MB, while it can hold the cache lock. +_PIP_TIMEOUT = 1800 # seconds +_IMPORT_TIMEOUT = 600 # seconds +# Installs stage in the cache under its lock; staging folders older than this are left by launches +# that were killed while installing. +_STALE_STAGING = 24 * 3600 # seconds +# Run by a fresh interpreter with the installed folder first on the Python path, like the workers. +_IMPORT_CHECK = "import deep_gemm; print(deep_gemm.__file__)" + + +def add_deepgemm_args(parser: argparse.ArgumentParser) -> None: + """Add ``--deepgemm-url`` and ``--deepgemm-version`` to the launcher's parser.""" + parser.add_argument( + "--deepgemm_url", + "--deepgemm-url", + default=None, + help=( + "Serve with the DeepGEMM wheel that pip selects from this package index " + "(https://.../simple), or with this wheel (a path or URL), installed into a cache " + "folder that is put first on the Python path before vLLM starts. Default: " + f"${DEEPGEMM_URL_ENV}." + ), + ) + parser.add_argument( + "--deepgemm_version", + "--deepgemm-version", + default=None, + help=( + "The version to install from the package index of --deepgemm-url instead of the " + "newest one, e.g. 2.8.0.post210 or, exactly, 2.8.0.post210+g6fec872c.torch2.13.cu130. " + f"Default: ${DEEPGEMM_VERSION_ENV}." + ), + ) + + +def resolve_deepgemm_args(args: argparse.Namespace) -> None: + """Install the DeepGEMM that ``--deepgemm-url`` or ``$DEEPGEMM_URL`` names, if any, and use it.""" + url = args.deepgemm_url or os.environ.get(DEEPGEMM_URL_ENV) + version = args.deepgemm_version or os.environ.get(DEEPGEMM_VERSION_ENV) + if not url: + if version: + raise ValueError( + f"{DEEPGEMM_VERSION_ENV} {version} pins a version of the package index that " + f"{DEEPGEMM_URL_ENV} names, which is not set." + ) + return + folder, installed = install_deepgemm(url, version) + prepend_python_path(folder) + print(f"[deepgemm] serving with deep_gemm {installed} from {url}: {folder}") + + +def install_deepgemm(url: str, version: str | None = None) -> tuple[Path, str]: + """Install the DeepGEMM that ``url`` names into a folder of the cache, unless it is there. + + Args: + url: A package index, or a wheel: a path, or an http(s) or file URL. + version: The version to install from the package index ``url``. + + Returns: + The folder, from which ``import deep_gemm`` loads the installed DeepGEMM with the installed + torch, and the installed version. + """ + name = unquote(urlsplit(url).path) if "://" in url else url + if name.endswith(".whl"): + if version is not None: + raise ValueError( + f"{DEEPGEMM_VERSION_ENV} pins a version of a package index, but {DEEPGEMM_URL_ENV} " + f"{url} is a wheel." + ) + installed, requirement = _wheel_version(Path(name).name), [url] + else: + pin = _parse_version(version) if version is not None else None + # A version with its local label is exact: once installed, the index is not needed. + installed = str(pin) if pin is not None and pin.local else _resolve(url, pin) + requirement = ["--index-url", url, f"deep_gemm=={installed}"] + wheels = deepgemm_cache_dir() / "wheels" + folder = wheels / _folder_name(installed, url) + if not folder.is_dir(): # complete, as published by a rename; a read-only cache works too + with _locked(wheels): + if not folder.is_dir(): + print(f"[deepgemm] installing deep_gemm {installed} from {url}", flush=True) + _pip_install(requirement, folder, url, installed) # checks the import first + return folder, installed + _check_import(folder, installed, url) # again: the environment can change + return folder, installed + + +def prepend_python_path(folder: Path) -> None: + """Put ``folder`` first on the Python path of this process and of the processes it starts.""" + path = str(folder) + if path not in sys.path: + sys.path.insert(0, path) + importlib.invalidate_caches() + python_path = os.environ.get("PYTHONPATH") + if not python_path: + os.environ["PYTHONPATH"] = path + elif python_path.split(os.pathsep)[0] != path: + os.environ["PYTHONPATH"] = f"{path}{os.pathsep}{python_path}" + + +def deepgemm_cache_dir() -> Path: + """The root of the cache of installed DeepGEMM wheels. + + ``$MODELOPT_DEEPGEMM_CACHE_DIR``, else ``modelopt/deepgemm`` under ``$VLLM_CACHE_ROOT``, else + ``~/.cache/modelopt/deepgemm``. + """ + if root := os.environ.get(DEEPGEMM_CACHE_ENV): + return Path(root).expanduser().absolute() + if vllm_root := os.environ.get("VLLM_CACHE_ROOT"): + return Path(vllm_root).expanduser().absolute() / "modelopt" / "deepgemm" + return Path.home() / ".cache" / "modelopt" / "deepgemm" + + +def _pip() -> list[str]: + # subprocess: pip has no supported Python API. Safe: an argument list without a shell, and pip + # installs only wheels, from the index or wheel that the user passes and must trust like any + # package source. + return [sys.executable, "-m", "pip", "install", *_PIP_OPTIONS] + + +def _run_pip(command: list[str], what: str, **kwargs) -> subprocess.CompletedProcess: + """Run the pip ``command`` for ``what`` within ``_PIP_TIMEOUT``.""" + try: + return subprocess.run(command, timeout=_PIP_TIMEOUT, check=False, **kwargs) + except subprocess.TimeoutExpired: + raise RuntimeError(f"pip did not finish {what} in {_PIP_TIMEOUT} s.") from None + + +def _parse_version(version: str) -> Version: + try: + return Version(version) + except InvalidVersion: + raise ValueError(f"{DEEPGEMM_VERSION_ENV} {version!r} is no version.") from None + + +def _wheel_version(name: str) -> str: + """The version of the deep_gemm wheel file ``name``.""" + try: + project, version, _, _ = parse_wheel_filename(name) + except InvalidWheelFilename as error: + raise ValueError(f"{name}: {error}") from None + if canonicalize_name(project) != _PROJECT: + raise ValueError(f"{name} is no deep_gemm wheel.") + return str(version) + + +def _resolve(index: str, pin: Version | None) -> str: + """The version of deep_gemm that pip selects on ``index``: the pinned or the newest one.""" + requirement = f"deep_gemm=={pin}" if pin is not None else "deep_gemm" + # pip prints its own errors, such as an http index on another host without --trusted-host. + result = _run_pip( + [*_pip(), "--dry-run", "--quiet", "--report", "-", "--index-url", index, requirement], + f"selecting deep_gemm on {index}", + stdout=subprocess.PIPE, + text=True, + ) + if result.returncode != 0: + raise RuntimeError( + f"pip finds no deep_gemm wheel for {_python_tag()} on {platform.machine()} on {index}, " + f"see its output above. {DEEPGEMM_URL_ENV} takes a package index (https://.../simple) " + "or a wheel; to serve a DeepGEMM source tree, install it into the environment with " + "`pip install --no-deps --no-build-isolation ` instead." + ) + selected = json.loads(result.stdout)["install"] + if len(selected) != 1: # --no-deps and --ignore-installed: deep_gemm only + raise RuntimeError(f"pip selected {len(selected)} packages for deep_gemm on {index}.") + return selected[0]["metadata"]["version"] + + +def _python_tag() -> str: + return f"{tags.interpreter_name()}{tags.interpreter_version()}" + + +def _folder_name(version: str, url: str) -> str: + """The cache folder of ``version`` from ``url``, which no other source of it shares.""" + return f"{version}-{_python_tag()}-{platform.machine()}-{_source_id(url)}" + + +def _source_id(url: str) -> str: + """A local wheel by its content, as a rebuild can keep its version; else the URL.""" + parts = urlsplit(url) + path = Path(unquote(parts.path) if parts.scheme == "file" else url) + if parts.scheme in ("", "file") and path.is_file(): + source = path.read_bytes() + else: + source = url.rstrip("/").encode() + return hashlib.sha256(source).hexdigest()[:12] + + +def _pip_install(requirement: list[str], folder: Path, url: str, version: str) -> None: + """``pip install --target`` into ``folder``: published complete, and only if it imports. + + Everyone can read it. + """ + staging = Path(tempfile.mkdtemp(prefix=".install-", dir=folder.parent)) + try: + command = [*_pip(), "--target", str(staging), *requirement] + if _run_pip(command, f"installing deep_gemm from {url}").returncode != 0: + raise RuntimeError(f"Installing deep_gemm from {url} failed, see the pip output above.") + for directory, _, files in os.walk(staging): + os.chmod(directory, 0o755) # mkdtemp's 0700 would hide it from other users + for path in (os.path.join(directory, name) for name in files): + if not os.path.islink(path): + os.chmod(path, os.stat(path).st_mode | 0o444) + _check_import(staging, version, url) + try: + staging.rename(folder) + except OSError: + if not folder.is_dir(): # else another process published it first + raise + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def _last_line(text: str) -> str: + return (text.strip().splitlines() or [""])[-1] + + +def _check_import(folder: Path, version: str, url: str) -> None: + """Check that ``import deep_gemm`` loads the one installed in ``folder``, with this torch. + + vLLM would fall back to its own copy without an error. + """ + python_path = os.environ.get("PYTHONPATH") + env = { + **os.environ, + "PYTHONPATH": f"{folder}{os.pathsep}{python_path}" if python_path else str(folder), + } + try: + # subprocess: a fresh interpreter, as a worker starts, runs the fixed _IMPORT_CHECK without + # a shell. + result = subprocess.run( + [sys.executable, "-c", _IMPORT_CHECK], + capture_output=True, + text=True, + env=env, + cwd=folder, + timeout=_IMPORT_TIMEOUT, + check=False, + ) + except subprocess.TimeoutExpired: + raise RuntimeError(f"import deep_gemm {version} from {url} timed out.") from None + if result.returncode != 0: + error = _last_line(result.stderr) or "no error message" + raise RuntimeError( + f"deep_gemm {version} from {url} does not import ({error}): it needs a build for this " + f"torch {torch.__version__} with CUDA {torch.version.cuda} and glibc " + f"{platform.libc_ver()[1] or '?'} or older. Pin one with --deepgemm-version or " + f"{DEEPGEMM_VERSION_ENV}." + ) + loaded = Path(_last_line(result.stdout)).resolve() + if not loaded.is_relative_to(folder.resolve()): + raise RuntimeError( + f"import deep_gemm loads {loaded} instead of deep_gemm {version} from {url}: " + "another deep_gemm comes first on the Python path, such as an editable install. " + "Uninstall it." + ) + + +@contextlib.contextmanager +def _locked(directory: Path) -> Iterator[None]: + """Hold an exclusive lock on the cache ``directory``, which it creates. + + Also removes the old staging folders that launches killed while installing left in it. + """ + directory.mkdir(parents=True, exist_ok=True) + with (directory / ".lock").open("a") as lock: + try: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + print(f"[deepgemm] waiting for another launch installing into {directory}", flush=True) + fcntl.flock(lock, fcntl.LOCK_EX) + for staging in directory.glob(".install-*"): + with contextlib.suppress(OSError): + if time.time() - staging.stat().st_mtime > _STALE_STAGING: + shutil.rmtree(staging, ignore_errors=True) + yield diff --git a/examples/vllm_serve/vllm_serve_fakequant.py b/examples/vllm_serve/vllm_serve_fakequant.py index 24e692d04ff..e2c334d9ab8 100644 --- a/examples/vllm_serve/vllm_serve_fakequant.py +++ b/examples/vllm_serve/vllm_serve_fakequant.py @@ -64,6 +64,7 @@ import vllm from packaging import version from vllm.entrypoints.cli.main import main as vllm_main +from vllm_deepgemm_utils import add_deepgemm_args, resolve_deepgemm_args from vllm_mlflow_utils import MLFLOW_ENV_VARS, add_mlflow_args, resolve_mlflow_args @@ -284,6 +285,7 @@ def main(): modelopt_parser = FlexibleArgumentParser(add_help=False) _add_fakequant_args(modelopt_parser) add_mlflow_args(modelopt_parser) + add_deepgemm_args(modelopt_parser) modelopt_args, rest_argv = modelopt_parser.parse_known_args(sys.argv[1:]) rest_argv = _default_to_serve(rest_argv) if rest_argv and rest_argv[0] != "serve": @@ -347,6 +349,8 @@ def main(): if not _has_flag(rest_argv, "--moe-backend"): rest_argv = [*rest_argv, "--moe-backend", "triton"] + # Before vLLM starts the workers, which import deep_gemm from the Python path it sets. + resolve_deepgemm_args(modelopt_args) _run_vllm_cli(rest_argv) diff --git a/tests/examples/vllm_serve/test_fakequant_example.py b/tests/examples/vllm_serve/test_fakequant_example.py index d5976cd9e21..2898798f09f 100644 --- a/tests/examples/vllm_serve/test_fakequant_example.py +++ b/tests/examples/vllm_serve/test_fakequant_example.py @@ -67,6 +67,8 @@ def clean_launcher_env(): "MODELOPT_MLFLOW_REQUIRED", "MODELOPT_MLFLOW_COMMAND", "MODELOPT_MLFLOW_RUN_NAME", + "DEEPGEMM_URL", + "DEEPGEMM_VERSION", ): os.environ.pop(key, None) yield @@ -156,6 +158,9 @@ def _launcher_import_modules(): add_mlflow_args=Mock(), resolve_mlflow_args=Mock(), ), + "vllm_deepgemm_utils": SimpleNamespace( + add_deepgemm_args=Mock(), resolve_deepgemm_args=Mock() + ), "vllm.entrypoints.openai.cli_args": SimpleNamespace( make_arg_parser=entrypoints.legacy_arg_parser ), @@ -204,6 +209,27 @@ def test_vllm_serve_main_disables_compile_cache(monkeypatch, clean_launcher_env) vllm_main.assert_called_once_with() +def test_vllm_serve_main_installs_deepgemm_before_vllm_starts(monkeypatch, clean_launcher_env): + """vLLM's workers import deep_gemm from the Python path that the install sets.""" + launcher = _load_fakequant_launcher(monkeypatch) + monkeypatch.setattr(launcher, "resolve_mlflow_args", Mock()) + vllm_main, _, _ = _stub_launcher_runtime(monkeypatch, launcher) + calls = [] + vllm_main.side_effect = lambda: calls.append("vllm") + monkeypatch.setattr( + launcher, "resolve_deepgemm_args", lambda args: calls.append(args.deepgemm_url) + ) + url = "https://example.com/simple" + monkeypatch.setattr( + sys, "argv", ["vllm_serve_fakequant.py", "/models/qwen", "--deepgemm-url", url] + ) + + launcher.main() + + assert calls == [url, "vllm"] + assert sys.argv == ["vllm", "serve", "/models/qwen"] + + def test_vllm_serve_entrypoint_dependency_error_propagates(monkeypatch): """Do not replace a missing parser dependency with a fallback import error.""" modules, _ = _launcher_import_modules() diff --git a/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py b/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py new file mode 100644 index 00000000000..c6d2e826097 --- /dev/null +++ b/tests/examples/vllm_serve/test_vllm_deepgemm_utils.py @@ -0,0 +1,567 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Installing the DeepGEMM that the vLLM fake-quant server imports, into a folder of a cache. + +``vllm_deepgemm_utils`` imports no vLLM. pip and the import check are stubbed, except in the last +tests, which install tiny pure-Python ``deep_gemm`` wheels with the real pip, from a local package +index or a file, and check the import in a real subprocess. +""" + +import argparse +import base64 +import fcntl +import hashlib +import importlib +import json +import os +import platform +import re +import stat +import subprocess +import sys +import time +import zipfile +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "vllm_serve" + +LABEL = "torch2.13.cu130" +NEWEST = f"2.8.0.post5+gbbbbbbb.{LABEL}" +ANY = "py3-none-any" # the tags of a wheel that every Python installs +INDEX = "https://pypi.example.com/simple" +GLIBC = platform.libc_ver()[1] or "?" +PIP = [sys.executable, "-m", "pip", "install", "--no-deps", "--no-input"] +PIP += ["--disable-pip-version-check", "--ignore-installed", "--only-binary", ":all:"] + + +def _isolated_module(monkeypatch, tmp_path): + """``vllm_deepgemm_utils`` with HOME under ``tmp_path``. + + The variables it reads are unset; ``PYTHONPATH`` and ``sys.path``, which it sets, are restored. + """ + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) # also restores sys.path + module = importlib.import_module("vllm_deepgemm_utils") + monkeypatch.setenv("HOME", str(tmp_path / "home")) + for name in ( + module.DEEPGEMM_URL_ENV, + module.DEEPGEMM_VERSION_ENV, + module.DEEPGEMM_CACHE_ENV, + "VLLM_CACHE_ROOT", + ): + monkeypatch.delenv(name, raising=False) + python_path = os.environ.get("PYTHONPATH") + monkeypatch.setenv("PYTHONPATH", python_path or "") + if python_path is None: + monkeypatch.delenv("PYTHONPATH") + return module + + +@pytest.fixture +def deepgemm(monkeypatch, tmp_path): + return _isolated_module(monkeypatch, tmp_path) + + +def _wheels(): + return Path.home() / ".cache/modelopt/deepgemm/wheels" + + +class _Subprocesses: + """Stubbed pip and import check. + + A pip dry run selects ``selected``, pip installs a fake deep_gemm, and the check loads it. + Records each command with its keyword arguments. + """ + + def __init__(self, deepgemm, monkeypatch, returncode=0): + self.calls, self.returncode, self.import_check = [], returncode, None + self.selected = NEWEST + monkeypatch.setattr(deepgemm.subprocess, "run", self) + + def __call__(self, command, **kwargs): + self.calls.append(SimpleNamespace(command=command, kwargs=kwargs)) + if command[1] == "-c": + folder = Path(kwargs["cwd"]) + assert kwargs["env"]["PYTHONPATH"].split(os.pathsep)[0] == str(folder) + if self.import_check is not None: + return self.import_check(command) + stdout = f"{folder / 'deep_gemm/__init__.py'}\n" + return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + if self.returncode != 0: + return subprocess.CompletedProcess(command, self.returncode, stdout="") + if "--dry-run" in command: + report = {"install": [{"metadata": {"name": "deep_gemm", "version": self.selected}}]} + return subprocess.CompletedProcess(command, 0, stdout=json.dumps(report)) + target = _target(command) + (target / "deep_gemm").mkdir() + (target / "deep_gemm/__init__.py").write_text("__version__ = '2.8.0'\n") + (target / "deep_gemm/_C.so").write_bytes(b"\x7fELF") + (target / "deep_gemm/_C.so").chmod(0o700) + return subprocess.CompletedProcess(command, 0) + + def pip(self, dry_run=False): + """The pip installs, or the pip dry runs.""" + return [ + call + for call in self.calls + if call.command[1:4] == ["-m", "pip", "install"] + and ("--dry-run" in call.command) == dry_run + ] + + +def _target(command): + return Path(command[command.index("--target") + 1]) + + +def test_index_installs_the_version_pip_selects_once(deepgemm, monkeypatch): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + folder, version = deepgemm.install_deepgemm(INDEX) + + assert (folder, version) == (_wheels() / deepgemm._folder_name(NEWEST, INDEX), NEWEST) + source = hashlib.sha256(INDEX.encode()).hexdigest()[:12] + assert folder.name == f"{NEWEST}-{deepgemm._python_tag()}-{os.uname().machine}-{source}" + (dry_run,) = subprocesses.pip(dry_run=True) + assert dry_run.command == [ + *PIP, + *["--dry-run", "--quiet", "--report", "-", "--index-url", INDEX, "deep_gemm"], + ] + # Only the report is captured: pip's own errors reach the terminal. + assert dry_run.kwargs["stdout"] == subprocess.PIPE and "stderr" not in dry_run.kwargs + assert dry_run.kwargs["timeout"] == deepgemm._PIP_TIMEOUT + (install,) = subprocesses.pip() + assert install.kwargs["timeout"] == deepgemm._PIP_TIMEOUT + staging = _target(install.command) + assert staging.parent == _wheels() and not staging.exists() # published by a rename + assert install.command == [ + *PIP, + *["--target", str(staging), "--index-url", INDEX, f"deep_gemm=={NEWEST}"], + ] + assert (folder / "deep_gemm/__init__.py").is_file() + assert subprocesses.calls[-1].command == [sys.executable, "-c", deepgemm._IMPORT_CHECK] + + # Installed: pip selects the version again, and only the import is checked. + subprocesses.calls.clear() + assert deepgemm.install_deepgemm(INDEX) == (folder, version) + assert len(subprocesses.pip(dry_run=True)) == 1 and subprocesses.pip() == [] + + +def test_index_version_can_be_pinned(deepgemm, monkeypatch): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + older = f"2.8.0.post3+gaaaaaaa.{LABEL}" + subprocesses.selected = older + # Without its local label: pip selects the build of that version. + assert deepgemm.install_deepgemm(INDEX, "2.8.0.post3")[1] == older + (dry_run,) = subprocesses.pip(dry_run=True) + assert dry_run.command[-1] == "deep_gemm==2.8.0.post3" + assert subprocesses.pip()[-1].command[-1] == f"deep_gemm=={older}" + + # With it: exact, so pip does not select, and once it is installed pip does not run. + subprocesses.calls.clear() + exact = "2.8.0.post6+gccccccc.torch2.14.cu130" + assert deepgemm.install_deepgemm(INDEX, exact.upper()) == ( # normalized + _wheels() / deepgemm._folder_name(exact, INDEX), + exact, + ) + assert subprocesses.pip(dry_run=True) == [] + assert subprocesses.pip()[-1].command[-1] == f"deep_gemm=={exact}" + subprocesses.calls.clear() + deepgemm.install_deepgemm(INDEX, exact) + assert [call.command[1] for call in subprocesses.calls] == ["-c"] + + +@pytest.mark.parametrize("form", ["path", "file", "https"]) +def test_wheel_is_installed_once_per_version(deepgemm, monkeypatch, tmp_path, form): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + wheel = tmp_path / f"deep_gemm-{NEWEST}-{ANY}.whl" + url = { + "path": str(wheel), + "file": wheel.as_uri(), + "https": f"https://files.example.com/{wheel.name.replace('+', '%2B')}", + }[form] + folder, version = deepgemm.install_deepgemm(url) + assert (folder, version) == (_wheels() / deepgemm._folder_name(NEWEST, url), NEWEST) + assert subprocesses.pip(dry_run=True) == [] # the file name has the version + (install,) = subprocesses.pip() + assert install.command == [*PIP, "--target", str(_target(install.command)), url] + assert deepgemm.install_deepgemm(url) == (folder, version) + assert len(subprocesses.pip()) == 1 + + +def test_sources_of_a_version_install_separately(deepgemm, monkeypatch, tmp_path): + """A rebuilt wheel or another index with the same version is not served from the cache.""" + subprocesses = _Subprocesses(deepgemm, monkeypatch) + wheel = tmp_path / f"deep_gemm-{NEWEST}-{ANY}.whl" + wheel.write_bytes(b"first build") + first, _ = deepgemm.install_deepgemm(str(wheel)) + assert deepgemm.install_deepgemm(wheel.as_uri())[0] == first # the same file + wheel.write_bytes(b"rebuilt with the same version") + rebuilt, _ = deepgemm.install_deepgemm(str(wheel)) + index, _ = deepgemm.install_deepgemm(INDEX) + other_index, _ = deepgemm.install_deepgemm("https://other.example.com/simple") + assert len({first, rebuilt, index, other_index}) == 4 + assert all(folder.is_dir() for folder in (first, rebuilt, index, other_index)) + assert len(subprocesses.pip()) == 4 + # Each source reuses its own folder. + assert deepgemm.install_deepgemm(str(wheel))[0] == rebuilt + assert deepgemm.install_deepgemm(f"{INDEX}/")[0] == index + assert len(subprocesses.pip()) == 4 + + +def test_a_build_that_does_not_import_is_not_cached(deepgemm, monkeypatch): + """The import is checked before the install is published, so a fixed build can be retried.""" + subprocesses = _Subprocesses(deepgemm, monkeypatch) + subprocesses.import_check = lambda command: subprocess.CompletedProcess( + command, 1, stdout="", stderr="ImportError: wrong torch\n" + ) + with pytest.raises(RuntimeError, match=r"does not import \(ImportError: wrong torch\)"): + deepgemm.install_deepgemm(INDEX, NEWEST) + assert [path.name for path in _wheels().iterdir()] == [".lock"] # no folder, no staging + subprocesses.import_check = None # fixed, e.g. the environment + folder, _ = deepgemm.install_deepgemm(INDEX, NEWEST) + assert folder.is_dir() and len(subprocesses.pip()) == 2 + + +def test_pip_timeouts_release_the_cache(deepgemm, monkeypatch): + """A pip run that hangs fails the launch, leaves no staging folder and releases the lock.""" + + def hang(command, **kwargs): + assert kwargs["timeout"] == deepgemm._PIP_TIMEOUT + raise subprocess.TimeoutExpired(command, kwargs["timeout"]) + + monkeypatch.setattr(deepgemm.subprocess, "run", hang) + with pytest.raises(RuntimeError, match=r"pip did not finish selecting deep_gemm on .* in"): + deepgemm.install_deepgemm(INDEX) + with pytest.raises(RuntimeError, match=r"pip did not finish installing deep_gemm from .* in"): + deepgemm.install_deepgemm(INDEX, NEWEST) # under the lock + assert [path.name for path in _wheels().iterdir()] == [".lock"] + with (_wheels() / ".lock").open() as lock: + fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) + + +@pytest.mark.parametrize( + ("url", "version", "match"), + [ + (INDEX, "latest", "DEEPGEMM_VERSION 'latest' is no version"), + (f"https://example.com/deep_gemm-2.8.0-{ANY}.whl", "2.8.0", r"pins .* but .* is a wheel"), + ( + f"https://example.com/other-1.0-{ANY}.whl", + None, + r"other-1.0-py3-none-any.whl is no deep_", + ), + ("dist/deep_gemm.whl", None, r"deep_gemm.whl: Invalid wheel filename"), + ], + ids=["version", "wheel-pin", "other-wheel", "wheel-name"], +) +def test_inputs_are_checked_before_pip_runs(deepgemm, monkeypatch, url, version, match): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + with pytest.raises(ValueError, match=match): + deepgemm.install_deepgemm(url, version) + assert subprocesses.calls == [] + + +def test_pip_failures_are_reported(deepgemm, monkeypatch): + """Nothing is built: a git URL or source tree finds no wheel, with a hint to install it.""" + _Subprocesses(deepgemm, monkeypatch, returncode=1) + git = "git+https://git.example.com/group/deepgemm.git@main" + with pytest.raises(RuntimeError) as error: + deepgemm.install_deepgemm(git) + assert str(error.value).startswith( + f"pip finds no deep_gemm wheel for {deepgemm._python_tag()} on {os.uname().machine} on " + f"{git}, see its output above." + ) + assert "install it into the environment with `pip install --no-deps" in str(error.value) + with pytest.raises(RuntimeError, match=re.escape(f"from {INDEX} failed, see the pip output")): + deepgemm.install_deepgemm(INDEX, NEWEST) + assert [path.name for path in _wheels().iterdir()] == [".lock"] # no folder, no staging + + +def test_version_needs_a_url(deepgemm): + with pytest.raises(ValueError, match=r"DEEPGEMM_VERSION 2.8.0 pins .* which is not set"): + deepgemm.resolve_deepgemm_args( + argparse.Namespace(deepgemm_url=None, deepgemm_version="2.8.0") + ) + + +def test_import_check_names_what_went_wrong(deepgemm, monkeypatch, tmp_path): + subprocesses = _Subprocesses(deepgemm, monkeypatch) + subprocesses.import_check = lambda command: subprocess.CompletedProcess( + command, + 1, + stdout="", + stderr="Traceback ...\nImportError: _C.so: undefined symbol: _ZN3c10\n", + ) + with pytest.raises(RuntimeError) as error: + deepgemm.install_deepgemm(INDEX) + assert str(error.value).startswith( + f"deep_gemm {NEWEST} from {INDEX} does not import (ImportError: _C.so: undefined symbol: " + "_ZN3c10): it needs a build for this torch " + ) + assert str(error.value).endswith( + f"and glibc {GLIBC} or older. Pin one with --deepgemm-version or DEEPGEMM_VERSION." + ) + + elsewhere = tmp_path / "src/deep_gemm/__init__.py" + subprocesses.import_check = lambda command: subprocess.CompletedProcess( + command, 0, stdout=f"Failed to load legacy kernels\n{elsewhere}\n", stderr="" + ) + with pytest.raises(RuntimeError, match=re.escape(f"loads {elsewhere} instead of deep_gemm")): + deepgemm.install_deepgemm(INDEX) + + +def test_install_holds_the_cache_lock(deepgemm, monkeypatch): + """Concurrent launches sharing the cache wait while one of them installs.""" + subprocesses, locked = _Subprocesses(deepgemm, monkeypatch), [] + + def run(command, **kwargs): + if command[1:3] == ["-m", "pip"]: + with (_wheels() / ".lock").open("a") as other: + try: + fcntl.flock(other, fcntl.LOCK_EX | fcntl.LOCK_NB) + fcntl.flock(other, fcntl.LOCK_UN) + locked.append(False) + except BlockingIOError: + locked.append(True) + return subprocesses(command, **kwargs) + + _wheels().mkdir(parents=True) + monkeypatch.setattr(deepgemm.subprocess, "run", run) + deepgemm.install_deepgemm(INDEX) + assert locked == [False, True] # the dry run, then the install + + +def test_install_rechecks_after_waiting_for_the_lock(deepgemm, monkeypatch, capsys): + """Another launch installed the same version while this one waited: no second install.""" + subprocesses, flock = _Subprocesses(deepgemm, monkeypatch), fcntl.flock + folder = _wheels() / deepgemm._folder_name(NEWEST, INDEX) + + def wait_for_other_launch(file, operation): + if operation & fcntl.LOCK_NB: + (folder / "deep_gemm").mkdir(parents=True) + raise BlockingIOError + return flock(file, operation) + + monkeypatch.setattr(deepgemm.fcntl, "flock", wait_for_other_launch) + assert deepgemm.install_deepgemm(INDEX) == (folder, NEWEST) + assert subprocesses.pip() == [] + assert "waiting for another launch" in capsys.readouterr().out + + +def test_installed_versions_need_no_lock(deepgemm, monkeypatch, tmp_path): + """A cache that this user cannot write, e.g. shared or read-only, serves what it holds.""" + subprocesses = _Subprocesses(deepgemm, monkeypatch) + wheel = str(tmp_path / f"deep_gemm-{NEWEST}-{ANY}.whl") + installed = [deepgemm.install_deepgemm(INDEX), deepgemm.install_deepgemm(wheel)] + + def unwritable(directory): + raise PermissionError(13, "Permission denied", str(directory / ".lock")) + + monkeypatch.setattr(deepgemm, "_locked", unwritable) + subprocesses.calls.clear() + assert [deepgemm.install_deepgemm(INDEX), deepgemm.install_deepgemm(wheel)] == installed + assert subprocesses.pip() == [] + + +def test_install_removes_what_killed_installs_left(deepgemm, monkeypatch): + """Installs stage under the lock; staging folders that old are left by killed launches.""" + _Subprocesses(deepgemm, monkeypatch) + stale, recent = _wheels() / ".install-stale", _wheels() / ".install-recent" + for staging in (stale, recent): + (staging / "deep_gemm").mkdir(parents=True) + day_ago = time.time() - deepgemm._STALE_STAGING - 60 + os.utime(stale, (day_ago, day_ago)) + deepgemm.install_deepgemm(INDEX) + assert not stale.exists() and recent.is_dir() + + +def test_installed_folder_is_readable_by_everyone(deepgemm, monkeypatch): + _Subprocesses(deepgemm, monkeypatch) + umask = os.umask(0o077) + try: + folder, _ = deepgemm.install_deepgemm(INDEX) + finally: + os.umask(umask) + modes = { + path.relative_to(folder).as_posix(): stat.S_IMODE(path.stat().st_mode) + for path in [folder, *folder.rglob("*")] + } + assert modes == { + ".": 0o755, + "deep_gemm": 0o755, + "deep_gemm/__init__.py": 0o644, + "deep_gemm/_C.so": 0o744, + } + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({}, "home/.cache/modelopt/deepgemm"), + ({"VLLM_CACHE_ROOT": "vllm"}, "vllm/modelopt/deepgemm"), + ({"VLLM_CACHE_ROOT": "vllm", "MODELOPT_DEEPGEMM_CACHE_DIR": "mine"}, "mine"), + ], + ids=["home", "vllm-cache-root", "override"], +) +def test_cache_dir_follows_the_environment(deepgemm, monkeypatch, tmp_path, env, expected): + _Subprocesses(deepgemm, monkeypatch) + for name, value in env.items(): + monkeypatch.setenv(name, str(tmp_path / value)) + assert deepgemm.deepgemm_cache_dir() == tmp_path / expected + folder, _ = deepgemm.install_deepgemm(INDEX) + assert folder.parent == tmp_path / expected / "wheels" + + +@pytest.mark.parametrize( + ("argv", "env", "expected"), + [ + (["--deepgemm-url", INDEX, "--deepgemm-version", "1"], {"DEEPGEMM_VERSION": "2"}, "1"), + (["--deepgemm_url", INDEX], {"DEEPGEMM_URL": "/other", "DEEPGEMM_VERSION": "2"}, "2"), + ([], {"DEEPGEMM_URL": INDEX}, None), + ([], {}, None), + ], + ids=["flags", "underscores", "env", "none"], +) +def test_flags_override_env_and_the_install_is_used(deepgemm, monkeypatch, argv, env, expected): + """The installed folder comes first on the Python path of the launcher and its processes.""" + for name, value in env.items(): + monkeypatch.setenv(name, value) + monkeypatch.setenv("PYTHONPATH", "/user/path") + installs, folder = [], Path("/cache/wheels/2.8.0.post5-cp312-x86_64") + + def install(url, version): + installs.append((url, version)) + return folder, "2.8.0.post5" + + monkeypatch.setattr(deepgemm, "install_deepgemm", install) + parser = argparse.ArgumentParser() + deepgemm.add_deepgemm_args(parser) + deepgemm.resolve_deepgemm_args(parser.parse_args(argv)) + if not argv and not env: + assert installs == [] and os.environ["PYTHONPATH"] == "/user/path" + assert str(folder) not in sys.path + return + assert installs == [(INDEX, expected)] + assert sys.path[0] == str(folder) + assert os.environ["PYTHONPATH"] == f"{folder}{os.pathsep}/user/path" + deepgemm.prepend_python_path(folder) # again: no change + assert sys.path.count(str(folder)) == 1 + assert os.environ["PYTHONPATH"] == f"{folder}{os.pathsep}/user/path" + + +def test_python_path_changes_do_not_outlive_the_test(tmp_path): + """The fixture's monkeypatch undoes the launcher's Python path changes.""" + environ, path = dict(os.environ), list(sys.path) + with pytest.MonkeyPatch.context() as monkeypatch: + deepgemm = _isolated_module(monkeypatch, tmp_path) + monkeypatch.setattr(deepgemm, "install_deepgemm", lambda url, version: (tmp_path, "2.8.0")) + deepgemm.resolve_deepgemm_args(argparse.Namespace(deepgemm_url="/x", deepgemm_version=None)) + assert sys.path[0] == str(tmp_path) + assert dict(os.environ) == environ and sys.path == path + + +# The real pip installs tiny pure-Python deep_gemm wheels from a file:// index or a file. + + +def _make_wheel(directory, version, tag=ANY): + """A pure-Python ``deep_gemm`` wheel of ``version`` whose package records its version.""" + dist_info = f"deep_gemm-{version}.dist-info" + files = { + "deep_gemm/__init__.py": f"__version__ = {version!r}\n", + f"{dist_info}/METADATA": f"Metadata-Version: 2.1\nName: deep_gemm\nVersion: {version}\n", + f"{dist_info}/WHEEL": ( + f"Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: {tag}\n" + ), + } + record = [] + for name, text in files.items(): + digest = base64.urlsafe_b64encode(hashlib.sha256(text.encode()).digest()).rstrip(b"=") + record.append(f"{name},sha256={digest.decode()},{len(text.encode())}") + files[f"{dist_info}/RECORD"] = "\n".join([*record, f"{dist_info}/RECORD,,", ""]) + path = directory / f"deep_gemm-{version}-{tag}.whl" + with zipfile.ZipFile(path, "w") as wheel: + for name, text in files.items(): + wheel.writestr(name, text) + return path + + +@pytest.fixture +def real_pip(deepgemm, monkeypatch, tmp_path): + """pip without the user's configuration; a local PEP 503 index of tiny wheels, its URL.""" + for name in [name for name in os.environ if name.startswith("PIP_")]: + monkeypatch.delenv(name) + monkeypatch.setenv("PIP_CONFIG_FILE", os.devnull) + files = tmp_path / "files" + files.mkdir() + wheels = [ + _make_wheel(files, f"2.8.0.post3+gaaaaaaa.{LABEL}"), + _make_wheel(files, NEWEST), + _make_wheel(files, f"2.8.0.post7+gddddddd.{LABEL}", tag="py3-none-win_amd64"), + ] + page = tmp_path / "simple" / "deep-gemm" + page.mkdir(parents=True) + links = "".join(f'{wheel.name}
\n' for wheel in wheels) + (page / "index.html").write_text(f"\n\n{links}\n") + return (tmp_path / "simple").as_uri() + + +def test_real_pip_installs_the_newest_wheel_for_this_machine(deepgemm, real_pip): + folder, version = deepgemm.install_deepgemm(real_pip) + assert version == NEWEST # not post7, which is for Windows + assert (folder / "deep_gemm/__init__.py").read_text() == f"__version__ = {NEWEST!r}\n" + assert (folder / f"deep_gemm-{NEWEST}.dist-info/METADATA").is_file() + # The import check ran in a fresh interpreter: an install that does not import fails it. + (folder / "deep_gemm/__init__.py").write_text("raise ImportError('wrong torch')\n") + with pytest.raises(RuntimeError, match=r"does not import \(ImportError: wrong torch\)"): + deepgemm.install_deepgemm(real_pip) + + +def test_real_pip_installs_a_pinned_version_and_the_launcher_uses_it(deepgemm, real_pip, capsys): + older = f"2.8.0.post3+gaaaaaaa.{LABEL}" + deepgemm.resolve_deepgemm_args( + argparse.Namespace(deepgemm_url=real_pip, deepgemm_version="2.8.0.post3") + ) + folder = _wheels() / deepgemm._folder_name(older, real_pip) + assert sys.path[0] == str(folder) and os.environ["PYTHONPATH"].startswith(str(folder)) + result = subprocess.run( + [sys.executable, "-c", "import deep_gemm; print(deep_gemm.__version__)"], + capture_output=True, + text=True, + check=True, + cwd=Path.home().parent, + ) + assert result.stdout == f"{older}\n" + assert f"serving with deep_gemm {older} from {real_pip}: {folder}" in capsys.readouterr().out + + +def test_real_pip_ignores_an_installed_deep_gemm(deepgemm, real_pip, monkeypatch, tmp_path): + """An installed deep_gemm, e.g. in site-packages, neither satisfies nor pins the selection.""" + installed = tmp_path / "site" + with zipfile.ZipFile(_make_wheel(tmp_path, f"2.8.0.post3+gaaaaaaa.{LABEL}")) as wheel: + wheel.extractall(installed) + monkeypatch.setenv("PYTHONPATH", str(installed)) + assert deepgemm.install_deepgemm(real_pip)[1] == NEWEST + folder, version = deepgemm.install_deepgemm(real_pip, "2.8.0.post3") # the installed version + assert version == f"2.8.0.post3+gaaaaaaa.{LABEL}" + assert (folder / "deep_gemm/__init__.py").is_file() + + +def test_real_pip_installs_a_wheel_file(deepgemm, real_pip, tmp_path): + wheel = _make_wheel(tmp_path, NEWEST) + for url in (str(wheel), wheel.as_uri()): + folder, version = deepgemm.install_deepgemm(url) + assert version == NEWEST + assert (folder / "deep_gemm/__init__.py").read_text() == f"__version__ = {NEWEST!r}\n"