MAINT: add threads/jobs primitives (#62)

Author: colinvwood

q2_diversity_lib/_util.py

@@ -104,7 +104,7 @@ def _validate_requested_cpus(wrapped_function, *args, **kwargs):
 
     cpus_requested = b_a_arguments[param_name]
 
-    if cpus_requested == 'auto':
+    if cpus_requested == 0:
         # mutate bound_arguments.arguments 'auto' to the requested # of cpus...
         b_a_arguments[param_name] = cpus_available
         # ...and update cpus requested to prevent TypeError

q2_diversity_lib/plugin_setup.py

@@ -7,7 +7,7 @@
 # ----------------------------------------------------------------------------
 
 from qiime2.plugin import (Plugin, Citations, Bool, Int, Range, Choices, Str,
-                           Float, List)
+                           Float, List, Threads)
 from q2_types.feature_table import (FeatureTable, Frequency, RelativeFrequency,
                                     PresenceAbsence)
 from q2_types.tree import Phylogeny, Rooted
@@ -56,7 +56,7 @@
     inputs={'table': FeatureTable[Frequency | RelativeFrequency
             | PresenceAbsence],
             'phylogeny': Phylogeny[Rooted]},
-    parameters={'threads': Int % Range(1, None) | Str % Choices(['auto'])},
+    parameters={'threads': Threads},
     outputs=[('vector', SampleData[AlphaDiversity])],
     input_descriptions={
         'table': "The feature table containing the samples for which Faith's "
@@ -146,7 +146,7 @@
 plugin.methods.register_function(
     function=beta.bray_curtis,
     inputs={'table': FeatureTable[Frequency | RelativeFrequency]},
-    parameters={'n_jobs': Int % Range(1, None) | Str % Choices(['auto'])},
+    parameters={'n_jobs': Threads},
     outputs=[('distance_matrix', DistanceMatrix)],
     input_descriptions={
         'table': "The feature table containing the samples for which "
@@ -174,7 +174,7 @@
     function=beta.jaccard,
     inputs={'table': FeatureTable[Frequency | RelativeFrequency
             | PresenceAbsence]},
-    parameters={'n_jobs': Int % Range(1, None) | Str % Choices(['auto'])},
+    parameters={'n_jobs': Threads},
     outputs=[('distance_matrix', DistanceMatrix)],
     input_descriptions={
         'table': "The feature table containing the samples for which "
@@ -202,7 +202,7 @@
     inputs={'table': FeatureTable[Frequency | RelativeFrequency
             | PresenceAbsence],
             'phylogeny': Phylogeny[Rooted]},
-    parameters={'threads': Int % Range(1, None) | Str % Choices(['auto']),
+    parameters={'threads': Threads,
                 'bypass_tips': Bool},
     outputs=[('distance_matrix', DistanceMatrix)],
     input_descriptions={
@@ -246,7 +246,7 @@
     function=beta.weighted_unifrac,
     inputs={'table': FeatureTable[Frequency | RelativeFrequency],
             'phylogeny': Phylogeny[Rooted]},
-    parameters={'threads': Int % Range(1, None) | Str % Choices(['auto']),
+    parameters={'threads': Threads,
                 'bypass_tips': Bool},
     outputs=[('distance_matrix', DistanceMatrix)],
     input_descriptions={
@@ -310,8 +310,7 @@
     function=beta.beta_passthrough,
     inputs={'table': FeatureTable[Frequency]},
     parameters={'metric': Str % Choices(beta.METRICS['NONPHYLO']['UNIMPL']),
-                'pseudocount': Int % Range(1, None),
-                'n_jobs': Int % Range(1, None) | Str % Choices(['auto'])},
+                'pseudocount': Int % Range(1, None), 'n_jobs': Threads},
     outputs=[('distance_matrix', DistanceMatrix)],
     input_descriptions={
         'table': 'The feature table containing the samples over which beta '
@@ -347,7 +346,7 @@
     inputs={'table': FeatureTable[Frequency],
             'phylogeny': Phylogeny[Rooted]},
     parameters={'metric': Str % Choices(beta.METRICS['PHYLO']['UNIMPL']),
-                'threads': Int % Range(1, None) | Str % Choices(['auto']),
+                'threads': Threads,
                 'variance_adjusted': Bool,
                 'alpha': Float % Range(0, 1, inclusive_end=True),
                 'bypass_tips': Bool},
@@ -412,13 +411,12 @@
     ]
 )
 
-
 plugin.methods.register_function(
     function=beta.beta_phylogenetic_meta_passthrough,
     inputs={'tables': List[FeatureTable[Frequency]],
             'phylogenies': List[Phylogeny[Rooted]]},
     parameters={'metric': Str % Choices(beta.METRICS['PHYLO']['UNIMPL']),
-                'threads': Int % Range(1, None) | Str % Choices(['auto']),
+                'threads': Threads,
                 'variance_adjusted': Bool,
                 'alpha': Float % Range(0, 1, inclusive_end=True),
                 'bypass_tips': Bool,

q2_diversity_lib/tests/test_util.py

@@ -234,13 +234,13 @@ def test_n_jobs_passed_as_default(self, mock_process):
         self.assertEqual(self.function_w_n_jobs_param(), 3)
 
     @mock.patch("q2_diversity_lib._util.psutil.Process")
-    def test_auto_passed_to_cpu_request(self, mock_process):
+    def test_zero_passed_to_cpu_request(self, mock_process):
         mock_process = psutil.Process()
         mock_process.cpu_affinity = mock.MagicMock(return_value=[0, 1, 2])
-        self.assertEqual(self.function_w_n_jobs_param('auto'), 3)
-        self.assertEqual(self.function_w_n_jobs_param(n_jobs='auto'), 3)
-        self.assertEqual(self.function_w_threads_param('auto'), 3)
-        self.assertEqual(self.function_w_threads_param(threads='auto'), 3)
+        self.assertEqual(self.function_w_n_jobs_param(0), 3)
+        self.assertEqual(self.function_w_n_jobs_param(n_jobs=0), 3)
+        self.assertEqual(self.function_w_threads_param(0), 3)
+        self.assertEqual(self.function_w_threads_param(threads=0), 3)
 
     @mock.patch("q2_diversity_lib._util.psutil.Process")
     def test_cpu_request_through_framework(self, mock_process):