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