gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] Default cache dir to /tmp/triton_$USER (#527)

Browse Source
This commit is contained in:
Bert Maher
2022-05-27 16:51:05 -04:00
committed by GitHub
parent c82a206684
commit 37037bb3be
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/triton/code_gen.py
View File

@@ -1150,6 +1150,11 @@ class DependenciesFinder(ast.NodeVisitor):
self.ret = hashlib.md5(self.ret).hexdigest() self.ret = hashlib.md5(self.ret).hexdigest()
def default_cache_dir():
import getpass
return f'/tmp/triton_{getpass.getuser()}'
class JITFunction: class JITFunction:
cache_hook = None cache_hook = None
@@ -1235,7 +1240,7 @@ class JITFunction:
hashed_key = hashlib.md5(key.encode("utf-8")).hexdigest() hashed_key = hashlib.md5(key.encode("utf-8")).hexdigest()
# create cache directory # create cache directory
cache_dir = os.environ.get('TRITON_CACHE_DIR', '/tmp/triton/') cache_dir = os.environ.get('TRITON_CACHE_DIR', default_cache_dir())
if cache_dir: if cache_dir:
os.makedirs(cache_dir, exist_ok=True) os.makedirs(cache_dir, exist_ok=True)