diff --git a/lib/cli/cli.rb b/lib/cli/cli.rb index 0425ee2..a759073 100644 --- a/lib/cli/cli.rb +++ b/lib/cli/cli.rb @@ -31,3 +31,4 @@ module CLI require_relative 'commands/gpu' require_relative 'commands/mpi' require_relative 'commands/array' +require_relative 'commands/modify' diff --git a/lib/cli/commands/modify.rb b/lib/cli/commands/modify.rb new file mode 100644 index 0000000..b85a85f --- /dev/null +++ b/lib/cli/commands/modify.rb @@ -0,0 +1,86 @@ +# frozen_string_literal: true + +require 'dry/cli' + +# Import subcommands like this +require_relative '../../services/modify_script' + +module AlcesJob + module CLI + module Commands + class Modify < Dry::CLI::Command + AlcesJob::CLI.register 'modify', self + desc 'This will modify a users script based on flags' + + argument :script, required: true, desc: 'The script to modify' + + option :job_name, type: :string, + desc: 'Sets the Slurm job name' + + option :nodes, type: :integer, + desc: 'Requests the number of compute nodes' + + option :ntasks, type: :integer, + desc: 'Specifies the total number of tasks' + + option :cpus_per_task, type: :integer, + desc: 'Specifies CPU cores per task' + + option :mem, type: :string, + desc: 'Sets the memory requirement for the job, e.g. 4G or 2000M' + + option :time, type: :string, + desc: 'Sets the job walltime limit, e.g. 02:00:00 or 1-00:00:00' + + option :partition, type: :string, + desc: 'Specifies the Slurm partition or queue to use' + + option :account, type: :string, + desc: 'Specifies the Slurm account to charge' + + option :gres, type: :string, + desc: 'Specifies generic resources such as GPUs, e.g. gpu:1' + + option :output, type: :string, + desc: 'Sets the Slurm stdout file path' + + option :error, type: :string, + desc: 'Sets the Slurm stderr file path' + + option :mail_user, type: :string, + desc: 'Sets the email address for Slurm notifications' + + option :mail_type, type: :string, + desc: 'Sets the Slurm mail notification type, e.g. BEGIN, END, FAIL' + + option :array, type: :string, + desc: 'Sets a Slurm array task specification' + + option :dependency, type: :string, + desc: 'Sets a Slurm dependency string' + + # option :modules, type: :array, default: [], + # desc: 'Loads one or more environment modules before running the job' + + # option :workdir, type: :string, + # desc: 'Changes to the specified working directory in the job script' + + option :command, type: :string, + desc: 'Specifies the shell command to execute in the script' + + # option :output_file, type: :string, + # desc: 'Writes the modified script to this output filename' + + # option :submit, type: :boolean, default: false, + # desc: 'Submits the script to Slurm automatically' + + def call(script:, **options) + AlcesJob::Services::ModifyScript.new( + script: script, + options: options + ).call + end + end + end + end +end diff --git a/lib/services/modify_script.rb b/lib/services/modify_script.rb new file mode 100644 index 0000000..14d3090 --- /dev/null +++ b/lib/services/modify_script.rb @@ -0,0 +1,142 @@ +# frozen_string_literal: true + +require_relative 'slurm_script_validator' + +module AlcesJob + module Services + class ModifyScript + def initialize(script:, options:) + @sbatch_options = { + job_name: 'job-name', + nodes: 'nodes', + ntasks: 'ntasks', + cpus_per_task: 'cpus-per-task', + mem: 'mem', + time: 'time', + partition: 'partition', + account: 'account', + gres: 'gres', + output: 'output', + error: 'error', + mail_user: 'mail-user', + mail_type: 'mail-type', + array: 'array', + dependency: 'dependency' + }.freeze + + @script = File.expand_path(script, Dir.pwd) + @options = options + return unless @options[:args]&.any? + + warn "ERROR: Unexpected arguments: #{@options[:args].join(' ')}" + warn 'Wrap the command in quotes, e.g. --command="python script.py"' + exit 1 + end + + def find_existing_job_name(lines) + job_line = lines.find { |line| line.start_with?('#SBATCH --job-name=') } + return nil unless job_line + + job_line.split('=', 2).last + end + + # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity + + def call + unless File.exist?(@script) + puts "Script not found: #{@script}" + return + end + + old_content = File.read(@script) + + lines = File.readlines(@script, chomp: true) + + edited_script = [] + + found_options = [] + + lines.each do |line| + if line.start_with?('#!') + edited_script << line + + elsif line.start_with?('#SBATCH') + parts = line.split[1] + + unless parts&.start_with?('--') && parts.include?('=') + edited_script << line + next + end + + name, _old_value = parts.split('=', 2) + + option_key = name.tr('-', '_').delete_prefix('__').to_sym + found_options << option_key + + if @options.key?(option_key) && !@options[option_key].nil? + new_value = @options[option_key] + edited_script << "#SBATCH #{name}=#{new_value}" + + else + edited_script << line + end + end + end + + puts found_options + + puts @options + + @options.each do |key, value| + next if found_options.include?(key) + next unless @sbatch_options.key?(key) + next if value.nil? + next if value == false + next if value.respond_to?(:empty?) && value.empty? + + sbatch_name = @sbatch_options[key] + edited_script << "#SBATCH --#{sbatch_name}=#{value}" + end + + job_name = @options[:job_name] || find_existing_job_name(lines) || 'slurm_job' + + if @options[:command] + puts 'hello' + edited_script << '' + edited_script << %(echo "Running job '#{job_name}'") if job_name + edited_script << '' + edited_script << @options[:command] + else + lines.each do |line| + edited_script << line if !line.start_with?('#!') && !line.start_with?('#SBATCH') + end + end + + puts edited_script + + File.write(@script, "#{edited_script.join("\n")}\n") + + validator = SlurmScriptValidator.new(@script) + + if validator.validate? + + puts 'Script updated successfully.' + + else + File.write(@script, old_content) + + puts 'Changes were invalid, so the script was reverted.' + + validator.errors.each do |error| + puts "ERROR: #{error}" + end + + end + validator.warnings.each do |warning| + puts "WARNING: #{warning}" + end + end + # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity + end + end +end