From 98469c3085e0041e6427743c5fed6f74fb5f4713 Mon Sep 17 00:00:00 2001
From: Joe Koning <koning@users.noreply.github.com>
Date: Tue, 7 Jun 2022 10:21:53 -0700
Subject: [PATCH] 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                                 | 3 +++
 docs/source/merlin_specification.rst         | 3 +++
 merlin/examples/workflows/flux/flux_par.yaml | 1 +
 merlin/spec/all_keys.py                      | 1 +
 merlin/study/batch.py                        | 2 +-
 5 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 587778da5..7dfe726a7 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/docs/source/merlin_specification.rst b/docs/source/merlin_specification.rst
index a84eae1a6..d01f6f83d 100644
--- a/docs/source/merlin_specification.rst
+++ b/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>
diff --git a/merlin/examples/workflows/flux/flux_par.yaml b/merlin/examples/workflows/flux/flux_par.yaml
index d401a0270..1fee4131e 100644
--- a/merlin/examples/workflows/flux/flux_par.yaml
+++ b/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:
diff --git a/merlin/spec/all_keys.py b/merlin/spec/all_keys.py
index 9af27d456..bc9d771ec 100644
--- a/merlin/spec/all_keys.py
+++ b/merlin/spec/all_keys.py
@@ -38,6 +38,7 @@
     "shell",
     "flux_path",
     "flux_start_opts",
+    "flux_exec",
     "flux_exec_workers",
     "launch_pre",
     "launch_args",
diff --git a/merlin/study/batch.py b/merlin/study/batch.py
index 69c68856a..88ab230a1 100644
--- a/merlin/study/batch.py
+++ b/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 += "/"