Added uses_gpu flag to Job Info class

lib/ood_core/job/adapters/slurm.rb

@@ -617,6 +617,12 @@ def get_state(st)
             STATE_MAP.fetch(st, :undetermined)
           end
 
+          def get_gpu_count(gres)
+            gpus = 0
+            gres.to_s.scan(/gpu:[^,]*(\d+)/).each { |val| gpus += val[0].to_i }
+            gpus
+          end
+
           # Parse hash describing Slurm job status
           def parse_job_info(v)
             allocated_nodes = parse_nodes(v[:node_list])
@@ -643,7 +649,8 @@ def parse_job_info(v)
               cpu_time: nil,
               submission_time: v[:submit_time] ? Time.parse(v[:submit_time]) : nil,
               dispatch_time: (v[:start_time].nil? || v[:start_time] == "N/A") ? nil : Time.parse(v[:start_time]),
-              native: v
+              native: v,
+              gpus: get_gpu_count(v[:gres])
             )
           end
 

lib/ood_core/job/info.rb

@@ -65,6 +65,10 @@ class Info
       # @return [Object] native info
       attr_reader :native
 
+      # Number of gpus allocated for job
+      # @return [Integer, nil] allocated total number of gpus
+      attr_reader :gpus
+
       # List of job array child task statuses
       # @note only relevant for job arrays
       # @return [Array<Task>] tasks
@@ -86,11 +90,12 @@ class Info
       # @param dispatch_time [#to_i, nil] dispatch time
       # @param tasks [Array<Hash>] tasks e.g. { id: '12345.owens-batch', status: :running }
       # @param native [Object] native info
+      # @param gpus [#to_i, 0] allocated total number of gpus
       def initialize(id:, status:, allocated_nodes: [], submit_host: nil,
                      job_name: nil, job_owner: nil, accounting_id: nil,
                      procs: nil, queue_name: nil, wallclock_time: nil,
                      wallclock_limit: nil, cpu_time: nil, submission_time: nil,
-                     dispatch_time: nil, native: nil, tasks: [],
+                     dispatch_time: nil, native: nil, gpus: 0, tasks: [],
                      **_)
         @id              = id.to_s
         @status          = Status.new(state: status.to_sym)
@@ -111,6 +116,7 @@ def initialize(id:, status:, allocated_nodes: [], submit_host: nil,
         @status = job_array_aggregate_status unless @tasks.empty?
 
         @native          = native
+        @gpus            = gpus            && gpus.to_i
       end
 
       # Create a new Info for a child task
@@ -147,10 +153,15 @@ def to_h
           submission_time: submission_time,
           dispatch_time:   dispatch_time,
           native:          native,
+          gpus:            gpus,
           tasks: tasks
         }
       end
 
+      def gpu?
+        gpus.positive?
+      end
+
       # The comparison operator
       # @param other [#to_h] object to compare against
       # @return [Boolean] whether objects are equivalent

spec/job/adapters/slurm_spec.rb

@@ -1195,4 +1195,29 @@ def job_info(opts = {})
       end
     end
   end
+
+  describe "#get_gpu_count" do
+    batch = OodCore::Job::Adapters::Slurm::Batch.new(cluster: "owens.osc.edu", conf: "/etc/slurm/conf/", bin: nil, bin_overrides: {}, submit_host: "owens.osc.edu", strict_host_checking: false)
+    adapter = OodCore::Job::Adapters::Slurm.new(slurm: batch)
+
+    context "when called" do
+      gres_cases = [
+        [nil, 0],
+        ["", 0],
+        ["gpu:v100-32g:2", 2],
+        ["gpu:v100-32g:2,pfsdir:1", 2],
+        ["third-thing:sub-thing:17,gpu:v100-32g:2,pfsdir:1", 2],
+        ["third-thing:sub-thing:17,pfsdir:1,gpu:v100-32g:2", 2],
+        ["gpu:v30-12g:2,gpu:v31-32g:1", 3],
+        ["gres:gpu:1", 1],
+        ["gres:pfsdir:ess", 0]
+      ]
+      gres_cases.each do |gc| 
+        it "does not return the correct number of gpus when gres=\"#{gc[0]}\"" do
+          gpus = adapter.send(:get_gpu_count, gc[0])
+          expect(gpus).to be(gc[1]);
+        end
+      end
+    end
+  end
 end

spec/job/info_spec.rb

@@ -32,6 +32,7 @@ def build_info(opts = {})
   it { is_expected.to respond_to(:submission_time) }
   it { is_expected.to respond_to(:dispatch_time) }
   it { is_expected.to respond_to(:native) }
+  it { is_expected.to respond_to(:gpus) }
   it { is_expected.to respond_to(:to_h) }
   it { is_expected.to respond_to(:tasks) }
 
@@ -151,6 +152,12 @@ def build_info(opts = {})
     it { is_expected.to eq("native") }
   end
 
+  describe "#gpus" do
+    subject { build_info(native: "gpus").native }
+
+    it { is_expected.to eq("gpus") }
+  end
+
   describe "#to_h" do
     subject { build_info.to_h }
 
@@ -170,6 +177,7 @@ def build_info(opts = {})
     it { is_expected.to have_key(:submission_time) }
     it { is_expected.to have_key(:dispatch_time) }
     it { is_expected.to have_key(:native) }
+    it { is_expected.to have_key(:gpus) }
   end
 
   describe "#==" do