Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
3.6 kB · 106 lines
Nix
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107{ lib, stdenv, buildPythonPackage, fetchPypi, addDriverRunpath, autoPatchelfHook, pypaInstallHook, wheelUnpackHook, cudaPackages, python, jaxlib,}:let inherit (jaxlib) version; inherit (cudaPackages) cudaAtLeast;
cudaLibPath = lib.makeLibraryPath ( with cudaPackages; [ (lib.getLib libcublas) # libcublas.so (lib.getLib cuda_cupti) # libcupti.so (lib.getLib cuda_cudart) # libcudart.so (lib.getLib cudnn) # libcudnn.so (lib.getLib libcufft) # libcufft.so (lib.getLib libcusolver) # libcusolver.so (lib.getLib libcusparse) # libcusparse.so (lib.getLib nccl) # libnccl.so (lib.getLib libnvjitlink) # libnvJitLink.so (lib.getLib addDriverRunpath.driverLink) # libcuda.so ] );
inbuildPythonPackage rec { pname = "jax-cuda12-pjrt"; inherit version; pyproject = false;
src = fetchPypi { pname = "jax_cuda12_pjrt"; inherit version; format = "wheel"; python = "py3"; dist = "py3"; platform = { x86_64-linux = "manylinux2014_x86_64"; aarch64-linux = "manylinux2014_aarch64"; } .${stdenv.hostPlatform.system}; hash = { x86_64-linux = "sha256-jNnq15SOosd4pQj+9dEVnot6v0/MxwN8P+Hb/NlQEtw="; aarch64-linux = "sha256-IvrwINLo98oeKRVjMkH333Z4tzxwePXwsvETJIM3994="; } .${stdenv.hostPlatform.system}; };
nativeBuildInputs = [ autoPatchelfHook pypaInstallHook wheelUnpackHook ];
# jax-cuda12-pjrt looks for ptxas, nvlink and nvvm 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 = '' mkdir -p $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin/ptxas ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/bin/nvlink ln -s ${cudaPackages.cuda_nvcc}/nvvm $out/${python.sitePackages}/jax_plugins/nvidia/cuda_nvcc/nvvm '';
# jax-cuda12-pjrt 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 = '' patchelf --add-rpath "${cudaLibPath}" $out/${python.sitePackages}/jax_plugins/xla_cuda12/xla_cuda_plugin.so '';
# FIXME: there are no tests, but we need to run preInstallCheck above doCheck = true;
pythonImportsCheck = [ "jax_plugins" ];
inherit cudaLibPath;
meta = { description = "JAX XLA PJRT Plugin for NVIDIA GPUs"; homepage = "https://github.com/jax-ml/jax/tree/main/jax_plugins/cuda"; sourceProvenance = [ lib.sourceTypes.binaryNativeCode ]; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ natsukium ]; platforms = lib.platforms.linux; # see CUDA compatibility matrix # https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder broken = !(cudaAtLeast "12.1") || !(lib.versionAtLeast cudaPackages.cudnn.version "9.1"); };}