Merge pull request #244 from alexhernandezgarcia/launch-time

Add --time to launch.py arguments
alexhernandezgarcia · Nov 13, 2023 · 0440190 · 0440190
2 parents 96f1eaf + 75413bb
commit 0440190
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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}
+python main.py {main_args}
diff --git 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}
+python main.py {main_args}