diff --git a/doc/changes/DM-53944.feature.rst b/doc/changes/DM-53944.feature.rst new file mode 100644 index 00000000..364a5946 --- /dev/null +++ b/doc/changes/DM-53944.feature.rst @@ -0,0 +1 @@ +Made improvements to how the environment specifications in submit yaml are handled. See ``Job Environment`` section in docs for details and examples. diff --git a/doc/lsst.ctrl.bps/quickstart.rst b/doc/lsst.ctrl.bps/quickstart.rst index c3705281..3e0bc8e6 100644 --- a/doc/lsst.ctrl.bps/quickstart.rst +++ b/doc/lsst.ctrl.bps/quickstart.rst @@ -1541,6 +1541,84 @@ invisible to the user. ``bps report`` will still show same labels and total counts as without ordering. ``cancel`` and ``restart`` will still work the same. + +.. _job-environment: + +Job Environment +--------------- + +One can specify environment values for jobs in the submit yaml. When creating +the job description for the WMS-plugins, BPS will search for environment +settings following normal scoping ordering but performs a special merge +instead of replacing the entire section. If an environment value has +another yaml environment variable, BPS will replace that yaml environment +variable with its value. This will allow normal path chaining +behavior across sections. + +.. note:: + + To not include an environment variable from another section, use the special + string ``BPS_NONE`` as the environment variable value (case-insensitive). + +Here are some examples, given the following BPS config yaml what +environment values BPS tells the WMS-plugin to set inside the job: + +.. code-block:: YAML + + var1: "root_val1" + environment: + VAR2: "root_val2" + VAR3: "root_val3" + VAR_PATH: "${PACKAGE_DIR}/root_dir:${VAR_PATH}" + TEST_VAR: "one {var1} three" + site: + site1: + var1: "site_val1" + environment: + VAR4: "site_val4" + VAR_PATH: "${PACKAGE_DIR}/site_dir:${VAR_PATH}" + cluster: + cl1: + var1: "cl1_val1" + environment: + VAR2: "BPS_NONE" + VAR3: "cl1_val3" + VAR4: "cl1_val4" + +#. BPS environment for a cl1 job for site1 : + + .. code-block:: + + VAR3="cl1_val3" + VAR4="cl1_val4" + VAR_PATH="${PACKAGE_DIR}/site_dir:${PACKAGE_DIR}/root_dir:${VAR_PATH}" + TEST_VAR="one cl1_val1 three" + +#. BPS environment for a non-cl1 job for site1 : + + .. code-block:: + + VAR2="root_val2" + VAR3="root_val3" + VAR4="site_val4" + VAR_PATH="${PACKAGE_DIR}/site_dir:${PACKAGE_DIR}/root_dir:${VAR_PATH}" + TEST_VAR="one site_val1 three" + +#. BPS environment for a non-cli job for some site other than site1: + + .. code-block:: + + VAR2="root_val2" + VAR3="root_val3" + VAR_PATH="${PACKAGE_DIR}/root_dir:${VAR_PATH}" + TEST_VAR="one root_val1 three" + +Different WMS-plugins create the job environment via different mechanisms +(e.g., the HTCondor plugin copies the submit environment to the job) or may +not support setting the environment via submit yaml. If the WMS-plugin +supports setting the environment via submit yaml, the job environment must +at least include what is in the submit yaml, but it may include more. + .. _bps-config-generation: Config Generation diff --git a/python/lsst/ctrl/bps/constants.py b/python/lsst/ctrl/bps/constants.py index 0353fe06..561474e5 100644 --- a/python/lsst/ctrl/bps/constants.py +++ b/python/lsst/ctrl/bps/constants.py @@ -29,11 +29,7 @@ from astropy import units as u -__all__ = [ - "DEFAULT_MEM_FMT", - "DEFAULT_MEM_RETRIES", - "DEFAULT_MEM_UNIT", -] +__all__ = ["BPS_NONE", "DEFAULT_MEM_FMT", "DEFAULT_MEM_RETRIES", "DEFAULT_MEM_UNIT"] DEFAULT_MEM_RETRIES = 5 @@ -47,3 +43,7 @@ DEFAULT_MEM_FMT = ".3f" """Default format specifier to use when reporting memory consumption. """ + +BPS_NONE = "BPS_NONE" +"""Special yaml value for a None value that can be read and written. +""" diff --git a/python/lsst/ctrl/bps/transform.py b/python/lsst/ctrl/bps/transform.py index cfb2538f..d1c12732 100644 --- a/python/lsst/ctrl/bps/transform.py +++ b/python/lsst/ctrl/bps/transform.py @@ -33,13 +33,17 @@ import math import os import re +from typing import Any from lsst.ctrl.bps import ClusteredQuantumGraph +from lsst.daf.butler import Config from lsst.pipe.base import QuantumGraph from lsst.utils.logging import VERBOSE from lsst.utils.timer import timeMethod from . import ( + BPS_NONE, + BPS_SEARCH_ORDER, DEFAULT_MEM_RETRIES, BpsConfig, GenericWorkflow, @@ -411,45 +415,7 @@ def _get_job_values(config, search_opt, cmd_line_key): else: job_values[attr] = getattr(default_gwjob, attr) - # Need to replace all config variables in environment values. - # Also change env vars in environment values to bash syntax. - # - # Note: Because job_values["environment"] is a BpsConfig and - # currently cannot have 2 search objects, for each environment - # setting, we have to get the setting string as is and then - # separately use the overall config to replace values inside - # the setting string. - tmp_job_env = job_values.get("environment", None) - if tmp_job_env: - _LOG.debug("_get_job_values: job_values['environment'] = %s", tmp_job_env) - - # Don't want to replace when getting environment setting string. - as_is_search_opt = { - "replaceVars": False, - "expandEnvVars": False, - "replaceEnvBps2Shell": False, - "replaceEnvShell2Bps": False, - } - - # When updating environment string, use given search options, - # but ensure making the environment string using bash syntax. - env_search_opt = copy.copy(search_opt) - env_search_opt["replaceVars"] = True # Replace bps config variables. - env_search_opt["replaceEnvBps2Shell"] = False # Replace bps syntax. - env_search_opt["replaceEnvShell2Bps"] = True # Do not replace shell env syntax. - env_search_opt["expandEnvVars"] = False # Do not replace with submission env value. - - job_env = {} # While replacing variables, convert to plain dict. - - for name in tmp_job_env: - # Get environment setting string as is. - value = tmp_job_env.search(name, as_is_search_opt)[1] - _LOG.debug("_get_job_values: as is value for %s = %s", name, value) - # Replace config vars and env placeholders - job_env[name] = config.modify_value(name, str(value), env_search_opt) - _LOG.debug("_get_job_values: new env value for %s = %s", name, job_env[name]) - # Save new dictionary back with other job values. - job_values["environment"] = job_env + job_values["environment"] = gather_job_environment(config, search_opt) # If the automatic memory scaling is enabled (i.e. the memory multiplier # is set and it is a positive number greater than 1.0), adjust number @@ -631,8 +597,6 @@ def create_generic_workflow( _, when_save = config.search("whenSaveJobQgraph", {"default": WhenToSaveQuantumGraphs.TRANSFORM.name}) save_qgraph_per_job = WhenToSaveQuantumGraphs[when_save.upper()] - search_opt = {"replaceVars": False, "expandEnvVars": False, "replaceEnvVars": True, "required": False} - generic_workflow = GenericWorkflow(name) # Save full run QuantumGraph for use by jobs @@ -664,13 +628,15 @@ def create_generic_workflow( gwjob = GenericWorkflowJob(cluster.name, cluster.label) # First get job values from cluster or cluster config - search_opt["curvals"] = {"curr_cluster": cluster.label} - found, value = config.search("computeSite", opt=search_opt) - if found: - search_opt["curvals"]["curr_site"] = value - found, value = config.search("computeCloud", opt=search_opt) - if found: - search_opt["curvals"]["curr_cloud"] = value + search_opt = config.get_search_opts(cluster.label) + search_opt.update( + { + "replaceVars": False, + "expandEnvVars": False, + "replaceEnvVars": True, + "required": False, + } + ) # If some config values are set for this cluster if cluster.label not in cached_job_values: @@ -694,6 +660,10 @@ def create_generic_workflow( _get_job_values(config["cluster"][cluster.label], search_opt, "runQuantumCommand") ) cluster_job_values = copy.copy(cached_job_values[cluster.label]) + # Environment is special because of the way it is merged. + # It needs to be set in the cached_pipetask_values, so + # don't include it here. + cluster_job_values.pop("environment", None) cluster_job_values["name"] = cluster.name cluster_job_values["label"] = cluster.label @@ -701,6 +671,12 @@ def create_generic_workflow( cluster_job_values["tags"] = cluster.tags _LOG.debug("cluster_job_values = %s", cluster_job_values) _handle_job_values(cluster_job_values, gwjob, cluster_job_values.keys()) + _LOG.debug( + "After _handle_job_values cluster %s: gwjob.environment = %s", cluster.label, gwjob.environment + ) + _LOG.debug( + "After _handle_job_values cluster %s: gwjob.arguments = %s", cluster.label, gwjob.arguments + ) # For purposes of whether to continue searching for a value is whether # the value evaluates to False. @@ -719,7 +695,19 @@ def create_generic_workflow( if task_label not in cached_pipetask_values: search_opt["curvals"]["curr_pipetask"] = task_label cached_pipetask_values[task_label] = _get_job_values(config, search_opt, "runQuantumCommand") + _LOG.debug( + "cached_pipetask_values[%s]['environment'] = %s", + task_label, + cached_pipetask_values[task_label].get("environment", None), + ) + _LOG.debug( + "cached_pipetask_values[%s]['arguments'] = %s", + task_label, + cached_pipetask_values[task_label].get("arguments", None), + ) _handle_job_values(cached_pipetask_values[task_label], gwjob, unset_attributes) + _LOG.debug("After _handle_job_values pipetask: gwjob.environment = %s", gwjob.environment) + _LOG.debug("After _handle_job_values pipetask: gwjob.arguments = %s", gwjob.arguments) # Update job with workflow attribute and profile values. qgraph_gwfile = _get_qgraph_gwfile( @@ -732,7 +720,10 @@ def create_generic_workflow( gwjob.cmdvals["qgraphNodeId"] = ",".join( sorted([f"{node_id}" for node_id in cluster.qgraph_node_ids]) ) + _LOG.debug("Before _enhance_command: gwjob.arguments = %s", gwjob.arguments) + _LOG.debug("Before _enhance_command: gwjob.cmdvals = %s", gwjob.cmdvals) _enhance_command(config, generic_workflow, gwjob, cached_job_values) + _LOG.debug("After _enhance_command: gwjob.environment = %s", gwjob.environment) # If writing per-job QuantumGraph files during TRANSFORM stage, # write it now while in memory. @@ -948,3 +939,91 @@ def add_final_job_as_sink(generic_workflow, final_job): generic_workflow.add_job(final_job) generic_workflow.add_job_relationships(gw_sinks, final_job.name) + + +def gather_job_environment(config: BpsConfig, search_opt: dict[str, Any]) -> dict[str, str]: + """Gather environment settings using given config search options. + + Parameters + ---------- + config : `lsst.ctrl.bps.BpsConfig` + Bps configuration. + search_opt : `dict` [`str`, `~typing.Any`] + Config search options. + + Returns + ------- + environment : `dict` [`str`, `str`] + Dictionary of environment variable names and values. + """ + # Don't want to replace when getting environment setting string. + as_is_search_opt = { + "replaceVars": False, + "expandEnvVars": False, + "replaceEnvBps2Shell": True, # Merge needs it to be shell syntax + "replaceEnvShell2Bps": False, + } + + def _update_env(env_updates: dict[str, str], environment: dict[str, str]): + """Update environment with values from given environment + section. + + Parameters + ---------- + env_updates : `dict` [`str`, `str`] or `~lsst.daf.butler.Config` + New values to use for update. + environment : `dict` [`str`, `str`] + Current environment values to update. Modified in place. + """ + for key, value in env_updates.items(): + value = config.modify_value(key, str(value), as_is_search_opt) + + for envkey in re.findall(r"\${([^}]+)}", value): + if envkey == key and envkey in environment: + oldval = environment[key] + value = re.sub(rf"\${{{envkey}}}", oldval, value) + environment[key] = value + + environment = {} + if ".environment" in config: + _, root_env = config.search("environment", as_is_search_opt) + + # Cast to Config to avoid replacing variables. + environment.update(Config(root_env)) + + curvals = search_opt.get("curvals", {}) + + # In order to do the concatenation correctly, must + # search in reverse order than normal config searches + for sect in reversed(BPS_SEARCH_ORDER): + sect_key = "curr_" + sect + if sect_key in curvals and sect in config and curvals[sect_key] in config[sect]: + search_sect = config[sect][curvals[sect_key]] + if "environment" in search_sect: + # Cast to Config to avoid replacing variables + _update_env(Config(search_sect["environment"]), environment) + + # Because not using normal config search, also have to check the + # search object if given. + if "searchobj" in search_opt and "environment" in search_opt["searchobj"]: + # Cast to Config to avoid replacing variables + env_sect = Config(search_opt["searchobj"]["environment"]) + _update_env(env_sect, environment) + + # Need to replace all config variables in environment values. + # Also remove any environment variable where value is "BPS_NONE". + new_environment = {} + if environment: + env_search_opt = copy.copy(search_opt) + env_search_opt["replaceVars"] = True # Replace bps config variables. + env_search_opt["replaceEnvBps2Shell"] = False # Keep bps syntax. + env_search_opt["replaceEnvShell2Bps"] = True # Replace shell env syntax. + env_search_opt["expandEnvVars"] = False # Do not replace with submission env values. + for name in environment: + # Ensure that environment values are strings + new_environment[name] = config.modify_value(name, str(environment[name]), env_search_opt) + if new_environment[name].upper() == BPS_NONE: + _LOG.debug("Removing %s from the merged environment", name) + del new_environment[name] + + return new_environment diff --git a/tests/cqg_test_utils.py b/tests/cqg_test_utils.py index 5a60969e..09035444 100644 --- a/tests/cqg_test_utils.py +++ b/tests/cqg_test_utils.py @@ -253,3 +253,84 @@ def get_add_quantum_args(key: str) -> tuple[uuid.UUID, str]: cqg.add_cluster(cluster) return qgraph, cqg + + +# T1(1,2) T1(1,4) T1(3,4) T5(1,2) T5(1,4) T5(3,4) +# | | | +# T2(1,2) T2(1,4) T2(3,4) +# | | | | | | +# | T2b(1,2) | T2b(1,4) | T2b(3,4) +# | | | +# T3(1,2) T3(1,4) T3(3,4) +# | | | +# T4(1,2) T4(1,4) T4(3,4) +def make_test_2_cluster_cqg(outdir): + """Make a ClusteredQuantumGraph with 2 sets of clusters for testing. + + Parameters + ---------- + outdir : `str` + Root used for the quantum graph filename stored + in the ClusteredQuantumGraph. The quantum graph is always saved to + this location. + + Returns + ------- + qgraph : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` + The fake QuantumGraph created for the test + ClusteredQuantumGraph returned separately. + cqg : `lsst.ctrl.bps.ClusteredQuantumGraph` + Clustered quantum graph. + """ + with make_test_helper() as helper: + qgc = helper.make_quantum_graph_builder(output_run="run").finish(attach_datastore_records=False) + qg_filename = f"{outdir}/test_file.qg" + # qgc.write(qg_filename) + qgraph = qgc.assemble() + cqg = ClusteredQuantumGraph("cqg2", qgraph, qg_filename) + + # since random hash ids, create mapping for tests + test_lookup = {} + for task_label, quanta_for_task in qgraph.quanta_by_task.items(): + for data_coordinate, quantum_id in quanta_for_task.items(): + data_id = dict(data_coordinate.required) + key = f"{task_label}_{data_id['D1']}_{data_id['D2']}" + test_lookup[key] = (quantum_id, cqg.qxgraph.nodes[quantum_id]) + + def get_add_quantum_args(key: str) -> tuple[uuid.UUID, str]: + quantum_id, quantum_info = test_lookup[key] + return quantum_id, quantum_info["task_label"] + + for dims in [(1, 2), (1, 4), (3, 4)]: + dims_str = f"{dims[0]}_{dims[1]}" + + # cluster T1,T2 + qc12 = QuantaCluster.from_quantum_info( + *test_lookup[f"T1_{dims_str}"], template=f"clusterT1T2_{dims_str}" + ) + qc12.add_quantum(*get_add_quantum_args(f"T2_{dims_str}")) + qc12.label = "clusterT1T2" # update label so doesnt look like only T1 + qc12.tags["label"] = qc12.label + + # cluster T3, T4 + qc34 = QuantaCluster.from_quantum_info( + *test_lookup[f"T3_{dims_str}"], template=f"clusterT3T4_{dims_str}" + ) + qc34.add_quantum(*get_add_quantum_args(f"T4_{dims_str}")) + qc34.label = "clusterT3T4" # update label so doesnt look like only T1 + qc34.tags["label"] = qc34.label + + # T1,T2 -> T3,T4 + cqg.add_cluster([qc34, qc12]) # reversed to check order is corrected in tests + cqg.add_dependency(qc12, qc34) + + # Add singleton dependency T2b + qc2b = QuantaCluster.from_quantum_info(*test_lookup[f"T2b_{dims_str}"], template=f"T2b_{dims_str}") + cqg.add_cluster(qc2b) + cqg.add_dependency(qc12, qc2b) + + # Add singleton independent + cluster = QuantaCluster.from_quantum_info(*test_lookup[f"T5_{dims_str}"], template=f"T5_{dims_str}") + cqg.add_cluster(cluster) + + return qgraph, cqg diff --git a/tests/data/config_for_transform.yaml b/tests/data/config_for_transform.yaml new file mode 100644 index 00000000..127c1c04 --- /dev/null +++ b/tests/data/config_for_transform.yaml @@ -0,0 +1,79 @@ +# For transform tests to be used with results of make_test_clustered_quantum_graph() +# Excessive environment settings for testing. + +computeSite: "root" +runQuantumCommand: "gexe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}" +var1: "root_val1" +environment: + VAR2: "root_val2" + VAR3: "root_val3" + VAR_PATH: "${PACKAGE_DIR}/root_dir:${VAR_PATH}" + TEST_VAR: "one {var1} three" + +clusterTemplate: "{D1}_{D2}" +cluster: + clusterT1T2: + requestMemory: 6144 + pipetasks: "T1, T2" + dimensions: "D1, D2" + var1: "cl12_val1" + environment: + VAR3: "cl12_val3" + VAR4: "cl12_val4" + VAR_PATH: "${PACKAGE_DIR}/cl12_dir:${VAR_PATH}" + clusterT3T4: + pipetasks: "T3, T4" + dimensions: "D1, D2" + var1: "cl34_val1" + environment: + VAR3: "cl34_val3" + VAR4: "cl34_val4" + VAR_PATH: "${PACKAGE_DIR}/cl34_dir:${VAR_PATH}" + +cloud: + cloud1: + runQuantumCommand: "c1exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}" + environment: + VAR3: "cloud1_val3" + VAR4: "cloud1_val4" + VAR_PATH: "${PACKAGE_DIR}/cloud1:${VAR_PATH}" +site: + site1: + var1: "site1_val1" + environment: + VAR4: "site1_val4" + VAR_PATH: "${PACKAGE_DIR}/site1_dir:${VAR_PATH}" + runQuantumCommand: "s1exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}" + +pipetask: + pipetaskInit: + var1: "init_val1" + requestMemory: 8096 + initPreCmdOpts: "--log-level=DEBUG" + T1: + var1: "T1_val1" + requestMemory: 1024 + T2: + requestMemory: 1024 + T2b: + var1: "T2b_val1" + environment: + VAR4: "T2b_val4" + VAR_PATH: "${PACKAGE_DIR}/t2b_dir:${VAR_PATH}" + T3: + requestMemory: 1024 + T4: + requestMemory: 4048 + T5: + var1: "T5_val1" + environment: + VAR4: "T5_val4" + VAR_PATH: "${PACKAGE_DIR}/t5_dir:${VAR_PATH}" + +finalJob: + command1: "/usr/bin/env" + var1: "final_val1" + environment: + VAR4: "final_val4" + VAR5: "final_val5" + VAR_PATH: "${PACKAGE_DIR}/final_dir:${VAR_PATH}" diff --git a/tests/test_transform.py b/tests/test_transform.py index c1785ed6..6c6c8800 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -26,15 +26,19 @@ # along with this program. If not, see . """Unit tests of transform.py.""" +import copy import dataclasses +import logging import os import shutil import tempfile import unittest -from cqg_test_utils import make_test_clustered_quantum_graph +from cqg_test_utils import make_test_2_cluster_cqg from lsst.ctrl.bps import ( + BPS_DEFAULTS, + BPS_NONE, BPS_SEARCH_ORDER, BpsConfig, GenericWorkflow, @@ -47,6 +51,7 @@ create_final_command, create_generic_workflow, create_generic_workflow_config, + gather_job_environment, ) TESTDIR = os.path.abspath(os.path.dirname(__file__)) @@ -70,63 +75,26 @@ class TestCreateGenericWorkflow(unittest.TestCase): """Tests of create_generic_workflow.""" def setUp(self): + logging.basicConfig(level=logging.WARNING) + logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO) self.tmpdir = tempfile.mkdtemp(dir=TESTDIR) - self.config = BpsConfig( - { - "runInit": True, - "computeSite": "global", - "runQuantumCommand": "gexe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}", - "clusterTemplate": "{D1}_{D2}", - "cluster": { - "cl1": {"pipetasks": "T1, T2", "dimensions": "D1, D2"}, - "cl2": {"pipetasks": "T3, T4", "dimensions": "D1, D2"}, - }, - "cloud": { - "cloud1": {"runQuantumCommand": "c1exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, - "cloud2": {"runQuantumCommand": "c2exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, - }, - "site": { - "site1": {"runQuantumCommand": "s1exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, - "site2": {"runQuantumCommand": "s2exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, - "global": {"runQuantumCommand": "s3exe -q {qgraphFile} --qgraph-node-id {qgraphNodeId}"}, - }, - # Needed because transform assumes they exist - "whenSaveJobQgraph": "NEVER", - "finalJob": {"whenRun": "ALWAYS", "command1": "/usr/bin/env"}, - }, - BPS_SEARCH_ORDER, + filename = os.path.join(TESTDIR, "data/config_for_transform.yaml") + self.orig_config = BpsConfig( + filename, BPS_SEARCH_ORDER, BPS_DEFAULTS, wms_service_class_fqn="wms_test_utils.WmsServiceSuccess" ) - _, self.cqg = make_test_clustered_quantum_graph(self.tmpdir) + _, self.cqg = make_test_2_cluster_cqg(self.tmpdir) def tearDown(self): shutil.rmtree(self.tmpdir, ignore_errors=True) - def testCreatingGenericWorkflowGlobal(self): - """Test creating a GenericWorkflow with global settings.""" - config = BpsConfig(self.config) - config["computeCloud"] = "cloud1" - config["computeSite"] = "site2" - config["queue"] = "global_queue" - print(config) - workflow = create_generic_workflow(config, self.cqg, "test_gw", self.tmpdir) - for jname in workflow: - gwjob = workflow.get_job(jname) - print(gwjob) - self.assertEqual(gwjob.compute_site, "site2") - self.assertEqual(gwjob.compute_cloud, "cloud1") - self.assertEqual(gwjob.executable.src_uri, "s2exe") - self.assertEqual(gwjob.queue, "global_queue") - final = workflow.get_final() - self.assertEqual(final.compute_site, "site2") - self.assertEqual(final.compute_cloud, "cloud1") - self.assertEqual(final.queue, "global_queue") - - def testCreatingQuantumGraphMixed(self): - """Test creating a GenericWorkflow with setting overrides.""" - config = BpsConfig(self.config) - config[".cluster.cl1.computeCloud"] = "cloud2" - config[".cluster.cl1.computeSite"] = "notthere" - config[".cluster.cl2.computeSite"] = "site1" + def testCreatingQuantumGraph(self): + """Test creating a GenericWorkflow with setting overrides. While + other tests exist to check get_job_valuei,gather_job_environment, + etc., going ahead and checking the integration of these by checking + results instead of using mocks. + """ + config = BpsConfig(self.orig_config) + config[".computeSite"] = "site1" config[".finalJob.queue"] = "special_final_queue" config[".finalJob.computeSite"] = "special_site" config[".finalJob.computeCloud"] = "special_cloud" @@ -134,29 +102,65 @@ def testCreatingQuantumGraphMixed(self): self.assertEqual(len(workflow) - 1, len(self.cqg)) # Don't count pipetaskInit for jname in workflow: gwjob = workflow.get_job(jname) - print(gwjob) - if jname.startswith("cl1"): - self.assertEqual(gwjob.compute_site, "notthere") - self.assertEqual(gwjob.compute_cloud, "cloud2") - self.assertEqual(gwjob.executable.src_uri, "c2exe") - elif jname.startswith("cl2"): - self.assertEqual(gwjob.compute_site, "site1") - self.assertIsNone(gwjob.compute_cloud) - self.assertEqual(gwjob.executable.src_uri, "s1exe") - elif jname.startswith("pipetask"): - self.assertEqual(gwjob.compute_site, "global") - self.assertIsNone(gwjob.compute_cloud) - self.assertEqual(gwjob.executable.src_uri, "s3exe") + self.assertEqual(gwjob.compute_site, "site1", f"failed for job {gwjob}") + self.assertIsNone(gwjob.compute_cloud, f"failed for job {gwjob}") + if gwjob.label == "pipetaskInit": + self.assertEqual( + gwjob.executable.src_uri, "/bin/pipetask", f"failed for job {gwjob}" + ) + else: + self.assertEqual(gwjob.executable.src_uri, "s1exe", f"failed for job {gwjob}") + self.assertIn("{qgraphNodeId}", gwjob.arguments) + + base_site_env_truth = { + "VAR2": "root_val2", + "VAR3": "root_val3", + "VAR4": "site1_val4", + "VAR_PATH": "/site1_dir:/root_dir:", + "TEST_VAR": "one site1_val1 three", + } + + match gwjob.label: + case "pipetaskInit": + env_truth = dict(base_site_env_truth) + env_truth["TEST_VAR"] = "one init_val1 three" + self.assertEqual(gwjob.environment, env_truth) + self.assertEqual(gwjob.cmdvals["initPreCmdOpts"], "--log-level=DEBUG") + # Taking pipetaskInit-defined requestMemory. + self.assertEqual(gwjob.request_memory, 8096) + case "clusterT1T2": + # Taking cluster-defined requestMemory. + self.assertEqual(gwjob.request_memory, 6144) + self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 2) + case "clusterT3T4": + # Taking max of the requestMemory for quanta in cluster. + self.assertEqual(gwjob.request_memory, 4048) + self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 2) + case "T2b": + # Taking default requestMemory from root section. + self.assertEqual(gwjob.request_memory, BPS_DEFAULTS["requestMemory"]) + self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 1) + case "T5": + # Taking default requestMemory from root section. + self.assertEqual(gwjob.request_memory, BPS_DEFAULTS["requestMemory"]) + self.assertEqual(len(gwjob.cmdvals["qgraphNodeId"].split(",")), 1) + case _: + # Should always have a label from above, but need to + # fail test if get different label. + self.fail(f"Invalid gwjob.label for job {gwjob}") # pragma: no cover final = workflow.get_final() - self.assertEqual(final.compute_site, "special_site") - self.assertEqual(final.compute_cloud, "special_cloud") - self.assertEqual(final.queue, "special_final_queue") + self.assertEqual(final.compute_site, "special_site", f"failed for final job {final}") + self.assertEqual(final.compute_cloud, "special_cloud", f"failed for final job {final}") + self.assertEqual(final.queue, "special_final_queue", f"failed for final job {final}") + self.assertEqual(final.request_memory, BPS_DEFAULTS["finalJob"]["requestMemory"]) class TestGetJobValues(unittest.TestCase): """Tests of _get_job_values.""" def setUp(self): + logging.basicConfig(level=logging.WARNING) + logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO) self.default_job = GenericWorkflowJob("default_job", "default_label") def testGettingDefaults(self): @@ -192,33 +196,166 @@ def testRetrievingCmdLine(self): self.assertEqual(job_values["executable"].src_uri, "/path/to/foo") self.assertEqual(job_values["arguments"], "bar.txt") - def testEnvironment(self): + @unittest.mock.patch("lsst.ctrl.bps.transform.gather_job_environment") + def testCallGatherJobEnvironmentNoSearchOpts(self, mock_gather): + # Test that _get_job_values passes right search options on + # to gather_job_environment function and didn't have side-effects. + env_truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"} + mock_gather.return_value = dict(env_truth) config = BpsConfig( { "var1": "two", "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} three"}, + "runQuantumCommand": "/path/to/foo bar.txt", + "pipetask": {"isr": {"requestMemory": 8096, "environment": {"ISR_VAR": "45"}}}, } ) - job_values = _get_job_values(config, {}, None) - truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"} - self.assertEqual(truth, job_values["environment"]) + search_opts = {} + config_copy = BpsConfig(config) + search_opts_copy = dict(search_opts) - def testEnvironmentOptions(self): + job_values = _get_job_values(config, search_opts, None) + env_truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"} + self.assertEqual(env_truth, job_values["environment"]) + mock_gather.assert_called_once_with(config, search_opts) + + # And didn't have side-effects that changed vars + self.assertEqual(config, config_copy) + self.assertEqual(search_opts, search_opts_copy) + + @unittest.mock.patch("lsst.ctrl.bps.transform.gather_job_environment") + def testCallGatherJobEnvironmentWithCurvals(self, mock_gather): + # Test that _get_job_values passes right search options on + # to gather_job_environment function and didn't have side-effects. + env_truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three", "ISR_VAR": "45"} + mock_gather.return_value = dict(env_truth) config = BpsConfig( { "var1": "two", "environment": {"TEST_INT": 1, "TEST_BOOL": False, "TEST_SPACES": "one {var1} three"}, - "finalJob": {"requestMemory": 8096, "command1": "/usr/bin/env"}, + "runQuantumCommand": "/path/to/foo bar.txt", + "pipetask": {"isr": {"requestMemory": 8096, "environment": {"ISR_VAR": "45"}}}, } ) - search_obj = config["finalJob"] - search_opts = {"replaceVars": False, "searchobj": search_obj} - job_values = _get_job_values(config, search_opts, None) - truth = {"TEST_INT": "1", "TEST_BOOL": "False", "TEST_SPACES": "one two three"} - self.assertEqual(truth, job_values["environment"]) + curvals = {"curr_pipetask": "isr"} + search_opts = {"replaceVars": False, "searchobj": {"curvals": curvals}} + + # Save copies to check no side-effects + config_copy = BpsConfig(config) + search_opts_copy = dict(search_opts) + + job_values = _get_job_values(config, search_opts, "runQuantumCommand") + mock_gather.assert_called_once_with(config, search_opts) + + self.assertEqual(job_values["environment"], env_truth) + self.assertEqual(job_values["executable"].src_uri, "/path/to/foo") + + # And didn't have side-effects that changed vars + self.assertEqual(config, config_copy) + self.assertEqual(search_opts, search_opts_copy) + + +class TestGatherJobEnvironment(unittest.TestCase): + """Tests for the gather_job_environment function.""" + + def setUp(self): + logging.basicConfig(level=logging.WARNING) + logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO) + # The directories don't match real ones, but are here to test + # environment variables in yaml environment section as well as + # appending values across sections. + filename = os.path.join(TESTDIR, "data/config_for_transform.yaml") + self.orig_config = BpsConfig( + filename, BPS_SEARCH_ORDER, BPS_DEFAULTS, wms_service_class_fqn="wms_test_utils.WmsServiceSuccess" + ) + self.orig_config[".finalJob.environment.VAR3"] = BPS_NONE + self.orig_config[".finalJob.environment.VAR4"] = BPS_NONE + # test case-insensitivity of BPS_NONE + self.orig_config[".cluster.clusterT1T2.environment.VAR2"] = BPS_NONE.lower() + + def testEnvironmentRootSiteCluster(self): + search_opts = { + "replaceVars": False, + "curvals": {"curr_cluster": "clusterT1T2", "curr_site": "site1"}, + } + job_env = gather_job_environment(self.orig_config, search_opts) + truth = { + "VAR3": "cl12_val3", + "VAR4": "cl12_val4", + "VAR_PATH": "/cl12_dir:/site1_dir:/root_dir" + ":", + "TEST_VAR": "one cl12_val1 three", + } + self.assertEqual(truth, job_env) self.assertEqual(search_opts["replaceVars"], False) - self.assertEqual(search_opts["searchobj"]["requestMemory"], 8096) - self.assertEqual(job_values["request_memory"], 8096) + + def testEnvironmentRootSite(self): + # Checking that doesn't pick up env from other cluster + search_opts = { + "replaceVars": False, + "curvals": {"curr_cluster": "notthere", "curr_site": "site1"}, + } + job_env = gather_job_environment(self.orig_config, search_opts) + truth = { + "VAR2": "root_val2", + "VAR3": "root_val3", + "VAR4": "site1_val4", + "VAR_PATH": "/site1_dir:/root_dir:", + "TEST_VAR": "one site1_val1 three", + } + self.assertEqual(truth, job_env) + self.assertEqual(search_opts["replaceVars"], False) + + def testEnvironmentRoot(self): + # Checking that doesn't pick up env from other cluster or site. + # Also check that doesn't modify our search_opts by setting opposites + # of what function uses. + orig_search_opts = { + "replaceVars": False, + "replaceEnvBps2Shell": True, + "replaceEnvShell2Bps": False, + "expandEnvVars": True, + "curvals": {"curr_cluster": "notthere", "curr_site": "notthere"}, + } + search_opts = copy.deepcopy(orig_search_opts) + + job_env = gather_job_environment(self.orig_config, search_opts) + truth = { + "VAR2": "root_val2", + "VAR3": "root_val3", + "VAR_PATH": "/root_dir:", + "TEST_VAR": "one root_val1 three", + } + self.assertEqual(truth, job_env) + self.assertEqual(orig_search_opts, search_opts) + + def testEnvironmentNoSearchOpts(self): + search_opts = {} + job_env = gather_job_environment(self.orig_config, search_opts) + truth = { + "VAR2": "root_val2", + "VAR_PATH": "/root_dir:", + "VAR3": "root_val3", + "TEST_VAR": "one root_val1 three", + } + self.assertEqual(truth, job_env) + self.assertEqual(search_opts, {}) + + def testSearchObj(self): + # Test that works with searchobj, like finalJob + search_opts = {"searchobj": self.orig_config["finalJob"], "curvals": {"curr_site": "site1"}} + copy_final = BpsConfig(self.orig_config["finalJob"]) + job_env = gather_job_environment(self.orig_config, search_opts) + # VAR3 and VAR4 removed in setUp + truth = { + "TEST_VAR": "one final_val1 three", + "VAR2": "root_val2", + "VAR5": "final_val5", + "VAR_PATH": "/final_dir:/site1_dir:/root_dir" + ":", + } + self.assertEqual(truth, job_env) + self.assertEqual(search_opts["searchobj"], copy_final) def testVarsInEnvironment(self): config = BpsConfig( @@ -236,6 +373,8 @@ class TestCreateFinalCommand(unittest.TestCase): """Tests for the create_final_command function.""" def setUp(self): + logging.basicConfig(level=logging.WARNING) + logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO) self.tmpdir = tempfile.TemporaryDirectory() self.script_beginning = [ "#!/bin/bash\n", @@ -357,6 +496,8 @@ class TestEnhanceCommand(unittest.TestCase): """Tests of _enhance_command function.""" def setUp(self): + logging.basicConfig(level=logging.WARNING) + logging.getLogger("lsst.ctrl.bps.bps_config").setLevel(logging.INFO) self.gw_exec = GenericWorkflowExec("test_exec", "/dummy/dir/pipetask") self.config = BpsConfig( {