Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168{ lib, stdenv, buildPythonPackage, fetchPypi, nix-update-script, python, cudaPackages,
# nativeBuildInputs autoPatchelfHook,
# dependencies cuda-bindings, cuda-pathfinder, numpy, pyelftools,}:
let inherit (stdenv.hostPlatform) system;
version = "13.4.0";
platforms = { x86_64-linux = "manylinux_2_24_x86_64.manylinux_2_28_x86_64"; aarch64-linux = "manylinux_2_24_aarch64.manylinux_2_28_aarch64"; };
hashes = { x86_64-linux = { cp311 = "sha256-xJC39QsYB0+hIZQPtzV8RNQqAkjuXKZLinpa7z2PUOY="; cp312 = "sha256-DWwobBkcjEZtzU3OqQ+g/RXxYohZK5A5rzo2fofjwLU="; cp313 = "sha256-8xeBWqCOTi8+kjmUXBOUALJBXffi8ZBae07IgXYxjGw="; cp314 = "sha256-OMr/i6qbr868gtGwSmr1jSwBlCCPhvavE/SkmPWP8mk="; }; aarch64-linux = { cp311 = "sha256-BYmPWvcPw+Hv6bSh/kCFE83N+wj5inl6X2myGyxDIlk="; cp312 = "sha256-EZh4dW3vNC+Q7t0CEhNE4AHRUCa/ehvixrmi2IK4Wuw="; cp313 = "sha256-S4L+oUNJwCYuYarQDj9mhs0diwJ9LavJhjFyIUGUgkA="; cp314 = "sha256-aDCxEjJTby+9salK8uX+hltDbRKVLhgFDeFGpki0jZI="; }; };
# Every wheel, keyed "<system>-<cpXY>" and exposed as passthru.srcs so the update script refreshes every hash. srcs = lib.concatMapAttrs ( wheelSystem: lib.mapAttrs' ( pyTag: hash: lib.nameValuePair "${wheelSystem}-${pyTag}" (fetchPypi { pname = "cupti_python"; inherit version hash; format = "wheel"; platform = platforms.${wheelSystem}; dist = pyTag; python = pyTag; abi = pyTag; }) ) ) hashes;
srcKey = "${system}-cp${lib.replaceStrings [ "." ] [ "" ] python.pythonVersion}";inbuildPythonPackage (finalAttrs: { pname = "cupti-python"; inherit version; format = "wheel"; __structuredAttrs = true;
src = srcs.${srcKey} or (throw "cupti-python: no wheel for ${srcKey}");
nativeBuildInputs = [ autoPatchelfHook ];
buildInputs = [ (lib.getLib cudaPackages.cuda_cupti) (lib.getLib stdenv.cc.cc) ];
pythonRemoveDeps = [ # Provided by cudaPackages.cuda_cupti instead of the nvidia-cuda-cupti wheel. "nvidia-cuda-cupti" ]; dependencies = [ cuda-bindings cuda-pathfinder numpy pyelftools ];
# The extension finds libcupti at run time with cuda.pathfinder's load_nvidia_dynamic_lib("cupti"), which # searches pip wheel and conda layouts, CUDA_HOME and the dynamic linker's default paths, but first returns a copy # that is already loaded. Link libcupti into the loader extension so autoPatchelfHook puts cuda_cupti on its # RUNPATH and it's loaded together with the module. preFixup = '' for so in "$out/${python.sitePackages}"/cupti/_internal/cupti.*.so; do patchelf --add-needed libcupti.so.${cudaPackages.cudaMajorVersion} "$so" done '';
pythonImportsCheck = [ "cupti" "cupti.cupti" ];
# No python tests doCheck = false;
passthru = { inherit srcs;
updateScript = nix-update-script { extraArgs = [ "--url=mirror://pypi/c/cupti-python/" # The wheels' fetchers carry no position, so name the file their hashes live in. "--override-filename=pkgs/development/python-modules/cupti-python/default.nix" "--no-src" ] ++ lib.concatMap (name: [ "--subpackage" "srcs.${name}" ]) (lib.attrNames srcs); };
# Runs the vector-add sample shipped in the wheel: it compiles a kernel with NVRTC, launches it under CUPTI activity # tracing, prints the activity records and verifies the result. testers.cuda.vectorAddDrv = cudaPackages.writeGpuTestPython { name = "cupti-python-vector-add-drv"; libraries = ps: [ ps.cupti-python ]; } '' import pathlib import subprocess import sys
import cupti
sample = pathlib.Path(cupti.__file__).parent.parent / "cupti-python-samples" / "CuptiVectorAddDrv.py" result = subprocess.run([sys.executable, sample, "--profile"], capture_output=True, text=True) print(result.stdout) print(result.stderr, file=sys.stderr)
assert result.returncode == 0, f"{sample.name} exited with {result.returncode}" assert "verify_result: PASS" in result.stdout, "vector add produced wrong results" kinds = {line.split()[0] for line in result.stdout.splitlines() if line.strip()} assert "CONCURRENT_KERNEL" in kinds, "CUPTI recorded no kernel activity" ''; tests.cuda.vectorAddDrv = finalAttrs.passthru.testers.cuda.vectorAddDrv.gpuCheck; };
meta = { description = "Python APIs for the NVIDIA CUDA Profiling Tools Interface (CUPTI)"; homepage = "https://docs.nvidia.com/cupti-python/"; changelog = "https://docs.nvidia.com/cupti-python/${finalAttrs.version}/release-notes/release-notes.html"; sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ]; license = lib.licenses.nvidiaCuptiPython; maintainers = with lib.maintainers; [ GaetanLepage connorbaker ]; teams = [ lib.teams.cuda ]; platforms = lib.attrNames platforms; # The wheels require nvidia-cuda-cupti~=13.4 and load libcupti.so.13. broken = cudaPackages.cudaOlder "13.4" || cudaPackages.cudaAtLeast "14"; };})