diff --git a/LAUNCH.md b/LAUNCH.md index aaadd7c2f..2d4b2358f 100644 --- a/LAUNCH.md +++ b/LAUNCH.md @@ -7,12 +7,12 @@ In the following, `$root` refers to the root of the current repository. ```sh usage: launch.py [-h] [--help-md] [--job_name JOB_NAME] [--outdir OUTDIR] [--cpus_per_task CPUS_PER_TASK] [--mem MEM] [--gres GRES] - [--partition PARTITION] [--modules MODULES] + [--partition PARTITION] [--time TIME] [--modules MODULES] [--conda_env CONDA_ENV] [--venv VENV] [--template TEMPLATE] [--code_dir CODE_DIR] [--git_checkout GIT_CHECKOUT] [--jobs JOBS] [--dry-run] [--verbose] [--force] -optional arguments: +options: -h, --help show this help message and exit --help-md Show an extended help message as markdown. Can be useful to overwrite LAUNCH.md with `$ python @@ -26,6 +26,9 @@ optional arguments: --gres GRES gres per node (e.g. gpu:1). Defaults to gpu:1 --partition PARTITION slurm partition to use for the job. Defaults to long + --time TIME wall clock time limit (e.g. 2-12:00:00). See: + https://slurm.schedmd.com/sbatch.html#OPT_time + Defaults to None --modules MODULES string after 'module load'. Defaults to anaconda/3 cuda/11.3 --conda_env CONDA_ENV @@ -69,6 +72,7 @@ modules : anaconda/3 cuda/11.3 outdir : $SCRATCH/gflownet/logs/slurm partition : long template : $root/mila/sbatch/template-conda.sh +time : None venv : None verbose : False ``` diff --git a/mila/launch.py b/mila/launch.py index dce36db8d..5e2f5379f 100644 --- a/mila/launch.py +++ b/mila/launch.py @@ -406,6 +406,7 @@ def code_dir_for_slurm_tmp_dir_checkout(git_checkout): "outdir": "$SCRATCH/gflownet/logs/slurm", "partition": "long", "template": "$root/mila/sbatch/template-conda.sh", + "time": None, "venv": None, "verbose": False, } @@ -454,6 +455,13 @@ def code_dir_for_slurm_tmp_dir_checkout(git_checkout): help="slurm partition to use for the job." + f" Defaults to {defaults['partition']}", ) + parser.add_argument( + "--time", + type=str, + help="wall clock time limit (e.g. 2-12:00:00). " + + "See: https://slurm.schedmd.com/sbatch.html#OPT_time" + + f" Defaults to {defaults['time']}", + ) parser.add_argument( "--modules", type=str, diff --git a/mila/sbatch/template-conda.sh b/mila/sbatch/template-conda.sh index f2ff3de33..983728b78 100644 --- a/mila/sbatch/template-conda.sh +++ b/mila/sbatch/template-conda.sh @@ -5,10 +5,11 @@ #SBATCH --mem={mem} #SBATCH --gres={gres} #SBATCH --partition={partition} +#SBATCH --time={time} module load {modules} conda activate {conda_env} cd {code_dir} -python main.py {main_args} \ No newline at end of file +python main.py {main_args} diff --git a/mila/sbatch/template-venv.sh b/mila/sbatch/template-venv.sh index d719ce58d..66cf0d694 100644 --- a/mila/sbatch/template-venv.sh +++ b/mila/sbatch/template-venv.sh @@ -5,10 +5,11 @@ #SBATCH --mem={mem} #SBATCH --gres={gres} #SBATCH --partition={partition} +#SBATCH --time={time} module load {modules} source {venv}/bin/activate cd {code_dir} -python main.py {main_args} \ No newline at end of file +python main.py {main_args}