6520: Assign weight to node based on price to preferentially schedule on
[arvados.git] / services / nodemanager / tests / test_jobqueue.py
1 #!/usr/bin/env python
2
3 from __future__ import absolute_import, print_function
4
5 import unittest
6 import mock
7
8 import arvnodeman.jobqueue as jobqueue
9 from . import testutil
10
11 class ServerCalculatorTestCase(unittest.TestCase):
12     def make_calculator(self, factors, **kwargs):
13         return jobqueue.ServerCalculator(
14             [(testutil.MockSize(n), {'cores': n}) for n in factors], **kwargs)
15
16     def calculate(self, servcalc, *constraints):
17         return servcalc.servers_for_queue(
18             [{'uuid': 'zzzzz-jjjjj-{:015x}'.format(index),
19               'runtime_constraints': cdict}
20              for index, cdict in enumerate(constraints)])
21
22     def test_empty_queue_needs_no_servers(self):
23         servcalc = self.make_calculator([1])
24         self.assertEqual([], servcalc.servers_for_queue([]))
25
26     def test_easy_server_count(self):
27         servcalc = self.make_calculator([1])
28         servlist = self.calculate(servcalc, {'min_nodes': 3})
29         self.assertEqual(3, len(servlist))
30
31     def test_implicit_server_count(self):
32         servcalc = self.make_calculator([1])
33         servlist = self.calculate(servcalc, {}, {'min_nodes': 3})
34         self.assertEqual(4, len(servlist))
35
36     def test_bad_min_nodes_override(self):
37         servcalc = self.make_calculator([1])
38         servlist = self.calculate(servcalc,
39                                   {'min_nodes': -2}, {'min_nodes': 'foo'})
40         self.assertEqual(2, len(servlist))
41
42     def test_ignore_unsatisfiable_jobs(self):
43         servcalc = self.make_calculator([1], max_nodes=9)
44         servlist = self.calculate(servcalc,
45                                   {'min_cores_per_node': 2},
46                                   {'min_ram_mb_per_node': 256},
47                                   {'min_nodes': 6},
48                                   {'min_nodes': 12},
49                                   {'min_scratch_mb_per_node': 200})
50         self.assertEqual(6, len(servlist))
51
52     def test_ignore_too_expensive_jobs(self):
53         servcalc = self.make_calculator([1, 2], max_nodes=12, max_price=6)
54         servlist = self.calculate(servcalc,
55                                   {'min_cores_per_node': 1, 'min_nodes': 6})
56         self.assertEqual(6, len(servlist))
57
58         servlist = self.calculate(servcalc,
59                                   {'min_cores_per_node': 2, 'min_nodes': 6})
60         self.assertEqual(0, len(servlist))
61
62     def test_job_requesting_max_nodes_accepted(self):
63         servcalc = self.make_calculator([1], max_nodes=4)
64         servlist = self.calculate(servcalc, {'min_nodes': 4})
65         self.assertEqual(4, len(servlist))
66
67     def test_cheapest_size(self):
68         servcalc = self.make_calculator([2, 4, 1, 3])
69         self.assertEqual(testutil.MockSize(1), servcalc.cheapest_size())
70
71     def test_next_biggest(self):
72         servcalc = self.make_calculator([1, 2, 4, 8])
73         servlist = self.calculate(servcalc,
74                                   {'min_cores_per_node': 3},
75                                   {'min_cores_per_node': 6})
76         self.assertEqual([servcalc.cloud_sizes[2].id,
77                           servcalc.cloud_sizes[3].id],
78                          [s.id for s in servlist])
79
80     def test_multiple_sizes(self):
81         servcalc = self.make_calculator([1, 2])
82         servlist = self.calculate(servcalc,
83                                   {'min_cores_per_node': 2},
84                                   {'min_cores_per_node': 1},
85                                   {'min_cores_per_node': 1})
86         self.assertEqual([servcalc.cloud_sizes[1].id,
87                           servcalc.cloud_sizes[0].id,
88                           servcalc.cloud_sizes[0].id],
89                          [s.id for s in servlist])
90
91         servlist = self.calculate(servcalc,
92                                   {'min_cores_per_node': 1},
93                                   {'min_cores_per_node': 2},
94                                   {'min_cores_per_node': 1})
95         self.assertEqual([servcalc.cloud_sizes[0].id,
96                           servcalc.cloud_sizes[1].id,
97                           servcalc.cloud_sizes[0].id],
98                          [s.id for s in servlist])
99
100         servlist = self.calculate(servcalc,
101                                   {'min_cores_per_node': 1},
102                                   {'min_cores_per_node': 1},
103                                   {'min_cores_per_node': 2})
104         self.assertEqual([servcalc.cloud_sizes[0].id,
105                           servcalc.cloud_sizes[0].id,
106                           servcalc.cloud_sizes[1].id],
107                          [s.id for s in servlist])
108
109
110
111 class JobQueueMonitorActorTestCase(testutil.RemotePollLoopActorTestMixin,
112                                    unittest.TestCase):
113     TEST_CLASS = jobqueue.JobQueueMonitorActor
114
115     class MockCalculator(object):
116         @staticmethod
117         def servers_for_queue(queue):
118             return [testutil.MockSize(n) for n in queue]
119
120
121     def build_monitor(self, side_effect, *args, **kwargs):
122         super(JobQueueMonitorActorTestCase, self).build_monitor(*args, **kwargs)
123         self.client.jobs().queue().execute.side_effect = side_effect
124
125     @mock.patch("subprocess.check_output")
126     def test_subscribers_get_server_lists(self, mock_squeue):
127         mock_squeue.return_value = ""
128
129         self.build_monitor([{'items': [1, 2]}], self.MockCalculator())
130         self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
131         self.stop_proxy(self.monitor)
132         self.subscriber.assert_called_with([testutil.MockSize(1),
133                                             testutil.MockSize(2)])
134
135     @mock.patch("subprocess.check_output")
136     def test_squeue_server_list(self, mock_squeue):
137         mock_squeue.return_value = """1 0 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzy
138 2 0 0 Resources zzzzz-zzzzz-zzzzzzzzzzzzzzz
139 """
140
141         super(JobQueueMonitorActorTestCase, self).build_monitor(jobqueue.ServerCalculator(
142             [(testutil.MockSize(n), {'cores': n, 'ram': n, 'scratch': n}) for n in range(1, 3)]))
143         self.monitor.subscribe(self.subscriber).get(self.TIMEOUT)
144         self.stop_proxy(self.monitor)
145         self.subscriber.assert_called_with([testutil.MockSize(1),
146                                             testutil.MockSize(2)])
147
148
149 if __name__ == '__main__':
150     unittest.main()