1
|
|
|
import re |
2
|
|
|
from collections import OrderedDict |
3
|
|
|
|
4
|
|
|
regex_walltime = re.compile("(\d+:){1,4}") |
5
|
|
|
regex_resource_nodes = re.compile("[a-zA-Z0-9]+(:ppn=\d+)?(:gpus=\d+)?(:[a-zA-Z0-9]+)*") |
6
|
|
|
regex_resource_pmem = re.compile("[0-9]+(b|kb|mb|gb|tb)?") |
7
|
|
|
|
8
|
|
|
|
9
|
|
|
class PBS(object): |
10
|
|
|
""" Offers functionalities to manage a PBS file. |
11
|
|
|
|
12
|
|
|
For more information about the PBS file format see: |
13
|
|
|
`http://docs.adaptivecomputing.com/suite/8-0/basic/help.htm#topics/torque/2-jobs/requestingRes.htm?TocPath=TORQUE Resource Manager|Submitting and managing jobs|Job submission|_____3` |
14
|
|
|
|
15
|
|
|
Parameters |
16
|
|
|
---------- |
17
|
|
|
queue_name : str |
18
|
|
|
name of the queue on which commands will be executed |
19
|
|
|
walltime : str |
20
|
|
|
maximum time allocated to execute every commands (DD:HH:MM:SS) |
21
|
|
|
""" |
22
|
|
|
def __init__(self, queue_name, walltime): |
23
|
|
|
if queue_name is None or len(queue_name) == 0: |
24
|
|
|
raise ValueError("Queue's name must be provided.") |
25
|
|
|
|
26
|
|
|
self.queue_name = queue_name |
27
|
|
|
self.modules = [] |
28
|
|
|
self.prolog = [] |
29
|
|
|
self.commands = [] |
30
|
|
|
self.epilog = [] |
31
|
|
|
|
32
|
|
|
self.resources = OrderedDict() |
33
|
|
|
self.add_resources(walltime=walltime) |
34
|
|
|
|
35
|
|
|
self.options = OrderedDict() |
36
|
|
|
self.add_options(q=queue_name) |
37
|
|
|
|
38
|
|
|
# Declares that all environment variables in the qsub command's environment are to be exported to the batch job. |
39
|
|
|
self.add_options(V="") |
40
|
|
|
|
41
|
|
|
def add_options(self, **options): |
42
|
|
|
""" Adds options to this PBS file. |
43
|
|
|
|
44
|
|
|
Parameters |
45
|
|
|
---------- |
46
|
|
|
**options : dict |
47
|
|
|
each key is the name of a PBS option (see `Options`) |
48
|
|
|
|
49
|
|
|
Options |
50
|
|
|
------- |
51
|
|
|
*A* : account_string |
52
|
|
|
Defines the account string associated with the job. |
53
|
|
|
*N* : name (up to 64 characters) |
54
|
|
|
Declares a name for the job. It must consist of printable, |
55
|
|
|
non white space characters with the first character alphabetic. |
56
|
|
|
""" |
57
|
|
|
for option_name, option_value in options.items(): |
58
|
|
|
# If known option, validate it. |
59
|
|
|
if option_name.strip('-') == 'N': |
60
|
|
|
if len(option_name) > 64: |
61
|
|
|
raise ValueError("Maximum number of characters for the name is: 64") |
62
|
|
|
|
63
|
|
|
self.options["-" + option_name] = option_value |
64
|
|
|
|
65
|
|
|
def add_resources(self, **resources): |
66
|
|
|
""" Adds resources to this PBS file. |
67
|
|
|
|
68
|
|
|
Parameters |
69
|
|
|
---------- |
70
|
|
|
**resources : dict |
71
|
|
|
each key is the name of a PBS resource (see `Resources`) |
72
|
|
|
|
73
|
|
|
Resources |
74
|
|
|
--------- |
75
|
|
|
*nodes* : nodes={<node_count>|<hostname>}[:ppn=<ppn>][:gpus=<gpu>][:<property>[:<property>]...] |
76
|
|
|
Specifies how many and what type of nodes to use |
77
|
|
|
**nodes={<node_count>|<hostname>}**: type of nodes |
78
|
|
|
**ppn=#**: Number of process per node requested for this job |
79
|
|
|
**gpus=#**: Number of process per node requested for this job |
80
|
|
|
**property**: A string specifying a node's feature |
81
|
|
|
*pmem*: pmem=[0-9]+(b|kb|mb|gb|tb) |
82
|
|
|
Specifies the maximum amount of physical memory used by any single process of the job. |
83
|
|
|
""" |
84
|
|
|
for resource_name, resource_value in resources.items(): |
85
|
|
|
# If known ressource, validate it. |
86
|
|
|
if resource_name == 'nodes': |
87
|
|
|
if re.match(regex_resource_nodes, str(resource_value)) is None: |
88
|
|
|
raise ValueError("Unknown format for PBS resource: nodes") |
89
|
|
|
elif resource_name == 'pmem': |
90
|
|
|
if re.match(regex_resource_pmem, str(resource_value)) is None: |
91
|
|
|
raise ValueError("Unknown format for PBS resource: pmem") |
92
|
|
|
elif resource_name == 'walltime': |
93
|
|
|
if re.match(regex_walltime, str(resource_value)) is None: |
94
|
|
|
raise ValueError("Unknown format for PBS resource: walltime (dd:hh:mm:ss)") |
95
|
|
|
|
96
|
|
|
self.resources[resource_name] = resource_value |
97
|
|
|
|
98
|
|
|
def add_modules_to_load(self, *modules): |
99
|
|
|
""" Adds modules to load prior to execute the job on a node. |
100
|
|
|
|
101
|
|
|
Parameters |
102
|
|
|
---------- |
103
|
|
|
*modules : list of str |
104
|
|
|
each string represents the name of the module to load |
105
|
|
|
""" |
106
|
|
|
self.modules += modules |
107
|
|
|
|
108
|
|
|
def add_to_prolog(self, *code): |
109
|
|
|
""" Adds the code to be executed before the commands. |
110
|
|
|
|
111
|
|
|
Parameters |
112
|
|
|
---------- |
113
|
|
|
*code : list of str |
114
|
|
|
Each string holds the code to be executed before the commands |
115
|
|
|
""" |
116
|
|
|
self.prolog += code |
117
|
|
|
|
118
|
|
|
def add_commands(self, *commands): |
119
|
|
|
""" Sets commands to execute on a node. |
120
|
|
|
|
121
|
|
|
Parameters |
122
|
|
|
---------- |
123
|
|
|
*commands : list of str |
124
|
|
|
each string represents a command that is part of this job |
125
|
|
|
""" |
126
|
|
|
self.commands += commands |
127
|
|
|
|
128
|
|
|
def add_to_epilog(self, *code): |
129
|
|
|
""" Adds the code to be executed after the commands. |
130
|
|
|
|
131
|
|
|
Parameters |
132
|
|
|
---------- |
133
|
|
|
*code : list of str |
134
|
|
|
Each string holds the code to be executed after the commands |
135
|
|
|
""" |
136
|
|
|
self.epilog += code |
137
|
|
|
|
138
|
|
|
def save(self, filename): |
139
|
|
|
""" Saves this PBS job to a file. |
140
|
|
|
|
141
|
|
|
Parameters |
142
|
|
|
---------- |
143
|
|
|
filename : str |
144
|
|
|
specified where to save this PBS file |
145
|
|
|
""" |
146
|
|
|
with open(filename, 'w') as pbs_file: |
147
|
|
|
pbs_file.write(str(self)) |
148
|
|
|
|
149
|
|
|
def __str__(self): |
150
|
|
|
pbs = [] |
151
|
|
|
pbs += ["#!/bin/bash"] |
152
|
|
|
|
153
|
|
|
for option_name, option_value in self.options.items(): |
154
|
|
|
if option_value == "": |
155
|
|
|
pbs += ["#PBS {0}".format(option_name)] |
156
|
|
|
else: |
157
|
|
|
pbs += ["#PBS {0} {1}".format(option_name, option_value)] |
158
|
|
|
|
159
|
|
|
for resource_name, resource_value in self.resources.items(): |
160
|
|
|
pbs += ["#PBS -l {0}={1}".format(resource_name, resource_value)] |
161
|
|
|
|
162
|
|
|
pbs += ["\n# Modules #"] |
163
|
|
|
for module in self.modules: |
164
|
|
|
pbs += ["module load " + module] |
165
|
|
|
|
166
|
|
|
pbs += ["\n# Prolog #"] |
167
|
|
|
pbs += self.prolog |
168
|
|
|
|
169
|
|
|
pbs += ["\n# Commands #"] |
170
|
|
|
pbs += ["{command}".format(command=command) for command in self.commands] |
171
|
|
|
|
172
|
|
|
pbs += ["\n# Epilog #"] |
173
|
|
|
pbs += self.epilog |
174
|
|
|
|
175
|
|
|
return "\n".join(pbs) |
176
|
|
|
|