Merge pull request #6822 from grondo/add-brokers

add method to request extra brokers on node 0

Author: mergify[bot]

src/bindings/python/flux/cli/alloc.py

@@ -19,11 +19,9 @@
 from flux.uri import JobURI
 
 
-class AllocCmd(base.MiniCmd):
+class AllocCmd(base.BatchAllocCmd):
     def __init__(self, prog, usage=None, description=None):
-        self.t0 = None
         super().__init__(prog, usage, description, exclude_io=True)
-        base.add_batch_alloc_args(self.parser)
         self.parser.add_argument(
             "-v",
             "--verbose",
@@ -44,24 +42,14 @@ def __init__(self, prog, usage=None, description=None):
         )
 
     def init_jobspec(self, args):
-        #  If number of slots not specified, then set it to node count
-        #   if set, otherwise raise an error.
-        if not args.nslots:
-            if not args.nodes:
-                raise ValueError("Number of slots to allocate must be specified")
-            args.nslots = args.nodes
-            args.exclusive = True
+        self.init_common(args)
 
         #  For --bg, do not run an rc2 (initial program) unless
         #    the user explicitly specified COMMAND:
         if args.bg and not args.COMMAND:
             args.broker_opts = args.broker_opts or []
             args.broker_opts.append("-Sbroker.rc2_none=1")
 
-        if args.dump:
-            args.broker_opts = args.broker_opts or []
-            args.broker_opts.append("-Scontent.dump=" + args.dump)
-
         jobspec = flux.job.JobspecV1.from_nest_command(
             command=args.COMMAND,
             num_slots=args.nslots,
@@ -73,6 +61,8 @@ def init_jobspec(self, args):
             conf=args.conf.config,
         )
 
+        self.update_jobspec_common(args, jobspec)
+
         #  For --bg, always allocate a pty, but not interactive,
         #   since an interactive pty causes the job shell to hang around
         #   until a pty client attaches, which may never happen.

src/bindings/python/flux/cli/base.py

@@ -40,7 +40,7 @@
 from flux.idset import IDset
 from flux.job import JobspecV1, JobWatcher
 from flux.progress import ProgressBar
-from flux.util import dict_merge, set_treedict
+from flux.util import dict_merge, get_treedict, set_treedict
 
 LOGGER = logging.getLogger("flux")
 
@@ -496,6 +496,12 @@ def update(self, value):
             return self.update_file(value, extension)
         return self.update_named_config(value)
 
+    def get(self, key, default=None):
+        """
+        Get a value from the current config using dotted key form
+        """
+        return get_treedict(self.config or {}, key, default=default)
+
 
 class ConfAction(argparse.Action):
     """Handle batch/alloc --conf option"""
@@ -547,7 +553,6 @@ class Xcmd:
         "flags": "--flags=",
         "begin_time": "--begin-time=",
         "signal": "--signal=",
-        "taskmap": "--taskmap=",
     }
 
     class Xinput:
@@ -1652,71 +1657,120 @@ def main(self, args):
         self.run_and_exit()
 
 
-def add_batch_alloc_args(parser):
-    """
-    Add "batch"-specific resource allocation arguments to parser object
-    which deal in slots instead of tasks.
-    """
-    parser.add_argument(
-        "--conf",
-        metavar="CONF",
-        default=BatchConfig(),
-        action=ConfAction,
-        help="Set configuration for a child Flux instance. CONF may be a "
-        + "multiline string in JSON or TOML, a configuration key=value, a "
-        + "path to a JSON or TOML file, or a configuration loaded by name "
-        + "from a standard search path. This option may specified multiple "
-        + "times, in which case the config is iteratively updated.",
-    )
-    parser.add_argument(
-        "--broker-opts",
-        metavar="OPTS",
-        default=None,
-        action="append",
-        help="Pass options to flux brokers",
-    )
-    parser.add_argument(
-        "--dump",
-        nargs="?",
-        const="flux-{{jobid}}-dump.tgz",
-        metavar="FILE",
-        help="Archive KVS on exit",
-    )
-    parser.add_argument(
-        "-n",
-        "--nslots",
-        type=int,
-        metavar="N",
-        help="Number of total resource slots requested."
-        + " The size of a resource slot may be specified via the"
-        + " -c, --cores-per-slot and -g, --gpus-per-slot options."
-        + " The default slot size is 1 core.",
-    )
-    parser.add_argument(
-        "-c",
-        "--cores-per-slot",
-        type=int,
-        metavar="N",
-        default=1,
-        help="Number of cores to allocate per slot",
-    )
-    parser.add_argument(
-        "-g",
-        "--gpus-per-slot",
-        type=int,
-        metavar="N",
-        help="Number of GPUs to allocate per slot",
-    )
-    parser.add_argument(
-        "-N",
-        "--nodes",
-        type=int,
-        metavar="N",
-        help="Distribute allocated resource slots across N individual nodes",
-    )
-    parser.add_argument(
-        "-x",
-        "--exclusive",
-        action="store_true",
-        help="With -N, --nodes, allocate nodes exclusively",
-    )
+class BatchAllocCmd(MiniCmd):
+    def __init__(self, prog, usage=None, description=None, exclude_io=True):
+        self.t0 = None
+        super().__init__(prog, usage, description, exclude_io)
+        self.parser.add_argument(
+            "--add-brokers", default=0, type=int, help=argparse.SUPPRESS
+        )
+        self.parser.add_argument(
+            "--conf",
+            metavar="CONF",
+            default=BatchConfig(),
+            action=ConfAction,
+            help="Set configuration for a child Flux instance. CONF may be a "
+            + "multiline string in JSON or TOML, a configuration key=value, a "
+            + "path to a JSON or TOML file, or a configuration loaded by name "
+            + "from a standard search path. This option may specified multiple "
+            + "times, in which case the config is iteratively updated.",
+        )
+        self.parser.add_argument(
+            "--broker-opts",
+            metavar="OPTS",
+            default=None,
+            action="append",
+            help="Pass options to flux brokers",
+        )
+        self.parser.add_argument(
+            "--dump",
+            nargs="?",
+            const="flux-{{jobid}}-dump.tgz",
+            metavar="FILE",
+            help="Archive KVS on exit",
+        )
+        self.parser.add_argument(
+            "-n",
+            "--nslots",
+            type=int,
+            metavar="N",
+            help="Number of total resource slots requested."
+            + " The size of a resource slot may be specified via the"
+            + " -c, --cores-per-slot and -g, --gpus-per-slot options."
+            + " The default slot size is 1 core.",
+        )
+        self.parser.add_argument(
+            "-c",
+            "--cores-per-slot",
+            type=int,
+            metavar="N",
+            default=1,
+            help="Number of cores to allocate per slot",
+        )
+        self.parser.add_argument(
+            "-g",
+            "--gpus-per-slot",
+            type=int,
+            metavar="N",
+            help="Number of GPUs to allocate per slot",
+        )
+        self.parser.add_argument(
+            "-N",
+            "--nodes",
+            type=int,
+            metavar="N",
+            help="Distribute allocated resource slots across N individual nodes",
+        )
+        self.parser.add_argument(
+            "-x",
+            "--exclusive",
+            action="store_true",
+            help="With -N, --nodes, allocate nodes exclusively",
+        )
+
+    def init_common(self, args):
+        """Common initialization code for batch/alloc"""
+        #  If number of slots not specified, then set it to node count
+        #   if set, otherwise raise an error.
+        if not args.nslots:
+            if not args.nodes:
+                raise ValueError("Number of slots to allocate must be specified")
+            args.nslots = args.nodes
+            args.exclusive = True
+
+        if args.dump:
+            args.broker_opts = args.broker_opts or []
+            args.broker_opts.append("-Scontent.dump=" + args.dump)
+
+        if args.add_brokers > 0:
+            if not args.nodes:
+                raise ValueError(
+                    "--add-brokers may only be specified with -N, --nnodes"
+                )
+            nbrokers = args.add_brokers
+            nnodes = args.nodes
+
+            # Force update taskmap with extra ranks on nodeid 0:
+            args.taskmap = f"manual:[[0,1,{1+nbrokers},1],[1,{nnodes-1},1,1]]"
+
+            # Exclude the additional brokers via configuration. However,
+            # don't throw away any ranks already excluded bythe user.
+            # Note: raises an exception if user excluded by hostname (unlikely)
+            exclude = IDset(args.conf.get("resource.exclude", default="")).set(
+                1, nbrokers
+            )
+            args.conf.update(f'resource.exclude="{exclude}"')
+
+    def update_jobspec_common(self, args, jobspec):
+        """Common jobspec update code for batch/alloc"""
+        # If args.add_brokers is being used, update jobspec task count
+        # to accurately reflect the updated task count.
+        if args.add_brokers > 0:
+            # Note: args.nodes required with add_brokers already checked above
+            total_tasks = args.nodes + args.add_brokers
+
+            # Overwrite task count with new total_tasks:
+            jobspec.tasks[0]["count"] = {"total": total_tasks}
+
+            # remove per-resource shell option which is no longer necessary:
+            del jobspec.attributes["system"]["shell"]["options"]["per-resource"]

src/bindings/python/flux/cli/batch.py

@@ -20,15 +20,14 @@
 LOGGER = logging.getLogger("flux-batch")
 
 
-class BatchCmd(base.MiniCmd):
+class BatchCmd(base.BatchAllocCmd):
     def __init__(self, prog, usage=None, description=None):
         super().__init__(prog, usage, description)
         self.parser.add_argument(
             "--wrap",
             action="store_true",
             help="Wrap arguments or stdin in a /bin/sh script",
         )
-        base.add_batch_alloc_args(self.parser)
         self.parser.add_argument(
             "SCRIPT",
             nargs=argparse.REMAINDER,
@@ -87,21 +86,11 @@ def process_script(self, args):
         return self.parse_directive_args(name, batchscript)
 
     def init_jobspec(self, args):
+        self.init_common(args)
+
         if args.wrap:
             self.script = f"#!/bin/sh\n{self.script}"
 
-        #  If number of slots not specified, then set it to node count
-        #   if set, otherwise raise an error.
-        if not args.nslots:
-            if not args.nodes:
-                raise ValueError("Number of slots to allocate must be specified")
-            args.nslots = args.nodes
-            args.exclusive = True
-
-        if args.dump:
-            args.broker_opts = args.broker_opts or []
-            args.broker_opts.append("-Scontent.dump=" + args.dump)
-
         #  If job name is not explicitly set in args, use the script name
         #   if a script was provided, else the string "batch" to
         #   indicate the script was set on flux batch stdin.
@@ -124,6 +113,8 @@ def init_jobspec(self, args):
             conf=args.conf.config,
         )
 
+        self.update_jobspec_common(args, jobspec)
+
         # Default output is flux-{{jobid}}.out
         # overridden by either --output=none or --output=kvs
         if not args.output:

t/t2714-python-cli-batch.t

@@ -272,4 +272,32 @@ test_expect_success 'flux batch: multiline --add-file requires name=' '
 	test_must_fail flux batch --dry-run directives6.sh >d6.out 2>&1 &&
 	grep "file name missing" d6.out
 '
+test_expect_success 'flux batch: hidden --add-brokers option requires -N' '
+	test_must_fail flux batch -n4 --add-brokers=1 --dry-run --wrap true
+'
+test_expect_success 'flux batch: hidden --add-brokers option works' '
+	cat <<-EOF >add-brokers.sh &&
+	#!/bin/sh
+	# flux: -N4 --output=add-brokers.out --error=add-brokers.err
+	# flux: --conf=tbon.topo=kary:2
+	# flux: --add-brokers=2
+	#
+	flux resource list -s free -no {nnodes}
+	flux resource status -s exclude -no {ranks}
+	flux exec -lr 1-2 flux getattr tbon.parent-endpoint \
+	    | sed s#://.*## | sort
+	flux run -N4 true
+	EOF
+	cat <<-EOF >add-brokers.expected &&
+	4
+	1-2
+	1: ipc
+	2: ipc
+	EOF
+	id=$(flux batch add-brokers.sh) &&
+	flux job wait-event $id clean &&
+	test_debug "cat add-brokers.err" &&
+	test_debug "cat add-brokers.out" &&
+	test_cmp add-brokers.expected add-brokers.out
+'
 test_done