Fix CWD-relative CUTLASS C++ compilation cache (#3404)

python/cutlass_cppgen/__init__.py

@@ -81,7 +81,24 @@ def cuda_install_path():
         _CUDA_INSTALL_PATH = os.getenv("CUDA_INSTALL_PATH", _cuda_install_path_from_nvcc())
     return _CUDA_INSTALL_PATH
 
-CACHE_FILE = "compiled_cache.db"
+def _cache_path() -> str:
+    """Return the per-user path used for compiled operation artifacts."""
+    cache_root = os.environ.get("CUTLASS_CPPGEN_CACHE") or os.environ.get("XDG_CACHE_HOME")
+    if not cache_root:
+        cache_root = os.path.join(os.path.expanduser("~"), ".cache")
+
+    cache_dir = os.path.join(os.path.expanduser(cache_root), "cutlass_cppgen")
+    os.makedirs(cache_dir, mode=0o700, exist_ok=True)
+    try:
+        os.chmod(cache_dir, 0o700)
+    except OSError:
+        # Some platforms do not support POSIX permission bits.
+        pass
+
+    return os.path.join(cache_dir, "compiled_cache.db")
+
+
+CACHE_FILE = _cache_path()
 
 from cutlass_library import (
     DataType,

python/cutlass_cppgen/backend/compiler.py

@@ -124,12 +124,20 @@ def convertToBinaryData(filename):
 
 
 def CDLLBin(host_binary):
-    tempfile.tempdir = "./"
-    temp_so = tempfile.NamedTemporaryFile(prefix="host_func", suffix=".so", delete=True)
-    with open(temp_so.name, "wb") as file:
-        file.write(host_binary)
-    host_lib = ctypes.CDLL(temp_so.name)
-    return host_lib
+    with tempfile.NamedTemporaryFile(prefix="host_func", suffix=".so", delete=True) as temp_so:
+        with open(temp_so.name, "wb") as file:
+            file.write(host_binary)
+        return ctypes.CDLL(temp_so.name)
+
+
+def _connect_cache():
+    connection = sqlite3.connect(CACHE_FILE)
+    try:
+        os.chmod(CACHE_FILE, 0o600)
+    except OSError:
+        # Some platforms do not support POSIX permission bits.
+        pass
+    return connection
 
 
 class ArtifactManager:
@@ -138,7 +146,7 @@ class ArtifactManager:
     """
 
     def __init__(self) -> None:
-        connection = sqlite3.connect(CACHE_FILE)
+        connection = _connect_cache()
         cursor = connection.cursor()
         # Create the table if it does not already exist
         sqlite_create_table_query = """
@@ -151,6 +159,7 @@ def __init__(self) -> None:
         cursor.execute(sqlite_create_table_query)
         connection.commit()
         cursor.close()
+        connection.close()
 
         self._nvrtc_compile_options = ["-std=c++17", "-default-device"]
         self._nvcc_compile_options = [
@@ -171,7 +180,7 @@ def nvcc(self):
         self.default_compile_options = self._nvcc_compile_options
 
     def insert_operation(self, op_key, cubin, hostfile, op_name, op_attrs):
-        connection = sqlite3.connect(CACHE_FILE)
+        connection = _connect_cache()
         cursor = connection.cursor()
         sqlite_insert_blob_query = """ INSERT OR IGNORE INTO compiled_operations (op_key, cubin, hostbin, op_name, op_attrs) VALUES (?, ?, ?, ?, ?)"""
 
@@ -182,14 +191,17 @@ def insert_operation(self, op_key, cubin, hostfile, op_name, op_attrs):
         cursor.execute(sqlite_insert_blob_query, data_tuple)
         connection.commit()
         cursor.close()
+        connection.close()
 
     def load_operation(self, op_key, extra_funcs):
-        connection = sqlite3.connect(CACHE_FILE)
+        connection = _connect_cache()
         cursor = connection.cursor()
         sqlite_fetch_blob_query = """SELECT * from compiled_operations where op_key = ?"""
         cursor.execute(sqlite_fetch_blob_query, (op_key,))
         record = cursor.fetchall()
         if len(record) == 0:
+            cursor.close()
+            connection.close()
             return False
         for row in record:
             key, cubin_image, host_binary, operation_name, op_attr = row
@@ -225,6 +237,8 @@ def load_operation(self, op_key, extra_funcs):
                     compiled_host_fns[attr] = func
 
             self.compiled_cache_host[key] = compiled_host_fns
+        cursor.close()
+        connection.close()
         return True
 
     def emit_compile_(self, operation_list, compilation_options, host_compilation_options):
@@ -308,7 +322,6 @@ def emit_compile_(self, operation_list, compilation_options, host_compilation_op
 
         else:  # with nvcc backend
             # emit code
-            tempfile.tempdir = "./"
             temp_cu = tempfile.NamedTemporaryFile(
                 prefix="kernel", suffix=".cu", delete=True)
             temp_cubin = tempfile.NamedTemporaryFile(
@@ -331,7 +344,6 @@ def emit_compile_(self, operation_list, compilation_options, host_compilation_op
             with open(temp_cubin.name, "rb") as file:
                 cubin_image = file.read()
 
-        tempfile.tempdir = "./"
         temp_src = tempfile.NamedTemporaryFile(
             prefix="host_src", suffix=".cu", delete=True)
 

test/python/cutlass_cppgen/test_cache.py

@@ -0,0 +1,93 @@
+#################################################################################################
+#
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
+# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+#################################################################################################
+
+import ctypes
+import importlib
+import os
+from pathlib import Path
+import tempfile
+import unittest
+from unittest import mock
+
+import cutlass_cppgen
+
+
+compiler = importlib.import_module("cutlass_cppgen.backend.compiler")
+
+
+class TestCache(unittest.TestCase):
+    def test_cache_path_uses_configured_cache_root(self):
+        with tempfile.TemporaryDirectory() as cache_root:
+            with mock.patch.dict(os.environ, {"CUTLASS_CPPGEN_CACHE": cache_root}):
+                cache_file = Path(cutlass_cppgen._cache_path())
+
+            self.assertEqual(cache_file, Path(cache_root) / "cutlass_cppgen" / "compiled_cache.db")
+            self.assertTrue(cache_file.parent.is_dir())
+            if os.name == "posix":
+                self.assertEqual(cache_file.parent.stat().st_mode & 0o777, 0o700)
+
+    def test_cache_path_falls_back_to_xdg_cache_home(self):
+        with tempfile.TemporaryDirectory() as cache_root:
+            with mock.patch.dict(os.environ, {"XDG_CACHE_HOME": cache_root}, clear=True):
+                cache_file = Path(cutlass_cppgen._cache_path())
+
+            self.assertEqual(cache_file, Path(cache_root) / "cutlass_cppgen" / "compiled_cache.db")
+
+    def test_cache_file_is_created_with_user_only_permissions(self):
+        with tempfile.TemporaryDirectory() as cache_root:
+            cache_file = Path(cache_root) / "compiled_cache.db"
+            with mock.patch.object(compiler, "CACHE_FILE", str(cache_file)):
+                connection = compiler._connect_cache()
+                connection.close()
+
+            if os.name == "posix":
+                self.assertEqual(cache_file.stat().st_mode & 0o777, 0o600)
+
+    def test_cdll_bin_does_not_use_the_current_working_directory(self):
+        with tempfile.TemporaryDirectory() as temp_root, tempfile.TemporaryDirectory() as work_root:
+            with mock.patch.dict(os.environ, {"TMPDIR": temp_root}):
+                with mock.patch.object(tempfile, "tempdir", None):
+                    with mock.patch.object(ctypes, "CDLL", return_value=object()) as cdll:
+                        old_cwd = os.getcwd()
+                        os.chdir(work_root)
+                        try:
+                            compiler.CDLLBin(b"not a real shared library")
+                        finally:
+                            os.chdir(old_cwd)
+
+            loaded_path = Path(cdll.call_args.args[0])
+            self.assertEqual(loaded_path.parent, Path(temp_root))
+            self.assertNotEqual(loaded_path.parent, Path(work_root))
+
+
+if __name__ == "__main__":
+    unittest.main()