fix GPU regex (#837)

OSC · May 21, 2024 · 6c34e9f · 6c34e9f
1 parent 4b21ae7
commit 6c34e9f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/lib/ood_core/job/adapters/slurm.rb b/lib/ood_core/job/adapters/slurm.rb
@@ -41,7 +41,7 @@ class Slurm < Adapter
         # calculated from gres string
         # @return [Integer] the number of gpus in gres
         def self.gpus_from_gres(gres)
-          gres.to_s.scan(/gpu:[^,]*(\d+)/).flatten.map(&:to_i).sum
+          gres.to_s.scan(/gpu[^(,]*[:=](\d+)/).flatten.map(&:to_i).sum
         end
 
         # Object used for simplified communication with a Slurm batch server

diff --git a/spec/job/adapters/slurm_spec.rb b/spec/job/adapters/slurm_spec.rb
@@ -1216,7 +1216,11 @@ def job_info(opts = {})
         ["gres/gpu:v100-32g=2", 2],
         ["gres/gpu:v100-32g=2,gres/gpu:v100-32g=4", 6],
         ["gres/gpu:v100-32g=2,gres:gpu:1,gres/gpu:v100-32g=4", 7],
-        ["gres/gpu:v100-32g=2,gres:pfsdir:1", 2]
+        ["gres/gpu:v100-32g=2,gres:pfsdir:1", 2],
+        ["gpu:p100:1,nsight:no_consume:1", 1],
+        ["gpu:p100:1(IDX:0),mps:0", 1],
+        ["gpu:a100:4(S:0-15)", 4],
+        ["gpu:a100:3(IDX:0,2-3),mps:0", 3],
       ]
       gres_cases.each do |gc| 
         it "does not return the correct number of gpus when gres=\"#{gc[0]}\"" do