Merge pull request #286763 from CertainLach/openai-triton

python311Packages.openai-triton: remove ptxas --version call when built without cuda support

pkgs/development/python-modules/openai-triton/0001-ptxas-disable-version-key-for-non-cuda-targets.patch

@@ -0,0 +1,27 @@
+From 10f3d49aa6084d1b9b9624017cce7df106b9fb7e Mon Sep 17 00:00:00 2001
+From: Yaroslav Bolyukin <iam@lach.pw>
+Date: Tue, 6 Feb 2024 13:51:28 +0100
+Subject: [PATCH] ptxas: disable version key for non-cuda targets
+
+---
+ python/triton/runtime/jit.py | 4 ++--
+ 1 file changed, 2 insertions(+), 2 deletions(-)
+
+diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py
+index d55972b4b..bd875a701 100644
+--- a/python/triton/runtime/jit.py
++++ b/python/triton/runtime/jit.py
+@@ -117,8 +117,8 @@ def version_key():
+         with open(lib.module_finder.find_spec(lib.name).origin, "rb") as f:
+             contents += [hashlib.md5(f.read()).hexdigest()]
+     # ptxas version
+-    ptxas = path_to_ptxas()[0]
+-    ptxas_version = hashlib.md5(subprocess.check_output([ptxas, "--version"])).hexdigest()
++    # ptxas = path_to_ptxas()[0]
++    ptxas_version = "noptxas"
+     return '-'.join(TRITON_VERSION) + '-' + ptxas_version + '-' + '-'.join(contents)
+ 
+ 
+-- 
+2.43.0
+

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

@@ -19,6 +19,9 @@
 , filelock
 , torchWithRocm
 , python
+
+, runCommand
+
 , cudaPackages
 , cudaSupport ? config.cudaSupport
 }:
@@ -46,6 +49,9 @@ buildPythonPackage rec {
     })
   ] ++ lib.optionals (!cudaSupport) [
     ./0000-dont-download-ptxas.patch
+    # openai-triton wants to get ptxas version even if ptxas is not
+    # used, resulting in ptxas not found error.
+    ./0001-ptxas-disable-version-key-for-non-cuda-targets.patch
   ];
 
   nativeBuildInputs = [
@@ -71,7 +77,12 @@ buildPythonPackage rec {
     zlib
   ];
 
-  propagatedBuildInputs = [ filelock ];
+  propagatedBuildInputs = [
+    filelock
+    # openai-triton uses setuptools at runtime:
+    # https://github.com/NixOS/nixpkgs/pull/286763/#discussion_r1480392652
+    setuptools
+  ];
 
   postPatch = let
     # Bash was getting weird without linting,
@@ -150,7 +161,18 @@ buildPythonPackage rec {
   # ];
 
   # Ultimately, torch is our test suite:
-  passthru.tests = { inherit torchWithRocm; };
+  passthru.tests = {
+    inherit torchWithRocm;
+    # Implemented as alternative to pythonImportsCheck, in case if circular dependency on torch occurs again,
+    # and pythonImportsCheck is commented back.
+    import-triton = runCommand "import-triton" { nativeBuildInputs = [(python.withPackages (ps: [ps.openai-triton]))]; } ''
+      python << \EOF
+      import triton
+      import triton.language
+      EOF
+      touch "$out"
+    '';
+  };
 
   pythonRemoveDeps = [
     # Circular dependency, cf. https://github.com/openai/triton/issues/1374