cudaPackages.writeGpuTestPython: add gpuCheckArgs parameter (#449057)

pkgs/development/cuda-modules/packages/writeGpuTestPython.nix

@@ -3,11 +3,13 @@
   runCommand,
   python3Packages,
   makeWrapper,
+  writableTmpDirAsHomeHook,
 }:
 {
   feature ? "cuda",
   name ? if feature == null then "cpu" else feature,
   libraries ? [ ], # [PythonPackage] | (PackageSet -> [PythonPackage])
+  gpuCheckArgs ? { },
   ...
 }@args:
 
@@ -28,6 +30,7 @@ let
     runCommand "tester-${name}"
       (
         lib.removeAttrs args [
+          "gpuCheckArgs"
           "libraries"
           "name"
         ]
@@ -52,10 +55,18 @@ let
   tester' = tester.overrideAttrs (oldAttrs: {
     passthru.gpuCheck =
       runCommand "test-${name}"
-        {
-          nativeBuildInputs = [ tester' ];
-          requiredSystemFeatures = lib.optionals (feature != null) [ feature ];
-        }
+        (
+          gpuCheckArgs
+          // {
+            nativeBuildInputs = [
+              tester'
+            ]
+            ++ gpuCheckArgs.nativeBuildInputs or [ ];
+
+            requiredSystemFeatures =
+              lib.optionals (feature != null) [ feature ] ++ gpuCheckArgs.requiredSystemFeatures or [ ];
+          }
+        )
         ''
           set -e
           ${tester.meta.mainProgram or (lib.getName tester')}

pkgs/development/python-modules/torch/tests/mk-torch-compile-check.nix

@@ -5,6 +5,7 @@
   libraries,
   name ? if feature == null then "torch-compile-cpu" else "torch-compile-${feature}",
   stdenv,
+  writableTmpDirAsHomeHook,
 }:
 let
   deviceStr = if feature == null then "" else '', device="cuda"'';
@@ -12,6 +13,12 @@ in
 cudaPackages.writeGpuTestPython
   {
     inherit name feature libraries;
+
+    # This could be accomplished with propagatedBuildInputs instead of introducing
+    gpuCheckArgs.nativeBuildInputs = [
+      # torch._inductor.exc.InductorError: PermissionError: [Errno 13] Permission denied: '/homeless-shelter'
+      writableTmpDirAsHomeHook
+    ];
     makeWrapperArgs = [
       "--suffix"
       "PATH"