diff --git a/pkgs/development/python-modules/triton/bin.nix b/pkgs/development/python-modules/triton/bin.nix new file mode 100644 index 000000000000000..ba8722c449bb38c --- /dev/null +++ b/pkgs/development/python-modules/triton/bin.nix @@ -0,0 +1,65 @@ +{ lib +, stdenv +, buildPythonPackage +, fetchurl +, isPy37 +, isPy38 +, isPy39 +, isPy310 +, isPy311 +, python +, pythonOlder +, pythonAtLeast +, filelock +, lit +, pythonRelaxDepsHook +, zlib +}: + +buildPythonPackage rec { + pname = "triton"; + version = "2.0.0"; + format = "wheel"; + + src = + let pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion; + unsupported = throw "Unsupported system"; + srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported; + in fetchurl srcs; + + disabled = !(isPy38 || isPy39 || isPy310 || isPy311); + + pythonRemoveDeps = [ "cmake" "torch" ]; + + nativeBuildInputs = [ + pythonRelaxDepsHook # torch and triton refer to each other so this hook is included to mitigate that. + ]; + + propagatedBuildInputs = [ + filelock + lit + zlib + ]; + + dontStrip = true; + + postFixup = '' + chmod +x "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas" + patchelf --set-interpreter "$(cat $NIX_CC/nix-support/dynamic-linker)" \ + "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas" + patchelf --set-rpath "${lib.makeLibraryPath [ stdenv.cc.cc.lib ]}" \ + "$out/${python.sitePackages}/triton/_C/libtriton.so" + patchelf --add-needed ${zlib.out}/lib/libz.so \ + "$out/${python.sitePackages}/triton/_C/libtriton.so" + ''; + + meta = with lib; { + description = "A language and compiler for custom Deep Learning operations"; + homepage = "https://github.com/openai/triton/"; + changelog = "https://github.com/openai/triton/releases/tag/v${version}"; + license = licenses.mit; + sourceProvenance = with sourceTypes; [ binaryNativeCode ]; + platforms = [ "x86_64-linux" ]; + maintainers = with maintainers; [ junjihashimoto ]; + }; +} diff --git a/pkgs/development/python-modules/triton/binary-hashes.nix b/pkgs/development/python-modules/triton/binary-hashes.nix new file mode 100644 index 000000000000000..79af94ac7804fa1 --- /dev/null +++ b/pkgs/development/python-modules/triton/binary-hashes.nix @@ -0,0 +1,31 @@ +# Warning: Need to update at the same time as torch-bin +# +# Precompiled wheels can be found at: +# https://download.pytorch.org/whl/torch_stable.html + +# To add a new version, run "prefetch.sh 'new-version'" to paste the generated file as follows. + +version : builtins.getAttr version { + "2.0.0" = { + x86_64-linux-38 = { + name = "triton-2.0.0-cp38-cp38-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; + hash = "sha256-nUl4KYt0/PWadf5x5TXAkrAjCIkzsvHfkz7DJhXkvu8="; + }; + x86_64-linux-39 = { + name = "triton-2.0.0-cp39-cp39-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; + hash = "sha256-dPEYwStDf7LKJeGgR1kXO1F1gvz0x74RkTMWx2QhNlY="; + }; + x86_64-linux-310 = { + name = "triton-2.0.0-cp310-cp310-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; + hash = "sha256-OIBu6WY/Sw981keQ6WxXk3QInlj0mqxKZggSGqVeJQU="; + }; + x86_64-linux-311 = { + name = "triton-2.0.0-cp311-cp311-linux_x86_64.whl"; + url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; + hash = "sha256-ImlBx7hZUhnd71mh/bgh6MdEKJoTJBXd1YT6zt60dbE="; + }; + }; +} diff --git a/pkgs/development/python-modules/triton/prefetch.sh b/pkgs/development/python-modules/triton/prefetch.sh new file mode 100755 index 000000000000000..a5172be7e3cef9a --- /dev/null +++ b/pkgs/development/python-modules/triton/prefetch.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env nix-shell +#!nix-shell -i bash -p nix-prefetch-scripts + +set -eou pipefail + +version=$1 + +linux_bucket="https://download.pytorch.org/whl" + +url_and_key_list=( + "x86_64-linux-38 $linux_bucket/triton-${version}-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp38-cp38-linux_x86_64.whl" + "x86_64-linux-39 $linux_bucket/triton-${version}-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp39-cp39-linux_x86_64.whl" + "x86_64-linux-310 $linux_bucket/triton-${version}-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp310-cp310-linux_x86_64.whl" + "x86_64-linux-311 $linux_bucket/triton-${version}-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp311-cp311-linux_x86_64.whl" +) + +hashfile=binary-hashes-"$version".nix +echo " \"$version\" = {" >> $hashfile + +for url_and_key in "${url_and_key_list[@]}"; do + key=$(echo "$url_and_key" | cut -d' ' -f1) + url=$(echo "$url_and_key" | cut -d' ' -f2) + name=$(echo "$url_and_key" | cut -d' ' -f3) + + echo "prefetching ${url}..." + hash=$(nix hash to-sri --type sha256 `nix-prefetch-url "$url" --name "$name"`) + + echo " $key = {" >> $hashfile + echo " name = \"$name\";" >> $hashfile + echo " url = \"$url\";" >> $hashfile + echo " hash = \"$hash\";" >> $hashfile + echo " };" >> $hashfile + + echo +done + +echo " };" >> $hashfile +echo "done." diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 797d29ed14db84b..aab62f74fd80dff 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11890,6 +11890,8 @@ self: super: with self; { trove-classifiers = callPackage ../development/python-modules/trove-classifiers { }; + triton-bin = callPackage ../development/python-modules/triton/bin.nix { }; + trueskill = callPackage ../development/python-modules/trueskill { }; trustme = callPackage ../development/python-modules/trustme { };