Coverage for python / lsst / ctrl / bps / parsl / job.py: 22%
91 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:19 -0800
1# This file is part of ctrl_bps_parsl.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org) and the LSST DESC (https://www.lsstdesc.org/).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <https://www.gnu.org/licenses/>.
28import os
29import re
30import subprocess
31from collections import defaultdict
32from collections.abc import Sequence
33from functools import partial
34from textwrap import dedent
35from typing import Any
37from parsl.app.bash import BashApp
38from parsl.app.futures import Future
40from lsst.ctrl.bps import BpsConfig, GenericWorkflow, GenericWorkflowJob
42from .configuration import get_bps_config_value
44__all__ = ("ParslJob", "get_file_paths")
46_env_regex = re.compile(r"<ENV:(\S+)>") # Regex for replacing <ENV:WHATEVER> in BPS job command-lines
47_file_regex = re.compile(r"<FILE:(\S+)>") # Regex for replacing <FILE:WHATEVER> in BPS job command-lines
50def run_command(
51 command_line: str,
52 inputs: Sequence[Future] = (),
53 stdout: str | None = None,
54 stderr: str | None = None,
55 parsl_resource_specification: dict[str, Any] | None = None,
56) -> str:
57 """Run a command.
59 This function exists to get information into parsl, through the ``inputs``,
60 ``stdout`` and ``stderr`` parameters. It needs to be wrapped by a parsl
61 ``bash_app`` decorator before use, after which it will return a `Future`.
63 Parameters
64 ----------
65 command_line : `str`
66 Command-line to have parsl run.
67 inputs : list of `Future`
68 Other commands that must have run before this.
69 stdout, stderr : `str`, optional
70 Filenames for stdout and stderr.
71 parsl_resource_specification : `dict`, optional
72 Resources required for job.
74 Returns
75 -------
76 command_line : `str`
77 Command-line to have parsl run.
78 """
79 return command_line
82def get_file_paths(workflow: GenericWorkflow, name: str) -> dict[str, str]:
83 """Extract file paths for a job.
85 Parameters
86 ----------
87 workflow : `GenericWorkflow`
88 BPS workflow that knows the file paths.
89 name : `str`
90 Job name.
92 Returns
93 -------
94 paths : `dict` mapping `str` to `str`
95 File paths for job, indexed by symbolic name.
96 """
97 return {ff.name: ff.src_uri for ff in workflow.get_job_inputs(name)}
100class ParslJob:
101 """Job to execute with parsl.
103 Parameters
104 ----------
105 generic : `GenericWorkflowJob`
106 BPS job information.
107 config : `BpsConfig`
108 BPS configuration.
109 file_paths : `dict` mapping `str` to `str`
110 File paths for job, indexed by symbolic name.
111 """
113 def __init__(
114 self,
115 generic: GenericWorkflowJob,
116 config: BpsConfig,
117 file_paths: dict[str, str],
118 ):
119 self.generic = generic
120 self.name = generic.name
121 self.config = config
122 self.file_paths = file_paths
123 self.future = None
124 self.done = False
126 # Determine directory for job stdout and stderr
127 log_dir = os.path.join(get_bps_config_value(self.config, "submitPath", str, required=True), "logs")
128 _, template = self.config.search(
129 "subDirTemplate",
130 opt={
131 "curvals": {"curr_site": self.config["computeSite"], "curr_cluster": self.generic.label},
132 "replaceVars": False,
133 "default": "",
134 },
135 )
136 job_vals = defaultdict(str)
137 job_vals["label"] = self.generic.label
138 if self.generic.tags:
139 job_vals.update(self.generic.tags)
140 subdir = template.format_map(job_vals)
141 # Call normpath just to make paths easier to read as templates tend
142 # to have variables that aren't used by every job. Avoid calling on
143 # empty string because changes it to dot.
144 same_part = os.path.normpath(os.path.join(log_dir, subdir, self.name))
145 self.stdout = same_part + ".stdout"
146 self.stderr = same_part + ".stderr"
148 def __reduce__(self):
149 """Recipe for pickling."""
150 return type(self), (self.generic, self.config, self.file_paths)
152 def get_command_line(self, allow_stage=True) -> str:
153 """Get the bash command-line to run to execute this job.
155 Parameters
156 ----------
157 allow_stage : `bool`
158 Allow staging of execution butler? This is not appropriate for the
159 initial or final jobs that run on the local nodes.
161 Returns
162 -------
163 command : `str`
164 Command-line to execute for job.
165 """
166 command: str = self.generic.executable.src_uri + " " + self.generic.arguments
167 if not allow_stage:
168 return command
169 exec_butler_dir = get_bps_config_value(self.config, "executionButlerDir", str)
170 if not exec_butler_dir or not os.path.isdir(exec_butler_dir):
171 # We're not using the execution butler
172 return command
174 # Add commands to copy the execution butler.
175 # This keeps workers from overloading the sqlite database.
176 # The copy can be deleted once we're done, because the original
177 # execution butler contains everything that's required.
178 job_dir = os.path.join(os.path.dirname(exec_butler_dir), self.name)
179 # Set the butlerConfig field to the location of the job-specific copy.
180 command = command.replace("<FILE:butlerConfig>", job_dir)
181 return dedent(
182 f"""
183 if [[ ! -d {job_dir} ]]; then mkdir -p {job_dir}; fi
184 cp {exec_butler_dir}/* {job_dir}
185 {command}
186 retcode=$?
187 rm -rf {job_dir}
188 exit $retcode
189 """
190 )
192 def evaluate_command_line(self, command: str) -> str:
193 """Evaluate the bash command-line.
195 BPS provides a command-line with symbolic names for BPS variables,
196 environment variables and files. Here, we replace those symbolic names
197 with the actual values, to provide a concrete command that can be
198 executed.
200 In replacing file paths, we are implicitly assuming that we are working
201 on a shared file system, i.e., that workers can see the butler
202 directory, and that files do not need to be staged to the worker.
204 Parameters
205 ----------
206 command : `str`
207 Command-line to execute, from BPS.
209 Returns
210 -------
211 command : `str`
212 Command ready for execution on a worker.
213 """
214 command = command.format(**self.generic.cmdvals) # BPS variables
216 # Make sure *all* symbolic names are resolved.
217 #
218 # In general, actual values for some symbolic names may contain other
219 # symbolic names. As a result, more than one iteration may be required
220 # to resolve all symbolic names. For example, an actual value for
221 # a filename may contain a symbolic name for an environment variable.
222 prev_command = command
223 while True:
224 command = re.sub(_env_regex, r"${\g<1>}", command) # Environment variables
225 command = re.sub(_file_regex, lambda match: self.file_paths[match.group(1)], command) # Files
226 if prev_command == command:
227 break
228 prev_command = command
230 return command
232 def get_resources(self) -> dict[str, Any]:
233 """Return what resources are required for executing this job."""
234 resources = {}
235 for bps_name, parsl_name, scale in (
236 ("request_memory", "memory", None), # Both BPS and WorkQueueExecutor use MB
237 ("request_cpus", "cores", None),
238 ("request_disk", "disk", None), # Both are MB
239 ("request_walltime", "running_time_min", None), # Both are minutes
240 ("priority", "priority", None),
241 ):
242 value = getattr(self.generic, bps_name)
243 if value is not None and scale is not None:
244 value *= scale
245 # Parsl's `HighThroughputExecutor` cannot have
246 # `priority=None`, but it can be omitted. By contrast,
247 # `WorkQueueExecutor` needs to have the other resource
248 # requests provided, but `priority` can be omitted, so we
249 # need special handling for `priority`.
250 if (parsl_name == "priority" and value is not None) or parsl_name != "priority":
251 resources[parsl_name] = value
252 return resources
254 def get_future(
255 self,
256 app: BashApp,
257 inputs: list[Future],
258 command_prefix: str | None = None,
259 resource_list: list | None = None,
260 ) -> Future | None:
261 """Get the parsl app future for the job.
263 This effectively queues the job for execution by a worker, subject to
264 dependencies.
266 Parameters
267 ----------
268 app : callable
269 A parsl bash_app decorator to use.
270 inputs : list of `Future`
271 Dependencies to be satisfied before executing this job.
272 command_prefix : `str`, optional
273 Bash commands to execute before the job command, e.g., for setting
274 the environment.
275 resource_list : `list`, optional
276 List of resource specifications to pass to the Parsl executor.
278 Returns
279 -------
280 future : `Future` or `None`
281 A `Future` object linked to the execution of the job, or `None` if
282 the job has already been done (e.g., by ``run_local``).
283 """
284 if self.done:
285 return None # Nothing to do
286 if not self.future:
287 command = self.get_command_line()
288 command = self.evaluate_command_line(command)
289 if command_prefix:
290 command = command_prefix + "\n" + command
291 resources = (
292 {k: v for k, v in self.get_resources().items() if k in resource_list}
293 if resource_list is not None
294 else {}
295 )
297 # Add a layer of indirection to which we can add a useful name.
298 # This name is used by parsl for tracking workflow status.
299 func = partial(run_command)
300 func.__name__ = self.generic.label # type: ignore
302 self.future = app(func)(
303 command,
304 inputs=inputs,
305 stdout=self.stdout,
306 stderr=self.stderr,
307 parsl_resource_specification=resources,
308 )
309 return self.future
311 def run_local(self):
312 """Run the command locally.
314 This is intended to support jobs that should not be done by a
315 worker.
316 """
317 if self.done: # Nothing to do
318 return
319 command = self.get_command_line(False)
320 command = self.evaluate_command_line(command)
321 os.makedirs(os.path.dirname(self.stdout), exist_ok=True)
322 os.makedirs(os.path.dirname(self.stderr), exist_ok=True)
323 with open(self.stdout, "w") as stdout, open(self.stderr, "w") as stderr:
324 subprocess.check_call(command, shell=True, executable="/bin/bash", stdout=stdout, stderr=stderr)
325 self.done = True