From 13973b5cf77c85e7fe2f6ec62df4906f88e4f1de Mon Sep 17 00:00:00 2001
From: Martin Weinelt <hexa@darmstadt.ccc.de>
Date: Fri, 16 Feb 2024 03:37:42 +0100
Subject: [PATCH] python311Packages.accelerate: fix build on linux

---
 pkgs/development/python-modules/accelerate/default.nix | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/accelerate/default.nix b/pkgs/development/python-modules/accelerate/default.nix
index f9a7be69f561..5f80a44ed7df 100644
--- a/pkgs/development/python-modules/accelerate/default.nix
+++ b/pkgs/development/python-modules/accelerate/default.nix
@@ -2,8 +2,6 @@
 , lib
 , buildPythonPackage
 , fetchFromGitHub
-, fetchpatch
-, pythonAtLeast
 , pythonOlder
 , pytestCheckHook
 , setuptools
@@ -13,6 +11,7 @@
 , pyyaml
 , safetensors
 , torch
+, cudatoolkit
 , evaluate
 , parameterized
 , transformers
@@ -52,6 +51,8 @@ buildPythonPackage rec {
   preCheck = ''
     export HOME=$(mktemp -d)
     export PATH=$out/bin:$PATH
+  '' + lib.optionalString (lib.meta.availableOn stdenv.hostPlatform cudatoolkit) ''
+    export TRITON_PTXAS_PATH="${cudatoolkit}/bin/ptxas"
   '';
   pytestFlagsArray = [ "tests" ];
   disabledTests = [
@@ -74,6 +75,8 @@ buildPythonPackage rec {
   ] ++ lib.optionals (stdenv.isLinux && stdenv.isAarch64) [
     # usual aarch64-linux RuntimeError: DataLoader worker (pid(s) <...>) exited unexpectedly
     "CheckpointTest"
+    # requires ptxas from cudatoolkit, which is unavailable on aarch64-linux
+    "test_dynamo_extract_model"
   ] ++ lib.optionals (stdenv.isDarwin && stdenv.isx86_64) [
     # RuntimeError: torch_shm_manager: execl failed: Permission denied
     "CheckpointTest"