Allow for flux exec arguments to limit the number of celery workers. (#…

…366)

* Added the flux_exec batch argument to allow for flux exec arguments,
e.g. flux_exec: flux exec -r "0-1" to run celery workers only on
ranks 0 and 1 of a multi-rank allocation.

* Remove period.

CHANGELOG.md

@@ -14,6 +14,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added redis.conf for default redis configuration for merlin in server/redis.conf
 - Added default configurations for merlin server command in merlin/server/*.yaml
 - Added documentation page docs/merlin_server.rst, docs/modules/server/configuration.rst, and docs/modules/server/commands.rst
+- Added the flux_exec batch argument to allow for flux exec arguments,
+  e.g. flux_exec: flux exec -r "0-1" to run celery workers only on
+  ranks 0 and 1 of a multi-rank allocation
 ### Changed
 - Rename lgtm.yml to .lgtm.yml
 ### Fixed

docs/source/merlin_specification.rst

@@ -48,6 +48,9 @@ see :doc:`./merlin_variables`.
      queue: pbatch
      flux_path: <optional path to flux bin>
      flux_start_opts: <optional flux start options>
+     flux_exec: <optional, flux exec command to launch workers on 
+                         all nodes if using flux and flux_exec_workers is True
+                         (flux exec) > 
      flux_exec_workers: <optional, flux argument to launch workers on 
                          all nodes. (True)> 
      launch_pre: <Any configuration needed before the srun or jsrun launch>

merlin/examples/workflows/flux/flux_par.yaml

@@ -6,6 +6,7 @@ batch:
   type: flux
   nodes: 1
   queue: pbatch
+  flux_exec: flux exec -r "0-1"
   flux_start_opts: -o,-S,log-filename=flux_par.out
 
 env:

merlin/spec/all_keys.py

@@ -38,6 +38,7 @@
     "shell",
     "flux_path",
     "flux_start_opts",
+    "flux_exec",
     "flux_exec_workers",
     "launch_pre",
     "launch_args",

merlin/study/batch.py

@@ -175,7 +175,7 @@ def batch_worker_launch(
 
         flux_exec: str = ""
         if flux_exec_workers:
-            flux_exec = "flux exec"
+            flux_exec = get_yaml_var(batch, "flux_exec", "flux exec")
 
         if "/" in flux_path:
             flux_path += "/"