Clush/Nodeset: add --pick N option

These options will pick N nodes at random from the nodeset before
doing the usual processing.

Closes #311.

Change-Id: I3a1fc87b7adc8bc4fe0d1a4041b02fd36529b1b5

Author: Dominique Martinet

lib/ClusterShell/CLI/Clush.py

@@ -53,6 +53,7 @@
 import signal
 import time
 import threading
+import random
 
 from ClusterShell.Defaults import DEFAULTS, _load_workerclass
 from ClusterShell.CLI.Config import ClushConfig, ClushConfigError
@@ -891,6 +892,15 @@ def main():
     if len(nodeset_base) < 1:
         parser.error('No node to run on.')
 
+    if options.pick and options.pick < len(nodeset_base):
+        # convert to string for sample as nsiter() is slower for big
+        # nodesets; and we assume options.pick will remain small-ish
+        keep = random.sample(nodeset_base, options.pick)
+        nodeset_base.intersection_update(','.join(keep))
+        if config.verbosity >= VERB_VERB:
+            msg = "Picked random nodes: %s" % nodeset_base
+            print Display.COLOR_RESULT_FMT % msg
+
     # Set open files limit.
     set_fdlimit(config.fd_max, display)
 

lib/ClusterShell/CLI/Nodeset.py

@@ -42,6 +42,7 @@
 import logging
 import math
 import sys
+import random
 
 from ClusterShell.CLI.Error import GENERIC_ERRORS, handle_generic_error
 from ClusterShell.CLI.OptionParser import OptionParser
@@ -290,6 +291,14 @@ def nodeset():
             # negative axis index (only single number supported)
             xset.fold_axis = [int(options.axis)]
 
+    if options.pick and options.pick < len(xset):
+        # convert to string for sample as nsiter() is slower for big
+        # nodesets; and we assume options.pick will remain small-ish
+        keep = random.sample(xset, options.pick)
+        # explicit class_set creation and str() convertion for RangeSet
+        keep = class_set(','.join([str(x) for x in keep]))
+        xset.intersection_update(keep)
+
     fmt = options.output_format # default to '%s'
 
     # Display result according to command choice

lib/ClusterShell/CLI/OptionParser.py

@@ -128,6 +128,9 @@ def install_nodes_options(self):
                           default=None, metavar='FILE',
                           help="topology configuration file to use for tree "
                                "mode")
+        optgrp.add_option("--pick", action="store", dest="pick",
+                          metavar="N", type="int",
+                          help="pick N node(s) at random in nodeset")
         self.add_option_group(optgrp)
 
     def install_display_options(self,
@@ -335,4 +338,7 @@ def install_nodeset_options(self):
         optgrp.add_option("--axis", action="store", dest="axis",
                           metavar="RANGESET", help="fold along these axis only "
                                                    "(axis 1..n for nD nodeset)")
+        optgrp.add_option("--pick", action="store", dest="pick",
+                          metavar="N", type="int",
+                          help="pick N node(s) at random in nodeset")
         self.add_option_group(optgrp)

tests/CLIClushTest.py

@@ -480,6 +480,15 @@ def test_033_worker_pdsh_tty(self):
         finally:
             delattr(ClusterShell.CLI.Clush, '_f_user_interaction')
 
+    def test_034_pick(self):
+        """test clush --pick"""
+        self._clush_t(["-w", "%s,localhost" % HOSTNAME, "--pick", "1",
+                      "echo foo"], None,
+                      re.compile(r"^(localhost|%s): foo\n$" % HOSTNAME))
+        self._clush_t(["-w", "%s,localhost" % HOSTNAME, "--pick", "2",
+                      "echo foo"], None,
+                      re.compile(r"^((localhost|%s): foo\n){2}$" % HOSTNAME))
+
 
 class CLIClushTest_B_StdinFailure(unittest.TestCase):
     """Unit test class for testing CLI/Clush.py and stdin failure"""

tests/CLINodesetTest.py

@@ -536,6 +536,15 @@ def test_024_axis_stdin(self):
         self._nodeset_t(["--axis=2","-f"], '\n'.join(ndnodes) + '\n',
                          ','.join(exp_result) + '\n')
 
+    def test_025_pick(self):
+        """test nodeset --pick"""
+        for num in range(1, 100):
+            self._nodeset_t(["--count", "--pick", str(num), "foo[1-100]"],
+                            None, "%s\n" % num)
+            self._nodeset_t(["--count", "--pick", str(num), "-R", "1-100"],
+                            None, "%s\n" % num)
+
+
 
 class CLINodesetGroupResolverTest1(CLINodesetTestBase):
     """Unit test class for testing CLI/Nodeset.py with custom Group Resolver"""