Coverage for tests/test_transform.py: 100%
250 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 08:59 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 08:59 +0000
1# This file is part of ctrl_bps.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.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/>.
27"""Unit tests of transform.py."""
29import copy
30import dataclasses
31import logging
32import os
33import shutil
34import tempfile
35import unittest
37from cqg_test_utils import make_test_2_cluster_cqg
39from lsst.ctrl.bps import (
40 BPS_DEFAULTS,
41 BPS_SEARCH_ORDER,
42 BpsConfig,
43 GenericWorkflow,
44 GenericWorkflowExec,
45 GenericWorkflowJob,
46)
47from lsst.ctrl.bps.transform import (
48 _enhance_command,
49 _get_job_values,
50 create_final_command,
51 create_generic_workflow,
52 create_generic_workflow_config,
53 gather_job_environment,
54)
56TESTDIR = os.path.abspath(os.path.dirname(__file__))
59class TestCreateGenericWorkflowConfig(unittest.TestCase):
60 """Tests of create_generic_workflow_config."""
62 def testCreate(self):
63 """Test successful creation of the config."""
64 config = BpsConfig({"a": 1, "b": 2, "uniqProcName": "testCreate"})
65 wf_config = create_generic_workflow_config(config, "/test/create/prefix")
66 self.assertIsInstance(wf_config, BpsConfig)
67 for key in config:
68 self.assertEqual(wf_config[key], config[key])
69 self.assertEqual(wf_config["workflowName"], "testCreate")
70 self.assertEqual(wf_config["workflowPath"], "/test/create/prefix")
73class TestCreateGenericWorkflow(unittest.TestCase):
74 """Tests of create_generic_workflow."""
76 def setUp(self):
77 logging.basicConfig(level=logging.WARNING)
78 logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO)
79 self.tmpdir = tempfile.mkdtemp(dir=TESTDIR)
80 filename = os.path.join(TESTDIR, "data/config_for_transform.yaml")
81 self.orig_config = BpsConfig(
82 filename, BPS_SEARCH_ORDER, BPS_DEFAULTS, wms_service_class_fqn="wms_test_utils.WmsServiceSuccess"
83 )
84 _, self.cqg = make_test_2_cluster_cqg(self.tmpdir)
86 def tearDown(self):
87 shutil.rmtree(self.tmpdir, ignore_errors=True)
89 def testCreatingQuantumGraph(self):
90 """Test creating a GenericWorkflow with setting overrides. While
91 other tests exist to check get_job_valuei,gather_job_environment,
92 etc., going ahead and checking the integration of these by checking
93 results instead of using mocks.
94 """
95 config = BpsConfig(self.orig_config)
96 config[".computeSite"] = "site1"
97 config[".finalJob.queue"] = "special_final_queue"
98 config[".finalJob.computeSite"] = "special_site"
99 config[".finalJob.computeCloud"] = "special_cloud"
100 workflow = create_generic_workflow(config, self.cqg, "test_gw", self.tmpdir)
101 self.assertEqual(len(workflow) - 1, len(self.cqg)) # Don't count pipetaskInit
102 for jname in workflow:
103 gwjob = workflow.get_job(jname)
104 self.assertEqual(gwjob.compute_site, "site1", f"failed for job {gwjob}")
105 self.assertIsNone(gwjob.compute_cloud, f"failed for job {gwjob}")
106 if gwjob.label == "pipetaskInit":
107 self.assertEqual(gwjob.executable.src_uri, "pipetask", f"failed for job {gwjob}")
108 else:
109 self.assertEqual(gwjob.executable.src_uri, "s1exe", f"failed for job {gwjob}")
110 self.assertIn("{qgraphNodeId}", gwjob.arguments)
112 base_site_env_truth = {
113 "VAR2": "root_val2",
114 "VAR3": "root_val3",
115 "VAR4": "site1_val4",
116 "VAR_PATH": "<ENV:PACKAGE_DIR>/site1_dir:<ENV:PACKAGE_DIR>/root_dir:<ENV:VAR_PATH>",
117 "TEST_VAR": "one site1_val1 three",
118 }
120 match gwjob.label:
121 case "pipetaskInit":
122 env_truth = dict(base_site_env_truth)
123 env_truth["TEST_VAR"] = "one init_val1 three"
124 self.assertEqual(gwjob.environment, env_truth)
125 self.assertEqual(gwjob.cmdvals["initPreCmdOpts"], "--log-level=DEBUG")
126 # Taking pipetaskInit-defined requestMemory.
127 self.assertEqual(gwjob.request_memory, 8096)
128 case "clusterT1T2":
129 # Taking cluster-defined requestMemory.
130 self.assertEqual(gwjob.request_memory, 6144)
131 self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 2)
132 case "clusterT3T4":
133 # Taking max of the requestMemory for quanta in cluster.
134 self.assertEqual(gwjob.request_memory, 4048)
135 self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 2)
136 case "T2b":
137 # Taking default requestMemory from root section.
138 self.assertEqual(gwjob.request_memory, BPS_DEFAULTS["requestMemory"])
139 self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 1)
140 case "T5":
141 # Taking default requestMemory from root section.
142 self.assertEqual(gwjob.request_memory, BPS_DEFAULTS["requestMemory"])
143 self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 1)
144 case _:
145 # Should always have a label from above, but need to
146 # fail test if get different label.
147 self.fail(f"Invalid gwjob.label for job {gwjob}") # pragma: no cover
148 final = workflow.get_final()
149 self.assertEqual(final.compute_site, "special_site", f"failed for final job {final}")
150 self.assertEqual(final.compute_cloud, "special_cloud", f"failed for final job {final}")
151 self.assertEqual(final.queue, "special_final_queue", f"failed for final job {final}")
152 self.assertEqual(final.request_memory, BPS_DEFAULTS["finalJob"]["requestMemory"])
155class TestGetJobValues(unittest.TestCase):
156 """Tests of _get_job_values."""
158 def setUp(self):
159 logging.basicConfig(level=logging.WARNING)
160 logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO)
161 self.default_job = GenericWorkflowJob("default_job", "default_label")
163 def testGettingDefaults(self):
164 """Test retrieving default values."""
165 config = BpsConfig({})
166 job_values = _get_job_values(config, {}, None)
167 self.assertTrue(
168 all(
169 getattr(self.default_job, field.name) == job_values[field.name]
170 for field in dataclasses.fields(self.default_job)
171 )
172 )
174 def testEnablingMemoryScaling(self):
175 """Test enabling the memory scaling mechanism."""
176 config = BpsConfig({"memoryMultiplier": 2.0})
177 job_values = _get_job_values(config, {}, None)
178 self.assertAlmostEqual(job_values["memory_multiplier"], 2.0)
179 self.assertEqual(job_values["number_of_retries"], 5)
181 def testDisablingMemoryScaling(self):
182 """Test disabling the memory scaling mechanism."""
183 config = BpsConfig({"memoryMultiplier": 0.5})
184 job_values = _get_job_values(config, {}, None)
185 self.assertIsNone(job_values["memory_multiplier"])
187 def testRetrievingCmdLine(self):
188 """Test retrieving the command line."""
189 cmd_line_key = "runQuantum"
190 config = BpsConfig({cmd_line_key: "/path/to/foo bar.txt"})
191 job_values = _get_job_values(config, {}, cmd_line_key)
192 self.assertEqual(job_values["executable"].name, "foo")
193 self.assertEqual(job_values["executable"].src_uri, "/path/to/foo")
194 self.assertEqual(job_values["arguments"], "bar.txt")
196 @unittest.mock.patch("lsst.ctrl.bps.transform.gather_job_environment")
197 def testCallGatherJobEnvironmentNoSearchOpts(self, mock_gather):
198 # Test that _get_job_values passes right search options on
199 # to gather_job_environment function and didn't have side-effects.
200 env_truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"}
201 mock_gather.return_value = dict(env_truth)
202 config = BpsConfig(
203 {
204 "var1": "two",
205 "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} three"},
206 "runQuantumCommand": "/path/to/foo bar.txt",
207 "pipetask": {"isr": {"requestMemory": 8096, "environment": {"ISR_VAR": "45"}}},
208 }
209 )
210 search_opts = {}
211 config_copy = BpsConfig(config)
212 search_opts_copy = dict(search_opts)
214 job_values = _get_job_values(config, search_opts, None)
215 env_truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"}
216 self.assertEqual(env_truth, job_values["environment"])
217 mock_gather.assert_called_once_with(config, search_opts)
219 # And didn't have side-effects that changed vars
220 self.assertEqual(config, config_copy)
221 self.assertEqual(search_opts, search_opts_copy)
223 @unittest.mock.patch("lsst.ctrl.bps.transform.gather_job_environment")
224 def testCallGatherJobEnvironmentWithCurvals(self, mock_gather):
225 # Test that _get_job_values passes right search options on
226 # to gather_job_environment function and didn't have side-effects.
227 env_truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three", "ISR_VAR": "45"}
228 mock_gather.return_value = dict(env_truth)
229 config = BpsConfig(
230 {
231 "var1": "two",
232 "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} three"},
233 "runQuantumCommand": "/path/to/foo bar.txt",
234 "pipetask": {"isr": {"requestMemory": 8096, "environment": {"ISR_VAR": "45"}}},
235 }
236 )
237 curvals = {"curr_pipetask": "isr"}
238 search_opts = {"replaceVars": False, "searchobj": {"curvals": curvals}}
240 # Save copies to check no side-effects
241 config_copy = BpsConfig(config)
242 search_opts_copy = dict(search_opts)
244 job_values = _get_job_values(config, search_opts, "runQuantumCommand")
245 mock_gather.assert_called_once_with(config, search_opts)
247 self.assertEqual(job_values["environment"], env_truth)
248 self.assertEqual(job_values["executable"].src_uri, "/path/to/foo")
250 # And didn't have side-effects that changed vars
251 self.assertEqual(config, config_copy)
252 self.assertEqual(search_opts, search_opts_copy)
255class TestGatherJobEnvironment(unittest.TestCase):
256 """Tests for the gather_job_environment function."""
258 def setUp(self):
259 logging.basicConfig(level=logging.WARNING)
260 logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO)
261 # The directories don't match real ones, but are here to test
262 # environment variables in yaml environment section as well as
263 # appending values across sections.
264 filename = os.path.join(TESTDIR, "data/config_for_transform.yaml")
265 self.orig_config = BpsConfig(
266 filename, BPS_SEARCH_ORDER, BPS_DEFAULTS, wms_service_class_fqn="wms_test_utils.WmsServiceSuccess"
267 )
269 def testEnvironmentRootSiteCluster(self):
270 search_opts = {
271 "replaceVars": False,
272 "curvals": {"curr_cluster": "clusterT1T2", "curr_site": "site1"},
273 }
274 job_env = gather_job_environment(self.orig_config, search_opts)
275 truth = {
276 "VAR3": "cl12_val3",
277 "VAR4": "cl12_val4",
278 "VAR_PATH": "<ENV:PACKAGE_DIR>/cl12_dir:<ENV:PACKAGE_DIR>/site1_dir:<ENV:PACKAGE_DIR>/root_dir"
279 ":<ENV:VAR_PATH>",
280 "TEST_VAR": "one cl12_val1 three",
281 }
282 self.assertEqual(truth, job_env)
283 self.assertEqual(search_opts["replaceVars"], False)
285 def testEnvironmentRootSite(self):
286 # Checking that doesn't pick up env from other cluster
287 search_opts = {
288 "replaceVars": False,
289 "curvals": {"curr_cluster": "notthere", "curr_site": "site1"},
290 }
291 job_env = gather_job_environment(self.orig_config, search_opts)
292 truth = {
293 "VAR2": "root_val2",
294 "VAR3": "root_val3",
295 "VAR4": "site1_val4",
296 "VAR_PATH": "<ENV:PACKAGE_DIR>/site1_dir:<ENV:PACKAGE_DIR>/root_dir:<ENV:VAR_PATH>",
297 "TEST_VAR": "one site1_val1 three",
298 }
299 self.assertEqual(truth, job_env)
300 self.assertEqual(search_opts["replaceVars"], False)
302 def testEnvironmentRoot(self):
303 # Checking that doesn't pick up env from other cluster or site.
304 # Also check that doesn't modify our search_opts by setting opposites
305 # of what function uses.
306 orig_search_opts = {
307 "replaceVars": False,
308 "replaceEnvBps2Shell": True,
309 "replaceEnvShell2Bps": False,
310 "expandEnvVars": True,
311 "curvals": {"curr_cluster": "notthere", "curr_site": "notthere"},
312 }
313 search_opts = copy.deepcopy(orig_search_opts)
315 job_env = gather_job_environment(self.orig_config, search_opts)
316 truth = {
317 "VAR2": "root_val2",
318 "VAR3": "root_val3",
319 "VAR_PATH": "<ENV:PACKAGE_DIR>/root_dir:<ENV:VAR_PATH>",
320 "TEST_VAR": "one root_val1 three",
321 }
322 self.assertEqual(truth, job_env)
323 self.assertEqual(orig_search_opts, search_opts)
325 def testEnvironmentNoSearchOpts(self):
326 search_opts = {}
327 job_env = gather_job_environment(self.orig_config, search_opts)
328 truth = {
329 "VAR2": "root_val2",
330 "VAR_PATH": "<ENV:PACKAGE_DIR>/root_dir:<ENV:VAR_PATH>",
331 "VAR3": "root_val3",
332 "TEST_VAR": "one root_val1 three",
333 }
334 self.assertEqual(truth, job_env)
335 self.assertEqual(search_opts, {})
337 def testSearchObj(self):
338 # Test that works with searchobj, like finalJob
339 search_opts = {"searchobj": self.orig_config["finalJob"], "curvals": {"curr_site": "site1"}}
340 copy_final = BpsConfig(self.orig_config["finalJob"])
341 job_env = gather_job_environment(self.orig_config, search_opts)
342 # VAR3 and VAR4 removed in setUp
343 truth = {
344 "TEST_VAR": "one final_val1 three",
345 "VAR2": "root_val2",
346 "VAR5": "final_val5",
347 "VAR_PATH": "<ENV:PACKAGE_DIR>/final_dir:<ENV:PACKAGE_DIR>/site1_dir:<ENV:PACKAGE_DIR>/root_dir"
348 ":<ENV:VAR_PATH>",
349 }
350 self.assertEqual(truth, job_env)
351 self.assertEqual(search_opts["searchobj"], copy_final)
353 def testVarsInEnvironment(self):
354 config = BpsConfig(
355 {
356 "var1": "two",
357 "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} <ENV:var3>"},
358 }
359 )
360 job_values = _get_job_values(config, {"replaceVars": True}, None)
361 truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two <ENV:var3>"}
362 self.assertEqual(truth, job_values["environment"])
365class TestCreateFinalCommand(unittest.TestCase):
366 """Tests for the create_final_command function."""
368 def setUp(self):
369 logging.basicConfig(level=logging.WARNING)
370 logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO)
371 self.tmpdir = tempfile.TemporaryDirectory()
372 self.script_beginning = [
373 "#!/bin/bash\n",
374 "\n",
375 "set -e\n",
376 "set -x\n",
377 "qgraphFile=$1\n",
378 "butlerConfig=$2\n",
379 ]
381 def tearDown(self):
382 self.tmpdir.cleanup()
384 def testSingleCommand(self):
385 """Test with single final job command."""
386 config_butler = f"{self.tmpdir.name}/test_repo"
387 config = BpsConfig(
388 {
389 "var1": "42a",
390 "var2": "42b",
391 "var3": "42c",
392 "butlerConfig": config_butler,
393 "finalJob": {"command1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}"},
394 }
395 )
396 gwf_exec, args = create_final_command(config, self.tmpdir.name)
397 self.assertEqual(args, f"<FILE:runQgraphFile> {config_butler}")
398 final_script = f"{self.tmpdir.name}/final_job.bash"
399 self.assertEqual(gwf_exec.src_uri, final_script)
400 with open(final_script) as infh:
401 lines = infh.readlines()
402 self.assertEqual(
403 lines, self.script_beginning + ["/usr/bin/echo 42a ${qgraphFile} 42b ${butlerConfig} 42c\n"]
404 )
406 def testMultipleCommands(self):
407 config_butler = f"{self.tmpdir.name}/test_repo"
408 config = BpsConfig(
409 {
410 "var1": "42a",
411 "var2": "42b",
412 "var3": "42c",
413 "butlerConfig": config_butler,
414 "finalJob": {
415 "command1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}",
416 "command2": "/usr/bin/uptime",
417 },
418 }
419 )
420 gwf_exec, args = create_final_command(config, self.tmpdir.name)
421 self.assertEqual(args, f"<FILE:runQgraphFile> {config_butler}")
422 final_script = f"{self.tmpdir.name}/final_job.bash"
423 self.assertEqual(gwf_exec.src_uri, final_script)
424 with open(final_script) as infh:
425 lines = infh.readlines()
426 self.assertEqual(
427 lines,
428 self.script_beginning
429 + ["/usr/bin/echo 42a ${qgraphFile} 42b ${butlerConfig} 42c\n", "/usr/bin/uptime\n"],
430 )
432 def testZeroCommands(self):
433 config_butler = f"{self.tmpdir.name}/test_repo"
434 config = BpsConfig(
435 {
436 "var1": "42a",
437 "var2": "42b",
438 "var3": "42c",
439 "butlerConfig": config_butler,
440 "finalJob": {
441 "cmd1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}",
442 "cmd2": "/usr/bin/uptime",
443 },
444 }
445 )
446 with self.assertRaisesRegex(RuntimeError, "finalJob.whenRun"):
447 _, _ = create_final_command(config, self.tmpdir.name)
449 def testWhiteSpaceOnlyCommand(self):
450 config_butler = f"{self.tmpdir.name}/test_repo"
451 config = BpsConfig(
452 {
453 "butlerConfig": config_butler,
454 "finalJob": {"command1": "", "command2": "\t \n"},
455 }
456 )
457 with self.assertRaisesRegex(RuntimeError, "finalJob.whenRun"):
458 _, _ = create_final_command(config, self.tmpdir.name)
460 def testSkipCommandUsingWhiteSpace(self):
461 config_butler = f"{self.tmpdir.name}/test_repo"
462 config = BpsConfig(
463 {
464 "var1": "42a",
465 "var2": "42b",
466 "var3": "42c",
467 "butlerConfig": config_butler,
468 "finalJob": {
469 "command1": "/usr/bin/echo {var1} {qgraphFile} {var2} {butlerConfig} {var3}",
470 "command2": "", # test skipping a command (i.e., overriding a default)
471 "command3": "/usr/bin/uptime",
472 },
473 }
474 )
475 gwf_exec, args = create_final_command(config, self.tmpdir.name)
476 self.assertEqual(args, f"<FILE:runQgraphFile> {config_butler}")
477 final_script = f"{self.tmpdir.name}/final_job.bash"
478 self.assertEqual(gwf_exec.src_uri, final_script)
479 with open(final_script) as infh:
480 lines = infh.readlines()
481 self.assertEqual(
482 lines,
483 self.script_beginning
484 + ["/usr/bin/echo 42a ${qgraphFile} 42b ${butlerConfig} 42c\n", "\n", "/usr/bin/uptime\n"],
485 )
488class TestEnhanceCommand(unittest.TestCase):
489 """Tests of _enhance_command function."""
491 def setUp(self):
492 logging.basicConfig(level=logging.WARNING)
493 logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO)
494 self.gw_exec = GenericWorkflowExec("test_exec", "/dummy/dir/pipetask")
495 self.config = BpsConfig(
496 {
497 # "profile": {},
498 "bpsUseShared": True,
499 "whenSaveJobQgraph": "NEVER",
500 "useLazyCommands": True,
501 # "memoryLimit": 32768,
502 "defOpts": "--long-log --log-file {submitPath}/{jobName}.{wmsAttemptNum}.json",
503 "submitPath": "/the/path",
504 }
505 )
506 self.cached_vals = {
507 "label1": {
508 "profile": {},
509 "bpsUseShared": True,
510 "whenSaveJobQgraph": "NEVER",
511 "useLazyCommands": True,
512 "memoryLimit": 32768,
513 "key1": "val1",
514 }
515 }
517 def testAttemptNum(self):
518 # test both in arguments as well as in variables in arguments
519 gwjob = GenericWorkflowJob("job1", "label1", executable=self.gw_exec)
520 gw = GenericWorkflow("test1")
521 gw.add_job(gwjob)
523 first_args = "{defOpts} run-qbb repo test.qg --summary {submitPath}/{jobName}-summary."
524 gwjob.arguments = first_args + "{wmsAttemptNum}.json"
526 new_arguments = first_args + "<WMS:attemptNum>.json"
527 new_opts = "--long-log --log-file /the/path/job1.<WMS:attemptNum>.json"
529 _enhance_command(self.config, gw, gwjob, {})
531 self.assertEqual(gwjob.arguments, new_arguments)
532 self.assertEqual(gwjob.cmdvals["defOpts"], new_opts)
534 def testKeyCachedCmdVal(self):
535 gwjob = GenericWorkflowJob("job1", "label1", executable=self.gw_exec)
536 gw = GenericWorkflow("test1")
537 gw.add_job(gwjob)
538 gwjob.arguments = "run-qbb repo test.qg -x {key1}"
539 self.assertNotIn("key1", gwjob.cmdvals)
540 _enhance_command(self.config, gw, gwjob, self.cached_vals)
541 self.assertEqual(gwjob.cmdvals["key1"], "val1")
543 def testS3Argument(self):
544 """Make sure s3 double slashes are not getting removed."""
545 gwjob = GenericWorkflowJob("job1", "label1", executable=self.gw_exec)
546 gw = GenericWorkflow("test1")
547 gw.add_job(gwjob)
548 s3 = "s3://user1@rubin-place-users/butler-pipeline1-processing.yaml"
549 gwjob.arguments = s3
550 _enhance_command(self.config, gw, gwjob, {})
551 self.assertEqual(gwjob.arguments, s3)
554if __name__ == "__main__":
555 unittest.main()