diff --git a/pkgs/development/python-modules/nccl4py/default.nix b/pkgs/development/python-modules/nccl4py/default.nix new file mode 100644 index 000000000000..8a1a9b4d2518 --- /dev/null +++ b/pkgs/development/python-modules/nccl4py/default.nix @@ -0,0 +1,102 @@ +{ + lib, + buildPythonPackage, + cudaPackages, + addDriverRunpath, + + # build-system + cython, + setuptools, + + # dependencies + cuda-core, + numpy, + packaging, + pythonOlder, + typing-extensions, +}: + +buildPythonPackage (finalAttrs: { + pname = "nccl4py"; + # `nccl4py` is versioned independently of `nccl` and should be the + # same as the contents of + # `${cudaPackages.nccl.src}/bindings/nccl4py/nccl/_version.py` + version = "0.3.0"; + pyproject = true; + __structuredAttrs = true; + + inherit (cudaPackages.nccl) src; + sourceRoot = "${finalAttrs.src.name}/bindings/nccl4py"; + + postPatch = + let + # taken from `cuda-bindings`'s `postPatch` + libCudaPath = + # Use cuda_compat to provide libcuda.so on pre-Thor Jetsons + if (cudaPackages.cuda_compat.meta.available or false) then + cudaPackages.cuda_compat + + # Else, use the host CUDA driver library + else + addDriverRunpath.driverLink; + in + '' + substituteInPlace nccl/bindings/_internal/nccl_linux.pyx \ + --replace-fail \ + "handle = dlopen('libcuda.so.1'" \ + "handle = dlopen('${libCudaPath}/lib/libcuda.so.1'" + + substituteInPlace nccl/bindings/_internal/nccl_linux.pyx \ + --replace-fail \ + 'cdef uintptr_t handle = load_nvidia_dynamic_lib("nccl")._handle_uint' \ + 'cdef uintptr_t handle = dlopen("${lib.getLib cudaPackages.nccl}/lib/libnccl.so.2", RTLD_NOW | RTLD_GLOBAL)' + ''; + + build-system = [ + cython + setuptools + ]; + + env = { + # `${sourceRoot}/setup.py` insists on reading only from $CUDA_HOME/include + CUDA_HOME = (lib.getInclude cudaPackages.cuda_cudart).outPath; + }; + + buildInputs = + lib.optionals (cudaPackages.cudaOlder "13.0") [ + cudaPackages.cuda_nvcc + ] + ++ lib.optionals (cudaPackages.cudaAtLeast "13.0") [ + cudaPackages.cuda_crt + ]; + + dependencies = [ + cuda-core + numpy + packaging + ] + ++ lib.optionals (pythonOlder "3.13") [ + typing-extensions + ]; + + pythonImportsCheck = [ + "nccl" + "nccl.bindings" + ]; + + # Upstream doesn't ship any tests. + doCheck = false; + + meta = { + description = "Python bindings for NCCL"; + homepage = "https://github.com/NVIDIA/nccl/blob/master/bindings/nccl4py/README.md"; + # `cudaPackages.nccl` is BSD3 but the bindings are licensed under + # Apache License 2.0 + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ + GaetanLepage + thefossguy + ]; + inherit (cudaPackages.nccl.meta) platforms badPlatforms; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 4e9d8d4746f3..16e02c0f3c1c 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11591,6 +11591,8 @@ self: super: with self; { nc-dnsapi = callPackage ../development/python-modules/nc-dnsapi { }; + nccl4py = callPackage ../development/python-modules/nccl4py { }; + ncclient = callPackage ../development/python-modules/ncclient { }; nclib = callPackage ../development/python-modules/nclib { };