From e712a04790a9f9e658c7d770e4ad514c38481b1a Mon Sep 17 00:00:00 2001
From: natsukium <tomoya.otabi@gmail.com>
Date: Sat, 30 Dec 2023 16:29:24 +0900
Subject: [PATCH] python311Packages.lightgbm: fix build on linux and with cuda

---
 .../python-modules/lightgbm/default.nix       | 23 +++++++++++++++----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/lightgbm/default.nix b/pkgs/development/python-modules/lightgbm/default.nix
index d2fc8cbc13a2..54ca047acdd9 100644
--- a/pkgs/development/python-modules/lightgbm/default.nix
+++ b/pkgs/development/python-modules/lightgbm/default.nix
@@ -1,4 +1,5 @@
 { lib
+, config
 , stdenv
 , buildPythonPackage
 , fetchPypi
@@ -19,12 +20,16 @@
 
 # optionals: gpu
 , boost
-, cudatoolkit
 , ocl-icd
 , opencl-headers
-, gpuSupport ? stdenv.isLinux
+, gpuSupport ? stdenv.isLinux && !cudaSupport
+, cudaSupport ? config.cudaSupport
+, cudaPackages
 }:
 
+assert gpuSupport -> cudaSupport != true;
+assert cudaSupport -> gpuSupport != true;
+
 buildPythonPackage rec {
   pname = "lightgbm";
   version = "4.1.0";
@@ -43,6 +48,8 @@ buildPythonPackage rec {
     pathspec
     pyproject-metadata
     scikit-build-core
+  ] ++ lib.optionals cudaSupport [
+    cudaPackages.cuda_nvcc
   ];
 
   dontUseCmakeConfigure = true;
@@ -51,10 +58,12 @@ buildPythonPackage rec {
     llvmPackages.openmp
   ]) ++ (lib.optionals gpuSupport [
     boost
-    cudatoolkit
     ocl-icd
     opencl-headers
-  ]);
+  ]) ++ lib.optionals cudaSupport [
+    cudaPackages.cuda_nvcc
+    cudaPackages.cuda_cudart
+  ];
 
   propagatedBuildInputs = [
     numpy
@@ -62,7 +71,11 @@ buildPythonPackage rec {
     scikit-learn
   ];
 
-  pypaBuildFlags = lib.optionalString gpuSupport "--config-setting=cmake.define.USE_CUDA=ON";
+  pypaBuildFlags = lib.optionals gpuSupport [
+    "--config-setting=cmake.define.USE_GPU=ON"
+  ] ++ lib.optionals cudaSupport [
+    "--config-setting=cmake.define.USE_CUDA=ON"
+  ];
 
   postConfigure = ''
     export HOME=$(mktemp -d)