From 7e7ae52dd316e1a31c8db7005adfae9f993f8445 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Mon, 10 Aug 2026 07:37:03 +0000 Subject: [PATCH 1/3] python3Packages.cupy: 14.0.1 -> 14.1.1 Diff: https://github.com/cupy/cupy/compare/v14.0.1...v14.1.1 Changelog: https://github.com/cupy/cupy/releases/tag/v14.1.1 --- .../python-modules/cupy/default.nix | 178 +++++++++++------- .../cupy/link-static-libraries.patch | 30 +++ 2 files changed, 140 insertions(+), 68 deletions(-) create mode 100644 pkgs/development/python-modules/cupy/link-static-libraries.patch diff --git a/pkgs/development/python-modules/cupy/default.nix b/pkgs/development/python-modules/cupy/default.nix index 3db369191fe9..c1326581981d 100644 --- a/pkgs/development/python-modules/cupy/default.nix +++ b/pkgs/development/python-modules/cupy/default.nix @@ -1,85 +1,118 @@ { lib, buildPythonPackage, + cudaPackages, fetchFromGitHub, + fetchpatch, # build-system cython, setuptools, # nativeBuildInputs - cudaPackages, - symlinkJoin, - addDriverRunpath, + autoAddDriverRunpath, + writableTmpDirAsHomeHook, # dependencies numpy, cuda-pathfinder, - - # tests - pytest-mock, - pytestCheckHook, }: -let - shouldUsePkg = lib.mapNullable (pkg: if pkg.meta.available or true then pkg else null); - - # some packages are not available on all platforms - cuda_nvprof = shouldUsePkg (cudaPackages.nvprof or null); - libcutensor = shouldUsePkg (cudaPackages.libcutensor or null); - nccl = shouldUsePkg (cudaPackages.nccl or null); - - outpaths = lib.filter (outpath: outpath != null) ( - with cudaPackages; - [ - cccl # - cuda_cudart - cuda_nvcc # - cuda_nvprof - cuda_nvrtc - cuda_nvtx - cuda_profiler_api - libcublas - libcufft - libcurand - libcusolver - libcusparse - libcusparse_lt # cusparseLt.h - ] - ); - cudatoolkit-joined = symlinkJoin { - name = "cudatoolkit-joined-${cudaPackages.cudaMajorMinorVersion}"; - paths = - outpaths ++ lib.concatMap (outpath: lib.map (output: outpath.${output}) outpath.outputs) outpaths; - }; -in buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs: { pname = "cupy"; - version = "14.0.1"; + version = "14.1.1"; pyproject = true; + __structuredAttrs = true; src = fetchFromGitHub { owner = "cupy"; repo = "cupy"; tag = "v${finalAttrs.version}"; fetchSubmodules = true; - hash = "sha256-TaEJ0BveUCXCRrNq9L49Tfbu0334+cANcVm5qnSOE1Q="; + hash = "sha256-8jreEbQA24V7EAD87z3uFwlr3LPDoGRCvF5vbf2dKvI="; }; - postPatch = '' - substituteInPlace pyproject.toml \ - --replace-fail \ - "Cython>=3.1,<3.2" \ - "Cython" - ''; + patches = [ + # Let cupy_setup_build.py find static libraries (.a) from buildInputs. + # By default, it looks for them in `$CUDA_PATH/lib{64,}`. + ./link-static-libraries.patch + + # Fix test collection with pytest>=9 + # https://github.com/cupy/cupy/pull/10020 + (fetchpatch { + name = "pytest-9-compat.patch"; + url = "https://github.com/cupy/cupy/commit/6f52f1541e0e047cdb84286793d6110567bce5c8.patch"; + hash = "sha256-D4VjOL0XO8gnfItGC/2MjuWYCJ7eNJYFn8jIyM+EojM="; + }) + ]; + + postPatch = + # Inject absolute path to nvcc instead of relying on runtime discovery heuristics + '' + substituteInPlace cupy/_environment.py \ + --replace-fail \ + "nvcc_path = os.environ.get('NVCC', None)" \ + "return '${lib.getExe cudaPackages.cuda_nvcc}'" + '' + # Inject absolute path to CUDA libraries which get disovered at runtime + + '' + substituteInPlace cupy/_core/core.pyx \ + --replace-fail \ + "_cuda_include_dir = find_nvidia_header_directory('cudart')" \ + "_cuda_include_dir = '${lib.getInclude cudaPackages.cuda_cudart}/include'" + + substituteInPlace cupy/cuda/__init__.py \ + --replace-fail \ + "from cuda import pathfinder" \ + "from ctypes import CDLL" \ + --replace-fail \ + 'pathfinder.load_nvidia_dynamic_lib("cufft")' \ + 'CDLL("${lib.getLib cudaPackages.libcufft}/lib/libcufft.so")' + + substituteInPlace cupy/cuda/cufft.pyx \ + --replace-fail \ + "from cuda import pathfinder" \ + "from ctypes import CDLL" \ + --replace-fail \ + "loaded_dl = pathfinder.load_nvidia_dynamic_lib('cufft')" \ + "loaded_dl = CDLL('${lib.getLib cudaPackages.libcufft}/lib/libcufft.so')" \ + --replace-fail \ + "handle = loaded_dl._handle_uint" \ + "handle = loaded_dl._handle" + + substituteInPlace cupy_backends/cuda/api/_runtime_softlink.pxi \ + --replace-fail \ + "from cuda import pathfinder" \ + "from ctypes import CDLL" \ + --replace-fail \ + "loaded_dl = pathfinder.load_nvidia_dynamic_lib('cudart')" \ + "loaded_dl = CDLL('${lib.getLib cudaPackages.cuda_cudart}/lib/libcudart.so')" \ + --replace-fail \ + "handle = loaded_dl._handle_uint" \ + "handle = loaded_dl._handle" + + substituteInPlace cupy_backends/cuda/libs/_cnvrtc.pxi \ + --replace-fail \ + "from cuda import pathfinder" \ + "from ctypes import CDLL" \ + --replace-fail \ + "loaded_dl = pathfinder.load_nvidia_dynamic_lib('nvrtc')" \ + "loaded_dl = CDLL('${lib.getLib cudaPackages.cuda_nvrtc}/lib/libnvrtc.so')" \ + --replace-fail \ + "handle = loaded_dl._handle_uint" \ + "handle = loaded_dl._handle" + + substituteInPlace cupy/cuda/compiler.py \ + --replace-fail \ + "cudadevrt = get_cuda_path()" \ + "return '${lib.getOutput "static" cudaPackages.cuda_cudart}/lib/libcudadevrt.a'" + ''; env = { LDFLAGS = toString [ # Fake libcuda.so (the real one is deployed impurely) "-L${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs" ]; - # NVCC = "${lib.getExe cudaPackages.cuda_nvcc}"; # FIXME: splicing/buildPackages - CUDA_PATH = "${cudatoolkit-joined}"; }; # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Setting both @@ -92,20 +125,36 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs export CUPY_NUM_NVCC_THREADS="$NIX_BUILD_CORES" ''; + enableParallelBuilding = true; + build-system = [ cython setuptools ]; nativeBuildInputs = [ - addDriverRunpath - cudatoolkit-joined + autoAddDriverRunpath + cudaPackages.cuda_nvcc + + writableTmpDirAsHomeHook # Needed for pythonImportsCheck ]; - buildInputs = [ - cudatoolkit-joined - libcutensor - nccl + buildInputs = with cudaPackages; [ + cuda_cudart # cuda.h, cuda_runtime.h, libcudart_static.a + cuda_cuxxfilt # nv_decode.h + cuda_nvrtc # nvrtc.h + cuda_nvtx # nvToolsExt.h + cuda_profiler_api # cuda_profiler_api.h + libcublas # cublas_v2.h + libcufft # cufft.h + libcurand # curand.h + libcusolver # cusolverDn.h + libcusparse # cusparse.h + libcusparse_lt # cusparseLt.h + libcutensor # cutensor.h + nccl # nccl.h + + (lib.getOutput "static" cuda_cuxxfilt) # libcufilt.a ]; dependencies = [ @@ -113,23 +162,15 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs numpy ]; - nativeCheckInputs = [ - pytest-mock - pytestCheckHook + pythonImportsCheck = [ + "cupy" + "cupy_backends" + "cupyx" ]; - # Won't work with the GPU, whose drivers won't be accessible from the build - # sandbox + # Won't work with the GPU, whose drivers won't be accessible from the build sandbox doCheck = false; - postFixup = '' - find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do - addDriverRunpath "$lib" - done - ''; - - enableParallelBuilding = true; - meta = { description = "NumPy-compatible matrix library accelerated by CUDA"; homepage = "https://cupy.chainer.org/"; @@ -140,5 +181,6 @@ buildPythonPackage.override { stdenv = cudaPackages.backendStdenv; } (finalAttrs "x86_64-linux" ]; maintainers = with lib.maintainers; [ GaetanLepage ]; + teams = [ lib.teams.cuda ]; }; }) diff --git a/pkgs/development/python-modules/cupy/link-static-libraries.patch b/pkgs/development/python-modules/cupy/link-static-libraries.patch new file mode 100644 index 000000000000..205e2b890f0f --- /dev/null +++ b/pkgs/development/python-modules/cupy/link-static-libraries.patch @@ -0,0 +1,30 @@ +Instead of looking up static libraries by hand in `$CUDA_PATH/lib{64,}` (which does not work with +the split CUDA redistributables), link them like any other library, so that the compiler wrapper +finds them in the (nix) `buildInputs`. + +--- a/install/cupy_builder/cupy_setup_build.py ++++ b/install/cupy_builder/cupy_setup_build.py +@@ -423,10 +423,7 @@ + + s = copy.deepcopy(settings) + if not no_cuda: +- s['libraries'] = module.libraries +- s['extra_objects'] = [ +- _find_static_library(name) for name in module.static_libraries +- ] ++ s['libraries'] = module.static_libraries + module.libraries + + compile_args = s.setdefault('extra_compile_args', []) + link_args = s.setdefault('extra_link_args', []) +@@ -490,9 +487,8 @@ + if (not no_cuda and not use_hip + and PLATFORM_LINUX + and f == 'cupy.cuda.function'): +- s_file['extra_objects'] = ( +- s_file.get('extra_objects', []) +- + [_find_static_library('cufilt')]) ++ s_file['libraries'] = ( ++ ['cufilt'] + s_file.get('libraries', [])) + + sources = module_extension_sources(f, use_cython, no_cuda) + extension = setuptools.Extension(name, sources, **s_file) From 13615fe56f775625fcadfe0c8170547d0029bda0 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Mon, 10 Aug 2026 21:56:50 +0000 Subject: [PATCH 2/3] python3Packages.array-api-compat: set CUPY_CACHE_DIR to fix tests --- pkgs/development/python-modules/array-api-compat/default.nix | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pkgs/development/python-modules/array-api-compat/default.nix b/pkgs/development/python-modules/array-api-compat/default.nix index f976a35514ae..383dd9cc332b 100644 --- a/pkgs/development/python-modules/array-api-compat/default.nix +++ b/pkgs/development/python-modules/array-api-compat/default.nix @@ -52,6 +52,11 @@ buildPythonPackage (finalAttrs: { pythonImportsCheck = [ "array_api_compat" ]; + # Otherwise, cupy will try to write in $HOME (/homeless-shelter) + preCheck = '' + export CUPY_CACHE_DIR=$(mktemp -d) + ''; + disabledTests = [ # CUDA (used via cupy) is not available in the testing sandbox "cupy" From 8b1d9199ece32a862c3505c3ec9748c3827d4234 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Mon, 10 Aug 2026 22:07:20 +0000 Subject: [PATCH 3/3] python3Packages.deepspeed: 0.19.2 -> 0.19.5 Diff: https://github.com/deepspeedai/DeepSpeed/compare/v0.19.2...v0.19.5 Changelog: https://github.com/deepspeedai/DeepSpeed/releases/tag/v0.19.5 --- .../python-modules/deepspeed/default.nix | 171 +++++++++++++----- 1 file changed, 129 insertions(+), 42 deletions(-) diff --git a/pkgs/development/python-modules/deepspeed/default.nix b/pkgs/development/python-modules/deepspeed/default.nix index 33a7836d4ff4..d5611c3f13a4 100644 --- a/pkgs/development/python-modules/deepspeed/default.nix +++ b/pkgs/development/python-modules/deepspeed/default.nix @@ -2,45 +2,56 @@ lib, buildPythonPackage, fetchFromGitHub, + symlinkJoin, + cudaPackages, + + # build-system setuptools, + torch, + + # dependencies einops, hjson, msgpack, ninja, numpy, - cupy, - cutlass, packaging, psutil, py-cpuinfo, pydantic, - torch, tqdm, - nix-update-script, - cudaPackages, - symlinkJoin, + # cuda-only: + nvidia-ml-py, + # cupy, + # cutlass, + + # tests + accelerate, + addBinToPathHook, + pytest-xdist, + pytestCheckHook, + tabulate, + transformers, + writableTmpDirAsHomeHook, }: let - cudaVersion = cudaPackages.cudaMajorMinorVersion; - inherit (torch) cudaCapabilities cudaSupport; cuda-native-redist = symlinkJoin { - name = "cuda-native-redist-${cudaVersion}"; + name = "cuda-native-redist"; paths = with cudaPackages; [ (lib.getDev cuda_cudart) (lib.getLib cuda_cudart) (lib.getStatic cuda_cudart) - cuda_nvcc + (lib.getBin cuda_nvcc) ]; }; - in -buildPythonPackage (finalAttrs: { +buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { pname = "deepspeed"; - version = "0.19.2"; + version = "0.19.5"; pyproject = true; __structuredAttrs = true; @@ -48,14 +59,49 @@ buildPythonPackage (finalAttrs: { owner = "deepspeedai"; repo = "DeepSpeed"; tag = "v${finalAttrs.version}"; - hash = "sha256-Nw1rw65hdqhARFR7W+XmmRT/pLkCi5nTF+6R9L3bLyo="; fetchSubmodules = true; + hash = "sha256-OIy8AehiwKhbdWrWslwc3yimAH4QYobhgg/inkHDx/U="; }; + postPatch = '' + substituteInPlace deepspeed/ops/op_builder/builder.py \ + --replace-fail \ + 'import distutils' \ + 'import setuptools._distutils' + '' + # `setup.py` appends `+{git_hash}` (here, `+unknown`) to the version unless `build.txt` exists, + # in which case its content is appended instead. + + '' + touch build.txt + '' + + lib.optionalString cudaSupport ( + # Hardcode CUDA_HOME to nix store path for JIT op compilation at runtime + '' + substituteInPlace deepspeed/ops/op_builder/builder.py \ + --replace-fail \ + "cuda_home = torch.utils.cpp_extension.CUDA_HOME" \ + "cuda_home = '${cuda-native-redist}'" + '' + # Hardcode CUTLASS_PATH to nix store path + + '' + substituteInPlace deepspeed/ops/op_builder/evoformer_attn.py \ + --replace-fail \ + 'self.cutlass_path = os.environ.get("CUTLASS_PATH")' \ + 'self.cutlass_path = "${lib.getInclude cudaPackages.cutlass}/include/cutlass"' + '' + ); + build-system = [ setuptools + torch ]; + # `deepspeed.ops.transformer.inference.triton` creates its autotune cache directory + # (`$TRITON_HOME`, defaulting to `$HOME/.triton`) at import time. + preBuild = '' + export TRITON_HOME=$(mktemp -d) + ''; + dependencies = [ einops hjson @@ -66,53 +112,94 @@ buildPythonPackage (finalAttrs: { psutil py-cpuinfo pydantic - setuptools torch tqdm ] ++ lib.optionals cudaSupport [ - cutlass - cupy + nvidia-ml-py ]; - postPatch = '' - substituteInPlace deepspeed/ops/op_builder/builder.py \ - --replace-fail 'import distutils' 'import setuptools._distutils' - '' - + lib.optionalString cudaSupport '' - # Hardcode CUDA_HOME to nix store path for JIT op compilation at runtime - substituteInPlace deepspeed/ops/op_builder/builder.py \ - --replace-fail \ - "cuda_home = torch.utils.cpp_extension.CUDA_HOME" \ - "cuda_home = '${cuda-native-redist}'" - - # Hardcode CUTLASS_PATH to nix store path - substituteInPlace deepspeed/ops/op_builder/evoformer_attn.py \ - --replace-fail \ - "self.cutlass_path = os.environ.get(\"CUTLASS_PATH\")" \ - "self.cutlass_path = '${cutlass}'" - ''; - env = lib.optionalAttrs cudaSupport { TORCH_CUDA_ARCH_LIST = lib.concatStringsSep ";" cudaCapabilities; }; - preConfigure = '' - # setuptools writes to ~/.cache during builds - export HOME=$TMPDIR + preCheck = '' + rm -rf deepspeed + '' + # Tests JIT-compile CPU ops, and the op builder pairs the ISA it detects (`-D__AVX512__`) with + # `-march=native`, which the cc wrapper strips. Let it through, otherwise the intrinsics fail to + # compile: "inlining failed in call to 'always_inline'". + + '' + export NIX_ENFORCE_NO_NATIVE=0 ''; - pythonImportsCheck = [ - "deepspeed" + pythonImportsCheck = [ "deepspeed" ]; + + nativeCheckInputs = [ + accelerate + addBinToPathHook # Tests run the deepspeed CLI + ninja + pytest-xdist + pytestCheckHook + tabulate + transformers + writableTmpDirAsHomeHook ]; - passthru.updateScript = nix-update-script { }; + pytestFlags = [ + # Tests memory consumption grows significantly with the number of parallel processes + # -> Limit the number of parallel jobs to prevent OOMing + "--maxprocesses=16" + ]; + + enabledTestPaths = [ + # The test suite seems to be very resource intensive. We stick to unit tests for now. + "tests/unit/" + ]; + + disabledTestPaths = [ + # Require unpackaged `mup` + "tests/unit/runtime/test_mup_optimizers.py" + ]; + + disabledTests = [ + # AssertionError: assert 2 == 1 + # + where 2 = len(['[pinned-memory checkpoint] crossed 32.000 GB: 33.000 GB pinned across 1 ... + "test_checkpoint_emits_info" + + # Fixture handling is incompatible with pytest >= 8.4, upstream pins pytest < 8.4: + # https://github.com/deepspeedai/DeepSpeed/pull/7327 + "TestDistributedFixture" + "TestZeROElasticCheckpoint" + "TestZeROPPLoadCheckpoint" + "TestZeROUniversalCheckpointDP" + + # Fail since we patch the cutlass/cuda path logic + "test_include_paths_accepts_cutlass_include_dir_directly" + "test_include_paths_finds_cutlass_from_cmake_prefix_path" + "test_include_paths_finds_cutlass_from_compiler_include_path" + "test_include_paths_finds_python_package_candidate_without_env" + "test_include_paths_reports_missing_cutlass" + "test_include_paths_uses_cutlass_path_env" + "test_installed_cuda_version_reports_a_runtime_only_cuda_home_as_missing" + + # Too long + "TestWarmupCosineLR" + ]; + + # Some python tests are skipped unless a GPU is visible + passthru.gpuCheck = finalAttrs.finalPackage.overrideAttrs { + requiredSystemFeatures = [ "cuda" ]; + }; meta = { - description = "Deep learning optimization library that makes distributed training and inference easy, efficient, and effective."; + description = "Deep learning optimization library that makes distributed training and inference easy, efficient, and effective"; homepage = "https://www.deepspeed.ai/"; + downloadPage = "https://github.com/deepspeedai/deepspeed"; changelog = "https://github.com/deepspeedai/DeepSpeed/releases/tag/${finalAttrs.src.tag}"; license = lib.licenses.asl20; + mainProgram = "deepspeed"; maintainers = with lib.maintainers; [ jlesquembre ]; + teams = [ lib.teams.cuda ]; }; })