python: add shared class for flux-alloc and batch

Problem: There's duplicated code in the Python cli AllocCmd and
BatchCmd classes because they do not share a common parent class.

Add a shared BatchAllocCmd class for `flux batch` and `flux alloc`.
Move the addition of common arguments to the class initializer,
and add an init_common() method for both classes to perform common
initialization after argument parsing for both cli utilities. Also
add a currently empty update_jobspec_common() method for common
updates to jobspec after it is created.

Derive AllocCmd and BatchCmd from  BatchAllocCmd. Remove duplicated
code that has now moved into BatchAllocCmd.init_common(). Call
BatchAllocCmd.update_jobspec_common() from both classes after jobspec
creation.

Author: grondo

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

@@ -1657,71 +1657,88 @@ 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(
+            "--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)
+
+    def update_jobspec_common(self, args, jobspec):
+        """Common jobspec update code for batch/alloc"""
+        pass

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: