services/api/script/dispatch_jobs.rb

   1 #!/usr/bin/env ruby
   2
   3 include Process
   4
   5 $signal = {}
   6 %w{TERM INT}.each do |sig|
   7   signame = sig
   8   Signal.trap(sig) do
   9     $stderr.puts "Received #{signame} signal"
  10     $signal[:term] = true
  11   end
  12 end
  13
  14 ENV["RAILS_ENV"] = ARGV[0] || "development"
  15
  16 require File.dirname(__FILE__) + '/../config/boot'
  17 require File.dirname(__FILE__) + '/../config/environment'
  18 require 'open3'
  19
  20 class Dispatcher
  21   include ApplicationHelper
  22
  23   def sysuser
  24     return act_as_system_user
  25   end
  26
  27   def refresh_todo
  28     @todo = Job.queue
  29   end
  30
  31   def start_jobs
  32     if Server::Application.config.whjobmanager_wrapper.to_s.match /^slurm/
  33       @idle_slurm_nodes = 0
  34       begin
  35         `sinfo`.
  36           split("\n").
  37           collect { |line| line.match /(\d+) +idle/ }.
  38           each do |re|
  39           @idle_slurm_nodes = re[1].to_i if re
  40         end
  41       rescue
  42       end
  43     end
  44
  45     @todo.each do |job|
  46
  47       min_nodes = 1
  48       begin
  49         if job.resource_limits['min_nodes']
  50           min_nodes = begin job.resource_limits['min_nodes'].to_i rescue 1 end
  51         end
  52       end
  53       next if @idle_slurm_nodes and @idle_slurm_nodes < min_nodes
  54
  55       next if @running[job.uuid]
  56       next if !take(job)
  57
  58       cmd_args = nil
  59       case Server::Application.config.whjobmanager_wrapper
  60       when :none
  61         cmd_args = []
  62       when :slurm_immediate
  63         cmd_args = ["salloc",
  64                     "--immediate",
  65                     "--exclusive",
  66                     "--no-kill",
  67                     "--job-name=#{job.uuid}",
  68                     "--nodes=#{min_nodes}"]
  69       else
  70         raise "Unknown whjobmanager_wrapper: #{Server::Application.config.whjobmanager_wrapper}"
  71       end
  72
  73       cmd_args << 'whjobmanager'
  74       cmd_args << "id=#{job.uuid}"
  75       cmd_args << "mrfunction=#{job.script}"
  76       job.script_parameters.each do |k,v|
  77         k = k.to_s
  78         if k == 'input'
  79           k = 'inputkey'
  80         else
  81           k = k.upcase
  82         end
  83         cmd_args << "#{k}=#{v}"
  84       end
  85       cmd_args << "revision=#{job.script_version}"
  86
  87       begin
  88         cmd_args << "stepspernode=#{job.resource_limits['max_tasks_per_node'].to_i}"
  89       rescue
  90         # OK if limit is not specified. OK to ignore if not integer.
  91       end
  92
  93       $stderr.puts "dispatch: #{cmd_args.join ' '}"
  94
  95       begin
  96         i, o, e, t = Open3.popen3(*cmd_args)
  97       rescue
  98         $stderr.puts "dispatch: popen3: #{$!}"
  99         sleep 1
 100         untake(job)
 101         next
 102       end
 103       $stderr.puts "dispatch: job #{job.uuid} start"
 104       $stderr.puts "dispatch: child #{t.pid} start"
 105       @running[job.uuid] = {
 106         stdin: i,
 107         stdout: o,
 108         stderr: e,
 109         wait_thr: t,
 110         job: job,
 111         stderr_buf: '',
 112         started: false,
 113         sent_int: 0
 114       }
 115       i.close
 116     end
 117   end
 118
 119   def take(job)
 120     lock_ok = false
 121     ActiveRecord::Base.transaction do
 122       job.reload
 123       if job.is_locked_by.nil? and
 124           job.update_attributes(is_locked_by: sysuser.uuid)
 125         lock_ok = true
 126       end
 127     end
 128     lock_ok
 129   end
 130
 131   def untake(job)
 132     job.reload
 133     job.update_attributes is_locked_by: nil
 134   end
 135
 136   def read_pipes
 137     @running.each do |job_uuid, j|
 138       job = j[:job]
 139
 140       # Throw away child stdout
 141       begin
 142         j[:stdout].read_nonblock(2**20)
 143       rescue Errno::EAGAIN, EOFError
 144       end
 145
 146       # Read whatever is available from child stderr
 147       stderr_buf = false
 148       begin
 149         stderr_buf = j[:stderr].read_nonblock(2**20)
 150       rescue Errno::EAGAIN, EOFError
 151       end
 152
 153       if stderr_buf
 154         j[:stderr_buf] << stderr_buf
 155         if j[:stderr_buf].index "\n"
 156           lines = j[:stderr_buf].lines "\n"
 157           if j[:stderr_buf][-1] == "\n"
 158             j[:stderr_buf] = ''
 159           else
 160             j[:stderr_buf] = lines.pop
 161           end
 162           lines.each do |line|
 163             $stderr.print "#{job_uuid} ! " unless line.index(job_uuid)
 164             $stderr.puts line
 165             line.chomp!
 166             if (re = line.match(/#{job_uuid} (\d+) (\S*) (.*)/))
 167               ignorethis, whjmpid, taskid, message = re.to_a
 168               if taskid == '' and message == 'start'
 169                 $stderr.puts "dispatch: noticed #{job_uuid} started"
 170                 j[:started] = true
 171                 ActiveRecord::Base.transaction do
 172                   j[:job].reload
 173                   j[:job].update_attributes running: true, started_at: Time.now
 174                 end
 175               elsif taskid == '' and (re = message.match /^revision (\S+)$/)
 176                 $stderr.puts "dispatch: noticed #{job_uuid} version #{re[1]}"
 177                 ActiveRecord::Base.transaction do
 178                   j[:job].reload
 179                   j[:job].script_version = re[1]
 180                   j[:job].save
 181                 end
 182               elsif taskid == '' and (re = message.match /^outputkey (\S+)$/)
 183                 $stderr.puts "dispatch: noticed #{job_uuid} output #{re[1]}"
 184                 j[:output] = re[1]
 185               elsif taskid == '' and (re = message.match /^meta key is (\S+)$/)
 186                 $stderr.puts "dispatch: noticed #{job_uuid} log #{re[1]}"
 187                 j[:log] = re[1]
 188                 ActiveRecord::Base.transaction do
 189                   j[:job].reload
 190                   j[:job].update_attributes log: j[:log]
 191                 end
 192               elsif taskid.match(/^\d+/) and (re = message.match /^failure /)
 193                 $stderr.puts "dispatch: noticed #{job_uuid} task fail"
 194                 ActiveRecord::Base.transaction do
 195                   j[:job].reload
 196                   j[:job].tasks_summary ||= {}
 197                   j[:job].tasks_summary[:failed] ||= 0
 198                   j[:job].tasks_summary[:failed] += 1
 199                   j[:job].save
 200                 end
 201               elsif (re = message.match(/^status: (\d+) done, (\d+) running, (\d+) todo/))
 202                 $stderr.puts "dispatch: noticed #{job_uuid} #{message}"
 203                 ActiveRecord::Base.transaction do
 204                   j[:job].reload
 205                   j[:job].tasks_summary ||= {}
 206                   j[:job].tasks_summary[:done] = re[1].to_i
 207                   j[:job].tasks_summary[:running] = re[2].to_i
 208                   j[:job].tasks_summary[:todo] = re[3].to_i
 209                   j[:job].save
 210                 end
 211                 if re[2].to_i == 0 and re[3].to_i == 0
 212                   $stderr.puts "dispatch: noticed #{job_uuid} succeeded"
 213                   j[:success] = true
 214                 end
 215               end
 216             end
 217           end
 218         end
 219       end
 220     end
 221   end
 222
 223   def reap_children
 224     return if 0 == @running.size
 225     pid_done = nil
 226     j_done = nil
 227
 228     if false
 229       begin
 230         pid_done = waitpid(-1, Process::WNOHANG | Process::WUNTRACED)
 231         if pid_done
 232           j_done = @running.values.
 233             select { |j| j[:wait_thr].pid == pid_done }.
 234             first
 235         end
 236       rescue SystemCallError
 237         # I have @running processes but system reports I have no
 238         # children. This is likely to happen repeatedly if it happens at
 239         # all; I will log this no more than once per child process I
 240         # start.
 241         if 0 < @running.select { |uuid,j| j[:warned_waitpid_error].nil? }.size
 242           children = @running.values.collect { |j| j[:wait_thr].pid }.join ' '
 243           $stderr.puts "dispatch: IPC bug: waitpid() error (#{$!}), but I have children #{children}"
 244         end
 245         @running.each do |uuid,j| j[:warned_waitpid_error] = true end
 246       end
 247     else
 248       @running.each do |uuid, j|
 249         if j[:wait_thr].status == false
 250           pid_done = j[:wait_thr].pid
 251           j_done = j
 252         end
 253       end
 254     end
 255
 256     return if !pid_done
 257
 258     job_done = j_done[:job]
 259     $stderr.puts "dispatch: child #{pid_done} exit"
 260     $stderr.puts "dispatch: job #{job_done.uuid} end"
 261
 262     # Ensure every last drop of stdout and stderr is consumed
 263     read_pipes
 264     if j_done[:stderr_buf] and j_done[:stderr_buf] != ''
 265       $stderr.puts j_done[:stderr_buf] + "\n"
 266     end
 267
 268     j_done[:wait_thr].value          # wait the thread
 269
 270     if !j_done[:started]
 271       # If the job never really started (due to a scheduling
 272       # failure), just put it back in the queue
 273       untake(job_done)
 274       $stderr.puts "dispatch: job #{job_done.uuid} requeued"
 275     else
 276       # Otherwise, mark the job as finished
 277       ActiveRecord::Base.transaction do
 278         job_done.reload
 279         job_done.log = j_done[:log]
 280         job_done.output = j_done[:output]
 281         job_done.success = j_done[:success]
 282         job_done.assert_finished
 283       end
 284     end
 285     @running.delete job_done.uuid
 286   end
 287
 288   def run
 289     act_as_system_user
 290     @running ||= {}
 291     $stderr.puts "dispatch: ready"
 292     while !$signal[:term] or @running.size > 0
 293       read_pipes
 294       if $signal[:term]
 295         @running.each do |uuid, j|
 296           if !j[:started] and j[:sent_int] < 2
 297             begin
 298               Process.kill 'INT', j[:wait_thr].pid
 299             rescue Errno::ESRCH
 300               # No such pid = race condition + desired result is
 301               # already achieved
 302             end
 303             j[:sent_int] += 1
 304           end
 305         end
 306       else
 307         refresh_todo unless did_recently(:refresh_todo, 1.0)
 308         start_jobs unless @todo.empty? or did_recently(:start_jobs, 1.0)
 309       end
 310       reap_children
 311       select(@running.values.collect { |j| [j[:stdout], j[:stderr]] }.flatten,
 312              [], [], 1)
 313     end
 314   end
 315
 316   protected
 317
 318   def did_recently(thing, min_interval)
 319     @did_recently ||= {}
 320     if !@did_recently[thing] or @did_recently[thing] < Time.now - min_interval
 321       @did_recently[thing] = Time.now
 322       false
 323     else
 324       true
 325     end
 326   end
 327 end
 328
 329 Dispatcher.new.run