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