Change unit of cuda-cache-size to MiB and add cuda_cache_dir

Author: Flamefire

easybuild/framework/easyblock.py

@@ -1027,6 +1027,24 @@ def make_dir(self, dir_name, clean, dontcreateinstalldir=False):
 
         mkdir(dir_name, parents=True)
 
+    def setup_cuda_cache(self):
+        cuda_cache_maxsize = build_option('cuda_cache_maxsize')
+        if cuda_cache_maxsize is None:
+            cuda_cache_maxsize = 1 * 1024  # 1 GiB default value
+        else:
+            cuda_cache_maxsize = int(cuda_cache_maxsize)
+        if cuda_cache_maxsize == 0:
+            self.log.info('Disabling CUDA PTX cache as per request')
+            env.setvar('CUDA_CACHE_DISABLE', '1')
+        else:
+            cuda_cache_dir = build_option('cuda_cache_dir')
+            if not cuda_cache_dir:
+                cuda_cache_dir = os.path.join(self.builddir, 'eb-cuda-cache')
+            self.log.info('Enabling CUDA PTX cache of size %s MiB at %s', cuda_cache_maxsize, cuda_cache_dir)
+            env.setvar('CUDA_CACHE_DISABLE', '0')
+            env.setvar('CUDA_CACHE_PATH', cuda_cache_dir)
+            env.setvar('CUDA_CACHE_MAXSIZE', str(cuda_cache_maxsize * 1024 * 1024))
+
     #
     # MODULE UTILITY FUNCTIONS
     #
@@ -2152,19 +2170,7 @@ def prepare_step(self, start_dir=True, load_tc_deps_modules=True):
 
         # Setup CUDA cache if required. If we don't do this, CUDA will use the $HOME for its cache files
         if get_software_root('CUDA') or get_software_root('CUDAcore'):
-            cuda_cache_maxsize = build_option('cuda_cache_maxsize')
-            if cuda_cache_maxsize is None:
-                cuda_cache_maxsize = 1 * 1024 * 1024 * 1024  # 1 GB default value
-            if cuda_cache_maxsize == 0:
-                self.log.info('Disabling CUDA PTX cache as per request')
-                env.setvar('CUDA_CACHE_DISABLE', '1')
-            else:
-                cuda_cache_dir = os.path.join(self.builddir, 'eb-cuda-cache')
-                self.log.info('Enabling CUDA PTX cache of size %s MB at %s',
-                              cuda_cache_maxsize / 1024 / 1024, cuda_cache_dir)
-                env.setvar('CUDA_CACHE_DISABLE', '0')
-                env.setvar('CUDA_CACHE_PATH', cuda_cache_dir)
-                env.setvar('CUDA_CACHE_MAXSIZE', str(cuda_cache_maxsize))
+            self.setup_cuda_cache()
 
         # guess directory to start configure/build/install process in, and move there
         if start_dir:

easybuild/tools/config.py

@@ -169,6 +169,7 @@ def mk_full_default_path(name, prefix=DEFAULT_PREFIX):
         'container_image_name',
         'container_template_recipe',
         'container_tmpdir',
+        'cuda_cache_dir',
         'cuda_cache_maxsize',
         'cuda_compute_capabilities',
         'download_timeout',

easybuild/tools/options.py

@@ -357,7 +357,9 @@ def override_options(self):
             'consider-archived-easyconfigs': ("Also consider archived easyconfigs", None, 'store_true', False),
             'containerize': ("Generate container recipe/image", None, 'store_true', False, 'C'),
             'copy-ec': ("Copy specified easyconfig(s) to specified location", None, 'store_true', False),
-            'cuda-cache-maxsize': ("Maximum size of the CUDA cache (in bytes) used for JIT compilation of PTX code. "
+            'cuda-cache-dir': ("Path to CUDA cache dir to use if enabled. Defaults to a path inside the build dir.",
+                               str, 'store', None, {'metavar': "PATH"}),
+            'cuda-cache-maxsize': ("Maximum size of the CUDA cache (in MiB) used for JIT compilation of PTX code. "
                                    "Leave value empty to let EasyBuild choose a value or '0' to disable the cache",
                                    int, 'store_or_None', None),
             'cuda-compute-capabilities': ("List of CUDA compute capabilities to use when building GPU software; "

test/framework/easyblock.py

@@ -1762,7 +1762,7 @@ def test_prepare_step_hmns(self):
         self.assertEqual(loaded_modules[0]['mod_name'], 'GCC/6.4.0-2.28')
 
     def test_prepare_step_cuda_cache(self):
-        """Test handling cuda-cache-maxsize option."""
+        """Test handling cuda-cache-* options."""
 
         init_config(build_options={'cuda_cache_maxsize': None})  # Automatic mode
 
@@ -1803,14 +1803,17 @@ def test_prepare_step_cuda_cache(self):
         self.assertNotIn('Enabling CUDA PTX cache', logtxt)
         self.assertEqual(os.environ['CUDA_CACHE_DISABLE'], '1')
 
-        init_config(build_options={'cuda_cache_maxsize': 1234567890})  # Specified size
+        # Specified size and location
+        cuda_cache_dir = os.path.join(self.test_prefix, 'custom-cuda-cache')
+        init_config(build_options={'cuda_cache_maxsize': 1234, 'cuda_cache_dir': cuda_cache_dir})
         write_file(eb.logfile, '')
         eb.prepare_step(start_dir=False)
         logtxt = read_file(eb.logfile)
         self.assertNotIn('Disabling CUDA PTX cache', logtxt)
         self.assertIn('Enabling CUDA PTX cache', logtxt)
         self.assertEqual(os.environ['CUDA_CACHE_DISABLE'], '0')
-        self.assertEqual(os.environ['CUDA_CACHE_MAXSIZE'], '1234567890')
+        self.assertEqual(os.environ['CUDA_CACHE_MAXSIZE'], str(1234 * 1024 * 1024))
+        self.assertEqual(os.environ['CUDA_CACHE_PATH'], cuda_cache_dir)
 
     def test_checksum_step(self):
         """Test checksum step"""