Support submitit Hydra launcher, to launch runs on SLURM clusters

.env.example

@@ -10,5 +10,8 @@ CAMUS_DATA_PATH="path/to/camus"
 ### API keys ###
 COMET_API_KEY="<your-comet-api-key>"
 
+### SLURM config ###
+SLURM_MAIL_USER="<mail-address-to-notify>"
+
 ### Error Flags ###
 # HYDRA_FULL_ERROR=1

pyproject.toml

@@ -28,6 +28,7 @@ torch = "~1.12.0"
 torchvision = "~0.13.0"
 pytorch-lightning = "~1.8.0"
 hydra-core = "~1.2.0"
+hydra-submitit-launcher = "*"
 torchmetrics = "*"
 torchinfo = "*"
 pathos = "*"

vital/config/hydra/launcher/alliancecan.yaml

@@ -0,0 +1,12 @@
+defaults:
+  - submitit_slurm
+
+timeout_min: ${oc.select:run_time_min,60}
+setup:
+  - "module load httpproxy" # load module allowing to connect to whitelisted domains
+  - "source ${venv}/bin/activate" # activate the pre-installed virtual environment
+  - "rsync -a ${data_dir_to_copy} $SLURM_TMPDIR" # copy the dataset to the compute node
+
+additional_parameters:
+  mail-user: ${oc.env:SLURM_MAIL_USER,null}
+  mail-type: ALL

vital/config/hydra/launcher/beluga.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - alliancecan
+
+cpus_per_task: 10
+gpus_per_node: 1
+mem_gb: 46

vital/config/launcher/alliancecan.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+
+# NOTE: This custom launcher (separate from the built-in 'hydra/launcher' node) is meant to override global options
+# that depend on the hardware resources available and for which the default options are for local runs
+# (e.g. GPU devices, number of dataloader workers, progress bar display, etc.)
+
+defaults:
+  - override /hydra/launcher: alliancecan
+
+trainer:
+  enable_progress_bar: False
+
+# Path to the root of the virtualenv to activate for the jobs
+venv: ???
+
+# Path of the data to copy from the shared filesystem to the compute node
+data_dir_to_copy: ???

vital/config/launcher/beluga.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+
+defaults:
+  - alliancecan
+  - override /hydra/launcher: beluga
+
+trainer:
+  devices: 1
+
+data:
+  num_workers: 9

vital/config/vital_default.yaml

@@ -6,6 +6,7 @@ defaults:
       - model_checkpoint
   - logger: comet/online
   - _self_
+  - optional launcher: null # List custom launcher after trainer/task/data, so that it can override their configs
 
 seed: null