add GPU

laszewsk · Oct 5, 2023 · 5fe3eca · 5fe3eca
1 parent 8c12f97
commit 5fe3eca
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/benchmarks/cloudmask/target/greene_v0.5/config-simple-rivanna.yaml b/benchmarks/cloudmask/target/greene_v0.5/config-simple-rivanna.yaml
@@ -40,7 +40,7 @@ submission:
 
 experiment:
 #  card_name: v100
-  card_name: v100
+  card_name: a100
   gpu_count: 1
   cpu_num: 1
   mem: "64GB"

diff --git a/benchmarks/cloudmask/target/greene_v0.5/experiment.py b/benchmarks/cloudmask/target/greene_v0.5/experiment.py
@@ -11,6 +11,7 @@
   --config=<config>       Specify the path to the config file [default: config-simple.yaml].
   --clean                 Clean up generated files.
   --out=<output>          Specify the output filename. If not provided, output will be written to stdout.
+  --gpu=<gpu>             the gpu
 """
 
 import os
@@ -26,6 +27,8 @@
 CONFIG = arguments['--config'] or "config-simple.yaml"
 CLEAN = arguments['--clean']
 OUTPUT = arguments['--out']
+GPU = arguments['--gpu'] or "a100"
+
 
 pprint(arguments)
 
@@ -134,6 +137,7 @@ def _print(content):
         replace_in_yaml(CONFIG_YAML, 'experiemnet.gpu_count', 1)
 
         # Modify the slurm script
+        replace_sbatch_in_slurm(SLURM_SCRIPT, 'gres', f"gpu:{GPU}:1")
         replace_sbatch_in_slurm(SLURM_SCRIPT, 'job-name', f"cloudmask-gpu-greene-epoch-{EXPERIMENT_ID}")
         replace_sbatch_in_slurm(SLURM_SCRIPT, 'time', timesArray[j])