From 1f8686373abe21bf6b1ce972f0ea55405b449329 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sun, 22 Aug 2021 20:37:42 +0000 Subject: [PATCH 1/7] python3Packages.jaxlib: init at 0.1.70 --- .../python-modules/jaxlib/default.nix | 45 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 47 insertions(+) create mode 100644 pkgs/development/python-modules/jaxlib/default.nix diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix new file mode 100644 index 000000000000..240c5a7d6d0e --- /dev/null +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -0,0 +1,45 @@ +# For the moment we only support the CPU backend of jaxlib. GPU and TPU backends +# require some additional work. Their wheels are not located on PyPI. +# * CPU/GPU: https://storage.googleapis.com/jax-releases/jax_releases.html +# * TPU: https://storage.googleapis.com/jax-releases/libtpu_releases.html + +{ autoPatchelfHook, buildPythonPackage, fetchPypi, isPy39, lib, stdenv +# propagatedBuildInputs +, absl-py, flatbuffers, scipy +}: + +buildPythonPackage rec { + pname = "jaxlib"; + version = "0.1.70"; + format = "wheel"; + + # At the time of writing (8/19/21), there are releases for 3.7-3.9. Supporting + # all of them is a pain, so we focus on 3.9, the current nixpkgs python3 + # version. + disabled = !isPy39; + + src = fetchPypi { + inherit pname version format; + dist = "cp39"; + python = "cp39"; + platform = "manylinux2010_x86_64"; + sha256 = "sha256-mytMTqoavpuRawj52MU5/iFj27SGlm8DaoQ5vd/3bss="; + }; + + # Prebuilt wheels are dynamically linked against things that nix can't find. + # Run `autoPatchelfHook` to automagically fix them. + nativeBuildInputs = [ autoPatchelfHook ]; + # Dynamic link dependencies + buildInputs = [ stdenv.cc.cc ]; + + # pip dependencies + propagatedBuildInputs = [ absl-py flatbuffers scipy ]; + + meta = with lib; { + description = "XLA library for JAX"; + homepage = "https://github.com/google/jax"; + license = licenses.asl20; + maintainers = with maintainers; [ samuela ]; + platforms = [ "x86_64-linux" ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 428913333280..fa17a8a7e7ee 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -3716,6 +3716,8 @@ in { javaproperties = callPackage ../development/python-modules/javaproperties { }; + jaxlib = callPackage ../development/python-modules/jaxlib { }; + JayDeBeApi = callPackage ../development/python-modules/JayDeBeApi { }; jc = callPackage ../development/python-modules/jc { }; From 426569a04174720817242de5e8d9157582595055 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sun, 22 Aug 2021 20:39:04 +0000 Subject: [PATCH 2/7] python3Packages.jax: init at 0.2.19 --- .../python-modules/jax/default.nix | 39 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 41 insertions(+) create mode 100644 pkgs/development/python-modules/jax/default.nix diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix new file mode 100644 index 000000000000..edf18bbe2ed0 --- /dev/null +++ b/pkgs/development/python-modules/jax/default.nix @@ -0,0 +1,39 @@ +{ buildPythonPackage, fetchFromGitHub, lib +# propagatedBuildInputs +, absl-py, numpy, opt-einsum +# checkInputs +, jaxlib, pytestCheckHook +}: + +buildPythonPackage rec { + pname = "jax"; + version = "0.2.19"; + + # Fetching from pypi doesn't allow us to run the test suite. See https://discourse.nixos.org/t/pythonremovetestsdir-hook-being-run-before-checkphase/14612/3. + src = fetchFromGitHub { + owner = "google"; + repo = pname; + rev = "jax-v${version}"; + sha256 = "sha256-pVn62G7pydR7ybkf7gSbu0FlEq2c0US6H2GTBAljup4="; + }; + + # jaxlib is _not_ included in propagatedBuildInputs because there are + # different versions of jaxlib depending on the desired target hardware. The + # JAX project ships separate wheels for CPU, GPU, and TPU. Currently only the + # CPU wheel is packaged. + propagatedBuildInputs = [ absl-py numpy opt-einsum ]; + + checkInputs = [ jaxlib pytestCheckHook ]; + # NOTE: Don't run the tests in the expiremental directory as they require flax + # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2. + # Not a big deal, this is how the JAX docs suggest running the test suite + # anyhow. + pytestFlagsArray = [ "-W ignore::DeprecationWarning" "tests/" ]; + + meta = with lib; { + description = "Differentiate, compile, and transform Numpy code."; + homepage = "https://github.com/google/jax"; + license = licenses.asl20; + maintainers = with maintainers; [ samuela ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index fa17a8a7e7ee..c844eea23dce 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -3716,6 +3716,8 @@ in { javaproperties = callPackage ../development/python-modules/javaproperties { }; + jax = callPackage ../development/python-modules/jax { }; + jaxlib = callPackage ../development/python-modules/jaxlib { }; JayDeBeApi = callPackage ../development/python-modules/JayDeBeApi { }; From a16117f1c042d4772d750077fa3c21421cf4a199 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sun, 22 Aug 2021 20:56:30 +0000 Subject: [PATCH 3/7] maintainer-list.nix: add samuela --- maintainers/maintainer-list.nix | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/maintainers/maintainer-list.nix b/maintainers/maintainer-list.nix index 37711fe0fa1e..9192c390d27b 100644 --- a/maintainers/maintainer-list.nix +++ b/maintainers/maintainer-list.nix @@ -9542,6 +9542,12 @@ githubId = 115821; name = "Sam Rose"; }; + samuela = { + email = "skainsworth@gmail.com"; + github = "samuela"; + githubId = 226872; + name = "Samuel Ainsworth"; + }; samueldr = { email = "samuel@dionne-riel.com"; github = "samueldr"; From 6f44416cf2655b87378d3752eda00dfccb39dea6 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Wed, 1 Sep 2021 21:04:02 +0000 Subject: [PATCH 4/7] python3Packages.jax: remove meta.description period --- pkgs/development/python-modules/jax/default.nix | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index edf18bbe2ed0..f0b4d6e26324 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -31,7 +31,7 @@ buildPythonPackage rec { pytestFlagsArray = [ "-W ignore::DeprecationWarning" "tests/" ]; meta = with lib; { - description = "Differentiate, compile, and transform Numpy code."; + description = "Differentiate, compile, and transform Numpy code"; homepage = "https://github.com/google/jax"; license = licenses.asl20; maintainers = with maintainers; [ samuela ]; From 0ff986154ea140f256e47a78233579ad27d8387a Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sat, 4 Sep 2021 05:47:09 +0000 Subject: [PATCH 5/7] python3Packages.jaxlib: add CUDA support --- .../python-modules/jaxlib/default.nix | 66 ++++++++++++++----- 1 file changed, 51 insertions(+), 15 deletions(-) diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix index 240c5a7d6d0e..039c3deb8fc8 100644 --- a/pkgs/development/python-modules/jaxlib/default.nix +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -1,16 +1,28 @@ -# For the moment we only support the CPU backend of jaxlib. GPU and TPU backends -# require some additional work. Their wheels are not located on PyPI. -# * CPU/GPU: https://storage.googleapis.com/jax-releases/jax_releases.html -# * TPU: https://storage.googleapis.com/jax-releases/libtpu_releases.html +# For the moment we only support the CPU and GPU backends of jaxlib. The TPU +# backend will require some additional work. Those wheels are located here: +# https://storage.googleapis.com/jax-releases/libtpu_releases.html. -{ autoPatchelfHook, buildPythonPackage, fetchPypi, isPy39, lib, stdenv +# For future reference, the easiest way to test that the gpu is being used is: +# NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib.override { cudaSupport = true; }" +# python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)" +# See https://github.com/google/jax/issues/971#issuecomment-508216439. + +{ addOpenGLRunpath, autoPatchelfHook, buildPythonPackage, config, fetchPypi +, fetchurl, isPy39, lib, stdenv # propagatedBuildInputs -, absl-py, flatbuffers, scipy +, absl-py, flatbuffers, scipy, cudatoolkit_11 +# Options: +, cudaSupport ? config.cudaSupport or false }: +assert cudaSupport -> lib.versionAtLeast cudatoolkit_11.version "11.1"; + +let + device = if cudaSupport then "gpu" else "cpu"; +in buildPythonPackage rec { pname = "jaxlib"; - version = "0.1.70"; + version = "0.1.71"; format = "wheel"; # At the time of writing (8/19/21), there are releases for 3.7-3.9. Supporting @@ -18,23 +30,47 @@ buildPythonPackage rec { # version. disabled = !isPy39; - src = fetchPypi { - inherit pname version format; - dist = "cp39"; - python = "cp39"; - platform = "manylinux2010_x86_64"; - sha256 = "sha256-mytMTqoavpuRawj52MU5/iFj27SGlm8DaoQ5vd/3bss="; - }; + src = { + cpu = fetchurl { + url = "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp39-none-manylinux2010_x86_64.whl"; + sha256 = "sha256:0rqhs6qabydizlv5d3rb20dbv6612rr7dqfniy9r6h4kazdinsn6"; + }; + gpu = fetchurl { + url = "https://storage.googleapis.com/jax-releases/cuda111/jaxlib-${version}+cuda111-cp39-none-manylinux2010_x86_64.whl"; + sha256 = "sha256:065kyzjsk9m84d138p99iymdiiicm1qz8a3iwxz8rspl43rwrw89"; + }; + }.${device}; # Prebuilt wheels are dynamically linked against things that nix can't find. # Run `autoPatchelfHook` to automagically fix them. - nativeBuildInputs = [ autoPatchelfHook ]; + nativeBuildInputs = [ autoPatchelfHook ] ++ lib.optional cudaSupport addOpenGLRunpath; # Dynamic link dependencies buildInputs = [ 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. Dependencies: + # * libcudart.so.11.0 -> cudatoolkit_11.lib + # * libcuda.so.1 -> opengl driver in /run/opengl-driver/lib + preInstallCheck = lib.optional cudaSupport '' + shopt -s globstar + + addOpenGLRunpath $out/**/*.so + + for file in $out/**/*.so; do + rpath=$(patchelf --print-rpath $file) + patchelf --set-rpath "$rpath:${lib.makeLibraryPath [ cudatoolkit_11.lib ]}" $file + done + ''; + # pip dependencies propagatedBuildInputs = [ absl-py flatbuffers scipy ]; + pythonImportsCheck = [ "jaxlib" ]; + meta = with lib; { description = "XLA library for JAX"; homepage = "https://github.com/google/jax"; From 3ea5dbdd720abec9a48665a69e69e5486bfe0fa0 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sat, 4 Sep 2021 20:57:32 +0000 Subject: [PATCH 6/7] python3Packages.jaxlib: add cudatoolkit_11 in propagatedBuildInputs --- pkgs/development/python-modules/jaxlib/default.nix | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix index 039c3deb8fc8..af0755bf9dd0 100644 --- a/pkgs/development/python-modules/jaxlib/default.nix +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -2,10 +2,13 @@ # backend will require some additional work. Those wheels are located here: # https://storage.googleapis.com/jax-releases/libtpu_releases.html. -# For future reference, the easiest way to test that the gpu is being used is: +# For future reference, the easiest way to test the GPU backend is to run # NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib.override { cudaSupport = true; }" -# python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)" -# See https://github.com/google/jax/issues/971#issuecomment-508216439. +# python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'" +# python -c "from jax import random; random.PRNGKey(0)" +# See https://github.com/google/jax/issues/971#issuecomment-508216439. There's +# no convenient way to test the GPU backend in the derivation since the nix +# build environment blocks access to the GPU. { addOpenGLRunpath, autoPatchelfHook, buildPythonPackage, config, fetchPypi , fetchurl, isPy39, lib, stdenv @@ -66,8 +69,8 @@ buildPythonPackage rec { done ''; - # pip dependencies - propagatedBuildInputs = [ absl-py flatbuffers scipy ]; + # pip dependencies and optionally cudatoolkit. + propagatedBuildInputs = [ absl-py flatbuffers scipy ] ++ lib.optional cudaSupport cudatoolkit_11; pythonImportsCheck = [ "jaxlib" ]; From 8f06bd48f26af655d9c112b27eb3d22a874685c4 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sat, 4 Sep 2021 21:41:40 +0000 Subject: [PATCH 7/7] python3Packages.jaxlib: add cudatoolkit_11 to rpaths --- pkgs/development/python-modules/jaxlib/default.nix | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix index af0755bf9dd0..7761c0372994 100644 --- a/pkgs/development/python-modules/jaxlib/default.nix +++ b/pkgs/development/python-modules/jaxlib/default.nix @@ -4,11 +4,14 @@ # For future reference, the easiest way to test the GPU backend is to run # NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib.override { cudaSupport = true; }" +# export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1 # python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'" # python -c "from jax import random; random.PRNGKey(0)" -# See https://github.com/google/jax/issues/971#issuecomment-508216439. There's -# no convenient way to test the GPU backend in the derivation since the nix -# build environment blocks access to the GPU. +# python -c "from jax import random; x = random.normal(random.PRNGKey(0), (100, 100)); x @ x" +# There's no convenient way to test the GPU backend in the derivation since the +# nix build environment blocks access to the GPU. See also: +# * https://github.com/google/jax/issues/971#issuecomment-508216439 +# * https://github.com/google/jax/issues/5723#issuecomment-913038780 { addOpenGLRunpath, autoPatchelfHook, buildPythonPackage, config, fetchPypi , fetchurl, isPy39, lib, stdenv @@ -57,6 +60,7 @@ buildPythonPackage rec { # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the # patchPhase. Dependencies: # * libcudart.so.11.0 -> cudatoolkit_11.lib + # * libcublas.so.11 -> cudatoolkit_11 # * libcuda.so.1 -> opengl driver in /run/opengl-driver/lib preInstallCheck = lib.optional cudaSupport '' shopt -s globstar @@ -65,7 +69,9 @@ buildPythonPackage rec { for file in $out/**/*.so; do rpath=$(patchelf --print-rpath $file) - patchelf --set-rpath "$rpath:${lib.makeLibraryPath [ cudatoolkit_11.lib ]}" $file + # For some reason `makeLibraryPath` on `cudatoolkit_11` maps to + # /lib which is different from /lib. + patchelf --set-rpath "$rpath:${cudatoolkit_11}/lib:${lib.makeLibraryPath [ cudatoolkit_11.lib ]}" $file done '';