changing tmpdir when runnign in slurm

facebookresearch · Oct 9, 2023 · 75426f4 · 75426f4
1 parent a2b9675
commit 75426f4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/audiocraft/train.py b/audiocraft/train.py
@@ -12,6 +12,7 @@
 import logging
 import multiprocessing
 import os
+from pathlib import Path
 import sys
 import typing as tp
 
@@ -119,6 +120,11 @@ def init_seed_and_system(cfg):
     logger.debug('Setting num threads to %d', cfg.num_threads)
     set_efficient_attention_backend(cfg.efficient_attention_backend)
     logger.debug('Setting efficient attention backend to %s', cfg.efficient_attention_backend)
+    if 'SLURM_JOB_ID' in os.environ:
+        tmpdir = Path('/scratch/slurm_tmpdir/' + os.environ['SLURM_JOB_ID'])
+        if tmpdir.exists():
+            logger.info("Changing tmpdir to %s", tmpdir)
+            os.environ['TMPDIR'] = str(tmpdir)
 
 
 @hydra_main(config_path='../config', config_name='config', version_base='1.1')