MAINT: add threads/jobs primitives

q2_phylogeny/_fasttree.py

@@ -38,7 +38,7 @@ def fasttree(alignment: AlignedDNAFASTAFormat,
         cmd = ['FastTree']
     else:
         env = os.environ.copy()
-        if n_threads == 'auto':
+        if n_threads == 0:
             env.pop('OMP_NUM_THREADS', 0)
         else:
             env.update({'OMP_NUM_THREADS': str(n_threads)})

q2_phylogeny/_iqtree.py

@@ -74,13 +74,15 @@ def _build_iqtree_command(
             '-m', str(substitution_model),
             '-pre', str(run_prefix)]
 
-    if n_cores == 'auto' and n_cores_max:
-        cmd += ['-nt', 'AUTO', '--threads-max', '%i' % n_cores_max]
-    elif n_cores == 'auto' and n_cores_max is None:
+    if n_cores == 0:
         cmd += ['-nt', 'AUTO']
     else:
+        print('n cores', n_cores)
         cmd += ['-nt', '%i' % n_cores]
 
+    if n_cores_max is not None and n_cores_max > 0:
+        cmd += ['--threads-max', '%i' % n_cores_max]
+
     if seed:
         cmd += ['-seed', '%i' % seed]
 
@@ -216,13 +218,14 @@ def _build_iqtree_ufbs_command(
             '-m', str(substitution_model),
             '-pre', str(run_prefix)]
 
-    if n_cores == 'auto' and n_cores_max:
-        cmd += ['-nt', 'AUTO', '--threads-max', '%i' % n_cores_max]
-    elif n_cores == 'auto' and n_cores_max is None:
+    if n_cores == 0:
         cmd += ['-nt', 'AUTO']
     else:
         cmd += ['-nt', '%i' % n_cores]
 
+    if n_cores_max is not None and n_cores_max > 0:
+        cmd += ['--threads-max', '%i' % n_cores_max]
+
     if seed:
         cmd += ['-seed', '%i' % seed]
 

q2_phylogeny/_raxml.py

@@ -14,6 +14,7 @@
 
 from q2_types.feature_data import AlignedDNAFASTAFormat
 from q2_types.tree import NewickFormat
+from qiime2.sdk.util import get_available_cores
 
 _raxml_versions = {
                    'Standard': '',
@@ -35,6 +36,9 @@ def run_command(cmd, verbose=True):
 
 
 def _set_raxml_version(raxml_version='Standard', n_threads=1):
+    if n_threads == 0:
+        n_threads = get_available_cores()
+
     if n_threads == 1:
         cmd = ['raxmlHPC' + _raxml_versions[raxml_version]]
         return cmd

q2_phylogeny/plugin_setup.py

@@ -7,7 +7,7 @@
 # ----------------------------------------------------------------------------
 
 from qiime2.plugin import (Plugin, Citations, Int, Range, Str, Choices, Bool,
-                           Float, List, TypeMatch, Metadata)
+                           Float, List, TypeMatch, Metadata, Threads)
 from q2_types.tree import Phylogeny, Unrooted, Rooted
 from q2_types.feature_data import FeatureData, AlignedSequence, Sequence
 from q2_types.feature_table import (FeatureTable, Frequency,
@@ -104,7 +104,7 @@
 plugin.methods.register_function(
     function=q2_phylogeny.fasttree,
     inputs={'alignment': FeatureData[AlignedSequence]},
-    parameters={'n_threads': Int % Range(1, None) | Str % Choices(['auto'])},
+    parameters={'n_threads': Threads},
     outputs=[('tree', Phylogeny[Unrooted])],
     input_descriptions={
         'alignment': ('Aligned sequences to be used for phylogenetic '
@@ -132,7 +132,7 @@
     parameters={
             'seed': Int,
             'n_searches': Int % Range(1, None),
-            'n_threads': Int % Range(1, None),
+            'n_threads': Threads,
             'substitution_model': Str % Choices(_RAXML_MODEL_OPT),
             'raxml_version': Str % Choices(_RAXML_VERSION_OPT)},
     outputs=[('tree', Phylogeny[Unrooted])],
@@ -171,7 +171,7 @@
             'seed': Int,
             'rapid_bootstrap_seed': Int,
             'bootstrap_replicates': Int % Range(10, None),
-            'n_threads': Int % Range(1, None),
+            'n_threads': Threads,
             'substitution_model': Str % Choices(_RAXML_MODEL_OPT),
             'raxml_version': Str % Choices(_RAXML_VERSION_OPT)},
     outputs=[('tree', Phylogeny[Unrooted])],
@@ -212,8 +212,8 @@
     inputs={'alignment': FeatureData[AlignedSequence]},
     parameters={
             'seed': Int,
-            'n_cores': Int % Range(1, None) | Str % Choices(['auto']),
-            'n_cores_max': Int % Range(2, None),
+            'n_cores': Threads,
+            'n_cores_max': Threads,
             'n_runs': Int % Range(1, None),
             'substitution_model': Str % Choices(_IQTREE_DNA_MODELS),
             'n_init_pars_trees': Int % Range(1, None),
@@ -297,8 +297,8 @@
     inputs={'alignment': FeatureData[AlignedSequence]},
     parameters={
             'seed': Int,
-            'n_cores': Int % Range(1, None) | Str % Choices(['auto']),
-            'n_cores_max': Int % Range(2, None),
+            'n_cores': Threads,
+            'n_cores_max': Threads,
             'n_runs': Int % Range(1, None),
             'substitution_model': Str % Choices(_IQTREE_DNA_MODELS),
             'n_init_pars_trees': Int % Range(1, None),
@@ -495,7 +495,7 @@
         'sequences': FeatureData[Sequence],
     },
     parameters={
-        'n_threads': Int % Range(1, None) | Str % Choices(['auto']),
+        'n_threads': Threads,
         'mask_max_gap_frequency': Float % Range(0, 1, inclusive_end=True),
         'mask_min_conservation': Float % Range(0, 1, inclusive_end=True),
         'parttree': Bool,
@@ -568,7 +568,7 @@
         'sequences': FeatureData[Sequence],
     },
     parameters={
-        'n_threads': Int % Range(1, None) | Str % Choices(['auto']),
+        'n_threads': Threads,
         'mask_max_gap_frequency': Float % Range(0, 1, inclusive_end=True),
         'mask_min_conservation': Float % Range(0, 1, inclusive_end=True),
         'seed': Int,
@@ -659,7 +659,7 @@
         'sequences': FeatureData[Sequence],
     },
     parameters={
-        'n_threads': Int % Range(1, None) | Str % Choices(['auto']),
+        'n_threads': Threads,
         'mask_max_gap_frequency': Float % Range(0, 1, inclusive_end=True),
         'mask_min_conservation': Float % Range(0, 1, inclusive_end=True),
         'parttree': Bool,

q2_phylogeny/tests/test_iqtree.py

@@ -103,7 +103,7 @@ def test_iqtree_n_cores_auto(self):
         input_sequences = AlignedDNAFASTAFormat(input_fp, mode='r')
 
         with redirected_stdio(stderr=os.devnull):
-            obs = iqtree(input_sequences, n_cores='auto')
+            obs = iqtree(input_sequences, n_cores=0)
         obs_tree = skbio.TreeNode.read(str(obs), convert_underscores=False)
 
         # load the resulting tree and test that it has the right number of
@@ -122,7 +122,7 @@ def test_iqtree_n_cores_auto_max(self):
         input_sequences = AlignedDNAFASTAFormat(input_fp, mode='r')
 
         with redirected_stdio(stderr=os.devnull):
-            obs = iqtree(input_sequences, n_cores='auto', n_cores_max=2)
+            obs = iqtree(input_sequences, n_cores=0, n_cores_max=2)
         obs_tree = skbio.TreeNode.read(str(obs), convert_underscores=False)
 
         # load the resulting tree and test that it has the right number of
@@ -197,7 +197,7 @@ def test_build_iqtree_command(self):
             with redirected_stdio(stderr=os.devnull):
                 obs = _build_iqtree_command(input_sequences,
                                             seed=1723,
-                                            n_cores='auto',
+                                            n_cores=0,
                                             n_cores_max=4,
                                             n_runs=2,
                                             substitution_model='MFP',
@@ -246,7 +246,7 @@ def test_build_iqtree_ufbs_command(self):
             with redirected_stdio(stderr=os.devnull):
                 obs = _build_iqtree_ufbs_command(input_sequences,
                                                  seed=1723,
-                                                 n_cores='auto',
+                                                 n_cores=0,
                                                  n_cores_max=6,
                                                  n_runs=5,
                                                  bootstrap_replicates=2000,
@@ -373,7 +373,7 @@ def test_iqtree_ultrafast_bootstrap_auto_threads(self):
         input_fp = self.get_data_path('aligned-dna-sequences-3.fasta')
         input_sequences = AlignedDNAFASTAFormat(input_fp, mode='r')
         with redirected_stdio(stderr=os.devnull):
-            obs = iqtree_ultrafast_bootstrap(input_sequences, n_cores='auto')
+            obs = iqtree_ultrafast_bootstrap(input_sequences, n_cores=0)
         obs_tree = skbio.TreeNode.read(str(obs))
         # load the resulting tree and test that it has the right number of
         # tips and the right tip ids
@@ -390,7 +390,7 @@ def test_iqtree_ultrafast_bootstrap_auto_threads_max(self):
         input_fp = self.get_data_path('aligned-dna-sequences-3.fasta')
         input_sequences = AlignedDNAFASTAFormat(input_fp, mode='r')
         with redirected_stdio(stderr=os.devnull):
-            obs = iqtree_ultrafast_bootstrap(input_sequences, n_cores='auto',
+            obs = iqtree_ultrafast_bootstrap(input_sequences, n_cores=0,
                                              n_cores_max=2)
         obs_tree = skbio.TreeNode.read(str(obs))
         # load the resulting tree and test that it has the right number of