add ability to specify gpus (#118)

OSC · Dec 26, 2024 · 92ca8dd · 92ca8dd
1 parent 4d3f98f
commit 92ca8dd
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/form.yml b/form.yml
@@ -12,6 +12,7 @@ form:
   - mode
   - bc_num_hours
   - node_type
+  - gpus
   - cuda_version
   - num_cores
   - auto_modules_app_jupyter
@@ -182,6 +183,7 @@ attributes:
           "any",     "any",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-owens: 1,
           data-max-num-cores-for-cluster-owens: 28,
           data-min-num-cores-for-cluster-pitzer: 1,
@@ -198,6 +200,7 @@ attributes:
           "40 core",     "any-40core",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-pitzer: 1,
           data-max-num-cores-for-cluster-pitzer: 40,
           data-option-for-cluster-owens: false,
@@ -211,6 +214,7 @@ attributes:
           "48 core",     "any-48core",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-pitzer: 1,
           data-max-num-cores-for-cluster-pitzer: 48,
           data-option-for-cluster-owens: false,
@@ -260,6 +264,7 @@ attributes:
           "largemem", "largemem",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-pitzer: 24,
           data-max-num-cores-for-cluster-pitzer: 48,
           data-option-for-cluster-owens: false,
@@ -273,6 +278,7 @@ attributes:
           "hugemem", "hugemem",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-owens: 4,
           data-max-num-cores-for-cluster-owens: 48,
           data-min-num-cores-for-cluster-pitzer: 20,
@@ -287,6 +293,7 @@ attributes:
           "debug",   "debug",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-owens: 1,
           data-max-num-cores-for-cluster-owens: 28,
           data-min-num-cores-for-cluster-pitzer: 1,
@@ -301,6 +308,7 @@ attributes:
           "pitzer",   "pitzer",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-kubernetes: 1,
           data-max-num-cores-for-cluster-kubernetes: 10,
           data-min-num-cores-for-cluster-kubernetes-test: 1,
@@ -316,6 +324,7 @@ attributes:
           "owens",   "owens",
           data-hide-cuda-version: true,
           data-set-cuda-version: '',
+          data-hide-gpus: true,
           data-min-num-cores-for-cluster-kubernetes: 1,
           data-max-num-cores-for-cluster-kubernetes: 10,
           data-min-num-cores-for-cluster-kubernetes-test: 1,
@@ -353,4 +362,8 @@ attributes:
           data-option-for-cluster-ascend: false,
           data-option-for-cluster-cardinal: false,
         ]
+  gpus:
+    widget: number_field
+    min: 0
+    max: 4
   classroom: 0
diff --git a/submit.yml.erb b/submit.yml.erb
@@ -40,8 +40,8 @@ batch_connect:
   conn_params:
     - jupyter_api
 script:
-<% if node_type =~ /gpu/ -%>
-  gpus_per_node: 1
+<% if node_type =~ /gpu/ && gpus.to_i.positive? -%>
+  gpus_per_node: <%= gpus %>
 <% end -%>
   native:
     <%- slurm_args.each do |arg| %>
@@ -70,7 +70,7 @@ script:
 ---
 script:
   wall_time: "<%= bc_num_hours.to_i * 3600 %>"
-  <%- if node_type =~ /gpu/ -%>
+  <%- if node_type =~ /gpu/ && gpus.to_i.positive? -%>
   gpus_per_node: 1
   <%- end -%>
   native: