pkgs/development/python-modules/openai-triton/bin.nix

   1 { lib
   2 , stdenv
   3 , addOpenGLRunpath
   4 , cudaPackages
   5 , buildPythonPackage
   6 , fetchurl
   7 , isPy38
   8 , isPy39
   9 , isPy310
  10 , isPy311
  11 , python
  12 , autoPatchelfHook
  13 , filelock
  14 , lit
  15 , pythonRelaxDepsHook
  16 , zlib
  17 }:
  18
  19 buildPythonPackage rec {
  20   pname = "triton";
  21   version = "2.1.0";
  22   format = "wheel";
  23
  24   src =
  25     let pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion;
  26         unsupported = throw "Unsupported system";
  27         srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported;
  28     in fetchurl srcs;
  29
  30   disabled = !(isPy38 || isPy39 || isPy310 || isPy311);
  31
  32   pythonRemoveDeps = [ "cmake" "torch" ];
  33
  34   buildInputs = [ zlib ];
  35
  36   nativeBuildInputs = [
  37     pythonRelaxDepsHook # torch and triton refer to each other so this hook is included to mitigate that.
  38     autoPatchelfHook
  39   ];
  40
  41   propagatedBuildInputs = [
  42     filelock
  43     lit
  44     zlib
  45   ];
  46
  47   dontStrip = true;
  48
  49   # If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"
  50   postFixup = ''
  51     chmod +x "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas"
  52   '' +
  53   (let
  54     # Bash was getting weird without linting,
  55     # but basically upstream contains [cc, ..., "-lcuda", ...]
  56     # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
  57     old = [ "-lcuda" ];
  58     new = [ "-lcuda" "-L${addOpenGLRunpath.driverLink}" "-L${cudaPackages.cuda_cudart}/lib/stubs/" ];
  59
  60     quote = x: ''"${x}"'';
  61     oldStr = lib.concatMapStringsSep ", " quote old;
  62     newStr = lib.concatMapStringsSep ", " quote new;
  63   in
  64     ''
  65       substituteInPlace $out/${python.sitePackages}/triton/common/build.py \
  66         --replace '${oldStr}' '${newStr}'
  67     '');
  68
  69   meta = with lib; {
  70     description = "A language and compiler for custom Deep Learning operations";
  71     homepage = "https://github.com/openai/triton/";
  72     changelog = "https://github.com/openai/triton/releases/tag/v${version}";
  73     # Includes NVIDIA's ptxas, but redistributions of the binary are not limited.
  74     # https://docs.nvidia.com/cuda/eula/index.html
  75     # triton's license is MIT.
  76     # openai-triton-bin includes ptxas binary, therefore unfreeRedistributable is set.
  77     license = with licenses; [ unfreeRedistributable mit ];
  78     sourceProvenance = with sourceTypes; [ binaryNativeCode ];
  79     platforms = [ "x86_64-linux" ];
  80     maintainers = with maintainers; [ junjihashimoto ];
  81   };
  82 }