Completed
Pull Request — master (#167)
by
unknown
28s
created

test_main_resume_by_expanding_pool()   B

Complexity

Conditions 6

Size

Total Lines 34

Duplication

Lines 0
Ratio 0 %

Importance

Changes 3
Bugs 0 Features 0
Metric Value
cc 6
c 3
b 0
f 0
dl 0
loc 34
rs 7.5384
1
import os
2
import unittest
3
import tempfile
4
import shutil
5
from os.path import join as pjoin, abspath
6
from mock import patch
7
from subprocess import call
8
import subprocess
9
from nose.tools import assert_true, assert_equal
10
from smartdispatch import smartdispatch_script
11
import six
12
import sys
13
14
class TestSmartdispatcher(unittest.TestCase):
15
16
    def setUp(self):
17
        self.testing_dir = tempfile.mkdtemp()
18
        self.logs_dir = os.path.join(self.testing_dir, 'SMART_DISPATCH_LOGS')
19
20
        self.folded_commands = 'echo "[1 2 3 4]" "[6 7 8]" "[9 0]"'
21
        self.commands = ["echo 1 6 9", "echo 1 6 0", "echo 1 7 9", "echo 1 7 0", "echo 1 8 9", "echo 1 8 0",
22
                         "echo 2 6 9", "echo 2 6 0", "echo 2 7 9", "echo 2 7 0", "echo 2 8 9", "echo 2 8 0",
23
                         "echo 3 6 9", "echo 3 6 0", "echo 3 7 9", "echo 3 7 0", "echo 3 8 9", "echo 3 8 0",
24
                         "echo 4 6 9", "echo 4 6 0", "echo 4 7 9", "echo 4 7 0", "echo 4 8 9", "echo 4 8 0"]
25
        self.nb_commands = len(self.commands)
26
27
        scripts_path = abspath(pjoin(os.path.dirname(__file__), os.pardir, "scripts"))
28
        self.smart_dispatch_command = '{} -C 1 -G 1 -q test -t 5:00 -x'.format(pjoin(scripts_path, 'smart-dispatch'))
29
        self.launch_command = "{0} launch {1}".format(self.smart_dispatch_command, self.folded_commands)
30
        self.resume_command = "{0} resume {{0}}".format(self.smart_dispatch_command)
31
32
        self.smart_dispatch_launcher_command = '{} -C 1 -G 1 -q test -t 5:00'.format(pjoin(scripts_path, 'smart-dispatch'))
33
        self.launcher_command = "{0} launch {1}".format(self.smart_dispatch_launcher_command, self.folded_commands)
34
35
        smart_dispatch_command_with_pool = '{} --pool 10 -C 1 -G 1 -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
36
        self.launch_command_with_pool = smart_dispatch_command_with_pool.format('launch ' + self.folded_commands)
37
        self.nb_workers = 10
38
39
        smart_dispatch_command_with_cores = '{} -C 1 -G 1 -c {{cores}} -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
40
        self.launch_command_with_cores = smart_dispatch_command_with_cores.format('launch ' + self.folded_commands, cores='{cores}')
41
42
        smart_dispatch_command_with_gpus = '{} -C 1 -G 1 -g {{gpus}} -q test -t 5:00 -x {{0}}'.format(pjoin(scripts_path, 'smart-dispatch'))
43
        self.launch_command_with_gpus = smart_dispatch_command_with_gpus.format('launch ' + self.folded_commands, gpus='{gpus}')
44
45
        self._cwd = os.getcwd()
46
        os.chdir(self.testing_dir)
47
48
    def tearDown(self):
49
        os.chdir(self._cwd)
50
        shutil.rmtree(self.testing_dir)
51
52
    def test_main_launch(self):
53
        # Actual test
54
        exit_status = call(self.launch_command, shell=True)
55
56
        # Test validation
57
        assert_equal(exit_status, 0)
58
        assert_true(os.path.isdir(self.logs_dir))
59
        assert_equal(len(os.listdir(self.logs_dir)), 1)
60
61
        batch_uid = os.listdir(self.logs_dir)[0]
62
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
63
        assert_equal(len(os.listdir(path_job_commands)), self.nb_commands + 1)
64
65
    def test_launch_using_commands_file(self):
66
        # Actual test
67
        commands_filename = "commands_to_run.txt"
68
        open(commands_filename, 'w').write("\n".join(self.commands))
69
70
        launch_command = self.smart_dispatch_command + " -f {0} launch".format(commands_filename)
71
        exit_status = call(launch_command, shell=True)
72
73
        # Test validation
74
        assert_equal(exit_status, 0)
75
        assert_true(os.path.isdir(self.logs_dir))
76
        assert_equal(len(os.listdir(self.logs_dir)), 1)
77
78
        batch_uid = os.listdir(self.logs_dir)[0]
79
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
80
        assert_equal(len(os.listdir(path_job_commands)), self.nb_commands + 1)
81
        assert_equal(open(pjoin(path_job_commands, 'commands.txt')).read(), "\n".join(self.commands) + "\n")
82
83
    def test_main_launch_with_pool_of_workers(self):
84
        # Actual test
85
        exit_status = call(self.launch_command_with_pool, shell=True)
86
87
        # Test validation
88
        assert_equal(exit_status, 0)
89
        assert_true(os.path.isdir(self.logs_dir))
90
        assert_equal(len(os.listdir(self.logs_dir)), 1)
91
92
        batch_uid = os.listdir(self.logs_dir)[0]
93
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
94
        assert_equal(len(os.listdir(path_job_commands)), self.nb_workers + 1)
95
96
    def test_main_launch_with_cores_command(self):
97
        # Actual test
98
        exit_status_0 = call(self.launch_command_with_cores.format(cores=0), shell=True)
99
        exit_status_100 = call(self.launch_command_with_cores.format(cores=100), shell=True)
100
101
        # Test validation
102
        assert_equal(exit_status_0, 2)
103
        assert_equal(exit_status_100, 2)        
104
        assert_true(os.path.isdir(self.logs_dir))
105
106
    def test_main_launch_with_gpus_command(self):
107
        # Actual test
108
        exit_status_0 = call(self.launch_command_with_gpus.format(gpus=0), shell=True)
109
        exit_status_100 = call(self.launch_command_with_gpus.format(gpus=100), shell=True)
110
111
        # Test validation
112
        assert_equal(exit_status_0, 0)
113
        assert_equal(exit_status_100, 2)
114
        assert_true(os.path.isdir(self.logs_dir))
115
116 View Code Duplication
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
117
    def test_gpu_check(self):
118
119
        argv = ['-x', '-g', '2', '-G', '1', '-C', '1', '-q', 'random', '-t', '00:00:10' ,'launch', 'echo', 'testing123']
120
121
        # Test if the check fail
122
        with self.assertRaises(SystemExit) as context:
123
            smartdispatch_script.main(argv=argv)
124
125
        self.assertTrue(context.exception.code, 2)
126
127
        # Test if the test pass
128
        argv[2] = '0'
129
130
        try:
131
            smartdispatch_script.main(argv=argv)
132
        except SystemExit as e:
133
            self.fail("The command failed the check, but it was supposed to pass.")
134
135
136
    def test_cpu_check(self):
137
138
        argv = ['-x', '-c', '2', '-C', '1', '-G', '1', '-t', '00:00:10', '-q', 'random', 'launch', 'echo', 'testing123']
139
140
        # Test if the check fail
141
        with self.assertRaises(SystemExit) as context:
142
            smartdispatch_script.main(argv=argv)
143
144
        self.assertTrue(context.exception.code, 2)
145
146
        # Test if the test pass
147
        argv[2] = '1'
148
149
        try:
150
            smartdispatch_script.main(argv=argv)
151
        except SystemExit as e:
152
            self.fail("The command failed the check, but it was supposed to pass.")
153
154
155
    @patch('subprocess.check_output')
156
    def test_launch_job_check(self, mock_check_output):
157
158
        #For this test, we won't call the script directly, since we want to mock subprocess.check_output
159
        argv = ['-t', '0:0:1', '-G', '1', '-C', '1', '-q', 'random', 'launch', 'echo', 'testing123']
160
161
        # Test if the test pass (i.e the script run normaly)
162
        mock_check_output.side_effect = None
163
        mock_check_output.return_value = ""
164
165
        try:
166
            smartdispatch_script.main(argv=argv)
167
        except SystemExit as e:
168
            self.fail("The launcher had no problem, but the script failed nonetheless.")
169
170
        # Test if the check fail
171
        mock_check_output.side_effect = subprocess.CalledProcessError(1, 1, "A wild error appeared!")
172
        
173
        try:
174
            with self.assertRaises(SystemExit) as context:
175
                smartdispatch_script.main(argv=argv)
176
177
                self.assertTrue(context.exception.code, 2)
178
        
179
        except subprocess.CalledProcessError:
180
            # Rerasing the exception
181
            orig_exc_type, orig_exc_value, orig_exc_traceback = sys.exc_info()
182
183
            new_exc = Exception("smartdispatch_script.main() raised subprocess.CalledProcessError unexpectedly")
184
            new_exc.reraised = True
185
            six.reraise(type(new_exc), new_exc, orig_exc_traceback)
186
187
    def test_main_resume(self):
188
        # Setup
189
        call(self.launch_command, shell=True)
190
        batch_uid = os.listdir(self.logs_dir)[0]
191
192
        # Simulate that some commands are in the running state.
193
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
194
        pending_commands_file = pjoin(path_job_commands, "commands.txt")
195
        running_commands_file = pjoin(path_job_commands, "running_commands.txt")
196
        commands = open(pending_commands_file).read().strip().split("\n")
197 View Code Duplication
        with open(running_commands_file, 'w') as running_commands:
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
198
            running_commands.write("\n".join(commands[::2]) + "\n")
199
        with open(pending_commands_file, 'w') as pending_commands:
200
            pending_commands.write("\n".join(commands[1::2]) + "\n")
201
202
        # Actual test (should move running commands back to pending).
203
        exit_status = call(self.resume_command.format(batch_uid), shell=True)
204
205
        # Test validation
206
        assert_equal(exit_status, 0)
207
        assert_true(os.path.isdir(self.logs_dir))
208
        assert_equal(len(os.listdir(self.logs_dir)), 1)
209
        assert_equal(len(open(running_commands_file).readlines()), 0)
210
        assert_equal(len(open(pending_commands_file).readlines()), len(commands))
211
212
        # Test when batch_uid is a path instead of a jobname.
213
        # Setup
214
        batch_uid = os.path.join(self.logs_dir, os.listdir(self.logs_dir)[0])
215
216
        # Simulate that some commands are in the running state.
217
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
218
        pending_commands_file = pjoin(path_job_commands, "commands.txt")
219
        running_commands_file = pjoin(path_job_commands, "running_commands.txt")
220
        commands = open(pending_commands_file).read().strip().split("\n")
221
        with open(running_commands_file, 'w') as running_commands:
222
            running_commands.write("\n".join(commands[::2]) + "\n")
223
        with open(pending_commands_file, 'w') as pending_commands:
224
            pending_commands.write("\n".join(commands[1::2]) + "\n")
225
226
        # Actual test (should move running commands back to pending).
227
        exit_status = call(self.resume_command.format(batch_uid), shell=True)
228
229
        # Test validation
230
        assert_equal(exit_status, 0)
231 View Code Duplication
        assert_true(os.path.isdir(self.logs_dir))
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
232
        assert_equal(len(os.listdir(self.logs_dir)), 1)
233
        assert_equal(len(open(running_commands_file).readlines()), 0)
234
        assert_equal(len(open(pending_commands_file).readlines()), len(commands))
235
236
    def test_main_resume_by_expanding_pool_default(self):
237
        # Create SMART_DISPATCH_LOGS structure.
238
        call(self.launch_command, shell=True)
239
        batch_uid = os.listdir(self.logs_dir)[0]
240
241
        # Simulate that some commands are in the running state.
242
        nb_commands_files = 2  # 'commands.txt' and 'running_commands.txt'
243
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
244
        pending_commands_file = pjoin(path_job_commands, "commands.txt")
245
        running_commands_file = pjoin(path_job_commands, "running_commands.txt")
246
        commands = open(pending_commands_file).read().strip().split("\n")
247
        with open(running_commands_file, 'w') as running_commands:
248
            running_commands.write("\n".join(commands[::2]) + "\n")
249
        with open(pending_commands_file, 'w') as pending_commands:
250
            pending_commands.write("\n".join(commands[1::2]) + "\n")
251
252
        # Remove PBS files so we can check that new ones are going to be created.
253
        for f in os.listdir(path_job_commands):
254
            if f.startswith('job_commands_') and f.endswith('.sh'):
255
                os.remove(pjoin(path_job_commands, f))
256
257
        # Should NOT move running commands back to pending but should add new workers.
258
        command_line = self.resume_command.format(batch_uid)
259
        command_line += " --expandPool"
260
        exit_status = call(command_line, shell=True)
261
262
        # Test validation
263
        assert_equal(exit_status, 0)
264
        assert_equal(len(open(running_commands_file).readlines()), len(commands[::2]))
265
        assert_equal(len(open(pending_commands_file).readlines()), len(commands[1::2]))
266
267
        nb_job_commands_files = len(os.listdir(path_job_commands))
268
        assert_equal(nb_job_commands_files-nb_commands_files, len(commands[1::2]))
269
270
    def test_main_resume_by_expanding_pool(self):
271
        # Create SMART_DISPATCH_LOGS structure.
272
        call(self.launch_command, shell=True)
273
        batch_uid = os.listdir(self.logs_dir)[0]
274
275
        # Simulate that some commands are in the running state.
276
        nb_commands_files = 2  # 'commands.txt' and 'running_commands.txt'
277
        path_job_commands = os.path.join(self.logs_dir, batch_uid, "commands")
278
        pending_commands_file = pjoin(path_job_commands, "commands.txt")
279
        running_commands_file = pjoin(path_job_commands, "running_commands.txt")
280
        commands = open(pending_commands_file).read().strip().split("\n")
281
        with open(running_commands_file, 'w') as running_commands:
282
            running_commands.write("\n".join(commands[::2]) + "\n")
283
        with open(pending_commands_file, 'w') as pending_commands:
284
            pending_commands.write("\n".join(commands[1::2]) + "\n")
285
286
        # Remove PBS files so we can check that new ones are going to be created.
287
        for f in os.listdir(path_job_commands):
288
            if f.startswith('job_commands_') and f.endswith('.sh'):
289
                os.remove(pjoin(path_job_commands, f))
290
291
        # Should NOT move running commands back to pending but should add new workers.
292
        nb_workers_to_add = 3
293
        command_line = self.resume_command.format(batch_uid)
294
        command_line += " --expandPool {}".format(nb_workers_to_add)
295
        exit_status = call(command_line, shell=True)
296
297
        # Test validation
298
        assert_equal(exit_status, 0)
299
        assert_equal(len(open(running_commands_file).readlines()), len(commands[::2]))
300
        assert_equal(len(open(pending_commands_file).readlines()), len(commands[1::2]))
301
302
        nb_job_commands_files = len(os.listdir(path_job_commands))
303
        assert_equal(nb_job_commands_files-nb_commands_files, nb_workers_to_add)
304