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