|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | + |
| 17 | +"""Nox session definitions for testing, linting, docs, and wheel builds. |
| 18 | +
|
| 19 | +Usage: |
| 20 | + python -m pip install nox uv # install nox and uv (once) |
| 21 | + nox -l # list all sessions |
| 22 | + nox -s gpu_megatron # run a GPU session (inside container) |
| 23 | + nox -s "unit-3.12(torch_211-tf_latest)" # run a specific unit test combination |
| 24 | + nox -s "unit-3.12(torch_211-tf_latest)" --no-venv # skip venv creation (reuse current env) |
| 25 | +""" |
| 26 | + |
| 27 | +import glob |
| 28 | +import os |
| 29 | +import shutil |
| 30 | + |
| 31 | +import nox |
| 32 | + |
| 33 | +nox.options.default_venv_backend = "uv" |
| 34 | +nox.options.envdir = "/tmp/.nox" |
| 35 | + |
| 36 | +TORCH_VERSIONS = { |
| 37 | + "torch_28": "torchvision~=0.23.0", |
| 38 | + "torch_29": "torchvision~=0.24.0", |
| 39 | + "torch_210": "torchvision~=0.25.0", |
| 40 | + "torch_211": "torchvision~=0.26.0", |
| 41 | +} |
| 42 | + |
| 43 | +TRANSFORMERS_VERSIONS = { |
| 44 | + "tf_latest": None, |
| 45 | + "tf_min": "transformers~=4.56.0", |
| 46 | +} |
| 47 | + |
| 48 | + |
| 49 | +def _cov_args(): |
| 50 | + """Return --cov when COVERAGE_PROCESS_START is set (CI only).""" |
| 51 | + return ["--cov"] if os.environ.get("COVERAGE_PROCESS_START") else [] |
| 52 | + |
| 53 | + |
| 54 | +# ─── CPU unit tests ─────────────────────────────────────────────────────────── |
| 55 | +@nox.session(python=["3.10", "3.11", "3.12", "3.13"]) |
| 56 | +@nox.parametrize("tf_ver", [nox.param(k, id=k) for k in TRANSFORMERS_VERSIONS]) |
| 57 | +@nox.parametrize("torch_ver", [nox.param(k, id=k) for k in TORCH_VERSIONS]) |
| 58 | +def unit(session, torch_ver, tf_ver): |
| 59 | + """Unit tests — parametrized over torch and transformers versions.""" |
| 60 | + session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]") |
| 61 | + tf_pin = TRANSFORMERS_VERSIONS[tf_ver] |
| 62 | + if tf_pin: |
| 63 | + session.install(tf_pin) |
| 64 | + session.run("python", "-m", "pytest", "tests/unit", *_cov_args()) |
| 65 | + |
| 66 | + |
| 67 | +@nox.session(python="3.12") |
| 68 | +@nox.parametrize("subset", ["onnx", "torch", "torch_deploy"]) |
| 69 | +def partial_unit(session, subset): |
| 70 | + """Unit tests with partial installs.""" |
| 71 | + if subset == "onnx": |
| 72 | + session.install("torchvision~=0.26.0", ".[onnx,dev-test]") |
| 73 | + session.run("python", "-m", "pytest", "tests/unit/onnx") |
| 74 | + elif subset == "torch": |
| 75 | + session.install("megatron-core", ".[dev-test]") |
| 76 | + session.run( |
| 77 | + "python", |
| 78 | + "-m", |
| 79 | + "pytest", |
| 80 | + "tests/unit/torch", |
| 81 | + "--ignore=tests/unit/torch/deploy", |
| 82 | + "--ignore=tests/unit/torch/puzzletron", |
| 83 | + ) |
| 84 | + else: # torch_deploy |
| 85 | + session.install(".[onnx,dev-test]") |
| 86 | + session.run("python", "-m", "pytest", "tests/unit/torch/deploy") |
| 87 | + |
| 88 | + |
| 89 | +# ─── GPU sessions (run inside containers — no new venv) ────────────────────── |
| 90 | +# `venv_backend="none"` skips creating a new venv so the session runs directly in the container's |
| 91 | +# existing Python environment (e.g. /opt/venv in NeMo) instead of an isolated one. |
| 92 | +# Use `python -m pip/pytest` to ensure the container's active venv Python is used, |
| 93 | +# not a stale PATH entry (e.g. NeMo container has pip → /usr/local/bin/pip but python → /opt/venv/bin/python). |
| 94 | +# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later |
| 95 | +@nox.session(venv_backend="none") |
| 96 | +def gpu(session): |
| 97 | + # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers |
| 98 | + # until https://github.com/tbenthompson/cppimport/pull/98 |
| 99 | + session.run( |
| 100 | + "python", |
| 101 | + "-m", |
| 102 | + "pip", |
| 103 | + "install", |
| 104 | + "--no-build-isolation", |
| 105 | + "git+https://github.com/Dao-AILab/fast-hadamard-transform.git", |
| 106 | + ) |
| 107 | + session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test]") |
| 108 | + session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x") |
| 109 | + session.run("python", "-m", "pip", "install", "cupy-cuda13x") |
| 110 | + session.run( |
| 111 | + "python", |
| 112 | + "-m", |
| 113 | + "pip", |
| 114 | + "install", |
| 115 | + "--no-build-isolation", |
| 116 | + "git+https://github.com/state-spaces/mamba.git", |
| 117 | + "git+https://github.com/Dao-AILab/causal-conv1d.git", |
| 118 | + ) |
| 119 | + session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) |
| 120 | + |
| 121 | + |
| 122 | +# Container: nvcr.io/nvidia/nemo:26.02 or later |
| 123 | +@nox.session(venv_backend="none") |
| 124 | +def gpu_megatron(session): |
| 125 | + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") |
| 126 | + session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) |
| 127 | + |
| 128 | + |
| 129 | +# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later |
| 130 | +@nox.session(venv_backend="none") |
| 131 | +def gpu_regression(session): |
| 132 | + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") |
| 133 | + session.run("python", "-m", "pytest", "tests/gpu_regression", *_cov_args()) |
| 134 | + |
| 135 | + |
| 136 | +# Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 or later |
| 137 | +@nox.session(venv_backend="none") |
| 138 | +def gpu_trtllm(session): |
| 139 | + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") |
| 140 | + session.run("python", "-m", "pytest", "tests/gpu_trtllm", *_cov_args()) |
| 141 | + |
| 142 | + |
| 143 | +# ─── Code quality ───────────────────────────────────────────────────────────── |
| 144 | +@nox.session |
| 145 | +def pre_commit_all(session): |
| 146 | + session.install("-e", ".[all,dev-lint]") |
| 147 | + session.run("pre-commit", "run", "--all-files", "--show-diff-on-failure") |
| 148 | + |
| 149 | + |
| 150 | +@nox.session |
| 151 | +def pre_commit_diff(session): |
| 152 | + session.install("-e", ".[all,dev-lint]") |
| 153 | + session.run("pre-commit", "run", "--from-ref", "origin/main", "--to-ref", "HEAD") |
| 154 | + |
| 155 | + |
| 156 | +# ─── Docs ───────────────────────────────────────────────────────────────────── |
| 157 | +@nox.session |
| 158 | +def docs(session): |
| 159 | + session.install("-e", ".[all,dev-docs]") |
| 160 | + shutil.rmtree("docs/build", ignore_errors=True) |
| 161 | + shutil.rmtree("docs/source/reference/generated", ignore_errors=True) |
| 162 | + with session.chdir("docs"): |
| 163 | + session.run( |
| 164 | + "sphinx-build", |
| 165 | + "source", |
| 166 | + "build/html", |
| 167 | + "--fail-on-warning", |
| 168 | + "--show-traceback", |
| 169 | + "--keep-going", |
| 170 | + ) |
| 171 | + |
| 172 | + |
| 173 | +@nox.session |
| 174 | +def docs_debug(session): |
| 175 | + session.install("-e", ".[all,dev-docs]") |
| 176 | + shutil.rmtree("docs/build", ignore_errors=True) |
| 177 | + shutil.rmtree("docs/source/reference/generated", ignore_errors=True) |
| 178 | + with session.chdir("docs"): |
| 179 | + session.run("sphinx-autobuild", "source", "build/html", "--host", "0.0.0.0") |
| 180 | + |
| 181 | + |
| 182 | +# ─── Wheel build ────────────────────────────────────────────────────────────── |
| 183 | +@nox.session |
| 184 | +def build_wheel(session): |
| 185 | + shutil.rmtree("build", ignore_errors=True) |
| 186 | + session.install("twine") |
| 187 | + session.run("pip", "wheel", "--no-deps", "--wheel-dir=dist", ".") |
| 188 | + wheels = glob.glob("dist/*.whl") |
| 189 | + session.run("twine", "check", *wheels) |
| 190 | + (modelopt_wheel,) = glob.glob("dist/nvidia_modelopt-*.whl") |
| 191 | + session.install(modelopt_wheel, "-f", "dist") |
| 192 | + with session.chdir("dist"): |
| 193 | + session.run("python", "-c", "import modelopt; print(modelopt.__version__)") |
0 commit comments