From 005ca586f1bac372ea4a4286d5d5d96518300aa7 Mon Sep 17 00:00:00 2001 From: Jeff Ohrstrom Date: Wed, 18 Jan 2023 11:09:18 -0500 Subject: [PATCH] get all available accounts for a user from the scheduler (#783) Start the `accounts` API on the adapter interface. Return account info objects for higher layers to stitch together. Co-authored-by: treydock --- lib/ood_core.rb | 1 + lib/ood_core/job/account_info.rb | 36 +++++++++ lib/ood_core/job/adapter.rb | 9 +++ lib/ood_core/job/adapters/helper.rb | 11 +++ lib/ood_core/job/adapters/slurm.rb | 31 +++++++- .../output/slurm/sacctmgr_show_accts.txt | 24 ++++++ spec/job/adapter_spec.rb | 8 +- spec/job/adapters/slurm_spec.rb | 78 +++++++++++++++++++ 8 files changed, 196 insertions(+), 2 deletions(-) create mode 100644 lib/ood_core/job/account_info.rb create mode 100644 spec/fixtures/output/slurm/sacctmgr_show_accts.txt diff --git a/lib/ood_core.rb b/lib/ood_core.rb index 019af8120..5660d76be 100644 --- a/lib/ood_core.rb +++ b/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" diff --git a/lib/ood_core/job/account_info.rb b/lib/ood_core/job/account_info.rb new file mode 100644 index 000000000..a35caff70 --- /dev/null +++ b/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 diff --git a/lib/ood_core/job/adapter.rb b/lib/ood_core/job/adapter.rb index c0a371de2..418c5295f 100644 --- a/lib/ood_core/job/adapter.rb +++ b/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] 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] def queues [] diff --git a/lib/ood_core/job/adapters/helper.rb b/lib/ood_core/job/adapters/helper.rb index f3ce31e70..5019f0175 100644 --- a/lib/ood_core/job/adapters/helper.rb +++ b/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 diff --git a/lib/ood_core/job/adapters/slurm.rb b/lib/ood_core/job/adapters/slurm.rb index 23f23c92f..1c9ba7ea1 100644 --- a/lib/ood_core/job/adapters/slurm.rb +++ b/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] 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] information describing submitted jobs diff --git a/spec/fixtures/output/slurm/sacctmgr_show_accts.txt b/spec/fixtures/output/slurm/sacctmgr_show_accts.txt new file mode 100644 index 000000000..2a37c43ff --- /dev/null +++ b/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 \ No newline at end of file diff --git a/spec/job/adapter_spec.rb b/spec/job/adapter_spec.rb index f08883e7c..79ee6ff2b 100644 --- a/spec/job/adapter_spec.rb +++ b/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 diff --git a/spec/job/adapters/slurm_spec.rb b/spec/job/adapters/slurm_spec.rb index 29937f267..af04a3eef 100644 --- a/spec/job/adapters/slurm_spec.rb +++ b/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 }