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