use slurm checkpoint dir

r2llab · Oct 31, 2024 · 441e409 · 441e409
1 parent 4197cc8
commit 441e409
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/wrangl/conf/wrangl_experiment.yaml b/wrangl/conf/wrangl_experiment.yaml
@@ -59,6 +59,7 @@ val_check_interval: 100
 flush_logs_every_n_steps: '${log_every_n_steps}'
 test_only: false
 autoresume: false
+use_slurm_checkpoint_dout: false  # set to True to use /checkpoint/$USER/$SLURM_JOB_ID for storage
 ckpt_path: 'latest.ckpt'
 val_sample_size: 100
 

diff --git a/wrangl/learn/supervised.py b/wrangl/learn/supervised.py
@@ -228,6 +228,8 @@ def run_train_test(cls, cfg: OmegaConf, train_dataset: torch.utils.data.Dataset,
         model_kwargs = model_kwargs or {}
         L.seed_everything(seed=cfg.seed, workers=True)
         dout = os.getcwd()
+        if cfg.use_slurm_checkpoint_dout:
+            dout = '/checkpoint/{}/{}'.format(os.environ['USER'], os.environ['SLURM_JOB_ID'])
 
         logger = logging.getLogger(name='{}:train_test'.format(cls.__name__))
         logger.info('Logging to {}'.format(dout))