From 530d5769d57e88bde975c5fefb9ecfcddba89742 Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Tue, 18 Jun 2024 22:03:08 +0000
Subject: [PATCH] python311Packages.torch: prevent libcuda_dirs() from running
 ldconfig

---
 .../python-modules/openai-triton/default.nix  | 39 ++++++++++++-------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/pkgs/development/python-modules/openai-triton/default.nix b/pkgs/development/python-modules/openai-triton/default.nix
index 6869704a1bdb..d1aaed52052d 100644
--- a/pkgs/development/python-modules/openai-triton/default.nix
+++ b/pkgs/development/python-modules/openai-triton/default.nix
@@ -1,10 +1,10 @@
 {
   lib,
   config,
+  addDriverRunpath,
   buildPythonPackage,
   fetchFromGitHub,
   fetchpatch,
-  addOpenGLRunpath,
   setuptools,
   pytestCheckHook,
   pythonRelaxDepsHook,
@@ -59,19 +59,23 @@ buildPythonPackage rec {
 
   postPatch =
     let
-      # Bash was getting weird without linting,
-      # but basically upstream contains [cc, ..., "-lcuda", ...]
-      # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
-      old = [ "-lcuda" ];
-      new = [
-        "-lcuda"
-        "-L${addOpenGLRunpath.driverLink}"
-        "-L${cudaPackages.cuda_cudart}/lib/stubs/"
-      ];
-
       quote = x: ''"${x}"'';
-      oldStr = lib.concatMapStringsSep ", " quote old;
-      newStr = lib.concatMapStringsSep ", " quote new;
+      subs.ldFlags =
+        let
+          # Bash was getting weird without linting,
+          # but basically upstream contains [cc, ..., "-lcuda", ...]
+          # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
+          old = [ "-lcuda" ];
+          new = [
+            "-lcuda"
+            "-L${addDriverRunpath.driverLink}"
+            "-L${cudaPackages.cuda_cudart}/lib/stubs/"
+          ];
+        in
+        {
+          oldStr = lib.concatMapStringsSep ", " quote old;
+          newStr = lib.concatMapStringsSep ", " quote new;
+        };
     in
     ''
       # Use our `cmakeFlags` instead and avoid downloading dependencies
@@ -86,11 +90,17 @@ buildPythonPackage rec {
       substituteInPlace unittest/CMakeLists.txt \
         --replace "include (\''${CMAKE_CURRENT_SOURCE_DIR}/googletest.cmake)" ""\
         --replace "include(GoogleTest)" "find_package(GTest REQUIRED)"
+
+      cat << \EOF > python/triton/common/build.py
+
+      def libcuda_dirs():
+          return [ "${addDriverRunpath.driverLink}/lib" ]
+      EOF
     ''
     + lib.optionalString cudaSupport ''
       # Use our linker flags
       substituteInPlace python/triton/common/build.py \
-        --replace '${oldStr}' '${newStr}'
+        --replace '${subs.ldFlags.oldStr}' '${subs.ldFlags.newStr}'
     '';
 
   nativeBuildInputs = [
@@ -123,7 +133,6 @@ buildPythonPackage rec {
     setuptools
   ];
 
-
   # Avoid GLIBCXX mismatch with other cuda-enabled python packages
   preConfigure =
     ''