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] || "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.crunch_job_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.crunch_job_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 crunch_job_wrapper: #{Server::Application.config.crunch_job_wrapper}"
  71       end
  72
  73       job_auth = ApiClientAuthorization.
  74         new(user: User.where('uuid=?', job.modified_by_user).first,
  75             api_client_id: 0)
  76       job_auth.save
  77
  78       cmd_args << 'crunch-job'
  79       cmd_args << '--job-api-token'
  80       cmd_args << job_auth.api_token
  81       cmd_args << '--job'
  82       cmd_args << job.uuid
  83
  84       commit = Commit.where(sha1: job.script_version).first
  85       if commit
  86         cmd_args << '--git-dir'
  87         cmd_args << File.
  88           join(Rails.configuration.git_repositories_dir,
  89                commit.repository_name,
  90                '.git')
  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         job_auth: job_auth
 115       }
 116       i.close
 117     end
 118   end
 119
 120   def take(job)
 121     # no-op -- let crunch-job take care of locking.
 122     true
 123   end
 124
 125   def untake(job)
 126     # no-op -- let crunch-job take care of locking.
 127     true
 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").to_a
 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           end
 160         end
 161       end
 162     end
 163   end
 164
 165   def reap_children
 166     return if 0 == @running.size
 167     pid_done = nil
 168     j_done = nil
 169
 170     if false
 171       begin
 172         pid_done = waitpid(-1, Process::WNOHANG | Process::WUNTRACED)
 173         if pid_done
 174           j_done = @running.values.
 175             select { |j| j[:wait_thr].pid == pid_done }.
 176             first
 177         end
 178       rescue SystemCallError
 179         # I have @running processes but system reports I have no
 180         # children. This is likely to happen repeatedly if it happens at
 181         # all; I will log this no more than once per child process I
 182         # start.
 183         if 0 < @running.select { |uuid,j| j[:warned_waitpid_error].nil? }.size
 184           children = @running.values.collect { |j| j[:wait_thr].pid }.join ' '
 185           $stderr.puts "dispatch: IPC bug: waitpid() error (#{$!}), but I have children #{children}"
 186         end
 187         @running.each do |uuid,j| j[:warned_waitpid_error] = true end
 188       end
 189     else
 190       @running.each do |uuid, j|
 191         if j[:wait_thr].status == false
 192           pid_done = j[:wait_thr].pid
 193           j_done = j
 194         end
 195       end
 196     end
 197
 198     return if !pid_done
 199
 200     job_done = j_done[:job]
 201     $stderr.puts "dispatch: child #{pid_done} exit"
 202     $stderr.puts "dispatch: job #{job_done.uuid} end"
 203
 204     # Ensure every last drop of stdout and stderr is consumed
 205     read_pipes
 206     if j_done[:stderr_buf] and j_done[:stderr_buf] != ''
 207       $stderr.puts j_done[:stderr_buf] + "\n"
 208     end
 209
 210     # Wait the thread
 211     j_done[:wait_thr].value
 212
 213     # Invalidate the per-job auth token
 214     j_done[:job_auth].update_attributes expires_at: Time.now
 215
 216     @running.delete job_done.uuid
 217   end
 218
 219   def run
 220     act_as_system_user
 221     @running ||= {}
 222     $stderr.puts "dispatch: ready"
 223     while !$signal[:term] or @running.size > 0
 224       read_pipes
 225       if $signal[:term]
 226         @running.each do |uuid, j|
 227           if !j[:started] and j[:sent_int] < 2
 228             begin
 229               Process.kill 'INT', j[:wait_thr].pid
 230             rescue Errno::ESRCH
 231               # No such pid = race condition + desired result is
 232               # already achieved
 233             end
 234             j[:sent_int] += 1
 235           end
 236         end
 237       else
 238         refresh_todo unless did_recently(:refresh_todo, 1.0)
 239         start_jobs unless @todo.empty? or did_recently(:start_jobs, 1.0)
 240       end
 241       reap_children
 242       select(@running.values.collect { |j| [j[:stdout], j[:stderr]] }.flatten,
 243              [], [], 1)
 244     end
 245   end
 246
 247   protected
 248
 249   def did_recently(thing, min_interval)
 250     @did_recently ||= {}
 251     if !@did_recently[thing] or @did_recently[thing] < Time.now - min_interval
 252       @did_recently[thing] = Time.now
 253       false
 254     else
 255       true
 256     end
 257   end
 258 end
 259
 260 Dispatcher.new.run