pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchPypi,
6 addDriverRunpath,
7 autoPatchelfHook,
8 pypaInstallHook,
9 wheelUnpackHook,
10 cudaPackages,
11 python,
12 jaxlib,
13}:
14let
15 inherit (jaxlib) version;
16
17 cudaLibPath = lib.makeLibraryPath (
18 with cudaPackages;
19 [
20 (lib.getLib libcublas) # libcublas.so
21 (lib.getLib cuda_cupti) # libcupti.so
22 (lib.getLib cuda_cudart) # libcudart.so
23 (lib.getLib cudnn) # libcudnn.so
24 (lib.getLib libcufft) # libcufft.so
25 (lib.getLib libcusolver) # libcusolver.so
26 (lib.getLib libcusparse) # libcusparse.so
27 (lib.getLib nccl) # libnccl.so
28 (lib.getLib libnvjitlink) # libnvJitLink.so
29 (lib.getLib addDriverRunpath.driverLink) # libcuda.so
30 ]
31 );
32
33in
34buildPythonPackage rec {
35 pname = "jax-cuda12-pjrt";
36 inherit version;
37 pyproject = false;
38
39 src = fetchPypi {
40 pname = "jax_cuda12_pjrt";
41 inherit version;
42 format = "wheel";
43 python = "py3";
44 dist = "py3";
45 platform =
46 {
47 x86_64-linux = "manylinux_2_27_x86_64";
48 aarch64-linux = "manylinux_2_27_aarch64";
49 }
50 .${stdenv.hostPlatform.system};
51 hash =
52 {
53 x86_64-linux = "sha256-OXdyaiozKwvTSDG96ytWUzY0QvMBLCmW/IgICq9rO60=";
54 aarch64-linux = "sha256-2H1mbQxSP62q23GU58J03MWg5/j40dfig1NT7zK+8Bw=";
55 }
56 .${stdenv.hostPlatform.system};
57 };
58
59 nativeBuildInputs = [
60 autoPatchelfHook
61 pypaInstallHook
62 wheelUnpackHook
63 ];
64
65 # jax-cuda12-pjrt looks for ptxas, nvlink and nvvm at runtime, eg when running `jax.random.PRNGKey(0)`.
66 # Linking into $out is the least bad solution. See
67 # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
68 # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211
69 # for more info.
70 postInstall = ''
71 mkdir -p $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin
72 ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin/ptxas
73 ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin/nvlink
74 ln -s ${cudaPackages.cuda_nvcc}/nvvm $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/nvvm
75 '';
76
77 # jax-cuda12-pjrt contains shared libraries that open other shared libraries via dlopen
78 # and these implicit dependencies are not recognized by ldd or
79 # autoPatchelfHook. That means we need to sneak them into rpath. This step
80 # must be done after autoPatchelfHook and the automatic stripping of
81 # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
82 # patchPhase.
83 preInstallCheck = ''
84 patchelf --add-rpath "${cudaLibPath}" $out/${python.sitePackages}/jax_plugins/xla_cuda12/xla_cuda_plugin.so
85 '';
86
87 # FIXME: there are no tests, but we need to run preInstallCheck above
88 doCheck = true;
89
90 pythonImportsCheck = [ "jax_plugins" ];
91
92 inherit cudaLibPath;
93
94 meta = {
95 description = "JAX XLA PJRT Plugin for NVIDIA GPUs";
96 homepage = "https://github.com/jax-ml/jax/tree/main/jax_plugins/cuda";
97 sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
98 license = lib.licenses.asl20;
99 maintainers = with lib.maintainers; [ natsukium ];
100 platforms = lib.platforms.linux;
101 # see CUDA compatibility matrix
102 # https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
103 broken = !(lib.versionAtLeast cudaPackages.cudnn.version "9.1");
104 };
105}