Module: Marenostrum::SLURM

Defined in:

lib/rbbt/hpc.rb

Class Method Summary

• .follow_job(slurm_basedir, tail = true) ⇒ Object
• .issue_template(template, options = {}) ⇒ Object
• .run_job(job, options = {}) ⇒ Object
• .template(args, options = {}) ⇒ Object
• .wait_for_job(slurm_basedir, time = 1) ⇒ Object

Class Method Details

.follow_job(slurm_basedir, tail = true) ⇒ Object

# File 'lib/rbbt/hpc.rb', line 412

def self.follow_job(slurm_basedir, tail = true)
  fjob = File.join(slurm_basedir, 'job.id')
  fout = File.join(slurm_basedir, 'std.out')
  ferr = File.join(slurm_basedir, 'std.err')
  fstatus = File.join(slurm_basedir, 'job.status')

  job = Open.read(fjob).strip if Open.exists?(fjob)

  if job
    status_txt = CMD.cmd("squeue --job #{job}").read
    STDERR.puts Log.color(:magenta, "Status [#{job.to_i}]:")
    STDERR.puts status_txt
    lines = status_txt.split("\n").length
  end

  if tail
    Log.severity = 10
    while ! File.exists? fout
      if job
        STDERR.puts
        Log.clear_line(STDERR)
        STDERR.write Log.color(:magenta, "Waiting for Output")
        3.times do
          STDERR.write Log.color(:magenta, ".")
          sleep 1
        end
        status_txt = CMD.cmd("squeue --job #{job}").read
        lines.times do
          Log.clear_line(STDERR)
        end
        Log.clear_line(STDERR)
        STDERR.puts Log.color(:magenta, "Status [#{job.to_i}]:")
        STDERR.puts status_txt
        lines = status_txt.split("\n").length
      end
    end
    STDERR.puts
    Log.clear_line(STDERR)
    STDERR.puts Log.color(:magenta, "Output:")
    begin
      CMD.cmd("squeue --job #{job} > #{fstatus}")
      out = CMD.cmd("tail -f '#{fout}'", :pipe => true) if File.exists?(fout) and not tail == :STDERR
      err = CMD.cmd("tail -f '#{ferr}'", :pipe => true) if File.exists?(ferr)

      terr = Misc.consume_stream(err, true, STDERR) if err
      tout = Misc.consume_stream(out, true, STDOUT) if out

      sleep 3 while CMD.cmd("squeue --job #{job}").read.include? job.to_s
    rescue Aborted
    ensure
      begin
        terr.exit if terr
        tout.exit if tout
        err.close if err
        err.join if err
      rescue Exception
      end

      begin
        out.close if out
        out.join if out
      rescue Exception
      end
    end
  end
end

.issue_template(template, options = {}) ⇒ Object

# File 'lib/rbbt/hpc.rb', line 365

def self.issue_template(template, options = {})

  slurm_basedir = options[:slurm_basedir]
  Open.mkdir slurm_basedir

  dry_run = options.delete :dry_run

  fout  = File.join(slurm_basedir, 'std.out')
  ferr  = File.join(slurm_basedir, 'std.err')
  fjob  = File.join(slurm_basedir, 'job.id')
  fexit = File.join(slurm_basedir, 'exit.status')
  fsync = File.join(slurm_basedir, 'sync.log')
  fcmd  = File.join(slurm_basedir, 'command.slurm')

  job = nil
  if options[:clean_job]
    [fcmd, fjob, fout, ferr, fsync, fexit].each do |file|
      Open.rm file if Open.exists? file
    end
  end

  return if Open.exists?(fexit)

  STDERR.puts Log.color(:magenta, "Issuing SLURM file: #{fcmd}")
  STDERR.puts template

  Open.write(fcmd, template) unless File.exists? fcmd
  if File.exists?(fjob)
    job = Open.read(fjob).to_i
  else
    if File.exists?(fout)
      return
    elsif dry_run
      STDERR.puts Log.color(:magenta, "To execute run: ") + Log.color(:blue, "sbatch '#{slurm_basedir}/command.slurm'")
      STDERR.puts Log.color(:magenta, "To monitor progress run (needs local rbbt): ") + Log.color(:blue, "rbbt mn --tail -w '#{slurm_basedir}'")
      raise Marenostrum::SBATCH, slurm_basedir
    else
      Open.rm fsync
      Open.rm fexit
      Open.rm fout
      Open.rm ferr
      job = CMD.cmd("sbatch '#{fcmd}'").read.scan(/\d+/).first.to_i
      Open.write(fjob, job.to_s)
    end
  end
end

.run_job(job, options = {}) ⇒ Object

# File 'lib/rbbt/hpc.rb', line 490

def self.run_job(job, options = {})
  options = IndiferentHash.setup(options.dup)
    
  dry_run          = options.delete :dry_run
  tail             = options.delete :tail

  workflow = job.workflow
  task = job.task_name

  keep_slurm_basedir = options.delete :keep_SLURM_slurm_basedir 
  slurm_basedir = options.delete :SLURM_basedir
  slurm_basedir = "~/rbbt-slurm" if slurm_basedir.nil?
  TmpFile.with_file(nil, !keep_slurm_basedir, :tmpdir => slurm_basedir, :prefix => "SLURM_rbbt_job-") do |tmp_directory|
    options[:slurm_basedir] ||= tmp_directory
    slurm_basedir = options[:slurm_basedir]
    inputs_dir = File.join(tmp_directory, 'inputs_dir')
    saved = Step.save_job_inputs(job, inputs_dir, options)
    if saved
      options[:inputs_dir] = inputs_dir
      cmd = ['workflow', 'task', workflow.to_s, task.to_s, '-pf', '--load_inputs', inputs_dir, '--log', (options[:log] || Log.severity).to_s]
    else
      cmd = ['workflow', 'task', workflow.to_s, task.to_s, '-pf', '--log', (options[:log] || Log.severity).to_s]
    end


    template = self.template(cmd, options)
    self.issue_template(template, options.merge(:slurm_basedir => slurm_basedir, :dry_run => dry_run))

    return unless tail

    t_monitor = Thread.new do
      self.follow_job(slurm_basedir, :STDERR)
    end
    self.wait_for_job(slurm_basedir)
    t_monitor.raise Aborted
    return unless Open.read(File.join(slurm_basedir, 'exit.status')).strip == '0'
    path = Open.read(File.join(slurm_basedir, 'std.out')).strip
    if Open.exists?(path) && job.path != path
      Log.info "Path of SLURM job #{path} is different from original job #{job.path}. Stablishing link."
      Open.ln path, job.path
      Open.ln path + '.info', job.path + '.info'  if Open.exists?(path + '.info')
      Open.ln path + '.files', job.path + '.files' if Open.exists?(path + '.files')
    end
  end
end

.template(args, options = {}) ⇒ Object

# File 'lib/rbbt/hpc.rb', line 15

def self.template(args, options = {})

  development      = options.delete :drbbt
  singularity      = options.delete :singularity
  contain          = options.delete :contain
  sync             = options.delete :sync
  user_group       = options.delete :user_group
  contain_and_sync = options.delete :contain_and_sync
  wipe_container   = options.delete :wipe_container
  copy_image       = options.delete :copy_image
  exclusive        = options.delete :exclusive
  highmem          = options.delete :highmem

  queue            = options.delete(:queue) || 'bsc_ls'
  task_cpus        = options.delete(:task_cpus) || 1
  nodes            = options.delete(:nodes) || 1
  time             = options.delete(:time) || "0:00:10"

  inputs_dir       = options.delete :inputs_dir
  config_keys      = options.delete :config_keys

  user = ENV['USER'] || `whoami`.strip
  group = File.basename(File.dirname(ENV['HOME']))

  if contain_and_sync
    contain = "/scratch/tmp/rbbt-#{user}" if contain.nil?
    sync = "~/.rbbt/var/jobs" if sync.nil?
    wipe_container = "post" if wipe_container.nil?
  end

  contain = nil if contain == "" || contain == "none"
  sync = nil if sync == "" || sync == "none"

  contain = File.expand_path(contain) if contain

  name = options[:name] ||= Misc.obj2digest({:options => options.collect{|k,v| [k,v]}.sort_by{|k,v| k.to_s }, :args => args})
  options.delete(:name)
  slurm_basedir = options[:slurm_basedir] ||= File.expand_path(File.join('~/rbbt-slurm', name)) if slurm_basedir.nil?
  options.delete(:slurm_basedir)

  rbbt_cmd = args.reject{|e| e == '--' }.collect{|e| e.include?(" ")? '"' + e + '"' : e } * " "

  rbbt_cmd += " "  << options.collect do |o,v|
    o = o.to_s
    case v
    when TrueClass 
      '--' << o
    when FalseClass 
      '--' << o << "=false"
    else
      ['--' << o, "'#{v}'"] * " "
    end
  end * " "

  rbbt_cmd << " --config_keys='#{config_keys}'" if config_keys and not config_keys.empty?


  time = Misc.format_seconds Misc.timespan(time) unless time.include? ":"


  #{{{ PREPARE LOCAL LOGFILES

  Open.mkdir slurm_basedir

  fout = File.join(slurm_basedir, 'std.out')
  ferr = File.join(slurm_basedir, 'std.err')
  fjob = File.join(slurm_basedir, 'job.id')
  fexit = File.join(slurm_basedir, 'exit.status')
  fsync = File.join(slurm_basedir, 'sync.log')
  fcmd = File.join(slurm_basedir, 'command.slurm')

  #{{{ GENERATE TEMPLATE

  # HEADER
  header ="#!/bin/bash\n#SBATCH --qos=\"\#{queue}\"\n#SBATCH --job-name=\"\#{name}\"\n#SBATCH --workdir=\"\#{Dir.pwd}\"\n#SBATCH --output=\"\#{fout}\"\n#SBATCH --error=\"\#{ferr}\"\n#SBATCH --cpus-per-task=\"\#{task_cpus}\"\n#SBATCH --time=\"\#{time}\"\n#SBATCH --nodes=\"\#{nodes}\"\n  EOF\n\n  prep  = \"\"\n\n  if highmem\n    header +=<<-EOF\n#SBATCH --constraint=highmem \n    EOF\n  end\n\n  if exclusive\n    header +=<<-EOF\n#SBATCH --exclusive \n    EOF\n  end\n\n  header +=<<-EOF\n#CMD: \#{rbbt_cmd}\n  EOF\n\n  # ENV\n  env = \"\"\n  env +=<<-EOF\n# Prepare env\n[[ -f ~/config/load.sh ]] && source ~/config/load.sh\nmodule load java\n\n# Calculate max available memory\nlet \"MAX_MEMORY=$SLURM_MEM_PER_CPU * $SLURM_CPUS_PER_TASK\"\n  EOF\n\n\n  # RUN\n  run = \"\"\n  exec_cmd = %(env _JAVA_OPTIONS=\"-Xms1g -Xmx${MAX_MEMORY}m\")\n\n\n  if singularity\n    \#{{{ SINGULARITY\n    \n    singularity_exec = %(singularity exec -e -B $SINGULARITY_OPT_DIR:/singularity_opt/ -B /apps/)\n    \n    env +=<<-EOF\nmodule load intel/2018.1\nmodule load singularity\nPROJECTS_ROOT=\"/gpfs/projects/bsc26/\"\nSINGULARITY_IMG=\"$PROJECTS_ROOT/rbbt.singularity.img\"\nSINGULARITY_OPT_DIR=\"$PROJECTS_ROOT/singularity_opt/\"\nSINGULARITY_RUBY_INLINE=\"$HOME/.singularity_ruby_inline\"\nmkdir -p \"$SINGULARITY_RUBY_INLINE\"\n    EOF\n\n    if contain\n      scratch_group_dir = File.join('/gpfs/scratch/', group)\n      projects_group_dir = File.join('/gpfs/projects/', group)\n\n      prep +=<<-EOF\n\n# Prepare container dir\nCONTAINER_DIR=\"\#{contain}\"\nmkdir -p $CONTAINER_DIR/.rbbt/etc/\n\nfor dir in .ruby_inline git home; do\nmkdir -p $CONTAINER_DIR/$dir\ndone\n\nfor tmpd in persist_locks  produce_locks  R_sockets  sensiblewrite  sensiblewrite_locks  step_info_locks  tsv_open_locks; do\nmkdir -p $CONTAINER_DIR/.rbbt/tmp/$tmpd\ndone\n\n# Copy environment \ncp ~/.rbbt/etc/environment $CONTAINER_DIR/.rbbt/etc/\n\n# Set search_paths\necho \"singularity: /singularity_opt/{PKGDIR}/{TOPLEVEL}/{SUBPATH}\" > $CONTAINER_DIR/.rbbt/etc/search_paths\necho \"rbbt_user: /home/rbbt/.rbbt/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\necho \"outside_home: $CONTAINER_DIR/home/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\necho \"group_projects: \#{projects_group_dir}/{PKGDIR}/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\necho \"group_scratch: \#{scratch_group_dir}/{PKGDIR}/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\necho \"user_projects: \#{projects_group_dir}/\#{user}/{PKGDIR}/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\necho \"user_scratch: \#{scratch_group_dir}/\#{user}/{PKGDIR}/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\n      EOF\n\n      if user_group && group != user_group\n        prep +=<<-EOF\n\n# Add user_group search_path\necho \"\#{user_group}: /gpfs/projects/\#{user_group}/{PKGDIR}/{TOPLEVEL}/{SUBPATH}\" >> $CONTAINER_DIR/.rbbt/etc/search_paths\n        EOF\n      end\n\n      if inputs_dir\n        prep +=<<-EOF\n\n# Copy inputs\n[[ -d '\#{inputs_dir}' ]] && cp -R '\#{inputs_dir}' $CONTAINER_DIR/inputs\n        EOF\n        rbbt_cmd = rbbt_cmd.sub(inputs_dir, \"\#{contain}/inputs\")\n      end\n\n      if copy_image\n        prep +=<<EOF\n\n# Copy image\nrsync -avz \"$SINGULARITY_IMG\" \"$CONTAINER_DIR/rbbt.singularity.img\" 1>&2\nSINGULARITY_IMG=\"$CONTAINER_DIR/rbbt.singularity.img\"\n"
      end

      if  wipe_container == "pre" || wipe_container == "both"
        if singularity
          prep +="\n# Clean container pre\nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rm -Rfv .rbbt/var/jobs &>> \#{fsync}\nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rbbt system clean -f &>> \#{fsync}\nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rm -Rfv tmp/ &>> \#{fsync}\n"
        else
          prep = ""
        end
      end
    end

    if contain
      singularity_exec << %( -C -H "$CONTAINER_DIR" \
-B /scratch/tmp \
#{ group != user_group ? "-B /gpfs/projects/#{user_group}" : "" } \
-B #{scratch_group_dir} \
-B #{projects_group_dir} \
-B "$SINGULARITY_RUBY_INLINE":"$CONTAINER_DIR/.ruby_inline":rw  \
-B ~/git:"$CONTAINER_DIR/git":ro \
#{Open.exists?('~/.rbbt/software/opt/')? '-B ~/.rbbt/software/opt/:"/opt/":ro' : '' } \
-B ~/.rbbt:"$CONTAINER_DIR/home/":ro \
"$SINGULARITY_IMG")
      exec_cmd << ' TMPDIR="$CONTAINER_DIR/.rbbt/tmp" '
    else
      singularity_exec += %( -B "$SINGULARITY_RUBY_INLINE":"$HOME/.ruby_inline":rw "$SINGULARITY_IMG" )
    end

    if development
      exec_cmd += " rbbt --dev='#{development}'"
    else
      exec_cmd += ' rbbt'
    end

    exec_cmd = singularity_exec + " " + exec_cmd
  else
    if development
      exec_cmd << " " << %(~/git/rbbt-util/bin/rbbt --dev=#{development})
    else
      exec_cmd << " " << 'rbbt'
    end

    if contain
      rbbt_cmd << " " << %(--workdir_all='#{contain}')
    end
  end


  cmd ="\#{exec_cmd} \\\\\n\#{rbbt_cmd}\n"

  run +="\n# Run command\n\#{cmd}\n\n# Save exit status\nexit_status=$?\n\n"

  # CODA
  coda = ""
  if sync
    if singularity
      coda +="singularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rbbt system clean all -q &>> \#{fsync}\n"
    else
      coda +="rbbt system clean all -q &>> \#{fsync}\n"
    end

    if sync.include?("=>")
      source, _sep, sync = sync.partition("=>")
      source = source.strip
      sync = sync.strip
      source = File.join(File.expand_path(contain), source)
    else
      source = File.join(File.expand_path(contain), '.rbbt/var/jobs')
    end

    target = File.expand_path(sync)
    coda +="\n# Sync data to target location\nmkdir -p \"$(dirname '\#{target}')\"\nrsync -avztAXHP --copy-unsafe-links \"\#{source}/\" \"\#{target}/\" &>> \#{fsync} \nsync_es=\"$?\" \nfind '\#{target}' -type l -ls | awk '$13 ~ /^\#{target.gsub('/','\\/')}/ { sub(\"\#{source}\", \"\#{target}\", $13); print $11, $13 }' | while read A B; do rm $A; ln -s $B $A; done\n"

    if  contain && (wipe_container == "post" || wipe_container == "both")
      prep ="if ls -A '\#{contain}' &> /dev/null ; then\necho \"ERROR: Container directory not empty, refusing to wipe. \#{contain}\" &>> \#{fsync}\nfi\n      EOF\n      if singularity\n        coda +=<<-EOF\nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rm -v /dev/shm/sem.*.{in,out,process} /dev/shm/sem.Session-PID.*.sem 2> /dev/null >> \#{fsync}\n\n\n# Clean container directory\n#if [ $exit_status == '0' -a $sync_es == '0' ]; then \nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rbbt system clean -f &>> \#{fsync}\nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rm -Rfv .rbbt/var/jobs &>> \#{fsync}\nsingularity exec -e -C -H \"$CONTAINER_DIR\" \"$SINGULARITY_IMG\" rm -Rfv tmp/ &>> \#{fsync}\n#else\n#    echo \"ERROR: Process failed or results could not sync correctly. Contain directory not purged\" &>> \#{fsync}\n#fi\n" + prep
      else
        coda +="\#{exec_cmd} system clean\nif [ $exit_status == '0' -a $sync_es == '0' ]; then \nrm -Rfv \#{contain} &>> \#{fsync}\nelse\necho \"ERROR: Process failed or results could not sync correctly. Contain directory not purged\" &>> \#{fsync}\nfi\nunset sync_es\n"

      end
    end
  end
  coda +="\n# Write exit status to file\necho $exit_status > \#{fexit}\n"
  if sync
    coda +="if [ \"$sync_es\" == '0' ]; then \n  unset sync_es\n  exit $exit_status\nelse\n  exit $sync_es\nfi\n" 
  else
    coda +="exit $exit_status\n" 
  end

  template = [header, env, prep, run, coda] * "\n"

  template
end

.wait_for_job(slurm_basedir, time = 1) ⇒ Object

# File 'lib/rbbt/hpc.rb', line 479

def self.wait_for_job(slurm_basedir, time = 1)
  fexit = File.join(slurm_basedir, 'exit.status')
  fjob = File.join(slurm_basedir, 'job.id')
  job = Open.read(fjob) if Open.exists?(fjob)


  while ! Open.exists?(fexit)
    sleep time
  end
end