From 4f6d51e7a2efd0e197012cee41609af6c8d31df4 Mon Sep 17 00:00:00 2001
From: Jeff Ohrstrom <johrstrom@osc.edu>
Date: Wed, 31 Jan 2024 13:13:45 -0500
Subject: [PATCH 1/2] add gpu awareness to queue_info

---
 lib/ood_core/job/adapters/slurm.rb | 5 +++--
 lib/ood_core/job/queue_info.rb     | 8 ++++++++
 spec/job/adapters/slurm_spec.rb    | 5 +++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/lib/ood_core/job/adapters/slurm.rb b/lib/ood_core/job/adapters/slurm.rb
index 80d64bbdb..b238f9c8a 100644
--- a/lib/ood_core/job/adapters/slurm.rb
+++ b/lib/ood_core/job/adapters/slurm.rb
@@ -327,13 +327,13 @@ def queues
 
             [].tap do |ret_arr|
               info_raw.each_line do |line|
-                ret_arr << str_to_acct_info(line)
+                ret_arr << str_to_queue_info(line)
               end
             end
           end
 
           private
-            def str_to_acct_info(line)
+            def str_to_queue_info(line)
               hsh = line.split(' ').map do |token|
                 m = token.match(/^(?<key>\w+)=(?<value>.+)$/)
                 [m[:key], m[:value]]
@@ -349,6 +349,7 @@ def str_to_acct_info(line)
 
 
               hsh[:deny_accounts] = hsh[:DenyAccounts].nil? ? [] : hsh[:DenyAccounts].to_s.split(',')
+              hsh[:tres] = hsh[:TRES].nil? ? {} : hsh[:TRES].to_s.split(',').map { |str| str.split('=') }.to_h
 
               OodCore::Job::QueueInfo.new(**hsh)
             end
diff --git a/lib/ood_core/job/queue_info.rb b/lib/ood_core/job/queue_info.rb
index 7c20a2629..de1c910a7 100644
--- a/lib/ood_core/job/queue_info.rb
+++ b/lib/ood_core/job/queue_info.rb
@@ -20,9 +20,13 @@ class OodCore::Job::QueueInfo
   # The accounts that are not allowed to use this queue.
   attr_reader :deny_accounts
 
+  # An Hash of Trackable Resources and their values.
+  attr_reader :tres
+
   def initialize(**opts)
     @name = opts.fetch(:name, 'unknown')
     @qos = opts.fetch(:qos, [])
+    @tres = opts.fetch(:tres, {})
 
     allow_accounts = opts.fetch(:allow_accounts, nil)
     @allow_accounts = if allow_accounts.nil?
@@ -42,4 +46,8 @@ def to_h
       [name, send(name)]
     end.to_h
   end
+
+  def gpu?
+    tres.keys.any? { |name| name.to_s.match?(/gpu/i) }
+  end
 end
diff --git a/spec/job/adapters/slurm_spec.rb b/spec/job/adapters/slurm_spec.rb
index 4417b15c0..e1923df9b 100644
--- a/spec/job/adapters/slurm_spec.rb
+++ b/spec/job/adapters/slurm_spec.rb
@@ -1333,11 +1333,16 @@ def job_info(opts = {})
         expect(systems_queue.allow_accounts).to eq(['root', 'pzs0708', 'pzs0710', 'pzs0722'])
         expect(systems_queue.deny_accounts).to eq([])
         expect(systems_queue.qos).to eq([])
+        expect(systems_queue.gpu?).to eq(true)
 
         quick_queue = queues.select { |q| q.name == 'quick' }.first
         expect(quick_queue.allow_accounts).to eq(nil)
         expect(quick_queue.deny_accounts).to eq(quick_deny_accounts)
         expect(quick_queue.qos).to eq(['quick'])
+        expect(quick_queue.gpu?).to eq(false)
+
+        gpu_queue = queues.select { |q| q.name == 'gpuserial' }.first
+        expect(gpu_queue.gpu?).to eq(true)
       end
     end
 

From b71acb6f3ffe1bb0c725255b4151c89aa8c797d1 Mon Sep 17 00:00:00 2001
From: Jeff Ohrstrom <johrstrom@osc.edu>
Date: Thu, 1 Feb 2024 09:52:14 -0500
Subject: [PATCH 2/2] strengthen this regex

---
 lib/ood_core/job/queue_info.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/ood_core/job/queue_info.rb b/lib/ood_core/job/queue_info.rb
index de1c910a7..5578ea5d4 100644
--- a/lib/ood_core/job/queue_info.rb
+++ b/lib/ood_core/job/queue_info.rb
@@ -48,6 +48,6 @@ def to_h
   end
 
   def gpu?
-    tres.keys.any? { |name| name.to_s.match?(/gpu/i) }
+    tres.keys.any? { |name| name.to_s.match?(%r{^gres/gpu($|:)}i) }
   end
 end