diff --git a/nixos/doc/manual/release-notes/rl-2505.section.md b/nixos/doc/manual/release-notes/rl-2505.section.md index 5079ca741469..85281beb0e83 100644 --- a/nixos/doc/manual/release-notes/rl-2505.section.md +++ b/nixos/doc/manual/release-notes/rl-2505.section.md @@ -258,6 +258,11 @@ - `siduck76-st` has been renamed to `st-snazzy`, like the project's [flake](https://github.com/siduck/st/blob/main/flake.nix). +- `python3Packages.jax` now directly depends on `python3Packages.jaxlib`. + As a result, packages that depend on jax no longer need to include jaxlib to their dependencies. + There is also a breaking change in the handling of CUDA. Instead of using a CUDA compatible jaxlib + as before, you can use plugins like `python3Packages.jax-cuda12-plugin`. + diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 5b7fb2d88a93..dd82358f03b7 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -6,46 +6,24 @@ { absl-py, - autoAddDriverRunpath, autoPatchelfHook, buildPythonPackage, - config, fetchPypi, - fetchurl, flatbuffers, - jaxlib-build, lib, ml-dtypes, python, scipy, stdenv, - # Options: - cudaSupport ? config.cudaSupport, - cudaPackages, }: let - inherit (cudaPackages) cudaVersion; - - version = "0.4.28"; - + version = "0.4.38"; inherit (python) pythonVersion; - cudaLibPath = lib.makeLibraryPath ( - with cudaPackages; - [ - (lib.getLib cuda_cudart) # libcudart.so - (lib.getLib cuda_cupti) # libcupti.so - (lib.getLib cudnn) # libcudnn.so - (lib.getLib libcufft) # libcufft.so - (lib.getLib libcusolver) # libcusolver.so - (lib.getLib libcusparse) # libcusparse.so - ] - ); - # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the # official instructions recommend installing CPU-only versions via PyPI. - cpuSrcs = + srcs = let getSrcFromPypi = { @@ -68,182 +46,129 @@ let }; in { - "3.9-x86_64-linux" = getSrcFromPypi { - platform = "manylinux2014_x86_64"; - dist = "cp39"; - hash = "sha256-Slbr8FtKTBeRaZ2HTgcvP4CPCYa0AQsU+1SaackMqdw="; - }; - "3.9-aarch64-darwin" = getSrcFromPypi { - platform = "macosx_11_0_arm64"; - dist = "cp39"; - hash = "sha256-sBVi7IrXVxm30DiXUkiel+trTctMjBE75JFjTVKCrTw="; - }; - "3.9-x86_64-darwin" = getSrcFromPypi { - platform = "macosx_10_14_x86_64"; - dist = "cp39"; - hash = "sha256-T5jMg3srbG3P4Kt/+esQkxSSCUYRmqOvn6oTlxj/J4c="; - }; - "3.10-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp310"; - hash = "sha256-47zcb45g+FVPQVwU2TATTmAuPKM8OOVGJ0/VRfh1dps="; + hash = "sha256-Ya7MuaJ8Z/24RQ9jVyQAGc1FEcudYqROR2R1bThIU60="; + }; + "3.10-aarch64-linux" = getSrcFromPypi { + platform = "manylinux2014_aarch64"; + dist = "cp310"; + hash = "sha256-7hnBY6j98IOdTBi4il+/tOcxunxDdBbT5Ug+Vwu3ZOQ="; }; "3.10-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp310"; - hash = "sha256-8Djmi9ENGjVUcisLvjbmpEg4RDenWqnSg/aW8O2fjAk="; + hash = "sha256-MLL1LLUNdHNK8vR3wlM6elg+O7eyyKzes2Hud9lAV3o="; }; "3.10-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp310"; - hash = "sha256-pCHSN/jCXShQFm0zRgPGc925tsJvUrxJZwS4eCKXvWY="; + hash = "sha256-VcGbnT8zpvxZ9kSqWiH7oCY5zN13bLSptVJmJfV4Of8="; }; "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp311"; - hash = "sha256-Rc4PPIQM/4I2z/JsN/Jsn/B4aV+T4MFiwyDCgfUEEnU="; + hash = "sha256-J1H/cDfWqZfQvg53zEvjgcWp+buLMU7bdVwTpv2Wn0U="; + }; + "3.11-aarch64-linux" = getSrcFromPypi { + platform = "manylinux2014_aarch64"; + dist = "cp311"; + hash = "sha256-Q9tYxMQnYnKWNmpWwQMY4fAPUDaQ4X+Uu0NEKT4ZleA="; }; "3.11-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp311"; - hash = "sha256-eThX+vN/Nxyv51L+pfyBH0NeQ7j7S1AgWERKf17M+Ck="; + hash = "sha256-P7DqrnNpFXr+y+rVCq8p5z/936d6IzXXIb2XlPPFEOQ="; }; "3.11-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp311"; - hash = "sha256-L/gpDtx7ksfq5SUX9lSSYz4mey6QZ7rT5MMj0hPnfPU="; + hash = "sha256-tn/eq9bf7Qi3do873/tSEWAIX4MFZpvRl77vYdCN4Is="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp312"; - hash = "sha256-RqGqhX9P7uikP8upXA4Kti1AwmzJcwtsaWVZCLo1n40="; + hash = "sha256-2tbAqWVnwG0IPARp/sQPIBIQsJk2W9aYvjGm0uyI/Vk="; + }; + "3.12-aarch64-linux" = getSrcFromPypi { + platform = "manylinux2014_aarch64"; + dist = "cp312"; + hash = "sha256-SW9FsOABojQTCc0MdK8LZwU33O15wWjLIwz8x3PwqoY"; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; - hash = "sha256-jdi//jhTcC9jzZJNoO4lc0pNGc1ckmvgM9dyun0cF10="; + hash = "sha256-8zvK/jLJelYuz2iU18QWdMgMCs3t+lQj1Jr1EUcUmHQ="; }; "3.12-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp312"; - hash = "sha256-1sCaVFMpciRhrwVuc1FG0sjHTCKsdCaoRetp8ya096A="; + hash = "sha256-P+/qmF8EFYFvO7r9PwOkNwUCde+brJpywTFOFkSsV8E="; + }; + + "3.13-x86_64-linux" = getSrcFromPypi { + platform = "manylinux2014_x86_64"; + dist = "cp313"; + hash = "sha256-LOd7qM2pJZpLypevwcci5CkabEY6Y/jTcsbtyFEX1iU="; + }; + "3.13-aarch64-linux" = getSrcFromPypi { + platform = "manylinux2014_aarch64"; + dist = "cp313"; + hash = "sha256-JIzKN3Hr8ksHD0lwE2TOraM+YTlEWwbHgsylrFrZK/Q="; + }; + "3.13-aarch64-darwin" = getSrcFromPypi { + platform = "macosx_11_0_arm64"; + dist = "cp313"; + hash = "sha256-b+MmuK82Y4fdR8zzElg7Kxf+0ScSybdKZIsYoTy9ur8="; + }; + "3.13-x86_64-darwin" = getSrcFromPypi { + platform = "macosx_10_14_x86_64"; + dist = "cp313"; + hash = "sha256-QeVa5YGKiC5XiehI9vFmh6wTK8+7Wl+hFKXRi3jQXy0="; }; }; - - # Note that the prebuilt jaxlib binary requires specific version of CUDA to - # work. The cuda12 jaxlib binaries only works with CUDA 12.2, and cuda11 - # jaxlib binaries only works with CUDA 11.8. This is why we need to find a - # binary that matches the provided cudaVersion. - gpuSrcVersionString = "cuda${cudaVersion}-${pythonVersion}"; - - # Find new releases at https://storage.googleapis.com/jax-releases - # When upgrading, you can get these hashes from prefetch.sh. See - # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index. - gpuSrcs = { - "cuda12.2-3.9" = fetchurl { - url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl"; - hash = "sha256-d8LIl22gIvmWfoyKfXKElZJXicPQIZxdS4HumhwQGCw="; - }; - "cuda12.2-3.10" = fetchurl { - url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; - hash = "sha256-PXtWv+UEcMWF8LhWe6Z1UGkf14PG3dkJ0Iop0LiimnQ="; - }; - "cuda12.2-3.11" = fetchurl { - url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl"; - hash = "sha256-QO2WSOzmJ48VaCha596mELiOfPsAGLpGctmdzcCHE/o="; - }; - "cuda12.2-3.12" = fetchurl { - url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl"; - hash = "sha256-ixWMaIChy4Ammsn23/3cCoala0lFibuUxyUr3tjfFKU="; - }; - }; in buildPythonPackage { pname = "jaxlib"; inherit version; format = "wheel"; - disabled = - !( - pythonVersion == "3.9" - || pythonVersion == "3.10" - || pythonVersion == "3.11" - || pythonVersion == "3.12" - ); - # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6. - src = - if !cudaSupport then - (cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}" - or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}") - ) - else - gpuSrcs."${gpuSrcVersionString}" or (throw "jaxlib-bin: No gpuSrc for ${gpuSrcVersionString}"); + src = ( + srcs."${pythonVersion}-${stdenv.hostPlatform.system}" + or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}") + ); # Prebuilt wheels are dynamically linked against things that nix can't find. # Run `autoPatchelfHook` to automagically fix them. - nativeBuildInputs = - lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ] - ++ lib.optionals cudaSupport [ autoAddDriverRunpath ]; + nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ]; # Dynamic link dependencies buildInputs = [ (lib.getLib stdenv.cc.cc) ]; - # jaxlib contains shared libraries that open other shared libraries via dlopen - # and these implicit dependencies are not recognized by ldd or - # autoPatchelfHook. That means we need to sneak them into rpath. This step - # must be done after autoPatchelfHook and the automatic stripping of - # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the - # patchPhase. - preInstallCheck = lib.optional cudaSupport '' - shopt -s globstar - - for file in $out/**/*.so; do - patchelf --add-rpath "${cudaLibPath}" "$file" - done - ''; - - propagatedBuildInputs = [ + dependencies = [ absl-py flatbuffers ml-dtypes scipy ]; - # jaxlib looks for ptxas at runtime, eg when running `jax.random.PRNGKey(0)`. - # Linking into $out is the least bad solution. See - # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 - # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211 - # for more info. - postInstall = lib.optional cudaSupport '' - mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin - ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas - ''; + pythonImportsCheck = [ "jaxlib" ]; - inherit (jaxlib-build) pythonImportsCheck; - - meta = with lib; { + meta = { description = "Prebuilt jaxlib backend from PyPi"; homepage = "https://github.com/google/jax"; - sourceProvenance = with sourceTypes; [ binaryNativeCode ]; - license = licenses.asl20; - maintainers = with maintainers; [ samuela ]; - platforms = [ - "aarch64-darwin" - "x86_64-linux" - "x86_64-darwin" - ]; - broken = - !(cudaSupport -> lib.versionAtLeast cudaVersion "11.1") - || !(cudaSupport -> lib.versionAtLeast cudaPackages.cudnn.version "8.2") - || !(cudaSupport -> stdenv.hostPlatform.isLinux) - || !(cudaSupport -> (gpuSrcs ? "cuda${cudaVersion}-${pythonVersion}")) + sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ]; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ samuela ]; + badPlatforms = [ # Fails at pythonImportsCheckPhase: # ...-python-imports-check-hook.sh/nix-support/setup-hook: line 10: 28017 Illegal instruction: 4 # /nix/store/5qpssbvkzfh73xih07xgmpkj5r565975-python3-3.11.9/bin/python3.11 -c # 'import os; import importlib; list(map(lambda mod: importlib.import_module(mod), os.environ["pythonImportsCheck"].split()))' - || (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64); + "x86_64-darwin" + ]; }; } diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 09b548c199a2..e30c9620859c 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -6562,9 +6562,7 @@ self: super: with self; { jax-jumpy = callPackage ../development/python-modules/jax-jumpy { }; - jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix { - inherit (pkgs.config) cudaSupport; - }; + jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix { }; jaxlib-build = callPackage ../development/python-modules/jaxlib rec { # Some platforms don't have `cudaSupport` defined, hence the need for 'or false'. @@ -6574,14 +6572,6 @@ self: super: with self; { jaxlib = jaxlib-bin; - jaxlibWithCuda = self.jaxlib.override { - cudaSupport = true; - }; - - jaxlibWithoutCuda = self.jaxlib.override { - cudaSupport = false; - }; - jaxopt = callPackage ../development/python-modules/jaxopt { }; jaxtyping = callPackage ../development/python-modules/jaxtyping { };