From 75426f4e1758bd7e4f9c0cd99928a28063650b4d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexandre=20D=C3=A9fossez?= Date: Mon, 9 Oct 2023 14:49:29 +0200 Subject: [PATCH] changing tmpdir when runnign in slurm --- audiocraft/train.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/audiocraft/train.py b/audiocraft/train.py index 22dd1178..5851222c 100644 --- 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')