Add ability to specify computeSite on the CLI

lsst · Nov 18, 2022 · 681cabf · 681cabf
1 parent f7beb8b
commit 681cabf
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 4 deletions.
diff --git a/doc/changes/DM-37044.feature.rst b/doc/changes/DM-37044.feature.rst
@@ -0,0 +1 @@
+Add ability to specify `computeSite` via the command line.
diff --git a/python/lsst/ctrl/bps/cli/cmd/commands.py b/python/lsst/ctrl/bps/cli/cmd/commands.py
@@ -70,6 +70,7 @@ def transform(*args, **kwargs):
 @click.command(cls=BpsCommand)
 @opt.config_file_argument(required=True)
 @opt.wms_service_option()
+@opt.compute_site_option()
 @opt.submission_options()
 def prepare(*args, **kwargs):
     """Prepare a workflow for submission."""
@@ -79,6 +80,7 @@ def prepare(*args, **kwargs):
 @click.command(cls=BpsCommand)
 @opt.config_file_argument(required=True)
 @opt.wms_service_option()
+@opt.compute_site_option()
 @opt.submission_options()
 def submit(*args, **kwargs):
     """Submit a workflow for execution."""
@@ -87,6 +89,7 @@ def submit(*args, **kwargs):
 
 @click.command(cls=BpsCommand)
 @opt.wms_service_option()
+@opt.compute_site_option()
 @click.option("--id", "run_id", help="Run id of workflow to restart.")
 def restart(*args, **kwargs):
     """Restart a failed workflow."""
@@ -95,6 +98,7 @@ def restart(*args, **kwargs):
 
 @click.command(cls=BpsCommand)
 @opt.wms_service_option()
+@opt.compute_site_option()
 @click.option("--id", "run_id", help="Restrict report to specific WMS run id.")
 @click.option("--user", help="Restrict report to specific user.")
 @click.option("--hist", "hist_days", default=0.0, help="Search WMS history X days for completed info.")
@@ -112,6 +116,7 @@ def report(*args, **kwargs):
 
 @click.command(cls=BpsCommand)
 @opt.wms_service_option()
+@opt.compute_site_option()
 @click.option("--id", "run_id", help="Run id of workflow to cancel.")
 @click.option("--user", help="User for which to cancel all submitted workflows.")
 @click.option(
@@ -135,6 +140,7 @@ def cancel(*args, **kwargs):
 
 @click.command(cls=BpsCommand)
 @opt.wms_service_option()
+@opt.compute_site_option()
 @click.option("--pass-thru", "pass_thru", default=str(), help="Pass the given string to the WMS service.")
 def ping(*args, **kwargs):
     """Ping workflow services."""

diff --git a/python/lsst/ctrl/bps/cli/opt/options.py b/python/lsst/ctrl/bps/cli/opt/options.py
@@ -26,6 +26,7 @@
     "extra_init_option",
     "extra_run_quantum_option",
     "wms_service_option",
+    "compute_site_option",
 ]
 
 from lsst.daf.butler.cli.utils import MWOptionDecorator
@@ -48,3 +49,8 @@
     "environment variable BPS_WMS_SERVICE_CLASS, default "
     "('lsst.ctrl.bps.wms.htcondor.HTCondorService')",
 )
+compute_site_option = MWOptionDecorator(
+    "--compute-site",
+    "compute_site",
+    help="Specification of the compute site used to run the workflow.",
+)
diff --git a/python/lsst/ctrl/bps/drivers.py b/python/lsst/ctrl/bps/drivers.py
@@ -90,6 +90,7 @@ def _init_submission_driver(config_file, **kwargs):
         "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.

diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py
@@ -35,7 +35,7 @@ def testPingNoArgs(self):
             mock_driver.return_value = 0
             result = self.runner.invoke(bps.cli, ["ping"])
             self.assertEqual(result.exit_code, 0)
-            mock_driver.assert_called_with(wms_service=None, pass_thru="")
+            mock_driver.assert_called_with(wms_service=None, compute_site=None, pass_thru="")
 
     def testPingClass(self):
         with unittest.mock.patch("lsst.ctrl.bps.cli.cmd.commands.ping_driver") as mock_driver:
@@ -44,7 +44,9 @@ def testPingClass(self):
                 bps.cli, ["ping", "--wms-service-class", "wms_test_utils.WmsServiceSuccess"]
             )
             self.assertEqual(result.exit_code, 0)
-            mock_driver.assert_called_with(wms_service="wms_test_utils.WmsServiceSuccess", pass_thru="")
+            mock_driver.assert_called_with(
+                wms_service="wms_test_utils.WmsServiceSuccess", compute_site=None, pass_thru=""
+            )
 
     def testPingFailure(self):
         with unittest.mock.patch("lsst.ctrl.bps.cli.cmd.commands.ping_driver") as mock_driver:
@@ -53,7 +55,11 @@ def testPingFailure(self):
                 bps.cli, ["ping", "--wms-service-class", "wms_test_utils.WmsServiceFailure"]
             )
             self.assertEqual(result.exit_code, 64)
-            mock_driver.assert_called_with(wms_service="wms_test_utils.WmsServiceFailure", pass_thru="")
+            mock_driver.assert_called_with(
+                wms_service="wms_test_utils.WmsServiceFailure",
+                compute_site=None,
+                pass_thru="",
+            )
 
     def testPingPassthru(self):
         with unittest.mock.patch("lsst.ctrl.bps.cli.cmd.commands.ping_driver") as mock_driver:
@@ -64,13 +70,17 @@ def testPingPassthru(self):
                     "ping",
                     "--wms-service-class",
                     "wms_test_utils.WmsServicePassThru",
+                    "--compute-site",
+                    "MY_COMPUTE_SITE",
                     "--pass-thru",
                     "EXTRA_VALUES",
                 ],
             )
             self.assertEqual(result.exit_code, 0)
             mock_driver.assert_called_with(
-                wms_service="wms_test_utils.WmsServicePassThru", pass_thru="EXTRA_VALUES"
+                wms_service="wms_test_utils.WmsServicePassThru",
+                compute_site="MY_COMPUTE_SITE",
+                pass_thru="EXTRA_VALUES",
             )