|
1
|
|
|
from nose.tools import assert_true, assert_equal, assert_raises |
|
2
|
|
|
from numpy.testing import assert_array_equal |
|
3
|
|
|
|
|
4
|
|
|
from smartdispatch.pbs import PBS |
|
5
|
|
|
import unittest |
|
6
|
|
|
import tempfile |
|
7
|
|
|
import shutil |
|
8
|
|
|
import os |
|
9
|
|
|
|
|
10
|
|
|
|
|
11
|
|
|
class TestPBS(unittest.TestCase): |
|
12
|
|
|
|
|
13
|
|
|
def setUp(self): |
|
14
|
|
|
self.testing_dir = tempfile.mkdtemp() |
|
15
|
|
|
self.queue_name = "qtest@cluster" |
|
16
|
|
|
self.walltime = "10:42" |
|
17
|
|
|
self.pbs = PBS(self.queue_name, self.walltime) |
|
18
|
|
|
|
|
19
|
|
|
def tearDown(self): |
|
20
|
|
|
shutil.rmtree(self.testing_dir) |
|
21
|
|
|
|
|
22
|
|
|
def test_constructor(self): |
|
23
|
|
|
assert_raises(ValueError, PBS, self.queue_name, None) |
|
24
|
|
|
assert_raises(ValueError, PBS, None, self.walltime) |
|
25
|
|
|
assert_raises(ValueError, PBS, "", self.walltime) |
|
26
|
|
|
assert_raises(ValueError, PBS, self.queue_name, "Wrong walltime format") |
|
27
|
|
|
|
|
28
|
|
|
def test_add_options(self): |
|
29
|
|
|
# Default options |
|
30
|
|
|
assert_equal(len(self.pbs.options), 2) |
|
31
|
|
|
assert_true('-V' in self.pbs.options.keys()) |
|
32
|
|
|
assert_equal(self.pbs.options['-q'], self.queue_name) |
|
33
|
|
|
|
|
34
|
|
|
self.pbs.add_options(A="option1") |
|
35
|
|
|
assert_equal(self.pbs.options["-A"], "option1") |
|
36
|
|
|
self.pbs.add_options(A="option2", B="option3") |
|
37
|
|
|
assert_equal(self.pbs.options["-A"], "option2") |
|
38
|
|
|
assert_equal(self.pbs.options["-B"], "option3") |
|
39
|
|
|
|
|
40
|
|
|
def test_add_sbatch_options(self): |
|
41
|
|
|
self.pbs.add_sbatch_options(a="value1") |
|
42
|
|
|
assert_equal(self.pbs.sbatch_options["-a"], "value1") |
|
43
|
|
|
self.pbs.add_sbatch_options(option1="value2", option2="value3") |
|
44
|
|
|
assert_equal(self.pbs.sbatch_options["--option1"], "value2") |
|
45
|
|
|
assert_equal(self.pbs.sbatch_options["--option2"], "value3") |
|
46
|
|
|
|
|
47
|
|
|
def test_add_resources(self): |
|
48
|
|
|
assert_equal(len(self.pbs.resources), 1) |
|
49
|
|
|
assert_equal(self.pbs.resources["walltime"], self.walltime) |
|
50
|
|
|
|
|
51
|
|
|
self.pbs.add_resources(nodes="23:ppn=2") |
|
52
|
|
|
assert_equal(self.pbs.resources["nodes"], "23:ppn=2") |
|
53
|
|
|
|
|
54
|
|
|
self.pbs.add_resources(pmem=123, walltime="42:42") |
|
55
|
|
|
assert_equal(self.pbs.resources["pmem"], 123) |
|
56
|
|
|
assert_equal(self.pbs.resources["walltime"], "42:42") |
|
57
|
|
|
|
|
58
|
|
|
def test_add_modules_to_load(self): |
|
59
|
|
|
assert_equal(self.pbs.modules, []) |
|
60
|
|
|
|
|
61
|
|
|
modules = ["cuda", "gcc", "python/2.7"] |
|
62
|
|
|
self.pbs.add_modules_to_load(modules[0], modules[1]) |
|
63
|
|
|
assert_array_equal(self.pbs.modules, modules[:2]) |
|
64
|
|
|
|
|
65
|
|
|
self.pbs.add_modules_to_load(modules[-1]) |
|
66
|
|
|
assert_array_equal(self.pbs.modules, modules) |
|
67
|
|
|
|
|
68
|
|
|
def test_add_commands(self): |
|
69
|
|
|
assert_equal(self.pbs.commands, []) |
|
70
|
|
|
|
|
71
|
|
|
commands = ["python my_script.py", "echo 1234", "git push origin master"] |
|
72
|
|
|
self.pbs.add_commands(commands[0]) |
|
73
|
|
|
assert_array_equal(self.pbs.commands, commands[0]) |
|
74
|
|
|
|
|
75
|
|
|
self.pbs.add_commands(*commands[1:]) |
|
76
|
|
|
assert_array_equal(self.pbs.commands, commands) |
|
77
|
|
|
|
|
78
|
|
|
def test_str(self): |
|
79
|
|
|
# Create simple PBS file |
|
80
|
|
|
expected = """#!/bin/bash |
|
81
|
|
|
#PBS -q qtest@mp2 |
|
82
|
|
|
#PBS -V |
|
83
|
|
|
#PBS -l walltime=01:00:00 |
|
84
|
|
|
|
|
85
|
|
|
# Modules # |
|
86
|
|
|
|
|
87
|
|
|
# Prolog # |
|
88
|
|
|
|
|
89
|
|
|
# Commands # |
|
90
|
|
|
|
|
91
|
|
|
# Epilog #""" |
|
92
|
|
|
pbs = PBS(queue_name="qtest@mp2", walltime="01:00:00") |
|
93
|
|
|
assert_equal(str(pbs), expected) |
|
94
|
|
|
|
|
95
|
|
|
# Create a more complex PBS file |
|
96
|
|
|
prolog = ['echo "This is prolog1"', |
|
97
|
|
|
'echo "This is prolog2"'] |
|
98
|
|
|
commands = ["cd $HOME; echo 1 2 3 1>> cmd1.o 2>> cmd1.e", |
|
99
|
|
|
"cd $HOME; echo 3 2 1 1>> cmd1.o 2>> cmd1.e"] |
|
100
|
|
|
epilog = ['echo "This is epilog1"', |
|
101
|
|
|
'echo "This is epilog2"'] |
|
102
|
|
|
modules = ["CUDA_Toolkit/6.0", "python2.7"] |
|
103
|
|
|
|
|
104
|
|
|
expected = """\ |
|
105
|
|
|
#!/bin/bash |
|
106
|
|
|
#PBS -q qtest@mp2 |
|
107
|
|
|
#PBS -V |
|
108
|
|
|
#PBS -A xyz-123-ab |
|
109
|
|
|
#PBS -l walltime=01:00:00 |
|
110
|
|
|
#PBS -l nodes=2:ppn=3:gpus=1 |
|
111
|
|
|
|
|
112
|
|
|
# Modules # |
|
113
|
|
|
module load CUDA_Toolkit/6.0 |
|
114
|
|
|
module load python2.7 |
|
115
|
|
|
|
|
116
|
|
|
# Prolog # |
|
117
|
|
|
{prolog1} |
|
118
|
|
|
{prolog2} |
|
119
|
|
|
|
|
120
|
|
|
# Commands # |
|
121
|
|
|
{command1} |
|
122
|
|
|
{command2} |
|
123
|
|
|
|
|
124
|
|
|
# Epilog # |
|
125
|
|
|
{epilog1} |
|
126
|
|
|
{epilog2}""".format( |
|
127
|
|
|
prolog1=prolog[0], prolog2=prolog[1], |
|
128
|
|
|
command1=commands[0], command2=commands[1], |
|
129
|
|
|
epilog1=epilog[0], epilog2=epilog[1]) |
|
130
|
|
|
|
|
131
|
|
|
pbs = PBS(queue_name="qtest@mp2", walltime="01:00:00") |
|
132
|
|
|
pbs.add_resources(nodes="2:ppn=3:gpus=1") |
|
133
|
|
|
pbs.add_options(A="xyz-123-ab") |
|
134
|
|
|
pbs.add_modules_to_load(*modules) |
|
135
|
|
|
pbs.add_to_prolog(*prolog) |
|
136
|
|
|
pbs.add_commands(*commands) |
|
137
|
|
|
pbs.add_to_epilog(*epilog) |
|
138
|
|
|
|
|
139
|
|
|
assert_equal(str(pbs), expected) |
|
140
|
|
|
|
|
141
|
|
|
def test_save(self): |
|
142
|
|
|
pbs_filename = os.path.join(self.testing_dir, "pbs.sh") |
|
143
|
|
|
self.pbs.save(pbs_filename) |
|
144
|
|
|
assert_equal(str(self.pbs), open(pbs_filename).read()) |
|
145
|
|
|
|
SMART-Lab /
smartdispatch
