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