From 5221661653cbdd9b6f235b4a29200458c1ea6859 Mon Sep 17 00:00:00 2001 From: Michelle Gower Date: Fri, 20 Feb 2026 11:25:28 -0600 Subject: [PATCH 1/4] Add ability to run submit processes as batch jobs. Note: Limited to same cluster with shared filesystems. --- doc/changes/DM-53494.feature.rst | 1 + doc/lsst.ctrl.bps/quickstart.rst | 49 +++ python/lsst/ctrl/bps/batch_submit.py | 234 ++++++++++++++ python/lsst/ctrl/bps/cli/cmd/__init__.py | 4 + python/lsst/ctrl/bps/cli/cmd/commands.py | 20 ++ python/lsst/ctrl/bps/drivers.py | 177 +++++++--- python/lsst/ctrl/bps/etc/bps_defaults.yaml | 22 ++ python/lsst/ctrl/bps/generic_workflow.py | 18 +- python/lsst/ctrl/bps/initialize.py | 5 +- python/lsst/ctrl/bps/pre_transform.py | 29 +- python/lsst/ctrl/bps/tests/gw_test_utils.py | 42 +++ python/lsst/ctrl/bps/wms_service.py | 12 + tests/qg_test_utils.py | 7 +- tests/test_batch_submit.py | 338 ++++++++++++++++++++ tests/test_cli_commands.py | 28 ++ tests/test_drivers.py | 332 +++++++++++++++++-- tests/test_pre_transform.py | 118 ++++++- 17 files changed, 1340 insertions(+), 96 deletions(-) create mode 100644 doc/changes/DM-53494.feature.rst create mode 100644 python/lsst/ctrl/bps/batch_submit.py create mode 100644 tests/test_batch_submit.py diff --git a/doc/changes/DM-53494.feature.rst b/doc/changes/DM-53494.feature.rst new file mode 100644 index 00000000..605b107f --- /dev/null +++ b/doc/changes/DM-53494.feature.rst @@ -0,0 +1 @@ +Added ability to run submit processes as batch jobs on same cluster with shared filesystems. diff --git a/doc/lsst.ctrl.bps/quickstart.rst b/doc/lsst.ctrl.bps/quickstart.rst index 4370ba89..80a198b9 100644 --- a/doc/lsst.ctrl.bps/quickstart.rst +++ b/doc/lsst.ctrl.bps/quickstart.rst @@ -1613,6 +1613,55 @@ Parsl). responsibility to remove them once no longer needed. The removal should be done regularly to avoid too many in single directory. +.. _bps_submit_as_batch: + +Submit Stages as Batch Jobs +--------------------------- + +In cases where one cannot run ``bps submit`` interactively (e.g., needs +too much memory), BPS can run the submit processes as batch jobs after +which the payload workflow will start running. This also is necessary +when running at a remote site (i.e., a site where the butler repository +isn't directly accessible). + +.. note:: + + Currently only the HTCondor WMS plugin supports this feature as described + below. The PanDA WMS plugin uses a special PanDA feature to run + ``bps submit`` on a remote site. See PanDA documentation for details. + +To tell BPS to execute submit stages as batch jobs, the submit yaml must +set ``bpsBatchSubmission`` to ``true``. + +The interactive submission process, ``bps submit ``, +is much shorter. It will create a workflow with two jobs: + +- ``buildQuantumGraph`` which creates the quantum graph. +- ``preparePayloadWorkflow`` which does the rest of the submission stages + seen when running ``bps submit``. These include clustering, creation of + the payload workflow, and preparing the WMS-specific workflow. + +One can set runtime values specific to those jobs (e.g., ``requestMemory``) in +sections with corresponding names similar to ``finalJob``. Currently the +logging-related command-line arguments aren't passed from ``bps batch-submit`` +to these jobs. Instead, one can set ``bpsPreCommandOpts``, which has the +same default as the payload job. + +Even with the new jobs, there is only one output run collection, one submit +directory and one top level WMS ID to be used with BPS commands. + +``bps report --id `` will show these 2 new jobs same as the payload +jobs. They will be the only lines to appear in the report until the +``preparePayloadWorkflow`` has finished at which time the expected payload +lines should appear (from ``pipetaskInit`` through ``finalJob``). + +``bps cancel`` can be used to abort the run during these new jobs or later +during the running of the payload jobs. Note, the ``finalJob`` job won't +automatically run unless the payload workflow was successfully submitted. + +See the corresponding section in the WMS-plugin documentation for additional +information and yaml settings. + .. _bps-troubleshooting: Troubleshooting diff --git a/python/lsst/ctrl/bps/batch_submit.py b/python/lsst/ctrl/bps/batch_submit.py new file mode 100644 index 00000000..a9701630 --- /dev/null +++ b/python/lsst/ctrl/bps/batch_submit.py @@ -0,0 +1,234 @@ +# This file is part of ctrl_bps. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This software is dual licensed under the GNU General Public License and also +# under a 3-clause BSD license. Recipients may choose which of these licenses +# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, +# respectively. If you choose the GPL option then the following text applies +# (but note that there is still no warranty even if you opt for BSD instead): +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Driver to run submit stages as batch jobs.""" + +__all__ = ["batch_payload_prepare", "create_batch_stages"] + +import logging +import os + +from lsst.resources import ResourcePath, ResourcePathExpression +from lsst.utils.logging import VERBOSE +from lsst.utils.timer import time_this, timeMethod + +from . import ( + DEFAULT_MEM_FMT, + DEFAULT_MEM_UNIT, + BpsConfig, + GenericWorkflow, + GenericWorkflowJob, + GenericWorkflowLazyGroup, +) +from .pre_transform import cluster_quanta, read_quantum_graph +from .prepare import prepare +from .transform import _get_job_values, transform + +_LOG = logging.getLogger(__name__) + + +@timeMethod(logger=_LOG, logLevel=VERBOSE) +def create_batch_stages( + config: BpsConfig, prefix: ResourcePathExpression +) -> tuple[GenericWorkflow, BpsConfig]: + """Create a GenericWorkflow that performs the submit stages as a workflow. + + Parameters + ---------- + config : `lsst.ctrl.bps.BpsConfig` + BPS configuration. + prefix : `lsst.resources.ResourcePathExpression` + Root path for any output files. + + Returns + ------- + generic_workflow : `lsst.ctrl.bps.GenericWorkflow` + The generic workflow transformed from the clustered quantum graph. + generic_workflow_config : `lsst.ctrl.bps.BpsConfig` + Configuration to accompany GenericWorkflow. + """ + prefix = ResourcePath(prefix) + generic_workflow: GenericWorkflow = GenericWorkflow(name=f"{config['uniqProcName']}_ctrl") + cmd_line_key = "jobCommand" + + # build QuantumGraph job + search_opt = {} + if "buildQuantumGraph" in config: + search_opt["searchobj"] = config.get("buildQuantumGraph") + build_job = GenericWorkflowJob( + name="buildQuantumGraph", + label="buildQuantumGraph", + ) + job_values = _get_job_values(config, search_opt, cmd_line_key) + if not job_values["executable"]: + raise RuntimeError( + f"Missing executable for buildQuantumGraph. Double check submit yaml for {cmd_line_key}" + ) + for key, value in job_values.items(): + if key not in {"name", "label"}: + setattr(build_job, key, value) + + generic_workflow.add_job(build_job) + generic_workflow.run_attrs.update( + { + "bps_isjob": "True", + "bps_project": config["project"], + "bps_campaign": config["campaign"], + "bps_run": config["uniqProcName"], + "bps_operator": config["operator"], + "bps_payload": config["payloadName"], + "bps_runsite": config["computeSite"], + } + ) + + # cluster/transform/prepare job + search_opt = {} + if "preparePayloadWorkflow" in config: + search_opt["searchobj"] = config.get("preparePayloadWorkflow") + prepare_job = GenericWorkflowLazyGroup( + name="preparePayloadWorkflow", + label="preparePayloadWorkflow", + ) + job_values = _get_job_values(config, search_opt, cmd_line_key) + if not job_values["executable"]: + raise RuntimeError( + f"Missing executable for preparePayloadWorkflow. Double check submit yaml for {cmd_line_key}" + ) + for key, value in job_values.items(): + if key not in {"name", "label"}: + setattr(prepare_job, key, value) + + generic_workflow.add_job(prepare_job, parent_names=["buildQuantumGraph"]) + + _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False}) + if save_workflow: + with prefix.join("bps_stages_generic_workflow.pickle").open("wb") as outfh: + generic_workflow.save(outfh, "pickle") + + return generic_workflow, config + + +@timeMethod(logger=_LOG, logLevel=VERBOSE) +def batch_payload_prepare(config: BpsConfig, prefix: ResourcePathExpression) -> None: + """Create a GenericWorkflow that performs the submit stages as a workflow. + + Parameters + ---------- + config : `lsst.ctrl.bps.BpsConfig` + BPS configuration. + + prefix : `lsst.resources.ResourcePathExpression` + Root path for any output files. + + Returns + ------- + generic_workflow : `lsst.ctrl.bps.GenericWorkflow` + The generic workflow transformed from the clustered quantum graph. + generic_workflow_config : `lsst.ctrl.bps.BpsConfig` + Configuration to accompany GenericWorkflow. + """ + prefix = ResourcePath(prefix) + # Read existing QuantumGraph + qgraph_filename = prefix.join(config["qgraphFileTemplate"]) + qgraph = read_quantum_graph(qgraph_filename) + config[".bps_defined.runQgraphFile"] = str(qgraph_filename) + + # Cluster + _LOG.info("Starting cluster stage (grouping quanta into jobs)") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Cluster stage completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + clustered_qgraph = cluster_quanta(config, qgraph, config["uniqProcName"]) + + _LOG.info("ClusteredQuantumGraph contains %d cluster(s)", len(clustered_qgraph)) + + submit_path = config[".bps_defined.submitPath"] + _, save_clustered_qgraph = config.search("saveClusteredQgraph", opt={"default": False}) + if save_clustered_qgraph: + clustered_qgraph.save(os.path.join(submit_path, "bps_clustered_qgraph.pickle")) + _, save_dot = config.search("saveDot", opt={"default": False}) + if save_dot: + clustered_qgraph.draw(os.path.join(submit_path, "bps_clustered_qgraph.dot")) + + # Transform + _LOG.info("Starting transform stage (creating generic workflow)") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Transform stage completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + generic_workflow, generic_workflow_config = transform(config, clustered_qgraph, submit_path) + _LOG.info("Generic workflow name '%s'", generic_workflow.name) + + num_jobs = sum(generic_workflow.job_counts.values()) + _LOG.info("GenericWorkflow contains %d job(s) (including final)", num_jobs) + + _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False}) + if save_workflow: + with open(os.path.join(submit_path, "bps_generic_workflow.pickle"), "wb") as outfh: + generic_workflow.save(outfh, "pickle") + _, save_dot = config.search("saveDot", opt={"default": False}) + if save_dot: + with open(os.path.join(submit_path, "bps_generic_workflow.dot"), "w") as outfh: + generic_workflow.draw(outfh, "dot") + + # Prepare + _LOG.info("Starting prepare stage (creating specific implementation of workflow)") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Prepare stage completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path) + + # Add payload workflow to currently running workflow + _LOG.info("Starting update workflow") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Workflow update completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + # Assuming submit_path for ctrl workflow is visible by this job. + wms_workflow.add_to_parent_workflow(generic_workflow_config) diff --git a/python/lsst/ctrl/bps/cli/cmd/__init__.py b/python/lsst/ctrl/bps/cli/cmd/__init__.py index 49c97190..2047458d 100644 --- a/python/lsst/ctrl/bps/cli/cmd/__init__.py +++ b/python/lsst/ctrl/bps/cli/cmd/__init__.py @@ -27,6 +27,8 @@ __all__ = [ "acquire", + "batch_acquire", + "batch_prepare", "cluster", "transform", "prepare", @@ -41,6 +43,8 @@ from .commands import ( acquire, + batch_acquire, + batch_prepare, cancel, cluster, ping, diff --git a/python/lsst/ctrl/bps/cli/cmd/commands.py b/python/lsst/ctrl/bps/cli/cmd/commands.py index 46ffb32a..29f68d2a 100644 --- a/python/lsst/ctrl/bps/cli/cmd/commands.py +++ b/python/lsst/ctrl/bps/cli/cmd/commands.py @@ -36,6 +36,8 @@ from ... import BpsSubprocessError from ...drivers import ( acquire_qgraph_driver, + batch_acquire_driver, + batch_prepare_driver, cancel_driver, cluster_qgraph_driver, ping_driver, @@ -219,3 +221,21 @@ def ping(*args, **kwargs): def submitcmd(*args, **kwargs): """Submit a command for execution.""" submitcmd_driver(*args, **kwargs) + + +@click.command(cls=BpsCommand) +@opt.config_file_argument(required=True) +@opt.submission_options() +def batch_acquire(*args, **kwargs): + """Run inside a batch job to create a new quantum graph.""" + with catch_errors(): + batch_acquire_driver(*args, **kwargs) + + +@click.command(cls=BpsCommand) +@opt.config_file_argument(required=True) +@opt.submission_options() +def batch_prepare(*args, **kwargs): + """Run payload workflow preparation inside a batch job.""" + with catch_errors(): + batch_prepare_driver(*args, **kwargs) diff --git a/python/lsst/ctrl/bps/drivers.py b/python/lsst/ctrl/bps/drivers.py index 8debefdc..c705867d 100644 --- a/python/lsst/ctrl/bps/drivers.py +++ b/python/lsst/ctrl/bps/drivers.py @@ -33,6 +33,8 @@ __all__ = [ "acquire_qgraph_driver", + "batch_acquire_driver", + "batch_prepare_driver", "cancel_driver", "cluster_qgraph_driver", "ping_driver", @@ -49,12 +51,22 @@ import logging import os from pathlib import Path +from typing import Any from lsst.pipe.base.quantum_graph import PredictedQuantumGraph from lsst.utils.timer import time_this from lsst.utils.usage import get_peak_mem_usage -from . import BPS_DEFAULTS, BPS_SEARCH_ORDER, DEFAULT_MEM_FMT, DEFAULT_MEM_UNIT, BpsConfig +from . import ( + BPS_DEFAULTS, + BPS_SEARCH_ORDER, + DEFAULT_MEM_FMT, + DEFAULT_MEM_UNIT, + BpsConfig, + ClusteredQuantumGraph, + GenericWorkflow, +) +from .batch_submit import batch_payload_prepare, create_batch_stages from .bps_reports import compile_code_summary, compile_job_summary from .bps_utils import _dump_env_info, _dump_pkg_info, _make_id_link from .cancel import cancel @@ -67,7 +79,7 @@ submit_path_validator, ) from .ping import ping -from .pre_transform import acquire_quantum_graph, cluster_quanta +from .pre_transform import acquire_quantum_graph, cluster_quanta, read_quantum_graph from .prepare import prepare from .report import display_report, retrieve_report from .restart import restart @@ -142,14 +154,16 @@ def acquire_qgraph_driver(config_file: str, **kwargs) -> tuple[BpsConfig, Predic mem_unit=DEFAULT_MEM_UNIT, mem_fmt=DEFAULT_MEM_FMT, ): - qgraph_file, qgraph = acquire_quantum_graph(config, out_prefix=submit_path) + qgraph_file = acquire_quantum_graph(config, out_prefix=submit_path) + qgraph = read_quantum_graph(qgraph_file) + _log_mem_usage() config[".bps_defined.runQgraphFile"] = qgraph_file return config, qgraph -def cluster_qgraph_driver(config_file, **kwargs): +def cluster_qgraph_driver(config_file: str, **kwargs: Any) -> tuple[BpsConfig, ClusteredQuantumGraph]: """Group quanta into clusters. Parameters @@ -193,7 +207,7 @@ def cluster_qgraph_driver(config_file, **kwargs): return config, clustered_qgraph -def transform_driver(config_file, **kwargs): +def transform_driver(config_file: str, **kwargs: Any) -> tuple[BpsConfig, GenericWorkflow]: """Create a workflow for a specific workflow management system. Parameters @@ -207,7 +221,7 @@ def transform_driver(config_file, **kwargs): ------- generic_workflow_config : `lsst.ctrl.bps.BpsConfig` Configuration to use when creating the workflow. - generic_workflow : `lsst.ctrl.bps.BaseWmsWorkflow` + generic_workflow : `lsst.ctrl.bps.GenericWorkflow` Representation of the abstract/scientific workflow specific to a given workflow management system. """ @@ -298,37 +312,40 @@ def submit_driver(config_file, **kwargs): "not accurately reflect actual memory usage by the bps process." ) - remote_build = {} config = BpsConfig( config_file, search_order=BPS_SEARCH_ORDER, defaults=BPS_DEFAULTS, wms_service_class_fqn=kwargs.get("wms_service"), ) - _, remote_build = config.search("remoteBuild", opt={"default": {}}) - if remote_build: - if config["wmsServiceClass"] == "lsst.ctrl.bps.panda.PanDAService": - if not remote_build.search("enabled", opt={"default": False})[1]: - remote_build = {} - _LOG.info("The workflow is submitted to the local Data Facility.") - else: - _LOG.info( - "Remote submission is enabled. The workflow is submitted to a remote Data Facility." - ) - _LOG.info("Initializing execution environment") - with time_this( - log=_LOG, - level=logging.INFO, - prefix=None, - msg="Initializing execution environment completed", - mem_usage=True, - mem_unit=DEFAULT_MEM_UNIT, - mem_fmt=DEFAULT_MEM_FMT, - ): - config = _init_submission_driver(config_file, **kwargs) - kwargs["remote_build"] = remote_build - kwargs["config_file"] = config_file - wms_workflow = None + translate_command_line_values(config, **kwargs) + + wms_service_class = config["wmsServiceClass"] + search_opts = config.get_search_opts() + + # Initialization is normally called as part of submission stages. + # But if running submission stages as batch job(s), need to + # run initialization separately. + + # PanDA-specific original options to run at sites with own Butler. + search_opts["default"] = {} + remote_build = config.search("remoteBuild", opt=search_opts)[1] + remote_build_enabled = False + if remote_build: # remoteBuild is a section of yaml + search_opts["default"] = False + remote_build_enabled = remote_build.search("enabled", opt=search_opts)[1] + + # BPS option to turn on running submission stages as batch jobs + search_opts["default"] = False + batch_submission_enabled = config.search("bpsBatchSubmission", opt=search_opts)[1] + + if remote_build_enabled or batch_submission_enabled: + _LOG.info("Running submission stages as batch job(s) is enabled.") + config = _init_submission_driver(config_file, **kwargs) + + if wms_service_class == "lsst.ctrl.bps.panda.PanDAService": + kwargs["remote_build"] = remote_build + kwargs["config_file"] = config_file else: _LOG.info("The workflow is submitted to the local Data Facility.") @@ -342,31 +359,38 @@ def submit_driver(config_file, **kwargs): mem_unit=DEFAULT_MEM_UNIT, mem_fmt=DEFAULT_MEM_FMT, ): - if not remote_build: - wms_workflow_config, wms_workflow = prepare_driver(config_file, **kwargs) - else: + if batch_submission_enabled: wms_workflow_config = config - - _LOG.info("Starting submit stage") - with time_this( - log=_LOG, - level=logging.INFO, - prefix=None, - msg="Submit stage completed", - mem_usage=True, - mem_unit=DEFAULT_MEM_UNIT, - mem_fmt=DEFAULT_MEM_FMT, - ): - workflow = submit(wms_workflow_config, wms_workflow, **kwargs) - if not wms_workflow: - wms_workflow = workflow - _LOG.info("Run '%s' submitted for execution with id '%s'", wms_workflow.name, wms_workflow.run_id) + wms_workflow = batch_submit(config) + else: + if remote_build_enabled: + wms_workflow_config = config + wms_workflow = None + else: + wms_workflow_config, wms_workflow = prepare_driver(config_file, **kwargs) + + _LOG.info("Starting submit stage") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Submit stage completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + workflow = submit(wms_workflow_config, wms_workflow, **kwargs) + if not wms_workflow: + wms_workflow = workflow + _LOG.info( + "Run '%s' submitted for execution with id '%s'", wms_workflow.name, wms_workflow.run_id + ) _log_mem_usage() _make_id_link(wms_workflow_config, wms_workflow.run_id) print(f"Run Id: {wms_workflow.run_id}") - print(f"Run Name: {wms_workflow.name}") + print(f"Run Name: {wms_workflow_config['uniqProcName']}") def restart_driver(wms_service, run_id): @@ -687,3 +711,56 @@ def _log_mem_usage() -> None: "Peak memory usage for bps process %s (main), %s (largest child process)", *tuple(f"{val.to(DEFAULT_MEM_UNIT):{DEFAULT_MEM_FMT}}" for val in get_peak_mem_usage()), ) + + +def batch_acquire_driver(config_file: str, **kwargs: Any) -> None: + """Create a quantum graph from pipeline definition in a batch job. + + Parameters + ---------- + config_file : `str` + Name of the configuration file. + **kwargs : `~typing.Any` + Additional modifiers to the configuration. + """ + config = BpsConfig(config_file) + submit_path = config[".bps_defined.submitPath"] + + _LOG.info("Starting acquire stage (generating and/or reading quantum graph)") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Acquire stage completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + _ = acquire_quantum_graph(config, out_prefix=submit_path) + _log_mem_usage() + + +def batch_prepare_driver(config_file: str, **kwargs: Any) -> None: + """Run workflow preparation in a batch job for an existing QuantumGraph. + + Parameters + ---------- + config_file : `str` + Name of the configuration file. + **kwargs : `~typing.Any` + Additional modifiers to the configuration. + """ + config = BpsConfig(config_file) + submit_path = config[".bps_defined.submitPath"] + + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Batch preparation completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + batch_payload_prepare(config, prefix=submit_path) + _log_mem_usage() diff --git a/python/lsst/ctrl/bps/etc/bps_defaults.yaml b/python/lsst/ctrl/bps/etc/bps_defaults.yaml index 7b1e39a0..92d4b72a 100644 --- a/python/lsst/ctrl/bps/etc/bps_defaults.yaml +++ b/python/lsst/ctrl/bps/etc/bps_defaults.yaml @@ -150,3 +150,25 @@ memoryLimit: 491520 # Default values for making id soft link to submit directory makeIdLink: False idLinkPath: "${PWD}/bps_links" + + +# Running submit stages as batch jobs +bpsBatchSubmission: false +bpsPreCommandOpts: "{defaultPreCmdOpts}" +buildQuantumGraph: + jobCommand: "${CTRL_BPS_DIR}/bin/bps {bpsPreCommandOpts} batch-acquire {configFile} --save-qgraph {fileDistributionEndPoint}{qgraphFile}" + requestMemory: 32768 +preparePayloadWorkflow: + #bpsPreCommandOpts: "--log-level=lsst.ctrl.bps.htcondor.prepare_utils=DEBUG --long-log" + jobCommand: "${CTRL_BPS_DIR}/bin/bps {bpsPreCommandOpts} batch-prepare {configFile} -g {fileDistributionEndPoint}{qgraphFile}" + requestMemory: 32768 + + +# Whether bps jobs, like preparePayloadWorkflow, use +# the site run temp space for staging files. Set +# default for local submissions. +bpsUseRunTempSpace: false + +# Whether BPS makes the exact command to be executed or +# whether WMS is responsible for any additional job commands. +bpsMakeCommand: true diff --git a/python/lsst/ctrl/bps/generic_workflow.py b/python/lsst/ctrl/bps/generic_workflow.py index c3c583bc..10c0092b 100644 --- a/python/lsst/ctrl/bps/generic_workflow.py +++ b/python/lsst/ctrl/bps/generic_workflow.py @@ -33,6 +33,7 @@ "GenericWorkflowFile", "GenericWorkflowGroup", "GenericWorkflowJob", + "GenericWorkflowLazyGroup", "GenericWorkflowNode", "GenericWorkflowNodeType", "GenericWorkflowNoopJob", @@ -127,6 +128,9 @@ class GenericWorkflowNodeType(IntEnum): GROUP = auto() """A special group (subdag) of jobs.""" + LAZY_GROUP = auto() + """When run will generate sub-workflow of jobs.""" + @dataclasses.dataclass(slots=True) class GenericWorkflowNode: @@ -474,7 +478,7 @@ def add_job( super().add_node(job.name, job=job) self.add_job_relationships(parent_names, job.name) self.add_job_relationships(job.name, child_names) - if job.node_type == GenericWorkflowNodeType.PAYLOAD: + if job.node_type in [GenericWorkflowNodeType.PAYLOAD, GenericWorkflowNodeType.LAZY_GROUP]: job = cast(GenericWorkflowJob, job) self.add_executable(job.executable) self._job_labels.add_job( @@ -1336,6 +1340,18 @@ def __init__(self, name: str, label: str, blocking: bool = False) -> None: self.blocking = blocking +@dataclasses.dataclass(slots=True) +class GenericWorkflowLazyGroup(GenericWorkflowJob): + """Node representing a group of jobs to be generated when run.""" + + # Docstring inherited. + + @property + def node_type(self) -> GenericWorkflowNodeType: + """Indicate this is a lazy group job.""" + return GenericWorkflowNodeType.LAZY_GROUP + + class GenericWorkflowLabels: """Label-oriented representation of the GenericWorkflowJobs.""" diff --git a/python/lsst/ctrl/bps/initialize.py b/python/lsst/ctrl/bps/initialize.py index 166c9ac3..f4d466c9 100644 --- a/python/lsst/ctrl/bps/initialize.py +++ b/python/lsst/ctrl/bps/initialize.py @@ -140,7 +140,10 @@ def init_submission( # save copy of configs (orig and expanded config) shutil.copy2(config_file, submit_path) - with open(f"{submit_path}/{config['uniqProcName']}_config.yaml", "w") as fh: + expanded_config_file = f"{submit_path}/{config['uniqProcName']}_config.yaml" + config[".bps_defined.configFile"] = expanded_config_file + + with open(expanded_config_file, "w") as fh: config.dump(fh) # Dump information about runtime environment and software versions in use. diff --git a/python/lsst/ctrl/bps/pre_transform.py b/python/lsst/ctrl/bps/pre_transform.py index e4cef200..c712c9ba 100644 --- a/python/lsst/ctrl/bps/pre_transform.py +++ b/python/lsst/ctrl/bps/pre_transform.py @@ -41,7 +41,7 @@ from lsst.pipe.base import QuantumGraph from lsst.pipe.base.pipeline_graph import TaskImportMode from lsst.pipe.base.quantum_graph import PredictedQuantumGraph -from lsst.resources import ResourcePath +from lsst.resources import ResourcePath, ResourcePathExpression from lsst.utils import doImport from lsst.utils.logging import VERBOSE from lsst.utils.timer import time_this, timeMethod @@ -50,7 +50,7 @@ @timeMethod(logger=_LOG, logLevel=VERBOSE) -def acquire_quantum_graph(config: BpsConfig, out_prefix: str = "") -> tuple[str, PredictedQuantumGraph]: +def acquire_quantum_graph(config: BpsConfig, out_prefix: str = "") -> str: """Read a quantum graph from a file or create one from scratch. Parameters @@ -64,10 +64,7 @@ def acquire_quantum_graph(config: BpsConfig, out_prefix: str = "") -> tuple[str, Returns ------- qgraph_filename : `str` - Name of file containing QuantumGraph that was read into qgraph. - qgraph : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` - A quantum graph read in from pre-generated file or one that is the - result of running code that generates it. + Name of file containing the quantum graph. """ # Check to see if user provided pre-generated QuantumGraph. found, input_qgraph_filename = config.search("qgraphFile") @@ -91,6 +88,24 @@ def acquire_quantum_graph(config: BpsConfig, out_prefix: str = "") -> tuple[str, with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed creating quantum graph"): qgraph_filename = create_quantum_graph(config, out_prefix) + return qgraph_filename + + +@timeMethod(logger=_LOG, logLevel=VERBOSE) +def read_quantum_graph(qgraph_filename: ResourcePathExpression) -> PredictedQuantumGraph: + """Read a quantum graph from a file. + + Parameters + ---------- + qgraph_filename : `lsst.resources.ResourcePathExpression` + Name of file containing PredictedQuantumGraph to be read. + + Returns + ------- + qgraph : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` + A quantum graph read in from pre-generated file or one that is the + result of running code that generates it. + """ _LOG.info("Reading quantum graph from '%s'", qgraph_filename) with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed reading quantum graph"): qgraph_path = ResourcePath(qgraph_filename) @@ -102,7 +117,7 @@ def acquire_quantum_graph(config: BpsConfig, out_prefix: str = "") -> tuple[str, qgraph = PredictedQuantumGraph.from_old_quantum_graph(QuantumGraph.loadUri(qgraph_path)) else: raise ValueError(f"Unrecognized extension for quantum graph file: {qgraph_filename}.") - return qgraph_filename, qgraph + return qgraph def execute(command: str, filename: str, write_buffering: int = 1) -> int: diff --git a/python/lsst/ctrl/bps/tests/gw_test_utils.py b/python/lsst/ctrl/bps/tests/gw_test_utils.py index e3cdb278..5fb9dfd2 100644 --- a/python/lsst/ctrl/bps/tests/gw_test_utils.py +++ b/python/lsst/ctrl/bps/tests/gw_test_utils.py @@ -33,6 +33,7 @@ "make_5_label_workflow", "make_5_label_workflow_2_groups", "make_5_label_workflow_middle_groups", + "make_lazy_workflow", ] import logging @@ -44,6 +45,7 @@ GenericWorkflowExec, GenericWorkflowGroup, GenericWorkflowJob, + GenericWorkflowLazyGroup, GenericWorkflowNodeType, GenericWorkflowNoopJob, ) @@ -656,3 +658,43 @@ def compare_generic_workflows(gwf1: GenericWorkflow, gwf2: GenericWorkflow) -> b equal = False return equal + + +def make_lazy_workflow(workflow_name: str, final: bool) -> GenericWorkflow: # pragma: no cover + """Create a simple workflow with a lazy workflow node for WMS tests. + + Parameters + ---------- + workflow_name : `str` + Name of the test workflow. + final : `bool` + Whether to add a final job. + + Returns + ------- + gwf : `lsst.ctrl.bps.GenericWorkflow` + The test workflow. + """ + gwf = GenericWorkflow(workflow_name) + + # job 1 + gwexec1 = GenericWorkflowExec("exec1", "my_exec1.sh", False) + job1 = GenericWorkflowJob("job1", "label1", executable=gwexec1) + gwf.add_job(job1, None) + + # lazy workflow job + gwexec2 = GenericWorkflowExec("exec2", "${CTRL_BPS_DIR}/python/lsst/ctrl/bps/_make_workflow.sh", False) + job2 = GenericWorkflowLazyGroup("lazy2", "label2", executable=gwexec2) + gwf.add_job(job2, [job1.name]) + + # Job after + gwexec3 = GenericWorkflowExec("exec3", "my_exec3.sh", False) + job3 = GenericWorkflowJob("job3", "label3", executable=gwexec3) + gwf.add_job(job3, [job2.name]) + + if final: + gwexec = GenericWorkflowExec("finalJob.bash", "finalJob.bash", True) + job = GenericWorkflowJob("finalJob", label="finalJob", executable=gwexec) + gwf.add_final(job) + + return gwf diff --git a/python/lsst/ctrl/bps/wms_service.py b/python/lsst/ctrl/bps/wms_service.py index 7665526d..44f59fd0 100644 --- a/python/lsst/ctrl/bps/wms_service.py +++ b/python/lsst/ctrl/bps/wms_service.py @@ -43,6 +43,8 @@ from enum import Enum from typing import Any +from . import BpsConfig + _LOG = logging.getLogger(__name__) @@ -576,3 +578,13 @@ def write(self, out_prefix): as well as internal WMS files. """ raise NotImplementedError + + def add_to_parent_workflow(self, config: BpsConfig) -> None: + """Add self to parent workflow. + + Parameters + ---------- + config : `lsst.ctrl.bps.BpsConfig` + Configuration. + """ + raise NotImplementedError diff --git a/tests/qg_test_utils.py b/tests/qg_test_utils.py index 1e61514d..b57daf11 100644 --- a/tests/qg_test_utils.py +++ b/tests/qg_test_utils.py @@ -304,7 +304,7 @@ def make_test_helper() -> InMemoryRepo: return helper -def make_test_quantum_graph(run: str = "run", uneven=False): +def make_test_quantum_graph(run: str = "run", uneven=False, save_filename: str = None): """Create a quantum graph for unit tests. Parameters @@ -314,6 +314,8 @@ def make_test_quantum_graph(run: str = "run", uneven=False): uneven : `bool`, optional Whether some of the quanta for initial tasks are not included as if finished in previous run. + save_filename : `str`, optional + Save test quantum graph to file using given filename. Returns ------- @@ -345,4 +347,7 @@ def make_test_quantum_graph(run: str = "run", uneven=False): } qgc.set_thin_graph() qgc.set_header_counts() + + if save_filename: + qgc.write(save_filename) return qgc.assemble() diff --git a/tests/test_batch_submit.py b/tests/test_batch_submit.py new file mode 100644 index 00000000..33c302dd --- /dev/null +++ b/tests/test_batch_submit.py @@ -0,0 +1,338 @@ +# This file is part of ctrl_bps. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This software is dual licensed under the GNU General Public License and also +# under a 3-clause BSD license. Recipients may choose which of these licenses +# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, +# respectively. If you choose the GPL option then the following text applies +# (but note that there is still no warranty even if you opt for BSD instead): +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +"""Unit tests for batch_submit.py.""" + +import logging +import shutil +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from lsst.ctrl.bps import BpsConfig, batch_submit + + +class TestCreateBatchStages(unittest.TestCase): + """Tests for create_batch_stages function.""" + + def testMissingBuildCmd(self): + """Missing buildQuantumGraph jobCommand""" + config = BpsConfig({"uniqProcName": "uniq_proc_name"}) + with self.assertRaisesRegex( + RuntimeError, "Missing executable for buildQuantumGraph. Double check submit yaml for jobCommand" + ): + _ = batch_submit.create_batch_stages(config, "not_used_prefix") + + def testMissingPrepareCmd(self): + """Missing preparePayloadWorkflow jobCommand""" + config = BpsConfig( + { + "configFile": "not_used_configFile", + "uniqProcName": "uniq_proc_name", + "operator": "testuser", + "payload": {"payloadName": "testPayload"}, + "bpsPreCommandOpts": "--long-log --log-level=VERBOSE", + "buildQuantumGraph": {"jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}"}, + } + ) + with self.assertRaisesRegex( + RuntimeError, + "Missing executable for preparePayloadWorkflow. Double check submit yaml for jobCommand", + ): + _ = batch_submit.create_batch_stages(config, "not_used_prefix") + + def testSuccess(self): + # No saving of files + config = BpsConfig( + { + "configFile": "not_used_configFile", + "uniqProcName": "uniq_proc_name", + "operator": "testuser", + "payload": {"payloadName": "testPayload"}, + "bpsPreCommandOpts": "--long-log --log-level=VERBOSE", + "buildQuantumGraph": { + "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}", + "requestMemory": 16384, + }, + "preparePayloadWorkflow": { + "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-prepare {configFile}", + "requestMemory": 24576, + }, + } + ) + + with tempfile.TemporaryDirectory() as tmpdir: + gw, config = batch_submit.create_batch_stages(config, tmpdir) + self.assertIn("buildQuantumGraph", gw) + job = gw.get_job("buildQuantumGraph") + self.assertIn("batch-acquire", job.arguments) + self.assertEqual(job.request_memory, 16384) + self.assertIn("preparePayloadWorkflow", gw) + job = gw.get_job("preparePayloadWorkflow") + self.assertIn("batch-prepare", job.arguments) + self.assertEqual(job.request_memory, 24576) + + # Check we didn't make any files + self.assertEqual(list(Path(tmpdir).iterdir()), []) + + def testSaving(self): + config = BpsConfig( + { + "configFile": "not_used_configFile", + "uniqProcName": "uniq_proc_name", + "operator": "testuser", + "payload": {"payloadName": "testPayload"}, + "bpsPreCommandOpts": "--long-log --log-level=VERBOSE", + "buildQuantumGraph": { + "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}", + "requestMemory": 16384, + }, + "preparePayloadWorkflow": { + "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-prepare {configFile}", + "requestMemory": 24576, + }, + "saveGenericWorkflow": True, + } + ) + with tempfile.TemporaryDirectory() as tmpdir: + gw, config = batch_submit.create_batch_stages(config, tmpdir) + self.assertTrue((Path(tmpdir) / "bps_stages_generic_workflow.pickle").exists()) + + +class TestBatchPayloadPrepare(unittest.TestCase): + """Tests for batch_payload_prepare function.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.config_info = { + "runQgraphFile": "run.qgraph", + "uniqProcName": "uniq_proc_name", + "computeSite": "site1", + "qgraphFileTemplate": "template.qgraph", + "bps_defined": {"submitPath": self.tmpdir}, + } + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _make_mocks(self, mock_cluster, mock_transform, mock_prepare): + """Configure the standard set of dependency mocks.""" + clustered_qgraph = MagicMock() + clustered_qgraph.__len__.return_value = 3 + mock_cluster.return_value = clustered_qgraph + + generic_workflow = MagicMock() + generic_workflow.name = "test_workflow" + generic_workflow.job_counts = {"label1": 5} + gwfile = MagicMock() + generic_workflow.get_file.return_value = gwfile + generic_workflow_config = MagicMock() + mock_transform.return_value = (generic_workflow, generic_workflow_config) + + wms_workflow = MagicMock() + mock_prepare.return_value = wms_workflow + + return clustered_qgraph, generic_workflow, generic_workflow_config, gwfile, wms_workflow + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testSuccessBasic(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """Test success with all save flags off and no run temp space.""" + _, generic_workflow, generic_workflow_config, gwfile, wms_workflow = self._make_mocks( + mock_cluster, mock_transform, mock_prepare + ) + config = BpsConfig(self.config_info) + + batch_submit.batch_payload_prepare(config, self.tmpdir) + + mock_read.assert_called_once_with("run.qgraph") + mock_cluster.assert_called_once() + mock_transform.assert_called_once() + mock_prepare.assert_called_once() + # The runQgraphFile should be marked as not transferred by the WMS. + self.assertFalse(gwfile.wms_transfer) + # The payload workflow should be attached to the running workflow. + wms_workflow.add_to_parent_workflow.assert_called_once_with(generic_workflow_config) + # No files should be written with all save flags off. + self.assertEqual(list(Path(self.tmpdir).iterdir()), []) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testSaveClusteredQgraph(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """Test saving of the clustered quantum graph.""" + clustered_qgraph, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) + self.config_info["saveClusteredQgraph"] = True + config = BpsConfig(self.config_info) + + batch_submit.batch_payload_prepare(config, self.tmpdir) + + clustered_qgraph.save.assert_called_once() + self.assertIn("bps_clustered_qgraph.pickle", clustered_qgraph.save.call_args[0][0]) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testSaveDotClustered(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """Test writing of the dot file.""" + clustered_qgraph, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) + self.config_info["saveDot"] = True + config = BpsConfig(self.config_info) + + batch_submit.batch_payload_prepare(config, self.tmpdir) + + clustered_qgraph.draw.assert_called_once() + self.assertIn("bps_clustered_qgraph.dot", clustered_qgraph.draw.call_args[0][0]) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testSaveGenericWorkflow(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """Test writing of the GenericWorkflow to a file.""" + _, generic_workflow, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) + self.config_info["saveGenericWorkflow"] = True + config = BpsConfig(self.config_info) + + batch_submit.batch_payload_prepare(config, self.tmpdir) + + generic_workflow.save.assert_called_once() + self.assertTrue((Path(self.tmpdir) / "bps_generic_workflow.pickle").exists()) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testSaveDotGeneric(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """Test saving the generic workflow dot file.""" + _, generic_workflow, *_ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) + self.config_info["saveDot"] = True + config = BpsConfig(self.config_info) + + batch_submit.batch_payload_prepare(config, self.tmpdir) + + generic_workflow.draw.assert_called_once() + self.assertEqual(generic_workflow.draw.call_args[0][1], "dot") + self.assertTrue((Path(self.tmpdir) / "bps_generic_workflow.dot").exists()) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testUseRunTempSpaceFound(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """When run temp space is enabled and endpoint set, src_uri updates.""" + _, _, _, gwfile, _ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) + self.config_info["bpsUseRunTempSpace"] = True + self.config_info["fileDistributionEndpoint"] = "/run/temp/space" + config = BpsConfig(self.config_info) + + batch_submit.batch_payload_prepare(config, self.tmpdir) + + self.assertEqual(gwfile.src_uri, str(Path("/run/temp/space") / "template.qgraph")) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testUseRunTempSpaceMissingEndpoint(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """Run temp space enabled, missing endpoint should raise KeyError.""" + self._make_mocks(mock_cluster, mock_transform, mock_prepare) + self.config_info["bpsUseRunTempSpace"] = True + config = BpsConfig(self.config_info) + + with self.assertRaisesRegex(KeyError, "fileDistributionEndpoint"): + batch_submit.batch_payload_prepare(config, self.tmpdir) + + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.transform") + @patch("lsst.ctrl.bps.batch_submit.cluster_quanta") + @patch("lsst.ctrl.bps.batch_submit.read_quantum_graph") + def testUseRunTempSpaceNotFound(self, mock_read, mock_cluster, mock_transform, mock_prepare): + """When bpsUseRunTempSpace is absent, a debug message is logged.""" + _, _, _, gwfile, _ = self._make_mocks(mock_cluster, mock_transform, mock_prepare) + config = BpsConfig(self.config_info) + + with self.assertLogs("lsst.ctrl.bps.batch_submit", level=logging.DEBUG) as cm: + batch_submit.batch_payload_prepare(config, self.tmpdir) + + self.assertTrue(any("missing bpsUseRunTempSpace" in msg for msg in cm.output)) + + +class TestBatchSubmit(unittest.TestCase): + """Tests for batch_submit function.""" + + def setUp(self): + self.config_info = {"bps_defined": {"submitPath": "/the/path"}} + + @patch("lsst.ctrl.bps.batch_submit._make_id_link") + @patch("lsst.ctrl.bps.batch_submit.submit") + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.create_batch_stages") + def testSuccessSubmits(self, mock_create, mock_prepare, mock_submit, mock_make_id_link): + """Without dryRun the control workflow is prepared and submitted.""" + generic_workflow = MagicMock() + config = BpsConfig(self.config_info) + mock_create.return_value = (generic_workflow, config) + wms_workflow = MagicMock() + wms_workflow.run_id = "run123" + mock_prepare.return_value = wms_workflow + + result = batch_submit.batch_submit(config) + + mock_create.assert_called_once() + mock_prepare.assert_called_once() + mock_submit.assert_called_once() + mock_make_id_link.assert_called_once_with(config, "run123") + self.assertIs(result, wms_workflow) + + @patch("lsst.ctrl.bps.batch_submit._make_id_link") + @patch("lsst.ctrl.bps.batch_submit.submit") + @patch("lsst.ctrl.bps.batch_submit.prepare") + @patch("lsst.ctrl.bps.batch_submit.create_batch_stages") + def testDryRun(self, mock_create, mock_prepare, mock_submit, mock_make_id_link): + """With dryRun the workflow is not submitted but still returned.""" + generic_workflow = MagicMock() + self.config_info["dryRun"] = True + config = BpsConfig(self.config_info) + mock_create.return_value = (generic_workflow, config) + wms_workflow = MagicMock() + wms_workflow.run_id = "run123" + mock_prepare.return_value = wms_workflow + + result = batch_submit.batch_submit(config) + + mock_submit.assert_not_called() + mock_make_id_link.assert_called_once_with(config, "run123") + self.assertIs(result, wms_workflow) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index 3bef420d..eb24f869 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -185,5 +185,33 @@ def testStatusNonZeroStatus(self): ) +class TestCommandBatchAcquire(unittest.TestCase): + """Test executing the batch-acquire subcommand.""" + + def setUp(self): + self.runner = LogCliRunner() + + def testBatchAcquire(self): + with unittest.mock.patch("lsst.ctrl.bps.cli.cmd.commands.batch_acquire_driver") as mock_driver: + mock_driver.return_value = 0 + result = self.runner.invoke(bps.cli, ["batch-acquire", "test.yaml"]) + self.assertEqual(result.exit_code, 0) + mock_driver.assert_called_once() + + +class TestCommandBatchPrepare(unittest.TestCase): + """Test executing the batch-prepare subcommand.""" + + def setUp(self): + self.runner = LogCliRunner() + + def testBatchPrepare(self): + with unittest.mock.patch("lsst.ctrl.bps.cli.cmd.commands.batch_prepare_driver") as mock_driver: + mock_driver.return_value = 0 + result = self.runner.invoke(bps.cli, ["batch-prepare", "test.yaml"]) + self.assertEqual(result.exit_code, 0) + mock_driver.assert_called_once() + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_drivers.py b/tests/test_drivers.py index df851e2d..3eb6326f 100644 --- a/tests/test_drivers.py +++ b/tests/test_drivers.py @@ -31,12 +31,12 @@ import shutil import tempfile import unittest +from pathlib import Path import yaml -from lsst.ctrl.bps import WmsRunReport, WmsStates +from lsst.ctrl.bps import BaseWmsWorkflow, BpsConfig, WmsRunReport, WmsStates, drivers from lsst.ctrl.bps.bps_reports import compile_code_summary, compile_job_summary -from lsst.ctrl.bps.drivers import _init_submission_driver, ping_driver, report_driver, status_driver TESTDIR = os.path.abspath(os.path.dirname(__file__)) @@ -63,7 +63,7 @@ def testDeprecatedOutCollection(self): with tempfile.NamedTemporaryFile(mode="w+", suffix=".yaml") as file: yaml.dump(config, stream=file) with self.assertRaisesRegex(KeyError, "outCollection"): - _init_submission_driver(file.name) + drivers._init_submission_driver(file.name) @unittest.mock.patch("lsst.ctrl.bps.initialize.BPS_DEFAULTS", {}) def testMissingOutputRun(self): @@ -71,7 +71,7 @@ def testMissingOutputRun(self): with tempfile.NamedTemporaryFile(mode="w+", suffix=".yaml") as file: yaml.dump(config, stream=file) with self.assertRaisesRegex(KeyError, "outputRun"): - _init_submission_driver(file.name) + drivers._init_submission_driver(file.name) @unittest.mock.patch("lsst.ctrl.bps.initialize.BPS_DEFAULTS", {}) def testMissingSubmitPath(self): @@ -79,19 +79,19 @@ def testMissingSubmitPath(self): with tempfile.NamedTemporaryFile(mode="w+", suffix=".yaml") as file: yaml.dump(config, stream=file) with self.assertRaisesRegex(KeyError, "submitPath"): - _init_submission_driver(file.name) + drivers._init_submission_driver(file.name) class TestPingDriver(unittest.TestCase): """Test ping.""" def testWmsServiceSuccess(self): - retval = ping_driver("wms_test_utils.WmsServiceSuccess") + retval = drivers.ping_driver("wms_test_utils.WmsServiceSuccess") self.assertEqual(retval, 0) def testWmsServiceFailure(self): with self.assertLogs(level=logging.ERROR) as cm: - retval = ping_driver("wms_test_utils.WmsServiceFailure") + retval = drivers.ping_driver("wms_test_utils.WmsServiceFailure") self.assertNotEqual(retval, 0) self.assertEqual(cm.records[0].getMessage(), "Couldn't contact service X") @@ -99,7 +99,7 @@ def testWmsServiceEnvVar(self): with unittest.mock.patch.dict( os.environ, {"BPS_WMS_SERVICE_CLASS": "wms_test_utils.WmsServiceSuccess"} ): - retval = ping_driver() + retval = drivers.ping_driver() self.assertEqual(retval, 0) @unittest.mock.patch( @@ -108,13 +108,13 @@ def testWmsServiceEnvVar(self): def testWmsServiceNone(self): with unittest.mock.patch.dict(os.environ, {}): with self.assertLogs(level=logging.INFO) as cm: - retval = ping_driver() + retval = drivers.ping_driver() self.assertEqual(retval, 0) self.assertEqual(cm.records[0].getMessage(), "DEFAULT None") def testWmsServicePassThru(self): with self.assertLogs(level=logging.INFO) as cm: - retval = ping_driver("wms_test_utils.WmsServicePassThru", "EXTRA_VALUES") + retval = drivers.ping_driver("wms_test_utils.WmsServicePassThru", "EXTRA_VALUES") self.assertEqual(retval, 0) self.assertRegex(cm.output[0], "INFO.+EXTRA_VALUES") @@ -124,13 +124,17 @@ class TestStatusDriver(unittest.TestCase): def testWmsServiceSuccess(self): with self.assertLogs(level=logging.INFO) as cm: - retval = status_driver("wms_test_utils.WmsServiceSuccess", run_id="/dummy/path", hist_days=3) + retval = drivers.status_driver( + "wms_test_utils.WmsServiceSuccess", run_id="/dummy/path", hist_days=3 + ) self.assertEqual(retval, WmsStates.SUCCEEDED.value) self.assertEqual(cm.records[0].getMessage(), "status: SUCCEEDED") def testWmsServiceFailure(self): with self.assertLogs(level=logging.WARNING) as cm: - retval = status_driver("wms_test_utils.WmsServiceFailure", run_id="/dummy/path", hist_days=3) + retval = drivers.status_driver( + "wms_test_utils.WmsServiceFailure", run_id="/dummy/path", hist_days=3 + ) self.assertEqual(retval, WmsStates.FAILED.value) self.assertEqual(cm.records[0].getMessage(), "Dummy error message.") @@ -139,7 +143,7 @@ def testWmsServiceFailure(self): ) def testWmsServiceNone(self): with unittest.mock.patch.dict(os.environ, {}): - retval = status_driver(None, run_id="/dummy/path", hist_days=3) + retval = drivers.status_driver(None, run_id="/dummy/path", hist_days=3) self.assertEqual(retval, WmsStates.RUNNING.value) @@ -152,7 +156,7 @@ class TestReportDriver(unittest.TestCase): def testWmsServiceFromDefaults(self): # Should not raise an exception and use default from BPS_DEFAULTS. with unittest.mock.patch.dict(os.environ, {}, clear=True): - report_driver( + drivers.report_driver( wms_service=None, run_id=None, user=None, @@ -165,7 +169,7 @@ def testWmsServiceFromEnvVar(self): with unittest.mock.patch.dict( os.environ, {"BPS_WMS_SERVICE_CLASS": "wms_test_utils.WmsServiceSuccess"} ): - report_driver( + drivers.report_driver( wms_service=None, run_id=None, user=None, @@ -178,7 +182,7 @@ def testWmsServiceFromEnvVar(self): def testHistDefault(self, mock_display, mock_retrieve): mock_retrieve.return_value = ([], []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id="123", user=None, @@ -195,7 +199,7 @@ def testHistDefault(self, mock_display, mock_retrieve): def testHistCustom(self, mock_display, mock_retrieve): mock_retrieve.return_value = ([], []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id="123", user=None, @@ -212,7 +216,7 @@ def testHistCustom(self, mock_display, mock_retrieve): def testPostprocessorsWithoutExitCodes(self, mock_display, mock_retrieve): mock_retrieve.return_value = ([], []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id="123", user=None, @@ -231,7 +235,7 @@ def testPostprocessorsWithoutExitCodes(self, mock_display, mock_retrieve): def testPostprocessorsWithExitCodes(self, mock_display, mock_retrieve): mock_retrieve.return_value = ([], []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id="123", user=None, @@ -251,7 +255,7 @@ def testPostprocessorsWithExitCodes(self, mock_display, mock_retrieve): def testPostprocessorsNoRunId(self, mock_display, mock_retrieve): mock_retrieve.return_value = ([], []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id=None, user=None, @@ -269,7 +273,7 @@ def testDisplayCalledIfRuns(self, mock_display, mock_retrieve): mock_runs = [WmsRunReport(wms_id="1", state=WmsStates.SUCCEEDED)] mock_retrieve.return_value = (mock_runs, []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id=None, user=None, @@ -288,7 +292,7 @@ def testDisplayCalledIfMessages(self, mock_display, mock_retrieve): mock_messages = ["Warning message 1", "Warning message 2"] mock_retrieve.return_value = ([], mock_messages) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id=None, user=None, @@ -307,7 +311,7 @@ def testDisplayCalledIfMessages(self, mock_display, mock_retrieve): def testNoRecordsFoundMessage(self, mock_print, mock_display, mock_retrieve): mock_retrieve.return_value = ([], []) - report_driver( + drivers.report_driver( wms_service="wms_test_utils.WmsServiceSuccess", run_id="123", user=None, @@ -325,5 +329,287 @@ def testNoRecordsFoundMessage(self, mock_print, mock_display, mock_retrieve): self.assertIn("123", call_args) +class TestAcquireQgraphDriver(unittest.TestCase): + """Test acquire_qgraph_driver function.""" + + def setUp(self): + self.tmpdir = Path(tempfile.mkdtemp()) + self.config_file = str(self.tmpdir / "config.yaml") + config = BpsConfig({"bps_defined": {"submitPath": str(self.tmpdir)}}) + + with open(self.config_file, "w") as fh: + config.dump(fh) + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + @unittest.mock.patch("lsst.ctrl.bps.drivers.read_quantum_graph") + @unittest.mock.patch("lsst.ctrl.bps.drivers.acquire_quantum_graph") + @unittest.mock.patch("lsst.ctrl.bps.drivers._init_submission_driver") + def testSuccess(self, mock_init, mock_acquire, mock_read): + drivers.acquire_qgraph_driver(self.config_file) + mock_init.assert_called_once() + mock_acquire.assert_called_once() + mock_read.assert_called_once() + + +class TestBatchAcquireDriver(unittest.TestCase): + """Test batch_acquire_driver function.""" + + def setUp(self): + self.tmpdir = Path(tempfile.mkdtemp()) + self.config_file = str(self.tmpdir / "config.yaml") + config = BpsConfig({"bps_defined": {"submitPath": str(self.tmpdir)}}) + + with open(self.config_file, "w") as fh: + config.dump(fh) + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + @unittest.mock.patch("lsst.ctrl.bps.drivers.acquire_quantum_graph") + def testSuccess(self, mock_acquire): + drivers.batch_acquire_driver(self.config_file) + mock_acquire.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.drivers.acquire_quantum_graph") + def testSaveQgraph(self, mock_acquire): + config = BpsConfig( + { + "bps_defined": {"submitPath": str(self.tmpdir)}, + "computeSite": "site1", + "saveQgraph": "/some/run.qgraph", + } + ) + with open(self.config_file, "w") as fh: + config.dump(fh) + + drivers.batch_acquire_driver(self.config_file) + + mock_acquire.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.drivers.ResourcePath") + @unittest.mock.patch("lsst.ctrl.bps.drivers.acquire_quantum_graph") + def testUseRunTempSpaceFound(self, mock_acquire, mock_resource_path): + config = BpsConfig( + { + "bps_defined": {"submitPath": str(self.tmpdir), "runQgraphFile": "/local/run.qgraph"}, + "computeSite": "site1", + "qgraphFileTemplate": "template.qgraph", + "bpsUseRunTempSpace": True, + "fileDistributionEndpoint": "/run/temp/space", + } + ) + with open(self.config_file, "w") as fh: + config.dump(fh) + + drivers.batch_acquire_driver(self.config_file) + + # The quantum graph should have been transferred to the staging area. + dest = mock_resource_path.return_value.join.return_value + dest.transfer_from.assert_called_once() + self.assertEqual(dest.transfer_from.call_args.kwargs["transfer"], "copy") + + @unittest.mock.patch("lsst.ctrl.bps.drivers.acquire_quantum_graph") + def testUseRunTempSpaceMissingEndpoint(self, mock_acquire): + config = BpsConfig( + { + "bps_defined": {"submitPath": str(self.tmpdir), "runQgraphFile": "/local/run.qgraph"}, + "computeSite": "site1", + "qgraphFileTemplate": "template.qgraph", + "bpsUseRunTempSpace": True, + } + ) + with open(self.config_file, "w") as fh: + config.dump(fh) + + with self.assertRaisesRegex(KeyError, "fileDistributionEndpoint"): + drivers.batch_acquire_driver(self.config_file) + + @unittest.mock.patch("lsst.ctrl.bps.drivers.acquire_quantum_graph") + def testUseRunTempSpaceNotFound(self, mock_acquire): + # config from setUp has no bpsUseRunTempSpace key. + with self.assertLogs("lsst.ctrl.bps.drivers", level=logging.DEBUG) as cm: + drivers.batch_acquire_driver(self.config_file) + self.assertTrue(any("missing bpsUseRunTempSpace" in msg for msg in cm.output)) + + +class TestBatchPrepareDriver(unittest.TestCase): + """Test batch_prepare_driver function.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.config_file = f"{self.tmpdir}/config.yaml" + config = BpsConfig({"bps_defined": {"submitPath": str(self.tmpdir)}}) + + with open(self.config_file, "w") as fh: + config.dump(fh) + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + @unittest.mock.patch("lsst.ctrl.bps.drivers.batch_payload_prepare") + def testSuccess(self, mock_prepare): + drivers.batch_prepare_driver(self.config_file) + mock_prepare.assert_called_once() + + +class _TestWorkflow(BaseWmsWorkflow): + def __init__(self, name, config=None, run_id=None): + super().__init__(name, config) + self.run_id = run_id + + def write(self, out_prefix): + pass # pragma: no cover + + def add_to_parent_workflow(self, config): + pass # pragma: no cover + + +class TestSubmitDriver(unittest.TestCase): + """Test submit_driver function.""" + + def setUp(self): + self.tmpdir = Path(tempfile.mkdtemp()) + self.config_file = str(self.tmpdir / "config.yaml") + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _write_config(self, config_info): + config = BpsConfig(config_info) + with open(self.config_file, "w") as fh: + config.dump(fh) + + @unittest.mock.patch("lsst.ctrl.bps.drivers._make_id_link") + @unittest.mock.patch("lsst.ctrl.bps.drivers.submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.batch_submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.prepare_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers._init_submission_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers.translate_command_line_values") + def testLocalSubmit( + self, mock_translate, mock_init, mock_prepare, mock_batch, mock_submit, mock_make_id_link + ): + """Neither remote build nor batch submission enabled.""" + self._write_config({"wmsServiceClass": "wms_test_utils.WmsServiceSuccess", "uniqProcName": "run1"}) + workflow = _TestWorkflow("run1", run_id="id1") + mock_prepare.return_value = (BpsConfig(self.config_file), workflow) + mock_submit.return_value = workflow + + drivers.submit_driver(self.config_file) + + mock_prepare.assert_called_once() + mock_submit.assert_called_once() + mock_batch.assert_not_called() + mock_init.assert_not_called() + mock_make_id_link.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.drivers._make_id_link") + @unittest.mock.patch("lsst.ctrl.bps.drivers.submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.batch_submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.prepare_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers._init_submission_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers.translate_command_line_values") + def testBatchSubmission( + self, mock_translate, mock_init, mock_prepare, mock_batch, mock_submit, mock_make_id_link + ): + """Batch submission enabled routes through batch_submit.""" + self._write_config( + { + "wmsServiceClass": "wms_test_utils.WmsServiceSuccess", + "uniqProcName": "run1", + "bpsBatchSubmission": True, + } + ) + workflow = _TestWorkflow("run1", run_id="id1") + mock_batch.return_value = workflow + + drivers.submit_driver(self.config_file) + + mock_batch.assert_called_once() + mock_prepare.assert_not_called() + mock_submit.assert_not_called() + mock_init.assert_called_once() + mock_make_id_link.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.drivers._make_id_link") + @unittest.mock.patch("lsst.ctrl.bps.drivers.submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.batch_submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.prepare_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers._init_submission_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers.translate_command_line_values") + def testRemoteBuildEnabled( + self, mock_translate, mock_init, mock_prepare, mock_batch, mock_submit, mock_make_id_link + ): + """Test when remoteBuild enabled.""" + self._write_config( + { + "wmsServiceClass": "wms_test_utils.WmsServiceSuccess", + "uniqProcName": "run1", + "remoteBuild": {"enabled": True}, + } + ) + workflow = _TestWorkflow("run1", run_id="id1") + mock_submit.return_value = workflow + + drivers.submit_driver(self.config_file) + + mock_init.assert_called_once() + mock_prepare.assert_not_called() + mock_batch.assert_not_called() + mock_submit.assert_called_once() + # submit is called with a None workflow, so its return value is used. + self.assertIsNone(mock_submit.call_args[0][1]) + mock_make_id_link.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.drivers._make_id_link") + @unittest.mock.patch("lsst.ctrl.bps.drivers.submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.batch_submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.prepare_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers._init_submission_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers.translate_command_line_values") + def testRemoteBuildDisabled( + self, mock_translate, mock_init, mock_prepare, mock_batch, mock_submit, mock_make_id_link + ): + """Test remoteBuild present but disabled.""" + self._write_config( + { + "wmsServiceClass": "wms_test_utils.WmsServiceSuccess", + "uniqProcName": "run1", + "remoteBuild": {"enabled": False}, + } + ) + workflow = _TestWorkflow("run1", run_id="id1") + mock_prepare.return_value = (BpsConfig(self.config_file), workflow) + mock_submit.return_value = workflow + + drivers.submit_driver(self.config_file) + + mock_init.assert_not_called() + mock_prepare.assert_called_once() + mock_submit.assert_called_once() + mock_batch.assert_not_called() + + @unittest.mock.patch("builtins.print") + @unittest.mock.patch("lsst.ctrl.bps.drivers._make_id_link") + @unittest.mock.patch("lsst.ctrl.bps.drivers.submit") + @unittest.mock.patch("lsst.ctrl.bps.drivers.prepare_driver") + @unittest.mock.patch("lsst.ctrl.bps.drivers.translate_command_line_values") + def testPrintsRunInfo(self, mock_translate, mock_prepare, mock_submit, mock_make_id_link, mock_print): + """Test run info is printed on completion.""" + self._write_config({"wmsServiceClass": "wms_test_utils.WmsServiceSuccess", "uniqProcName": "run1"}) + workflow = _TestWorkflow("run1", run_id="id1") + mock_prepare.return_value = (BpsConfig(self.config_file), workflow) + mock_submit.return_value = workflow + + drivers.submit_driver(self.config_file) + + printed = " ".join(str(call.args[0]) for call in mock_print.call_args_list) + self.assertIn("Run Id:", printed) + self.assertIn("id1", printed) + self.assertIn("Run Name:", printed) + self.assertIn("run1", printed) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_pre_transform.py b/tests/test_pre_transform.py index 72515034..7125dd44 100644 --- a/tests/test_pre_transform.py +++ b/tests/test_pre_transform.py @@ -33,8 +33,11 @@ import unittest from pathlib import Path -from lsst.ctrl.bps import BpsConfig, BpsSubprocessError, ClusteredQuantumGraph -from lsst.ctrl.bps.pre_transform import cluster_quanta, create_quantum_graph, execute, update_quantum_graph +from qg_test_utils import make_test_quantum_graph + +from lsst.ctrl.bps import BpsConfig, BpsSubprocessError, ClusteredQuantumGraph, pre_transform +from lsst.pipe.base import QuantumGraph +from lsst.pipe.base.quantum_graph import PredictedQuantumGraph from lsst.pipe.base.tests.mocks import InMemoryRepo TESTDIR = os.path.abspath(os.path.dirname(__file__)) @@ -56,7 +59,7 @@ def testSuccessfulExecution(self): content = "Successful execution" command = f"{sys.executable} -c 'print(\"{content}\")'" with self.assertLogs(logger=self.logger, level="INFO") as cm: - status = execute(command, self.file.name) + status = pre_transform.execute(command, self.file.name) self.assertIn(content, cm.output[0]) self.file.seek(0) file_contents = self.file.read() @@ -66,7 +69,7 @@ def testSuccessfulExecution(self): def testFailingExecution(self): """Test exit status if command failed.""" - status = execute("false", self.file.name) + status = pre_transform.execute("false", self.file.name) self.assertIn("false", self.file.read()) self.assertNotEqual(status, 0) @@ -92,7 +95,7 @@ def testSuccess(self): """Test if a new quantum graph was created successfully.""" config = BpsConfig(self.settings, search_order=[]) with self.assertLogs(logger=self.logger, level="INFO") as cm: - qgraph_filename = create_quantum_graph(config, self.tmpdir) + qgraph_filename = pre_transform.create_quantum_graph(config, self.tmpdir) _, command = config.search("createQuantumGraph", opt={"curvals": {"qgraphFile": qgraph_filename}}) self.assertIn(command, cm.output[0]) self.assertTrue(os.path.exists(qgraph_filename)) @@ -102,14 +105,14 @@ def testCommandMissing(self): del self.settings["createQuantumGraph"] config = BpsConfig(self.settings, search_order=[]) with self.assertRaisesRegex(KeyError, "command.*not found"): - create_quantum_graph(config, self.tmpdir) + pre_transform.create_quantum_graph(config, self.tmpdir) def testFailure(self): """Test if error is caught when the quantum graph creation fails.""" self.settings["createQuantumGraph"] = "bash -c 'exit 2'" config = BpsConfig(self.settings, search_order=[]) with self.assertRaises(BpsSubprocessError) as cm: - create_quantum_graph(config, self.tmpdir) + pre_transform.create_quantum_graph(config, self.tmpdir) self.assertEqual(cm.exception.errno, errno.ENOENT) self.assertIn("non-zero exit code", str(cm.exception)) @@ -143,7 +146,7 @@ def testSuccess(self): """Test if the quantum graph was updated.""" config = BpsConfig(self.settings, search_order=[]) with self.assertLogs(logger=self.logger, level="INFO") as cm: - update_quantum_graph(config, str(self.src), self.tmpdir) + pre_transform.update_quantum_graph(config, str(self.src), self.tmpdir) _, command = config.search("updateQuantumGraph", opt={"curvals": {"qgraphFile": str(self.src)}}) self.assertIn("backing up", cm.output[0].lower()) self.assertIn("completed", cm.output[1].lower()) @@ -156,7 +159,7 @@ def testSuccessInPlace(self): """Test if a quantum graph was updated inplace.""" config = BpsConfig(self.settings, search_order=[]) with self.assertLogs(logger=self.logger, level="INFO") as cm: - update_quantum_graph(config, str(self.src), self.tmpdir, inplace=True) + pre_transform.update_quantum_graph(config, str(self.src), self.tmpdir, inplace=True) _, command = config.search("updateQuantumGraph", opt={"curvals": {"qgraphFile": str(self.src)}}) self.assertIn(command, cm.output[0]) self.assertTrue(self.src.read_text(), "bar\n") @@ -167,18 +170,107 @@ def testCommandMissing(self): del self.settings["updateQuantumGraph"] config = BpsConfig(self.settings, search_order=[]) with self.assertRaisesRegex(KeyError, "command.*not found"): - update_quantum_graph(config, str(self.src), self.tmpdir) + pre_transform.update_quantum_graph(config, str(self.src), self.tmpdir) def testFailure(self): """Test if error is caught when the command fails.""" self.settings["updateQuantumGraph"] = "bash -c 'exit 2'" config = BpsConfig(self.settings, search_order=[]) with self.assertRaises(BpsSubprocessError) as cm: - update_quantum_graph(config, str(self.src), self.tmpdir) + pre_transform.update_quantum_graph(config, str(self.src), self.tmpdir) self.assertEqual(cm.exception.errno, errno.ENOENT) self.assertRegex(str(cm.exception), "non-zero exit code") +class TestReadQuantumGraph(unittest.TestCase): + """Test read_quantum_graph method.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp(dir=TESTDIR) + self.logger = logging.getLogger("lsst.ctrl.bps") + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def testBadExtension(self): + with self.assertRaisesRegex(ValueError, "Unrecognized extension for quantum graph file"): + _ = pre_transform.read_quantum_graph("mygraph.badext") + + def testReadQG(self): + filename = f"{self.tmpdir}/testQG.qg" + self.qg = make_test_quantum_graph("test_read", save_filename=filename) + self.assertTrue(os.path.exists(filename)) + + qg = pre_transform.read_quantum_graph(filename) + self.assertEqual(len(qg), 18) + + @unittest.mock.patch.object(QuantumGraph, "loadUri") + @unittest.mock.patch.object(PredictedQuantumGraph, "from_old_quantum_graph") + def testReadOldQG(self, mock_from, mock_load): + # Instead of maintaining old format in ctrl_bps, just make sure bps + # function calls methods to read old format and convert old format. + # This test will fail if from_old_quantum_graph or QuantumGraph are + # removed. Those calls should also be removed from read_quantum_graph. + mock_from.return_value = "quantum_graph" + mock_load.return_value = "quantum_graph" + filename = f"{self.tmpdir}/testQG.qgraph" + _ = pre_transform.read_quantum_graph(filename) + mock_load.assert_called_once() + mock_from.assert_called_once() + + +class TestAcquireQuantumGraph(unittest.TestCase): + """Test acquire_quantum_graph method.""" + + def setUp(self): + self.tmpdir = tempfile.mkdtemp(dir=TESTDIR) + self.logger = logging.getLogger("lsst.ctrl.bps") + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + @unittest.mock.patch("lsst.ctrl.bps.pre_transform.create_quantum_graph") + def testCreateGraph(self, mock_create): + qgraph_filename = f"{self.tmpdir}/created.qg" + mock_create.return_value = qgraph_filename + results = pre_transform.acquire_quantum_graph(BpsConfig({}), out_prefix=self.tmpdir) + self.assertEqual(results, qgraph_filename) + mock_create.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.pre_transform.update_quantum_graph") + def testExistingGraphNoCopy(self, mock_update): + filename = "original.qg" + config = BpsConfig({"qgraphFile": filename, "finalJob": {"dummy_var": "dummy_val"}}) + + results = pre_transform.acquire_quantum_graph(config, out_prefix=None) + self.assertEqual(results, filename) + mock_update.assert_called_once() + + @unittest.mock.patch("lsst.ctrl.bps.pre_transform.update_quantum_graph") + def testExistingGraphNoCopyNoUpdate(self, mock_update): + filename = "original.qg" + config = BpsConfig({"qgraphFile": filename}) + + results = pre_transform.acquire_quantum_graph(config, out_prefix=None) + self.assertEqual(results, filename) + mock_update.assert_not_called() + + @unittest.mock.patch("lsst.ctrl.bps.pre_transform.update_quantum_graph") + def testExistingGraphCopy(self, mock_update): + filename = Path(self.tmpdir) / "original.qg" + with open(filename, "w") as fh: + fh.write("test file") + path = Path(self.tmpdir) / "run_dir" + path.mkdir(parents=True, exist_ok=True) + + config = BpsConfig({"qgraphFile": str(filename), "finalJob": {"dummy_var": "dummy_val"}}) + + results = pre_transform.acquire_quantum_graph(config, out_prefix=path) + self.assertEqual(results, str(path / filename.name)) + self.assertTrue(Path(results).exists()) + mock_update.assert_called_once() + + class TestClusterQuanta(unittest.TestCase): """Test cluster_quanta method. Other tests cover functions cluster_quanta calls so mocking them here. @@ -197,7 +289,7 @@ def testValidate(self, mock_validate): with InMemoryRepo() as repo: qgraph = repo.make_quantum_graph() with self.assertRaisesRegex(RuntimeError, "Fake error"): - _ = cluster_quanta(config, qgraph, "a_name") + _ = pre_transform.cluster_quanta(config, qgraph, "a_name") @unittest.mock.patch.object(ClusteredQuantumGraph, "validate") def testNoValidate(self, mock_validate): @@ -211,7 +303,7 @@ def testNoValidate(self, mock_validate): config = BpsConfig(settings, search_order=[]) with InMemoryRepo() as repo: qgraph = repo.make_quantum_graph() - _ = cluster_quanta(config, qgraph, "a_name") + _ = pre_transform.cluster_quanta(config, qgraph, "a_name") if __name__ == "__main__": From be47dea016bfac81392a6357647a0579efd5fbbe Mon Sep 17 00:00:00 2001 From: Michelle Gower Date: Thu, 16 Apr 2026 12:07:09 -0500 Subject: [PATCH 2/4] Update handling of environment and file transfers. --- python/lsst/ctrl/bps/batch_submit.py | 176 +++++++++++++++++++---- python/lsst/ctrl/bps/bps_config.py | 33 +++++ python/lsst/ctrl/bps/cli/cmd/commands.py | 2 + python/lsst/ctrl/bps/drivers.py | 49 +++++-- python/lsst/ctrl/bps/initialize.py | 62 +++++--- python/lsst/ctrl/bps/pre_transform.py | 6 +- python/lsst/ctrl/bps/transform.py | 36 +++-- tests/data/initialize_config.yaml | 13 ++ tests/data/initialize_config_truth.yaml | 13 ++ tests/test_batch_submit.py | 26 +++- tests/test_bpsconfig.py | 70 ++++++++- tests/test_drivers.py | 4 +- 12 files changed, 409 insertions(+), 81 deletions(-) diff --git a/python/lsst/ctrl/bps/batch_submit.py b/python/lsst/ctrl/bps/batch_submit.py index a9701630..4570faee 100644 --- a/python/lsst/ctrl/bps/batch_submit.py +++ b/python/lsst/ctrl/bps/batch_submit.py @@ -27,10 +27,11 @@ """Driver to run submit stages as batch jobs.""" -__all__ = ["batch_payload_prepare", "create_batch_stages"] +__all__ = ["batch_payload_prepare", "batch_submit", "create_batch_stages"] import logging import os +from pathlib import Path from lsst.resources import ResourcePath, ResourcePathExpression from lsst.utils.logging import VERBOSE @@ -41,12 +42,15 @@ DEFAULT_MEM_UNIT, BpsConfig, GenericWorkflow, + GenericWorkflowFile, GenericWorkflowJob, GenericWorkflowLazyGroup, ) +from .bps_utils import _make_id_link from .pre_transform import cluster_quanta, read_quantum_graph from .prepare import prepare -from .transform import _get_job_values, transform +from .submit import submit +from .transform import _enhance_command, _get_job_values, transform _LOG = logging.getLogger(__name__) @@ -73,16 +77,52 @@ def create_batch_stages( """ prefix = ResourcePath(prefix) generic_workflow: GenericWorkflow = GenericWorkflow(name=f"{config['uniqProcName']}_ctrl") - cmd_line_key = "jobCommand" + generic_workflow.run_attrs.update( + { + "bps_isjob": "True", + "bps_project": config["project"], + "bps_campaign": config["campaign"], + "bps_run": config["uniqProcName"], + "bps_operator": config["operator"], + "bps_payload": config["payloadName"], + } + ) - # build QuantumGraph job - search_opt = {} - if "buildQuantumGraph" in config: - search_opt["searchobj"] = config.get("buildQuantumGraph") + # Save full run QuantumGraph for use by jobs + qgraph_file = GenericWorkflowFile( + "runQgraphFile", + src_uri=config["runQgraphFile"], + wms_transfer=True, + job_access_remote=True, + job_shared=True, + ) + generic_workflow.add_file(qgraph_file) + + # Save config file for use by jobs + config_file = GenericWorkflowFile( + "configFile", + src_uri=config["configFile"], + wms_transfer=True, + job_access_remote=False, + job_shared=True, + ) + generic_workflow.add_file(config_file) + + # Build QuantumGraph job build_job = GenericWorkflowJob( name="buildQuantumGraph", label="buildQuantumGraph", ) + search_opt = config.get_search_opts(build_job.label) + search_opt.update( + { + "replaceVars": False, + "expandEnvVars": False, + "replaceEnvVars": True, + "required": False, + } + ) + cmd_line_key = "jobCommand" job_values = _get_job_values(config, search_opt, cmd_line_key) if not job_values["executable"]: raise RuntimeError( @@ -93,26 +133,30 @@ def create_batch_stages( setattr(build_job, key, value) generic_workflow.add_job(build_job) - generic_workflow.run_attrs.update( - { - "bps_isjob": "True", - "bps_project": config["project"], - "bps_campaign": config["campaign"], - "bps_run": config["uniqProcName"], - "bps_operator": config["operator"], - "bps_payload": config["payloadName"], - "bps_runsite": config["computeSite"], - } - ) + _LOG.debug("build job's arguments: %s", build_job.arguments) + + generic_workflow.add_file(config_file) + generic_workflow.add_job_outputs(build_job.name, [qgraph_file]) + generic_workflow.add_job_inputs(build_job.name, [config_file]) + _enhance_command(config, generic_workflow, build_job, job_values) + _LOG.debug("build job's arguments: %s", build_job.arguments) - # cluster/transform/prepare job - search_opt = {} - if "preparePayloadWorkflow" in config: - search_opt["searchobj"] = config.get("preparePayloadWorkflow") + # Build cluster/transform/prepare job prepare_job = GenericWorkflowLazyGroup( name="preparePayloadWorkflow", label="preparePayloadWorkflow", ) + search_opt = config.get_search_opts(prepare_job.label) + search_opt.update( + { + "replaceVars": False, + "expandEnvVars": False, + "replaceEnvVars": True, + "required": False, + } + ) + _LOG.debug("preparePayloadWorkflow search_opt = %s", search_opt) + cmd_line_key = "jobCommand" job_values = _get_job_values(config, search_opt, cmd_line_key) if not job_values["executable"]: raise RuntimeError( @@ -123,6 +167,8 @@ def create_batch_stages( setattr(prepare_job, key, value) generic_workflow.add_job(prepare_job, parent_names=["buildQuantumGraph"]) + generic_workflow.add_job_inputs(prepare_job.name, [qgraph_file, config_file]) + _enhance_command(config, generic_workflow, prepare_job, job_values) _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False}) if save_workflow: @@ -140,7 +186,6 @@ def batch_payload_prepare(config: BpsConfig, prefix: ResourcePathExpression) -> ---------- config : `lsst.ctrl.bps.BpsConfig` BPS configuration. - prefix : `lsst.resources.ResourcePathExpression` Root path for any output files. @@ -151,11 +196,9 @@ def batch_payload_prepare(config: BpsConfig, prefix: ResourcePathExpression) -> generic_workflow_config : `lsst.ctrl.bps.BpsConfig` Configuration to accompany GenericWorkflow. """ - prefix = ResourcePath(prefix) # Read existing QuantumGraph - qgraph_filename = prefix.join(config["qgraphFileTemplate"]) - qgraph = read_quantum_graph(qgraph_filename) - config[".bps_defined.runQgraphFile"] = str(qgraph_filename) + qgraph_uri = config["runQgraphFile"] + qgraph = read_quantum_graph(qgraph_uri) # Cluster _LOG.info("Starting cluster stage (grouping quanta into jobs)") @@ -197,6 +240,26 @@ def batch_payload_prepare(config: BpsConfig, prefix: ResourcePathExpression) -> num_jobs = sum(generic_workflow.job_counts.values()) _LOG.info("GenericWorkflow contains %d job(s) (including final)", num_jobs) + # Want to read quantum graph from temp space if told to use it. + gwfile = generic_workflow.get_file("runQgraphFile") + gwfile.wms_transfer = False + # Root computeSite is currently how to specify where the pipeline + # will actually be run. + found, use_run_temp_space = config.search( + "bpsUseRunTempSpace", opt={"curvals": {"curr_site": config[".computeSite"]}} + ) + if found and use_run_temp_space: + found, run_temp_space = config.search( + "fileDistributionEndpoint", opt={"curvals": {"curr_site": config[".computeSite"]}} + ) + if found: + _LOG.debug("run_temp_space = %s", run_temp_space) + gwfile.src_uri = str(Path(run_temp_space) / config["qgraphFileTemplate"]) + else: + raise KeyError("Config is missing fileDistributionEndpoint.") + elif not found: + _LOG.debug("Config is missing bpsUseRunTempSpace") + _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False}) if save_workflow: with open(os.path.join(submit_path, "bps_generic_workflow.pickle"), "wb") as outfh: @@ -232,3 +295,62 @@ def batch_payload_prepare(config: BpsConfig, prefix: ResourcePathExpression) -> ): # Assuming submit_path for ctrl workflow is visible by this job. wms_workflow.add_to_parent_workflow(generic_workflow_config) + + +def batch_submit(config: BpsConfig): + """Submit a workflow for execution with preparation done in batch jobs. + + Parameters + ---------- + config : `lsst.ctrl.bps.BpsConfig` + BPS configuration. + + Returns + ------- + wms_workflow : `lsst.ctrl.bps.BaseWmsWorkflow` + Submitted workflow. + """ + submit_path = config[".bps_defined.submitPath"] + + _LOG.info("Starting to create control workflow") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Creation completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + generic_workflow, config = create_batch_stages(config, submit_path) + + _LOG.info("Starting to prepare control workflow") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Preparation completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + wms_workflow = prepare(config, generic_workflow, submit_path) + + _, dry_run = config.search("dryRun", opt={"default": False}) + if not dry_run: + _LOG.info("Starting to submit control workflow") + with time_this( + log=_LOG, + level=logging.INFO, + prefix=None, + msg="Submission completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + submit(config, wms_workflow) + _LOG.info("Run '%s' submitted for execution with id '%s'", wms_workflow.name, wms_workflow.run_id) + + _make_id_link(config, wms_workflow.run_id) + + return wms_workflow diff --git a/python/lsst/ctrl/bps/bps_config.py b/python/lsst/ctrl/bps/bps_config.py index c2c561f0..d76295fa 100644 --- a/python/lsst/ctrl/bps/bps_config.py +++ b/python/lsst/ctrl/bps/bps_config.py @@ -497,3 +497,36 @@ def _recursive_generate_config(self, recursive_key: str, sub_config: Config) -> _LOG.debug("After config = %s", self) else: raise ValueError(f"Unparsable {genkey} value='{value}'") + + def get_search_opts(self, label: str | None = None) -> dict[str, Any]: + """Create base search options given a label. + + Parameters + ---------- + label : `str` or None, optional + If given, label for which to define BpsConfig search options. + + Returns + ------- + search_opts : `dict` [`str`, `~typing.Any`] + Base BpsConfig search options for given label. + """ + search_opts = {"curvals": {}} + if label: + search_opts["curvals"]["label"] = label + if label in self["cluster"]: + search_opts["curvals"]["curr_cluster"] = label + elif label in self["pipetask"]: + search_opts["curvals"]["curr_pipetask"] = label + elif label in self: + search_opts["searchobj"] = self[label] + + # Site/cloud can be defined per cluster/pipetask/special job + found, val = self.search("computeSite", search_opts) + if found: + search_opts["curvals"]["curr_site"] = val + found, val = self.search("computeCloud", search_opts) + if found: + search_opts["curvals"]["curr_cloud"] = val + + return search_opts diff --git a/python/lsst/ctrl/bps/cli/cmd/commands.py b/python/lsst/ctrl/bps/cli/cmd/commands.py index 29f68d2a..e31663f6 100644 --- a/python/lsst/ctrl/bps/cli/cmd/commands.py +++ b/python/lsst/ctrl/bps/cli/cmd/commands.py @@ -31,6 +31,7 @@ import click +from lsst.ctrl.mpexec.cli.opt import save_qgraph_option from lsst.daf.butler.cli.utils import MWCommand from ... import BpsSubprocessError @@ -226,6 +227,7 @@ def submitcmd(*args, **kwargs): @click.command(cls=BpsCommand) @opt.config_file_argument(required=True) @opt.submission_options() +@save_qgraph_option() def batch_acquire(*args, **kwargs): """Run inside a batch job to create a new quantum graph.""" with catch_errors(): diff --git a/python/lsst/ctrl/bps/drivers.py b/python/lsst/ctrl/bps/drivers.py index c705867d..96ce368f 100644 --- a/python/lsst/ctrl/bps/drivers.py +++ b/python/lsst/ctrl/bps/drivers.py @@ -54,6 +54,7 @@ from typing import Any from lsst.pipe.base.quantum_graph import PredictedQuantumGraph +from lsst.resources import ResourcePath from lsst.utils.timer import time_this from lsst.utils.usage import get_peak_mem_usage @@ -66,7 +67,7 @@ ClusteredQuantumGraph, GenericWorkflow, ) -from .batch_submit import batch_payload_prepare, create_batch_stages +from .batch_submit import batch_payload_prepare, batch_submit from .bps_reports import compile_code_summary, compile_job_summary from .bps_utils import _dump_env_info, _dump_pkg_info, _make_id_link from .cancel import cancel @@ -77,6 +78,7 @@ out_collection_validator, output_run_validator, submit_path_validator, + translate_command_line_values, ) from .ping import ping from .pre_transform import acquire_quantum_graph, cluster_quanta, read_quantum_graph @@ -97,7 +99,7 @@ def _init_submission_driver(config_file: str, **kwargs) -> BpsConfig: ---------- config_file : `str` Name of the configuration file. - **kwargs : `~typing.Any` + **kwargs : `dict` [`str`, `~typing.Any`] Additional modifiers to the configuration. Returns @@ -713,18 +715,22 @@ def _log_mem_usage() -> None: ) -def batch_acquire_driver(config_file: str, **kwargs: Any) -> None: +def batch_acquire_driver(config_file: str, **kwargs) -> None: """Create a quantum graph from pipeline definition in a batch job. Parameters ---------- config_file : `str` Name of the configuration file. - **kwargs : `~typing.Any` + **kwargs Additional modifiers to the configuration. """ config = BpsConfig(config_file) - submit_path = config[".bps_defined.submitPath"] + translate_command_line_values(config, **kwargs) + + found, val = config.search("saveQgraph") + if found: + config[".bps_defined.runQgraphFile"] = val _LOG.info("Starting acquire stage (generating and/or reading quantum graph)") with time_this( @@ -736,21 +742,48 @@ def batch_acquire_driver(config_file: str, **kwargs: Any) -> None: mem_unit=DEFAULT_MEM_UNIT, mem_fmt=DEFAULT_MEM_FMT, ): - _ = acquire_quantum_graph(config, out_prefix=submit_path) + _ = acquire_quantum_graph(config, out_prefix="") + + # Copy quantum graph to staging area + found, use_run_temp_space = config.search( + "bpsUseRunTempSpace", opt={"curvals": {"curr_site": config[".computeSite"]}} + ) + if found and use_run_temp_space: + found, run_temp_space = config.search( + "fileDistributionEndpoint", opt={"curvals": {"curr_site": config[".computeSite"]}} + ) + if found: + _LOG.debug("run_temp_space = %s", run_temp_space) + dest = ResourcePath(run_temp_space, forceDirectory=True).join(config["qgraphFileTemplate"]) + src = ResourcePath(config[".bps_defined.runQgraphFile"], forceDirectory=False) + # S3 clients explicitly instantiate here to overpass this + # https://stackoverflow.com/questions/52820971/is-boto3-client-thread-safe + dest.exists() + + _LOG.debug("Copying quantum graph from %s to %s", src, dest) + dest.transfer_from(src, transfer="copy") + else: + raise KeyError("Config is missing fileDistributionEndpoint.") + elif not found: + _LOG.debug("Config is missing bpsUseRunTempSpace") + _log_mem_usage() -def batch_prepare_driver(config_file: str, **kwargs: Any) -> None: +def batch_prepare_driver(config_file: str, **kwargs) -> None: """Run workflow preparation in a batch job for an existing QuantumGraph. Parameters ---------- config_file : `str` Name of the configuration file. - **kwargs : `~typing.Any` + **kwargs Additional modifiers to the configuration. """ config = BpsConfig(config_file) + translate_command_line_values(config, **kwargs) + + config[".bps_defined.runQgraphFile"] = kwargs["qgraph"] submit_path = config[".bps_defined.submitPath"] with time_this( diff --git a/python/lsst/ctrl/bps/initialize.py b/python/lsst/ctrl/bps/initialize.py index f4d466c9..3867f357 100644 --- a/python/lsst/ctrl/bps/initialize.py +++ b/python/lsst/ctrl/bps/initialize.py @@ -33,6 +33,7 @@ "out_collection_validator", "output_run_validator", "submit_path_validator", + "translate_command_line_values", ] import getpass @@ -40,6 +41,7 @@ import re import shutil from collections.abc import Callable, Iterable +from pathlib import Path from lsst.ctrl.bps import BPS_DEFAULTS, BPS_SEARCH_ORDER, BpsConfig from lsst.ctrl.bps.bps_utils import _dump_env_info, _dump_pkg_info, mkdir @@ -49,6 +51,39 @@ _LOG = logging.getLogger(__name__) +def translate_command_line_values(config: BpsConfig, **kwargs) -> None: + """Override config with command-line values. + + Parameters + ---------- + config : `lsst.ctrl.bps.BpsConfig` + BPS configuration. + **kwargs + Additional modifiers to the configuration from the command line. + """ + _LOG.debug("translate_command_line_values: kwargs = %s", kwargs) + # Handle diffs between pipetask argument names vs bps yaml + translation = { + "input": "inCollection", + "output_run": "outputRun", + "qgraph": "qgraphFile", + "pipeline": "pipelineYaml", + "wms_service": "wmsServiceClass", + "compute_site": "computeSite", + } + for key, value in kwargs.items(): + # Don't want to override config with None or empty string values. + if value: + # pipetask argument parser converts some values to list, + # but bps will want string. + if not isinstance(value, str) and isinstance(value, Iterable): + value = ",".join(value) + new_key = translation.get(key, re.sub(r"_(\S)", lambda match: match.group(1).upper(), key)) + config[f".bps_cmdline.{new_key}"] = value + + _LOG.debug("translate_command_line_values: .bps_cmdline = %s", config[".bps_cmdline"]) + + def init_submission( config_file: str, validators: Iterable[Callable[[BpsConfig], None]] = (), **kwargs ) -> BpsConfig: @@ -63,7 +98,7 @@ def init_submission( A list of functions performing checks on the given configuration. Each function should take a single argument, a BpsConfig object, and raise if the check fails. By default, no checks are performed. - **kwargs : `~typing.Any` + **kwargs Additional modifiers to the configuration. Returns @@ -78,25 +113,7 @@ def init_submission( wms_service_class_fqn=kwargs.get("wms_service"), ) - # Override config with command-line values. - # Handle diffs between pipetask argument names vs bps yaml - translation = { - "input": "inCollection", - "output_run": "outputRun", - "qgraph": "qgraphFile", - "pipeline": "pipelineYaml", - "wms_service": "wmsServiceClass", - "compute_site": "computeSite", - } - for key, value in kwargs.items(): - # Don't want to override config with None or empty string values. - if value: - # pipetask argument parser converts some values to list, - # but bps will want string. - if not isinstance(value, str) and isinstance(value, Iterable): - value = ",".join(value) - new_key = translation.get(key, re.sub(r"_(\S)", lambda match: match.group(1).upper(), key)) - config[f".bps_cmdline.{new_key}"] = value + translate_command_line_values(config, **kwargs) # Run validation tests on the given config if any. for validator in validators: @@ -138,6 +155,11 @@ def init_submission( submit_path = mkdir(config["submitPath"]) config[".bps_defined.submitPath"] = str(submit_path) + # Pre-make quantum graph filename + prefix = Path(submit_path) + qgraph_filename = prefix / config["qgraphFileTemplate"] + config[".bps_defined.runQgraphFile"] = str(qgraph_filename) + # save copy of configs (orig and expanded config) shutil.copy2(config_file, submit_path) expanded_config_file = f"{submit_path}/{config['uniqProcName']}_config.yaml" diff --git a/python/lsst/ctrl/bps/pre_transform.py b/python/lsst/ctrl/bps/pre_transform.py index c712c9ba..ccb708a5 100644 --- a/python/lsst/ctrl/bps/pre_transform.py +++ b/python/lsst/ctrl/bps/pre_transform.py @@ -190,8 +190,10 @@ def create_quantum_graph(config: BpsConfig, out_prefix: str = "") -> str: BpsSubprocessError Raised if the command for generating the QuantumGraph failed. """ - # Create name of file to store QuantumGraph. - qgraph_filename = os.path.join(out_prefix, config["qgraphFileTemplate"]) + found, qgraph_filename = config.search(".bps_defined.runQgraphFile") + if not found: + # Create name of file to store QuantumGraph. + qgraph_filename = os.path.join(out_prefix, config["qgraphFileTemplate"]) # Get QuantumGraph generation command. search_opt = {"curvals": {"qgraphFile": qgraph_filename}} diff --git a/python/lsst/ctrl/bps/transform.py b/python/lsst/ctrl/bps/transform.py index ea4e075e..2c1b6846 100644 --- a/python/lsst/ctrl/bps/transform.py +++ b/python/lsst/ctrl/bps/transform.py @@ -47,7 +47,11 @@ GenericWorkflowFile, GenericWorkflowJob, ) -from .bps_utils import WhenToSaveQuantumGraphs, create_job_quantum_graph_filename, save_qg_subgraph +from .bps_utils import ( + WhenToSaveQuantumGraphs, + create_job_quantum_graph_filename, + save_qg_subgraph, +) # All available job attributes. _ATTRS_ALL = frozenset([field.name for field in dataclasses.fields(GenericWorkflowJob)]) @@ -216,22 +220,21 @@ def _enhance_command(config, generic_workflow, gwjob, cached_job_values): """ _LOG.debug("gwjob given to _enhance_command: %s", gwjob) - curvals = { - "curr_pipetask": gwjob.label, - "curr_cluster": gwjob.label, - "jobName": gwjob.name, - "jobLabel": gwjob.label, - } + search_opt = config.get_search_opts(gwjob.label) + + search_opt["curvals"]["jobName"] = gwjob.name + search_opt["curvals"]["jobLabel"] = gwjob.label for key, value in gwjob.tags.items(): - curvals[key] = value + search_opt["curvals"][key] = value - search_opt = { - "curvals": curvals, - "replaceVars": False, - "expandEnvVars": False, - "replaceEnvVars": True, - "required": False, - } + search_opt.update( + { + "replaceVars": False, + "expandEnvVars": False, + "replaceEnvVars": True, + "required": False, + } + ) if gwjob.label not in cached_job_values: cached_job_values[gwjob.label] = {} @@ -269,13 +272,16 @@ def _enhance_command(config, generic_workflow, gwjob, cached_job_values): # (Be careful to not replace env variables as they may # be different in compute job.) search_opt["replaceVars"] = True + _LOG.debug("before cmdvals = %s (search_opt = %s)", gwjob.cmdvals, search_opt) for key in re.findall(r"{([^}]+)}", gwjob.arguments): + _LOG.debug("looking for %s in cmdvals", key) if key in gwjob.cmdvals: continue elif key in cached_job_values[gwjob.label]: gwjob.cmdvals[key] = cached_job_values[gwjob.label][key] else: _, gwjob.cmdvals[key] = config.search(key, opt=search_opt) + _LOG.debug("after cmdvals = %s", gwjob.cmdvals) # backwards compatibility if not cached_job_values[gwjob.label]["useLazyCommands"]: diff --git a/tests/data/initialize_config.yaml b/tests/data/initialize_config.yaml index a528ecef..98f07a3a 100644 --- a/tests/data/initialize_config.yaml +++ b/tests/data/initialize_config.yaml @@ -16,7 +16,20 @@ pipetask: genval1: "-c val1:0.1 bpsEval(lsst.ctrl.bps.tests.config_test_utils.generate_value_1, {p1})" genval2: "-c val1:bpsEval(max, [{p1}, {p3}])" +cluster: + cl1: + p3: 64 + p5: 55 + finalJob: p1: 9 p2: "unused final" bpsGenerateConfig: "lsst.ctrl.bps.tests.config_test_utils.generate_config_2({p1}, param3={p3})" + +site: + site1: + p1: 33 + +cloud: + cloud1: + p1: 34 diff --git a/tests/data/initialize_config_truth.yaml b/tests/data/initialize_config_truth.yaml index f6e501fb..e3119e16 100644 --- a/tests/data/initialize_config_truth.yaml +++ b/tests/data/initialize_config_truth.yaml @@ -20,6 +20,11 @@ pipetask: genval1: "-c val1:0.1 bpsEval(lsst.ctrl.bps.tests.config_test_utils.generate_value_1, {p1})" genval2: "-c val1:bpsEval(max, [{p1}, {p3}])" +cluster: + cl1: + p3: 64 + p5: 55 + finalJob: p1: 9 p2: "unused final" @@ -28,3 +33,11 @@ finalJob: gencfg_4: 9 gencfg_5: -3 gencfg_6: 16 + +site: + site1: + p1: 33 + +cloud: + cloud1: + p1: 34 diff --git a/tests/test_batch_submit.py b/tests/test_batch_submit.py index 33c302dd..7eed4fbe 100644 --- a/tests/test_batch_submit.py +++ b/tests/test_batch_submit.py @@ -39,9 +39,19 @@ class TestCreateBatchStages(unittest.TestCase): """Tests for create_batch_stages function.""" + def setUp(self): + self.common_config = { + "bpsUseShared": True, + "whenSaveJobQgraph": "NEVER", + "useLazyCommands": True, + "submitPath": "/the/path", + } + def testMissingBuildCmd(self): """Missing buildQuantumGraph jobCommand""" - config = BpsConfig({"uniqProcName": "uniq_proc_name"}) + config_info = dict(self.common_config) + config_info.update({"uniqProcName": "uniq_proc_name"}) + config = BpsConfig(config_info) with self.assertRaisesRegex( RuntimeError, "Missing executable for buildQuantumGraph. Double check submit yaml for jobCommand" ): @@ -49,16 +59,17 @@ def testMissingBuildCmd(self): def testMissingPrepareCmd(self): """Missing preparePayloadWorkflow jobCommand""" - config = BpsConfig( + config_info = dict(self.common_config) + config_info.update( { "configFile": "not_used_configFile", "uniqProcName": "uniq_proc_name", "operator": "testuser", "payload": {"payloadName": "testPayload"}, - "bpsPreCommandOpts": "--long-log --log-level=VERBOSE", "buildQuantumGraph": {"jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}"}, } ) + config = BpsConfig(config_info) with self.assertRaisesRegex( RuntimeError, "Missing executable for preparePayloadWorkflow. Double check submit yaml for jobCommand", @@ -67,13 +78,13 @@ def testMissingPrepareCmd(self): def testSuccess(self): # No saving of files - config = BpsConfig( + config_info = dict(self.common_config) + config_info.update( { "configFile": "not_used_configFile", "uniqProcName": "uniq_proc_name", "operator": "testuser", "payload": {"payloadName": "testPayload"}, - "bpsPreCommandOpts": "--long-log --log-level=VERBOSE", "buildQuantumGraph": { "jobCommand": "${CTRL_BPS_DIR}/bin/bps batch-acquire {configFile}", "requestMemory": 16384, @@ -84,6 +95,7 @@ def testSuccess(self): }, } ) + config = BpsConfig(config_info) with tempfile.TemporaryDirectory() as tmpdir: gw, config = batch_submit.create_batch_stages(config, tmpdir) @@ -100,7 +112,8 @@ def testSuccess(self): self.assertEqual(list(Path(tmpdir).iterdir()), []) def testSaving(self): - config = BpsConfig( + config_info = dict(self.common_config) + config_info.update( { "configFile": "not_used_configFile", "uniqProcName": "uniq_proc_name", @@ -118,6 +131,7 @@ def testSaving(self): "saveGenericWorkflow": True, } ) + config = BpsConfig(config_info) with tempfile.TemporaryDirectory() as tmpdir: gw, config = batch_submit.create_batch_stages(config, tmpdir) self.assertTrue((Path(tmpdir) / "bps_stages_generic_workflow.pickle").exists()) diff --git a/tests/test_bpsconfig.py b/tests/test_bpsconfig.py index 7eb69ac6..d063ea09 100644 --- a/tests/test_bpsconfig.py +++ b/tests/test_bpsconfig.py @@ -466,11 +466,12 @@ def testWithSearchOrder(self): # config ordering is used. # Ditto for finalJob (which isn't a search section). # Checking all in single function to ensure doesn't quit early. + self.maxDiff = None + self.config.generate_config() filename = os.path.join(TESTDIR, "data/initialize_config_truth.yaml") truth = BpsConfig(filename, BPS_SEARCH_ORDER, defaults={}) - self.assertEqual(self.config, truth) def testBpsEval(self): @@ -511,5 +512,72 @@ def testBpsEvalInvalid(self): _ = self.config.search("badkey1", opt=test_opt) +class TestBpsConfigGetSearchOpts(unittest.TestCase): + """Tests for BpsConfig.get_search_opts.""" + + def testPipetaskLabel(self): + filename = os.path.join(TESTDIR, "data/initialize_config_truth.yaml") + config = BpsConfig(filename, BPS_SEARCH_ORDER, defaults={}) + config["computeCloud"] = "cloud1" + config["computeSite"] = "site1" + results = config.get_search_opts("ptask1") + truth = { + "curvals": { + "label": "ptask1", + "curr_pipetask": "ptask1", + "curr_site": "site1", + "curr_cloud": "cloud1", + } + } + self.assertEqual(truth, results) + + def testClusterLabel(self): + filename = os.path.join(TESTDIR, "data/initialize_config_truth.yaml") + config = BpsConfig(filename, BPS_SEARCH_ORDER, defaults={}) + config["computeCloud"] = "cloud1" + config["computeSite"] = "site1" + results = config.get_search_opts("cl1") + truth = { + "curvals": {"label": "cl1", "curr_cluster": "cl1", "curr_site": "site1", "curr_cloud": "cloud1"} + } + self.assertEqual(truth, results) + + def testNoMatchingLabel(self): + filename = os.path.join(TESTDIR, "data/initialize_config_truth.yaml") + config = BpsConfig(filename, BPS_SEARCH_ORDER, defaults={}) + config["computeCloud"] = "cloud1" + config["computeSite"] = "site1" + results = config.get_search_opts("notthere") + truth = {"curvals": {"label": "notthere", "curr_site": "site1", "curr_cloud": "cloud1"}} + self.assertEqual(truth, results) + + def testNoMatchingAnything(self): + filename = os.path.join(TESTDIR, "data/initialize_config_truth.yaml") + config = BpsConfig(filename, BPS_SEARCH_ORDER, defaults={}) + config["computeCloud"] = "cloud_notthere" + config["computeSite"] = "site_notthere" + results = config.get_search_opts("notthere") + truth = { + "curvals": {"label": "notthere", "curr_site": "site_notthere", "curr_cloud": "cloud_notthere"} + } + self.assertEqual(truth, results) + + def testFinal(self): + filename = os.path.join(TESTDIR, "data/initialize_config_truth.yaml") + config = BpsConfig(filename, BPS_SEARCH_ORDER, defaults={}) + config["computeCloud"] = "cloud1" + config["computeSite"] = "site1" + results = config.get_search_opts("finalJob") + truth = { + "curvals": { + "label": "finalJob", + "curr_site": "site1", + "curr_cloud": "cloud1", + }, + "searchobj": config["finalJob"], + } + self.assertEqual(truth, results) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_drivers.py b/tests/test_drivers.py index 3eb6326f..425bd0c6 100644 --- a/tests/test_drivers.py +++ b/tests/test_drivers.py @@ -359,7 +359,7 @@ class TestBatchAcquireDriver(unittest.TestCase): def setUp(self): self.tmpdir = Path(tempfile.mkdtemp()) self.config_file = str(self.tmpdir / "config.yaml") - config = BpsConfig({"bps_defined": {"submitPath": str(self.tmpdir)}}) + config = BpsConfig({"bps_defined": {"submitPath": str(self.tmpdir)}, "computeSite": "site1"}) with open(self.config_file, "w") as fh: config.dump(fh) @@ -450,7 +450,7 @@ def tearDown(self): @unittest.mock.patch("lsst.ctrl.bps.drivers.batch_payload_prepare") def testSuccess(self, mock_prepare): - drivers.batch_prepare_driver(self.config_file) + drivers.batch_prepare_driver(self.config_file, qgraph="test.qg") mock_prepare.assert_called_once() From 3b65562c398a5a042303193112f2a20735e5b125 Mon Sep 17 00:00:00 2001 From: Michelle Gower Date: Mon, 29 Jun 2026 20:12:03 -0500 Subject: [PATCH 3/4] Add run site to report. --- python/lsst/ctrl/bps/batch_submit.py | 1 + python/lsst/ctrl/bps/bps_reports.py | 16 ++++++++++++++++ python/lsst/ctrl/bps/report.py | 16 ++-------------- python/lsst/ctrl/bps/wms_service.py | 3 +++ tests/test_bps_reports.py | 27 +++++++++++++++------------ tests/wms_test_utils.py | 1 + 6 files changed, 38 insertions(+), 26 deletions(-) diff --git a/python/lsst/ctrl/bps/batch_submit.py b/python/lsst/ctrl/bps/batch_submit.py index 4570faee..69f4b64d 100644 --- a/python/lsst/ctrl/bps/batch_submit.py +++ b/python/lsst/ctrl/bps/batch_submit.py @@ -83,6 +83,7 @@ def create_batch_stages( "bps_project": config["project"], "bps_campaign": config["campaign"], "bps_run": config["uniqProcName"], + "bps_runsite": config["computeSite"], "bps_operator": config["operator"], "bps_payload": config["payloadName"], } diff --git a/python/lsst/ctrl/bps/bps_reports.py b/python/lsst/ctrl/bps/bps_reports.py index 142e84d9..6243bf16 100644 --- a/python/lsst/ctrl/bps/bps_reports.py +++ b/python/lsst/ctrl/bps/bps_reports.py @@ -28,6 +28,7 @@ """Classes and functions used in reporting run status.""" __all__ = [ + "DISPLAY_SUMMARY_FIELDS", "BaseRunReport", "DetailedRunReport", "ExitCodesReport", @@ -46,6 +47,20 @@ _LOG = logging.getLogger(__name__) +DISPLAY_SUMMARY_FIELDS = [ + ("X", "S"), + ("STATE", "S"), + ("%S", "S"), + ("ID", "S"), + ("OPERATOR", "S"), + ("PROJECT", "S"), + ("CAMPAIGN", "S"), + ("SITE", "S"), + ("PAYLOAD", "S"), + ("RUN", "S"), +] + + class BaseRunReport(abc.ABC): """The base class representing a run report. @@ -175,6 +190,7 @@ def add(self, run_report, use_global_id=False): run_report.operator, run_report.project, run_report.campaign, + run_report.site, run_report.payload, run_report.run, ) diff --git a/python/lsst/ctrl/bps/report.py b/python/lsst/ctrl/bps/report.py index b9275a0c..93c40b1c 100644 --- a/python/lsst/ctrl/bps/report.py +++ b/python/lsst/ctrl/bps/report.py @@ -40,7 +40,7 @@ from lsst.utils import doImportType -from .bps_reports import DetailedRunReport, ExitCodesReport, SummaryRunReport +from .bps_reports import DISPLAY_SUMMARY_FIELDS, DetailedRunReport, ExitCodesReport, SummaryRunReport from .wms_service import BaseWmsService, WmsRunReport, WmsStates _LOG = logging.getLogger(__name__) @@ -83,19 +83,7 @@ def display_report( file : TextIO File or file-like object to write the output to. """ - run_brief = SummaryRunReport( - [ - ("X", "S"), - ("STATE", "S"), - ("%S", "S"), - ("ID", "S"), - ("OPERATOR", "S"), - ("PROJECT", "S"), - ("CAMPAIGN", "S"), - ("PAYLOAD", "S"), - ("RUN", "S"), - ] - ) + run_brief = SummaryRunReport(DISPLAY_SUMMARY_FIELDS) if is_detailed: fields = [(" ", "S")] + [(state.name, "i") for state in WmsStates] + [("EXPECTED", "i")] diff --git a/python/lsst/ctrl/bps/wms_service.py b/python/lsst/ctrl/bps/wms_service.py index 44f59fd0..7cc9a9b8 100644 --- a/python/lsst/ctrl/bps/wms_service.py +++ b/python/lsst/ctrl/bps/wms_service.py @@ -249,6 +249,9 @@ class WmsRunReport: operator: str | None = None """Username of the operator who submitted the run.""" + site: str | None = None + """Compute site for payload jobs.""" + run_summary: str | None = None """Job counts per label.""" diff --git a/tests/test_bps_reports.py b/tests/test_bps_reports.py index 07b53731..7ffa125b 100644 --- a/tests/test_bps_reports.py +++ b/tests/test_bps_reports.py @@ -35,6 +35,7 @@ from wms_test_utils import TEST_REPORT from lsst.ctrl.bps import ( + DISPLAY_SUMMARY_FIELDS, BaseRunReport, DetailedRunReport, ExitCodesReport, @@ -113,23 +114,14 @@ class SummaryRunReportTestCase(unittest.TestCase): """Test a summary run report.""" def setUp(self): - self.fields = [ - ("X", "S"), - ("STATE", "S"), - ("%S", "S"), - ("ID", "S"), - ("OPERATOR", "S"), - ("PROJECT", "S"), - ("CAMPAIGN", "S"), - ("PAYLOAD", "S"), - ("RUN", "S"), - ] + self.fields = DISPLAY_SUMMARY_FIELDS self.run = WmsRunReport( wms_id="1.0", global_wms_id="foo#1.0", path="/path/to/run", label="label", run="run", + site="SITE1", project="dev", campaign="testing", payload="test", @@ -146,7 +138,9 @@ def setUp(self): self.report = SummaryRunReport(self.fields) self.expected = Table(dtype=self.fields) - self.expected.add_row(["", "RUNNING", "50", "1.0", "tester", "dev", "testing", "test", "run"]) + self.expected.add_row( + ["", "RUNNING", "50", "1.0", "tester", "dev", "testing", "SITE1", "test", "run"] + ) self.expected_output = io.StringIO() self.actual_output = io.StringIO() @@ -157,6 +151,8 @@ def tearDown(self): def testAddWithNoFlag(self): """Test adding a report for a run with no issues.""" + self.maxDiff = None + print("\n".join(self.expected.pformat(max_lines=-1, max_width=-1)), file=self.expected_output) self.report.add(self.run) @@ -166,6 +162,8 @@ def testAddWithNoFlag(self): def testAddWithFailedFlag(self): """Test adding a run with a failed job.""" + self.maxDiff = None + self.expected["X"][0] = "F" print("\n".join(self.expected.pformat(max_lines=-1, max_width=-1)), file=self.expected_output) @@ -180,6 +178,8 @@ def testAddWithFailedFlag(self): def testAddWithHeldFlag(self): """Test adding a run with a held job.""" + self.maxDiff = None + self.expected["X"][0] = "H" print("\n".join(self.expected.pformat(max_lines=-1, max_width=-1)), file=self.expected_output) @@ -194,6 +194,8 @@ def testAddWithHeldFlag(self): def testAddWithDeletedFlag(self): """Test adding a run with a deleted job.""" + self.maxDiff = None + self.expected["X"][0] = "D" print("\n".join(self.expected.pformat(max_lines=-1, max_width=-1)), file=self.expected_output) @@ -229,6 +231,7 @@ def setUp(self): path="/path/to/run", label="label", run="run", + site="TEST", project="dev", campaign="testing", payload="test", diff --git a/tests/wms_test_utils.py b/tests/wms_test_utils.py index 7fd036f0..1ec54713 100644 --- a/tests/wms_test_utils.py +++ b/tests/wms_test_utils.py @@ -38,6 +38,7 @@ path="/path/to/run", label="label", run="run", + site="SITE1", project="dev", campaign="testing", payload="test", From 16b698bcd24d8271f883a9b054890310c318f0e5 Mon Sep 17 00:00:00 2001 From: Michelle Gower Date: Mon, 27 Jul 2026 17:33:47 -0700 Subject: [PATCH 4/4] Add DEBUG timing statements to init_submission. --- python/lsst/ctrl/bps/initialize.py | 114 +++++++++++++++++++++-------- 1 file changed, 83 insertions(+), 31 deletions(-) diff --git a/python/lsst/ctrl/bps/initialize.py b/python/lsst/ctrl/bps/initialize.py index 3867f357..f25cfe45 100644 --- a/python/lsst/ctrl/bps/initialize.py +++ b/python/lsst/ctrl/bps/initialize.py @@ -43,10 +43,17 @@ from collections.abc import Callable, Iterable from pathlib import Path -from lsst.ctrl.bps import BPS_DEFAULTS, BPS_SEARCH_ORDER, BpsConfig +from lsst.ctrl.bps import ( + BPS_DEFAULTS, + BPS_SEARCH_ORDER, + DEFAULT_MEM_FMT, + DEFAULT_MEM_UNIT, + BpsConfig, +) from lsst.ctrl.bps.bps_utils import _dump_env_info, _dump_pkg_info, mkdir from lsst.pipe.base import Instrument from lsst.utils import doImport +from lsst.utils.timer import time_this _LOG = logging.getLogger(__name__) @@ -113,11 +120,29 @@ def init_submission( wms_service_class_fqn=kwargs.get("wms_service"), ) - translate_command_line_values(config, **kwargs) - - # Run validation tests on the given config if any. - for validator in validators: - validator(config) + with time_this( + log=_LOG, + level=logging.DEBUG, + prefix=None, + msg="Translating command line values completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + translate_command_line_values(config, **kwargs) + + with time_this( + log=_LOG, + level=logging.DEBUG, + prefix=None, + msg="Validation tests completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + # Run validation tests on the given config if any. + for validator in validators: + validator(config) # Set some initial values config[".bps_defined.timestamp"] = Instrument.makeCollectionTimestamp() @@ -128,32 +153,50 @@ def init_submission( if "uniqProcName" not in config: config[".bps_defined.uniqProcName"] = config["outputRun"].replace("/", "_") - # If requested, run WMS plugin checks early in submission process to - # ensure WMS has what it will need for prepare() or submit(). - if kwargs.get("runWmsSubmissionChecks", False): - found, wms_class = config.search("wmsServiceClass") - if not found: - raise KeyError("Missing wmsServiceClass in bps config. Aborting.") - - # Check that can import wms service class. - wms_service_class = doImport(wms_class) - wms_service = wms_service_class(config) - - try: - wms_service.run_submission_checks() - except NotImplementedError: - # Allow various plugins to implement only when needed to do extra - # checks. - _LOG.debug("run_submission_checks is not implemented in %s.", wms_class) - else: - _LOG.debug("Skipping submission checks.") + with time_this( + log=_LOG, + level=logging.DEBUG, + prefix=None, + msg="Submission tests completed", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + # If requested, run WMS plugin checks early in submission process to + # ensure WMS has what it will need for prepare() or submit(). + if kwargs.get("runWmsSubmissionChecks", False): + found, wms_class = config.search("wmsServiceClass") + if not found: + raise KeyError("Missing wmsServiceClass in bps config. Aborting.") + + # Check that can import wms service class. + wms_service_class = doImport(wms_class) + wms_service = wms_service_class(config) + + try: + wms_service.run_submission_checks() + except NotImplementedError: + # Allow various plugins to implement only when needed to do + # extra checks. + _LOG.debug("run_submission_checks is not implemented in %s.", wms_class) + else: + _LOG.debug("Skipping submission checks.") # Replace all bpsGenerateConfig config.generate_config() - # Make submit directory to contain all outputs. - submit_path = mkdir(config["submitPath"]) - config[".bps_defined.submitPath"] = str(submit_path) + with time_this( + log=_LOG, + level=logging.DEBUG, + prefix=None, + msg="Submit path created", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + # Make submit directory to contain all outputs. + submit_path = mkdir(config["submitPath"]) + config[".bps_defined.submitPath"] = str(submit_path) # Pre-make quantum graph filename prefix = Path(submit_path) @@ -168,9 +211,18 @@ def init_submission( with open(expanded_config_file, "w") as fh: config.dump(fh) - # Dump information about runtime environment and software versions in use. - _dump_env_info(f"{submit_path}/{config['uniqProcName']}.env.info.yaml") - _dump_pkg_info(f"{submit_path}/{config['uniqProcName']}.pkg.info.yaml") + with time_this( + log=_LOG, + level=logging.DEBUG, + prefix=None, + msg="Saved environment and package information", + mem_usage=True, + mem_unit=DEFAULT_MEM_UNIT, + mem_fmt=DEFAULT_MEM_FMT, + ): + # Dump information about runtime environment and software versions. + _dump_env_info(f"{submit_path}/{config['uniqProcName']}.env.info.yaml") + _dump_pkg_info(f"{submit_path}/{config['uniqProcName']}.pkg.info.yaml") return config