From c2f57864ba797728789497fa6934e03051c2b9b2 Mon Sep 17 00:00:00 2001
From: "Mark A. Grondona" <mark.grondona@gmail.com>
Date: Tue, 29 Mar 2022 07:08:07 -0700
Subject: [PATCH] flux-mini: set ntasks/slots to nnodes when ntasks not set

Problem: It is inconvenient to require the specification of both ntasks
and nnodes when a user wants one task/slot per node. Until recently,
it was not possible to handle this in a coherent manner, though, so
the Python Jobspec class and flux-mini commands throw an error whenever
the node count is greater than the number of requested tasks/slots.

Now that node exclusivity can be set in the jobspec, though, it is
possible to set ntasks/slots to the number of nodes (when ntasks is
not explicitly set), by also defaulting the node exclusive flag to
True for this case.

This allows `flux mini run -N4 command` to work consistently regardless
of whether or not the enclosing instance defaults to node exclusive
allocation.

Fixes #4228
---
 src/cmd/flux-mini.py | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/cmd/flux-mini.py b/src/cmd/flux-mini.py
index 4080654722d9..fd13f5474c69 100755
--- a/src/cmd/flux-mini.py
+++ b/src/cmd/flux-mini.py
@@ -709,7 +709,6 @@ def __init__(self):
             "-n",
             "--ntasks",
             metavar="N",
-            default="1",
             help="Number of tasks to start",
         )
         self.parser.add_argument(
@@ -742,6 +741,15 @@ def init_jobspec(self, args):
         if not args.command:
             raise ValueError("job command and arguments are missing")
 
+        #  If ntasks not set, then set it to either node count, with
+        #   exclusive flag enabled, or to 1 (the default).
+        if not args.ntasks:
+            if args.nodes:
+                args.ntasks = args.nodes
+                args.exclusive = True
+            else:
+                args.ntasks = 1
+
         #  Ensure integer args are converted to int() here.
         #  This is done because we do not use type=int in argparse in order
         #   to allow these options to be mutable for bulksubmit:
@@ -1510,8 +1518,14 @@ def read_script(args):
 
     def init_jobspec(self, args):
         # If no script (reading from stdin), then use "flux" as arg[0]
+
+        #  If number of slots not specified, then set it to node count
+        #   if set, otherwise raise an error.
         if not args.nslots:
-            raise ValueError("Number of slots to allocate must be specified")
+            if not args.nodes:
+                raise ValueError("Number of slots to allocate must be specified")
+            args.nslots = args.nodes
+            args.exclusive = True
 
         jobspec = JobspecV1.from_batch_command(
             script=self.read_script(args),
@@ -1556,8 +1570,13 @@ def __init__(self):
 
     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:
-            raise ValueError("Number of slots to allocate must be specified")
+            if not args.nodes:
+                raise ValueError("Number of slots to allocate must be specified")
+            args.nslots = args.nodes
+            args.exclusive = True
 
         jobspec = JobspecV1.from_nest_command(
             command=args.COMMAND,