get all available accounts for a user from the scheduler

lib/ood_core.rb

@@ -12,6 +12,7 @@ module Job
     require "ood_core/job/script"
     require "ood_core/job/info"
     require "ood_core/job/cluster_info"
+    require "ood_core/job/account_info"
     require "ood_core/job/queue_info"
     require "ood_core/job/status"
     require "ood_core/job/adapter"

lib/ood_core/job/account_info.rb

@@ -0,0 +1,36 @@
+module OodCore
+  module Job
+
+    class AccountInfo
+
+      # The name of the account.
+      attr_reader :name
+      alias to_s name
+
+      # The QoS values this account can use.
+      attr_reader :qos
+
+      # The cluster this account is associated with.
+      attr_reader :cluster
+
+      # The queue this account can use. nil means there is no queue info
+      # for this account.
+      attr_reader :queue
+
+      def initialize(**opts)
+        orig_name = opts.fetch(:name, 'unknown')
+        @name = OodCore::Job::Adapters::Helper.upcase_accounts? ? orig_name.upcase : orig_name
+        @qos = opts.fetch(:qos, [])
+        @cluster = opts.fetch(:cluster, nil)
+        @queue = opts.fetch(:queue, nil)
+      end
+
+      def to_h
+        instance_variables.map do |var|
+          name = var.to_s.gsub('@', '').to_sym
+          [name, send(name)]
+        end.to_h
+      end
+    end
+  end
+end

lib/ood_core/job/adapter.rb

@@ -198,8 +198,17 @@ def job_name_illegal_chars
         ENV["OOD_JOB_NAME_ILLEGAL_CHARS"].to_s
       end
 
+      # Retrieve the accounts available to use  for the current user.
+      #
+      # Subclasses that do not implement this will return empty arrays.
+      # @return [Array<AccountInfo>] the accounts available to the user.
+      def accounts
+        []
+      end
+
       # Return the list of queues for this scheduler.
       #
+      # Subclasses that do not implement this will return empty arrays.
       # @return [Array<QueueInfo>]
       def queues
         []

lib/ood_core/job/adapters/helper.rb

@@ -31,6 +31,17 @@ def self.ssh_wrap(submit_host, cmd, cmd_args, strict_host_checking = true, env =
 
           return 'ssh', args + [cmd] + cmd_args
         end
+
+        # Determine whether to upcase account strings when returning adapter#accounts
+        def self.upcase_accounts?
+          env_var = ENV['OOD_UPCASE_ACCOUNTS']
+
+          if env_var.nil? || env_var.to_s.downcase == 'false'
+            false
+          else
+            true
+          end
+        end
       end
     end
   end

lib/ood_core/job/adapters/slurm.rb

@@ -1,4 +1,5 @@
 require "time"
+require 'etc'
 require "ood_core/refinements/hash_extensions"
 require "ood_core/refinements/array_extensions"
 require "ood_core/job/adapters/helper"
@@ -178,6 +179,27 @@ def get_jobs(id: "", owner: nil, attrs: nil)
             return [{ id: id, state: 'undetermined' }]
           end
 
+          def accounts
+            user = Etc.getlogin
+            args = ['-nP', 'show', 'users', 'withassoc', 'format=account,cluster,partition,qos', 'where', "user=#{user}"]
+
+            [].tap do |accts|
+              call('sacctmgr', *args).each_line do |line|
+                acct, cluster, queue, qos = line.split('|')
+                next if acct.nil?
+
+                args = {
+                  name: acct,
+                  qos: qos.to_s.chomp.split(','),
+                  cluster: cluster,
+                  queue: queue.empty? ? nil : queue
+                }
+                info = OodCore::Job::AccountInfo.new(**args) unless acct.nil?
+                accts << info unless acct.nil?
+              end
+            end
+          end
+
           def squeue_fields(attrs)
             if attrs.nil?
               all_squeue_fields
@@ -355,7 +377,7 @@ def call(cmd, *args, env: {}, stdin: "")
               cmd = OodCore::Job::Adapters::Helper.bin_path(cmd, bin, bin_overrides)
 
               args  = args.map(&:to_s)
-              args.concat ["-M", cluster] if cluster
+              args.concat ["-M", cluster] if cluster && cmd != 'sacctmgr'
 
               env = env.to_h
               env["SLURM_CONF"] = conf.to_s if conf
@@ -513,6 +535,13 @@ def cluster_info
           @slurm.get_cluster_info
         end
 
+        # Retrieve the accounts available to use  for the current user.
+        #
+        # @return [Array<String>] the accounts available to the user.
+        def accounts
+          @slurm.accounts
+        end
+
         # Retrieve info for all jobs from the resource manager
         # @raise [JobAdapterError] if something goes wrong getting job info
         # @return [Array<Info>] information describing submitted jobs

spec/fixtures/output/slurm/sacctmgr_show_accts.txt

@@ -0,0 +1,24 @@
+pzs0715|ascend|partition_a|ascend-default
+pzs0714|ascend|partition_b|ascend-default
+pzs1124|owens||owens-default,staff,phoenix,geophys,hal,gpt
+pzs1118|owens||owens-default
+pzs1117|owens||owens-default
+pzs1010|owens||owens-default
+pzs0715|owens||owens-default
+pzs0714|owens||owens-default
+pde0006|owens||owens-default
+pas2051|owens||owens-default
+pas1871|owens||owens-default
+pas1754|owens||owens-default
+pas1604|owens||owens-default
+pzs1124|pitzer||pitzer-default
+pzs1118|pitzer||pitzer-default
+pzs1117|pitzer||pitzer-default
+pzs1010|pitzer||pitzer-default
+pzs0715|pitzer||pitzer-default
+pzs0714|pitzer||pitzer-default
+pde0006|pitzer||pitzer-default
+pas2051|pitzer||pitzer-default
+pas1871|pitzer||pitzer-default
+pas1754|pitzer||pitzer-default
+pas1604|pitzer||pitzer-default

spec/job/adapter_spec.rb

@@ -209,8 +209,14 @@
     end
   end
 
+  describe '#accounts' do
+    it 'returns an empty array by default' do
+      expect(subject.accounts).to eq([])
+    end
+  end
+
   describe '#queues' do
-    it 'returns an empty array' do
+    it 'returns an empty array by default' do
       expect(adapter.queues).to eq([])
     end
   end

spec/job/adapters/slurm_spec.rb

@@ -1227,6 +1227,84 @@ def job_info(opts = {})
     end
   end
 
+  describe '#accounts' do
+    context 'when sacctmgr returns successfully' do
+      let(:slurm) { OodCore::Job::Adapters::Slurm::Batch.new }
+      let(:expected_accounts) {["pzs0715", "pzs0714", "pzs1124", "pzs1118", "pzs1117", "pzs1010", "pde0006", "pas2051", "pas1871", "pas1754", "pas1604"]}
+
+      it 'returns the correct accounts names' do
+        allow(Etc).to receive(:getlogin).and_return('me')
+        allow(Open3).to receive(:capture3)
+                          .with({}, 'sacctmgr', '-nP', 'show', 'users', 'withassoc', 'format=account,cluster,partition,qos', 'where', 'user=me', {stdin_data: ''})
+                          .and_return([File.read('spec/fixtures/output/slurm/sacctmgr_show_accts.txt'), '',  double("success?" => true)])
+
+        expect(subject.accounts.map(&:to_s).uniq).to eq(expected_accounts)
+      end
+
+      # TODO test for qos & cluster once the API solidifies
+      it 'parses qos correctly' do
+        allow(Etc).to receive(:getlogin).and_return('me')
+        allow(Open3).to receive(:capture3)
+                          .with({}, 'sacctmgr', '-nP', 'show', 'users', 'withassoc', 'format=account,cluster,partition,qos', 'where', 'user=me', {stdin_data: ''})
+                          .and_return([File.read('spec/fixtures/output/slurm/sacctmgr_show_accts.txt'), '',  double("success?" => true)])
+
+        accts = subject.accounts
+        acct_w_qos = accts.select { |a| a.name == 'pzs1124' && a.cluster == 'owens' }.first
+        expect(acct_w_qos.qos).to eq(['owens-default', 'staff', 'phoenix', 'geophys', 'hal', 'gpt'])
+
+        other_accts = accts - [acct_w_qos]
+        other_accts.each do |acct|
+          expect(acct.qos).to eq(["#{acct.cluster}-default"])
+        end
+      end
+
+      it 'parses partition correctly' do
+        allow(Etc).to receive(:getlogin).and_return('me')
+        allow(Open3).to receive(:capture3)
+                          .with({}, 'sacctmgr', '-nP', 'show', 'users', 'withassoc', 'format=account,cluster,partition,qos', 'where', 'user=me', {stdin_data: ''})
+                          .and_return([File.read('spec/fixtures/output/slurm/sacctmgr_show_accts.txt'), '',  double("success?" => true)])
+
+        accts = subject.accounts
+        acct_w_partitions = accts.select { |a| a.cluster == 'ascend' }
+        acct_w_no_partitions = accts.select { |a| a.queue.nil? }
+
+        expect(acct_w_partitions.size).to eq(2)
+        expect(accts - acct_w_no_partitions).to eq(acct_w_partitions)
+        expect(acct_w_partitions.select {|a| a.name == 'pzs0715'}.first.queue).to eq('partition_a')
+        expect(acct_w_partitions.select {|a| a.name == 'pzs0714'}.first.queue).to eq('partition_b')
+      end
+    end
+
+    context 'when sacctmgr fails' do
+      let(:slurm) { OodCore::Job::Adapters::Slurm::Batch.new }
+
+      it 'raises the error' do
+        allow(Etc).to receive(:getlogin).and_return('me')
+        allow(Open3).to receive(:capture3)
+                          .with({}, 'sacctmgr', '-nP', 'show', 'users', 'withassoc', 'format=account,cluster,partition,qos', 'where', 'user=me', {stdin_data: ''})
+                          .and_return(['', 'the error message',  double("success?" => false)])
+
+        expect { subject.accounts }.to raise_error(OodCore::Job::Adapters::Slurm::Batch::Error, 'the error message')
+      end
+    end
+
+    context 'when OOD_UPCASE_ACCOUNTS is set' do
+      let(:slurm) { OodCore::Job::Adapters::Slurm::Batch.new }
+      let(:expected_accounts) {["PZS0715", "PZS0714", "PZS1124", "PZS1118", "PZS1117", "PZS1010", "PDE0006", "PAS2051", "PAS1871", "PAS1754", "PAS1604"]}
+
+      it 'returns the correct accounts' do
+        allow(Etc).to receive(:getlogin).and_return('me')
+        allow(Open3).to receive(:capture3)
+                          .with({}, 'sacctmgr', '-nP', 'show', 'users', 'withassoc', 'format=account,cluster,partition,qos', 'where', 'user=me', {stdin_data: ''})
+                          .and_return([File.read('spec/fixtures/output/slurm/sacctmgr_show_accts.txt'), '',  double("success?" => true)])
+
+        with_modified_env({ OOD_UPCASE_ACCOUNTS: 'true'}) do
+          expect(subject.accounts.map(&:to_s).uniq).to eq(expected_accounts)
+        end
+      end
+    end
+  end
+
   describe '#queues' do
     context 'when scontrol returns successfully' do
       let(:slurm) { OodCore::Job::Adapters::Slurm::Batch.new }