services/api/test/unit/crunch_dispatch_test.rb

   1 require 'test_helper'
   2 require 'crunch_dispatch'
   3
   4 class CrunchDispatchTest < ActiveSupport::TestCase
   5   test 'choose cheaper nodes first' do
   6     act_as_system_user do
   7       # Replace test fixtures with a set suitable for testing dispatch
   8       Node.destroy_all
   9
  10       # Idle nodes with different prices
  11       [['compute1', 3.20, 32],
  12        ['compute2', 1.60, 16],
  13        ['compute3', 0.80, 8]].each do |hostname, price, cores|
  14         Node.create!(hostname: hostname,
  15                      info: {
  16                        'slurm_state' => 'idle',
  17                      },
  18                      properties: {
  19                        'cloud_node' => {
  20                          'price' => price,
  21                        },
  22                        'total_cpu_cores' => cores,
  23                        'total_ram_mb' => cores*1024,
  24                        'total_scratch_mb' => cores*10000,
  25                      })
  26       end
  27
  28       # Node with no price information
  29       Node.create!(hostname: 'compute4',
  30                    info: {
  31                      'slurm_state' => 'idle',
  32                    },
  33                    properties: {
  34                      'total_cpu_cores' => 8,
  35                      'total_ram_mb' => 8192,
  36                      'total_scratch_mb' => 80000,
  37                    })
  38
  39       # Cheap but busy node
  40       Node.create!(hostname: 'compute5',
  41                    info: {
  42                      'slurm_state' => 'alloc',
  43                    },
  44                    properties: {
  45                      'cloud_node' => {
  46                        'price' => 0.10,
  47                      },
  48                      'total_cpu_cores' => 32,
  49                      'total_ram_mb' => 32768,
  50                      'total_scratch_mb' => 320000,
  51                    })
  52     end
  53
  54     dispatch = CrunchDispatch.new
  55     [[1, 16384, ['compute2']],
  56      [2, 16384, ['compute2', 'compute1']],
  57      [2, 8000, ['compute4', 'compute3']],
  58     ].each do |min_nodes, min_ram, expect_nodes|
  59       job = Job.new(runtime_constraints: {
  60                       'min_nodes' => min_nodes,
  61                       'min_ram_mb_per_node' => min_ram,
  62                     })
  63       nodes = dispatch.nodes_available_for_job_now job
  64       assert_equal expect_nodes, nodes
  65     end
  66   end
  67
  68   test 'respond to TERM' do
  69     lockfile = Rails.root.join 'tmp', 'dispatch.lock'
  70     ENV['CRUNCH_DISPATCH_LOCKFILE'] = lockfile.to_s
  71     begin
  72       pid = Process.fork do
  73         begin
  74           # Abandon database connections inherited from parent
  75           # process.  Credit to
  76           # https://github.com/kstephens/rails_is_forked
  77           ActiveRecord::Base.connection_handler.connection_pools.each_value do |pool|
  78             pool.instance_eval do
  79               @reserved_connections = {}
  80               @connections = []
  81             end
  82           end
  83           ActiveRecord::Base.establish_connection
  84
  85           dispatch = CrunchDispatch.new
  86           dispatch.stubs(:did_recently).returns true
  87           dispatch.run []
  88         ensure
  89           Process.exit!
  90         end
  91       end
  92       assert_with_timeout 5, "Dispatch did not lock #{lockfile}" do
  93         !can_lock(lockfile)
  94       end
  95     ensure
  96       Process.kill("TERM", pid)
  97     end
  98     assert_with_timeout 20, "Dispatch did not unlock #{lockfile}" do
  99       can_lock(lockfile)
 100     end
 101   end
 102
 103   def assert_with_timeout timeout, message
 104     t = 0
 105     while (t += 0.1) < timeout
 106       if yield
 107         return
 108       end
 109       sleep 0.1
 110     end
 111     assert false, message + " (waited #{timeout} seconds)"
 112   end
 113
 114   def can_lock lockfile
 115     lockfile.open(File::RDWR|File::CREAT, 0644) do |f|
 116       return f.flock(File::LOCK_EX|File::LOCK_NB)
 117     end
 118   end
 119 end