diff --git a/pkgs/development/python-modules/openai-triton/0001-ptxas-disable-version-key-for-non-cuda-targets.patch b/pkgs/development/python-modules/openai-triton/0001-ptxas-disable-version-key-for-non-cuda-targets.patch new file mode 100644 index 000000000000..3941d54b8b37 --- /dev/null +++ b/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 +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 + diff --git a/pkgs/development/python-modules/openai-triton/default.nix b/pkgs/development/python-modules/openai-triton/default.nix index 46f1e6de5684..018852bdc743 100644 --- a/pkgs/development/python-modules/openai-triton/default.nix +++ b/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