Merge pull request #5219 from grondo/issue#5143

flux-resource: add `-i, --include` option to filter hosts for `status` and `drain` commands

Author: mergify[bot]

doc/man1/flux-resource.rst

@@ -68,7 +68,7 @@ COMMANDS
    states as with ``flux resource list``. By default, the *STATE* reported
    by ``flux resource info`` is "all".
 
-**status**  [-n] [-o FORMAT] [-s STATE,...] [--skip-empty]
+**status**  [-n] [-o FORMAT] [-s STATE,...] [-i TARGETS] [--skip-empty]
    Show system view of resources. This command queries both the resource
    service and scheduler to identify resources that are available,
    excluded by configuration, or administratively drained or draining.
@@ -85,6 +85,12 @@ COMMANDS
    valid states include "avail", "exclude", "draining", "drained", and "all".
    The special "drain" state is shorthand for "drained,draining".
 
+   With *-i, --include=TARGETS*, the results are filtered to only include
+   resources matching **TARGETS**, which may be specified either as an idset
+   of broker ranks or list of hosts in hostlist form. It is not an error to
+   specify ranks or hosts which do not exist, the result will be filtered
+   to include only those ranks or hosts that are present in *TARGETS*.
+
    The *-o,--format=FORMAT* option customizes output formatting (See the
    OUTPUT FORMAT section below for details).
 
@@ -94,13 +100,19 @@ COMMANDS
    unless the ``-s, --states`` option is used. Suppression of empty lines
    can may be forced with the ``--skip-empty`` option.
 
-**drain** [-n] [-o FORMAT] [-f] [-u] [targets] [reason ...]
+**drain** [-n] [-o FORMAT] [-i TARGETS] [-f] [-u] [targets] [reason ...]
    If specified without arguments, list drained nodes. In this mode,
    *-n,--no-header* suppresses header from output and *-o,--format=FORMAT*
    customizes output formatting (see below).  The *targets* argument is an
    IDSET or HOSTLIST specifying nodes to drain.  Any remaining arguments
    are assumed to be a reason to be recorded with the drain event.
 
+   With *-i, --include=TARGETS*, **drain** output is filtered to only include
+   resources matching **TARGETS**, which may be specified either as an idset
+   of broker ranks or list of hosts in hostlist form. It is not an error to
+   specify ranks or hosts which do not exist, the result will be filtered
+   to include only those ranks or hosts that are present in *TARGETS*.
+
    By default, **flux resource drain** will fail if any of the *targets*
    are already drained. To change this behavior, use either of the
    *-f, --force* or *-u, --update* options. With *--force*, the *reason* for

src/bindings/python/flux/hostlist.py

@@ -160,7 +160,7 @@ def __iter__(self):
     def __contains__(self, name):
         """Test if a hostname is in a Hostlist"""
         try:
-            self.pimpl.find(name)
+            self.find(name)
         except FileNotFoundError:
             return False
         return True
@@ -218,6 +218,34 @@ def copy(self):
         """Copy a Hostlist object"""
         return Hostlist(handle=self.pimpl.copy())
 
+    def find(self, host):
+        """Return the position of a host in a Hostlist"""
+        return self.pimpl.find(host)
+
+    def index(self, hosts, ignore_nomatch=False):
+        """
+        Return a list of integers corresponding to the indices of ``hosts``
+        in the current Hostlist.
+        Args:
+            hosts (str, Hostlist): List of hosts to find
+            ignore_nomatch (bool): Ignore hosts in ``hosts`` that are not
+             present in Hostlist. Otherwise, FileNotFound error is raised
+             with the missing hosts.
+        """
+        if not isinstance(hosts, Hostlist):
+            hosts = Hostlist(hosts)
+        ids = []
+        notfound = Hostlist()
+        for host in hosts:
+            try:
+                ids.append(self.find(host))
+            except FileNotFoundError:
+                notfound.append(host)
+        if notfound and not ignore_nomatch:
+            suffix = "s" if len(notfound) > 1 else ""
+            raise FileNotFoundError(f"host{suffix} '{notfound}' not found")
+        return ids
+
 
 def decode(arg):
     """

src/bindings/python/flux/resource/status.py

@@ -54,21 +54,52 @@ def __init__(self, rstatus=None, allocated_ranks=None):
         if allocated_ranks is None:
             allocated_ranks = IDset()
 
+        self.rstatus = rstatus
         self.rset = ResourceSet(rstatus["R"])
+        self.nodelist = self.rset.nodelist
+        self.allocated_ranks = allocated_ranks
+
+        self._recalculate()
+
+    def filter(self, include):
+        """
+        Filter the reported resources in a ResourceStatus object
+        Args:
+            include (str, IDset, Hostlist): restrict the current set of
+             reported ranks to the given ranks or hosts.
+        """
+        try:
+            include_ranks = IDset(include)
+        except ValueError:
+            include_ranks = self.nodelist.index(include, ignore_nomatch=True)
+        self._recalculate(include_ranks)
 
-        # get idset of all ranks and nodelist:
+    def _recalculate(self, include_ranks=None):
+        """
+        Recalculate derived idsets and drain_info, only including ranks
+        in the IDset 'include_ranks' if given.
+        Args:
+            include_ranks (IDset): restrict the current set of reported ranks.
+        """
+        # get idset of all ranks:
         self.all = self.rset.ranks
-        self.nodelist = self.rset.nodelist
 
         # offline/online
-        self.offline = IDset(rstatus["offline"])
-        self.online = IDset(rstatus["online"])
+        self.offline = IDset(self.rstatus["offline"])
+        self.online = IDset(self.rstatus["online"])
 
         # excluded: excluded by configuration
-        self.exclude = IDset(rstatus["exclude"])
+        self.exclude = IDset(self.rstatus["exclude"])
 
         # allocated: online and allocated by scheduler
-        self.allocated = allocated_ranks
+        self.allocated = self.allocated_ranks
+
+        # If include_ranks was provided, filter all idsets to only those
+        # that intersect the provided idset
+        if include_ranks is not None:
+            for name in ("all", "offline", "online", "exclude", "allocated"):
+                result = getattr(self, name).intersect(include_ranks)
+                setattr(self, name, result)
 
         # drained: free+drain
         self.drained = IDset()
@@ -77,8 +108,10 @@ def __init__(self, rstatus=None, allocated_ranks=None):
 
         # drain_info: ranks, timestamp, reason tuples for all drained resources
         self.drain_info = []
-        for drain_ranks, entry in rstatus["drain"].items():
+        for drain_ranks, entry in self.rstatus["drain"].items():
             ranks = IDset(drain_ranks)
+            if include_ranks is not None:
+                ranks = ranks.intersect(include_ranks)
             self.drained += ranks - self.allocated
             self.draining += ranks - self.drained
             info = DrainInfo(ranks, entry["timestamp"], entry["reason"])

src/cmd/flux-resource.py

@@ -366,6 +366,12 @@ def status(args):
     else:
         rstatus = resource_status(flux.Flux()).get()
 
+    if args.include:
+        try:
+            rstatus.filter(include=args.include)
+        except (ValueError, TypeError) as exc:
+            raise ValueError(f"--include: {exc}") from None
+
     formatter = flux.util.OutputFormat(fmt, headings=headings)
 
     # Remove any `{color*}` fields if color is off
@@ -582,6 +588,13 @@ def main():
         help="Update only. Do not return an error if one or more targets "
         + "are already drained. Do not overwrite any existing drain reason.",
     )
+    drain_parser.add_argument(
+        "-i",
+        "--include",
+        metavar="TARGETS",
+        help="Include only specified targets in output set. TARGETS may be "
+        + "provided as an idset or hostlist.",
+    )
     drain_parser.add_argument(
         "-o",
         "--format",
@@ -634,6 +647,13 @@ def main():
         metavar="STATE,...",
         help="Output resources in given states",
     )
+    status_parser.add_argument(
+        "-i",
+        "--include",
+        metavar="TARGETS",
+        help="Include only specified targets in output set. TARGETS may be "
+        + "provided as an idset or hostlist.",
+    )
     status_parser.add_argument(
         "-n", "--no-header", action="store_true", help="Suppress header output"
     )

src/common/libhostlist/hostlist.c

@@ -456,7 +456,7 @@ static int append_range_list_with_suffix (struct hostlist *hl,
     unsigned long j;
 
     /* compute max buffer size for this set of hosts */
-    int size = strlen (pfx) + strlen (sfx) + 20 + rng->width; 
+    int size = strlen (pfx) + strlen (sfx) + 20 + rng->width;
 
 
     for (i = 0; i < n; i++) {
@@ -854,8 +854,8 @@ static void hostlist_coalesce (struct hostlist *hl)
              */
             if (new->hi < hprev->hi)
                 hnext->hi = hprev->hi;
-    
-            /*  
+
+            /*
              *  The duplicated range will inserted piecemeal below,
              *   e.g. [5-7,6-8] -> [5-6,6-7,7-8]
              *  Therefore adjust the end of hprev to new->lo (the

src/common/libhostlist/hostlist.h

@@ -92,7 +92,7 @@ void hostlist_sort (struct hostlist * hl);
  */
 void hostlist_uniq (struct hostlist *hl);
 
-/* 
+/*
  *  Return the host at the head of hostlist 'hl', or NULL if list is empty.
  *  Leaves internal cursor pointing at the head item.
  *
@@ -102,7 +102,7 @@ void hostlist_uniq (struct hostlist *hl);
  */
 const char * hostlist_first (struct hostlist *hl);
 
-/* 
+/*
  *  Return the host at the tail of hostlist 'hl', or NULL if list is empty.
  *  Leaves internal cursor pointing at the last item.
  *
@@ -112,7 +112,7 @@ const char * hostlist_first (struct hostlist *hl);
  */
 const char * hostlist_last (struct hostlist *hl);
 
-/* 
+/*
  *  Advance the internal cursor and return the next host in 'hl' or NULL
  *   if list is empty or the end of the list has been reached.
  *

t/python/t0020-hostlist.py

@@ -160,6 +160,24 @@ def test_copy(self):
         hl.delete("foo[0-3]")
         self.assertEqual(str(cp), "foo[0-3]")
 
+    def test_find(self):
+        hl = hostlist.decode("foo[0-3]")
+        self.assertEqual(hl.find("foo1"), 1)
+        self.assertEqual(hl.find("foo3"), 3)
+        with self.assertRaises(FileNotFoundError):
+            hl.find("foo4")
+
+    def test_index_method(self):
+        hl = hostlist.decode("foo[0-10]")
+        self.assertListEqual(hl.index("foo[1,5,7]"), [1, 5, 7])
+        self.assertListEqual(hl.index(hostlist.decode("foo1")), [1])
+        self.assertListEqual(hl.index("foo[9-11]", ignore_nomatch=True), [9, 10])
+        self.assertListEqual(hl.index("foo11", ignore_nomatch=True), [])
+        with self.assertRaises(FileNotFoundError):
+            hl.index("foo[9-11]")
+        with self.assertRaises(FileNotFoundError):
+            hl.index("foo11")
+
 
 if __name__ == "__main__":
     unittest.main(testRunner=TAPTestRunner())

t/t2311-resource-drain.t

@@ -318,6 +318,12 @@ test_expect_success 'flux resource drain differentiates drain/draining' '
 	test $(flux resource status -s drain -no {nnodes}) -eq ${SIZE}
 '
 
+test_expect_success 'flux resource drain supports --include' '
+	flux resource drain -ni 0 >drain-include.output &&
+	test_debug "cat drain-include.output" &&
+	test $(wc -l <drain-include.output) -eq 1
+'
+
 test_expect_success 'flux resource drain works without scheduler loaded' '
 	flux module unload sched-simple &&
 	flux resource drain &&

t/t2351-resource-status-input.t

@@ -26,6 +26,21 @@ for input in ${INPUTDIR}/*.json; do
     '
 done
 
+#  Ensure all tested inputs can also work with --include
+#  We simply restrict to rank 0 and then ensure {ranks} returns only 0
+for input in ${INPUTDIR}/*.json; do
+    name=$(basename ${input%%.json})
+    test_expect_success "flux-resource status input --include check: $name" '
+        base=${input%%.json} &&
+        name=$(basename $base)-i &&
+        flux resource status -o "{ranks} {nodelist}" --include=0 \
+            --from-stdin < $input > $name.output 2>&1 &&
+        test_debug "cat $name.output" &&
+	grep "^0[^,-]" $name.output
+    '
+done
+
+
 test_expect_success 'flux-resource status: header included with all formats' '
 	cat <<-EOF >headers.expected &&
 	state==STATE
@@ -192,4 +207,24 @@ test_expect_success 'flux-resource status: lines are combined based on format' '
 	EOF
 	test_cmp ts-detailed2.expected ts-detailed2.out
 '
+test_expect_success 'flux-resource status: --include works with ranks' '
+	INPUT=${INPUTDIR}/drain.json &&
+	flux resource status --include=1,3 --from-stdin \
+		-no "{nnodes}" <$INPUT >drain-include.out &&
+	test_debug "cat drain-include.out" &&
+	test "$(cat drain-include.out)" = "2"
+'
+test_expect_success 'flux-resource status: --include works with hostnames' '
+	INPUT=${INPUTDIR}/drain.json &&
+	flux resource status --include=foo[1,3] --from-stdin \
+		-no "{nodelist}" <$INPUT >drain-include-host.out &&
+	test_debug "cat drain-include-host.out" &&
+	test "$(cat drain-include-host.out)" = "foo[1,3]"
+'
+test_expect_success 'flux-resource status: --include works with invalid host' '
+	INPUT=${INPUTDIR}/drain.json &&
+	flux resource status --include=foo7 --from-stdin \
+		-no "{nodelist}" <$INPUT >drain-empty.out 2>&1 &&
+	test_must_be_empty drain-fail.out
+'
 test_done