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 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       if Server::Application.config.crunch_job_user
  74         cmd_args.unshift("sudo", "-u",
  75                          Server::Application.config.crunch_job_user)
  76       end
  77
  78       job_auth = ApiClientAuthorization.
  79         new(user: User.where('uuid=?', job.modified_by_user).first,
  80             api_client_id: 0)
  81       job_auth.save
  82
  83       cmd_args << 'crunch-job'
  84       cmd_args << '--job-api-token'
  85       cmd_args << job_auth.api_token
  86       cmd_args << '--job'
  87       cmd_args << job.uuid
  88
  89       commit = Commit.where(sha1: job.script_version).first
  90       if commit
  91         cmd_args << '--git-dir'
  92         cmd_args << File.
  93           join(Rails.configuration.git_repositories_dir,
  94                commit.repository_name,
  95                '.git')
  96       end
  97
  98       $stderr.puts "dispatch: #{cmd_args.join ' '}"
  99
 100       begin
 101         i, o, e, t = Open3.popen3(*cmd_args)
 102       rescue
 103         $stderr.puts "dispatch: popen3: #{$!}"
 104         sleep 1
 105         untake(job)
 106         next
 107       end
 108       $stderr.puts "dispatch: job #{job.uuid} start"
 109       $stderr.puts "dispatch: child #{t.pid} start"
 110       @running[job.uuid] = {
 111         stdin: i,
 112         stdout: o,
 113         stderr: e,
 114         wait_thr: t,
 115         job: job,
 116         stderr_buf: '',
 117         started: false,
 118         sent_int: 0,
 119         job_auth: job_auth
 120       }
 121       i.close
 122     end
 123   end
 124
 125   def take(job)
 126     # no-op -- let crunch-job take care of locking.
 127     true
 128   end
 129
 130   def untake(job)
 131     # no-op -- let crunch-job take care of locking.
 132     true
 133   end
 134
 135   def read_pipes
 136     @running.each do |job_uuid, j|
 137       job = j[:job]
 138
 139       # Throw away child stdout
 140       begin
 141         j[:stdout].read_nonblock(2**20)
 142       rescue Errno::EAGAIN, EOFError
 143       end
 144
 145       # Read whatever is available from child stderr
 146       stderr_buf = false
 147       begin
 148         stderr_buf = j[:stderr].read_nonblock(2**20)
 149       rescue Errno::EAGAIN, EOFError
 150       end
 151
 152       if stderr_buf
 153         j[:stderr_buf] << stderr_buf
 154         if j[:stderr_buf].index "\n"
 155           lines = j[:stderr_buf].lines("\n").to_a
 156           if j[:stderr_buf][-1] == "\n"
 157             j[:stderr_buf] = ''
 158           else
 159             j[:stderr_buf] = lines.pop
 160           end
 161           lines.each do |line|
 162             $stderr.print "#{job_uuid} ! " unless line.index(job_uuid)
 163             $stderr.puts line
 164           end
 165         end
 166       end
 167     end
 168   end
 169
 170   def reap_children
 171     return if 0 == @running.size
 172     pid_done = nil
 173     j_done = nil
 174
 175     if false
 176       begin
 177         pid_done = waitpid(-1, Process::WNOHANG | Process::WUNTRACED)
 178         if pid_done
 179           j_done = @running.values.
 180             select { |j| j[:wait_thr].pid == pid_done }.
 181             first
 182         end
 183       rescue SystemCallError
 184         # I have @running processes but system reports I have no
 185         # children. This is likely to happen repeatedly if it happens at
 186         # all; I will log this no more than once per child process I
 187         # start.
 188         if 0 < @running.select { |uuid,j| j[:warned_waitpid_error].nil? }.size
 189           children = @running.values.collect { |j| j[:wait_thr].pid }.join ' '
 190           $stderr.puts "dispatch: IPC bug: waitpid() error (#{$!}), but I have children #{children}"
 191         end
 192         @running.each do |uuid,j| j[:warned_waitpid_error] = true end
 193       end
 194     else
 195       @running.each do |uuid, j|
 196         if j[:wait_thr].status == false
 197           pid_done = j[:wait_thr].pid
 198           j_done = j
 199         end
 200       end
 201     end
 202
 203     return if !pid_done
 204
 205     job_done = j_done[:job]
 206     $stderr.puts "dispatch: child #{pid_done} exit"
 207     $stderr.puts "dispatch: job #{job_done.uuid} end"
 208
 209     # Ensure every last drop of stdout and stderr is consumed
 210     read_pipes
 211     if j_done[:stderr_buf] and j_done[:stderr_buf] != ''
 212       $stderr.puts j_done[:stderr_buf] + "\n"
 213     end
 214
 215     # Wait the thread
 216     j_done[:wait_thr].value
 217
 218     # Invalidate the per-job auth token
 219     j_done[:job_auth].update_attributes expires_at: Time.now
 220
 221     @running.delete job_done.uuid
 222   end
 223
 224   def run
 225     act_as_system_user
 226     @running ||= {}
 227     $stderr.puts "dispatch: ready"
 228     while !$signal[:term] or @running.size > 0
 229       read_pipes
 230       if $signal[:term]
 231         @running.each do |uuid, j|
 232           if !j[:started] and j[:sent_int] < 2
 233             begin
 234               Process.kill 'INT', j[:wait_thr].pid
 235             rescue Errno::ESRCH
 236               # No such pid = race condition + desired result is
 237               # already achieved
 238             end
 239             j[:sent_int] += 1
 240           end
 241         end
 242       else
 243         refresh_todo unless did_recently(:refresh_todo, 1.0)
 244         start_jobs unless @todo.empty? or did_recently(:start_jobs, 1.0)
 245       end
 246       reap_children
 247       select(@running.values.collect { |j| [j[:stdout], j[:stderr]] }.flatten,
 248              [], [], 1)
 249     end
 250   end
 251
 252   protected
 253
 254   def did_recently(thing, min_interval)
 255     @did_recently ||= {}
 256     if !@did_recently[thing] or @did_recently[thing] < Time.now - min_interval
 257       @did_recently[thing] = Time.now
 258       false
 259     else
 260       true
 261     end
 262   end
 263 end
 264
 265 Dispatcher.new.run