Feat/m5 testing

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'

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

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