services/api/script/crunch-dispatch.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] || ENV["RAILS_ENV"] || "development"
  15
  16 require File.dirname(__FILE__) + '/../config/boot'
  17 require File.dirname(__FILE__) + '/../config/environment'
  18 require 'open3'
  19
  20 $redis ||= Redis.new
  21
  22 class Dispatcher
  23   include ApplicationHelper
  24
  25   def sysuser
  26     return act_as_system_user
  27   end
  28
  29   def refresh_todo
  30     @todo = Job.queue
  31   end
  32
  33   def start_jobs
  34     if Server::Application.config.crunch_job_wrapper.to_s.match /^slurm/
  35       @idle_slurm_nodes = 0
  36       begin
  37         `sinfo`.
  38           split("\n").
  39           collect { |line| line.match /(\d+) +idle/ }.
  40           each do |re|
  41           @idle_slurm_nodes = re[1].to_i if re
  42         end
  43       rescue
  44       end
  45     end
  46
  47     @todo.each do |job|
  48
  49       min_nodes = 1
  50       begin
  51         if job.resource_limits['min_nodes']
  52           min_nodes = begin job.resource_limits['min_nodes'].to_i rescue 1 end
  53         end
  54       end
  55       next if @idle_slurm_nodes and @idle_slurm_nodes < min_nodes
  56
  57       next if @running[job.uuid]
  58       next if !take(job)
  59
  60       cmd_args = nil
  61       case Server::Application.config.crunch_job_wrapper
  62       when :none
  63         cmd_args = []
  64       when :slurm_immediate
  65         cmd_args = ["salloc",
  66                     "--chdir=/",
  67                     "--immediate",
  68                     "--exclusive",
  69                     "--no-kill",
  70                     "--job-name=#{job.uuid}",
  71                     "--nodes=#{min_nodes}"]
  72       else
  73         raise "Unknown crunch_job_wrapper: #{Server::Application.config.crunch_job_wrapper}"
  74       end
  75
  76       if Server::Application.config.crunch_job_user
  77         cmd_args.unshift("sudo", "-E", "-u",
  78                          Server::Application.config.crunch_job_user,
  79                          "PERLLIB=#{ENV['PERLLIB']}")
  80       end
  81
  82       job_auth = ApiClientAuthorization.
  83         new(user: User.where('uuid=?', job.modified_by_user_uuid).first,
  84             api_client_id: 0)
  85       job_auth.save
  86
  87       cmd_args << (ENV['CRUNCH_JOB_BIN'] || `which crunch-job`.strip)
  88       cmd_args << '--job-api-token'
  89       cmd_args << job_auth.api_token
  90       cmd_args << '--job'
  91       cmd_args << job.uuid
  92
  93       commit = Commit.where(sha1: job.script_version).first
  94       if commit
  95         cmd_args << '--git-dir'
  96         if File.exists?(File.
  97                         join(Rails.configuration.git_repositories_dir,
  98                              commit.repository_name + '.git'))
  99           cmd_args << File.
 100             join(Rails.configuration.git_repositories_dir,
 101                  commit.repository_name + '.git')
 102         else
 103           cmd_args << File.
 104             join(Rails.configuration.git_repositories_dir,
 105                  commit.repository_name, '.git')
 106         end
 107       end
 108
 109       $stderr.puts "dispatch: #{cmd_args.join ' '}"
 110
 111       begin
 112         i, o, e, t = Open3.popen3(*cmd_args)
 113       rescue
 114         $stderr.puts "dispatch: popen3: #{$!}"
 115         sleep 1
 116         untake(job)
 117         next
 118       end
 119       $stderr.puts "dispatch: job #{job.uuid} start"
 120       $stderr.puts "dispatch: child #{t.pid} start"
 121       @running[job.uuid] = {
 122         stdin: i,
 123         stdout: o,
 124         stderr: e,
 125         wait_thr: t,
 126         job: job,
 127         stderr_buf: '',
 128         started: false,
 129         sent_int: 0,
 130         job_auth: job_auth
 131       }
 132       i.close
 133     end
 134   end
 135
 136   def take(job)
 137     # no-op -- let crunch-job take care of locking.
 138     true
 139   end
 140
 141   def untake(job)
 142     # no-op -- let crunch-job take care of locking.
 143     true
 144   end
 145
 146   def read_pipes
 147     @running.each do |job_uuid, j|
 148       job = j[:job]
 149
 150       # Throw away child stdout
 151       begin
 152         j[:stdout].read_nonblock(2**20)
 153       rescue Errno::EAGAIN, EOFError
 154       end
 155
 156       # Read whatever is available from child stderr
 157       stderr_buf = false
 158       begin
 159         stderr_buf = j[:stderr].read_nonblock(2**20)
 160       rescue Errno::EAGAIN, EOFError
 161       end
 162
 163       if stderr_buf
 164         j[:stderr_buf] << stderr_buf
 165         if j[:stderr_buf].index "\n"
 166           lines = j[:stderr_buf].lines("\n").to_a
 167           if j[:stderr_buf][-1] == "\n"
 168             j[:stderr_buf] = ''
 169           else
 170             j[:stderr_buf] = lines.pop
 171           end
 172           lines.each do |line|
 173             $stderr.print "#{job_uuid} ! " unless line.index(job_uuid)
 174             $stderr.puts line
 175             $redis.publish job_uuid, "#{Time.now.ctime.to_s} #{line}"
 176           end
 177         end
 178       end
 179     end
 180   end
 181
 182   def reap_children
 183     return if 0 == @running.size
 184     pid_done = nil
 185     j_done = nil
 186
 187     if false
 188       begin
 189         pid_done = waitpid(-1, Process::WNOHANG | Process::WUNTRACED)
 190         if pid_done
 191           j_done = @running.values.
 192             select { |j| j[:wait_thr].pid == pid_done }.
 193             first
 194         end
 195       rescue SystemCallError
 196         # I have @running processes but system reports I have no
 197         # children. This is likely to happen repeatedly if it happens at
 198         # all; I will log this no more than once per child process I
 199         # start.
 200         if 0 < @running.select { |uuid,j| j[:warned_waitpid_error].nil? }.size
 201           children = @running.values.collect { |j| j[:wait_thr].pid }.join ' '
 202           $stderr.puts "dispatch: IPC bug: waitpid() error (#{$!}), but I have children #{children}"
 203         end
 204         @running.each do |uuid,j| j[:warned_waitpid_error] = true end
 205       end
 206     else
 207       @running.each do |uuid, j|
 208         if j[:wait_thr].status == false
 209           pid_done = j[:wait_thr].pid
 210           j_done = j
 211         end
 212       end
 213     end
 214
 215     return if !pid_done
 216
 217     job_done = j_done[:job]
 218     $stderr.puts "dispatch: child #{pid_done} exit"
 219     $stderr.puts "dispatch: job #{job_done.uuid} end"
 220     $redis.publish job_done.uuid, "end"
 221
 222     # Ensure every last drop of stdout and stderr is consumed
 223     read_pipes
 224     if j_done[:stderr_buf] and j_done[:stderr_buf] != ''
 225       $stderr.puts j_done[:stderr_buf] + "\n"
 226     end
 227
 228     # Wait the thread
 229     j_done[:wait_thr].value
 230
 231     # Invalidate the per-job auth token
 232     j_done[:job_auth].update_attributes expires_at: Time.now
 233
 234     @running.delete job_done.uuid
 235   end
 236
 237   def run
 238     act_as_system_user
 239     @running ||= {}
 240     $stderr.puts "dispatch: ready"
 241     while !$signal[:term] or @running.size > 0
 242       read_pipes
 243       if $signal[:term]
 244         @running.each do |uuid, j|
 245           if !j[:started] and j[:sent_int] < 2
 246             begin
 247               Process.kill 'INT', j[:wait_thr].pid
 248             rescue Errno::ESRCH
 249               # No such pid = race condition + desired result is
 250               # already achieved
 251             end
 252             j[:sent_int] += 1
 253           end
 254         end
 255       else
 256         refresh_todo unless did_recently(:refresh_todo, 1.0)
 257         start_jobs unless @todo.empty? or did_recently(:start_jobs, 1.0)
 258       end
 259       reap_children
 260       select(@running.values.collect { |j| [j[:stdout], j[:stderr]] }.flatten,
 261              [], [], 1)
 262     end
 263   end
 264
 265   protected
 266
 267   def did_recently(thing, min_interval)
 268     @did_recently ||= {}
 269     if !@did_recently[thing] or @did_recently[thing] < Time.now - min_interval
 270       @did_recently[thing] = Time.now
 271       false
 272     else
 273       true
 274     end
 275   end
 276 end
 277
 278 Dispatcher.new.run