adds @_safely_count_cpus, tests

q2_diversity_lib/_util.py

@@ -9,6 +9,7 @@
 import numpy as np
 from functools import wraps
 from inspect import signature
+import psutil
 
 
 def _drop_undefined_samples(counts: np.ndarray, sample_ids: np.ndarray,
@@ -36,3 +37,31 @@ def wrapper(*args, **kwargs):
                 raise ValueError("The provided table object is empty")
         return some_function(*args, **kwargs)
     return wrapper
+
+
+def _safely_count_cpus(some_function):
+    @wraps(some_function)
+    def wrapper(*args, **kwargs):
+        # https://psutil.readthedocs.io/en/latest/index.html#psutil.cpu_count
+        # `Process.cpu_affinity` may not be available on all systems. If not,
+        # fall back to the original cpu counting mechanism.
+        sig = signature(wrapper)
+        try:
+            cpus = len(psutil.Process().cpu_affinity())
+            print(sig.parameters)
+            if 'system_cpus' in sig.parameters:
+                kwargs['system_cpus'] = cpus
+            else:
+                raise AttributeError("The _safely_count_cpus decorator may "
+                                     "not be applied to callables without "
+                                     "'system_cpus' parameter.")
+        except AttributeError:
+            cpus = psutil.cpu_count(logical=False)
+            if 'system_cpus' in sig.parameters:
+                kwargs['system_cpus'] = cpus
+            else:
+                raise AttributeError("The _safely_count_cpus decorator may "
+                                     "not be applied to callables without "
+                                     "'system_cpus' parameter.")
+        return some_function(*args, **kwargs)
+    return wrapper

q2_diversity_lib/tests/test_util.py

@@ -7,9 +7,11 @@
 # ----------------------------------------------------------------------------
 
 from qiime2.plugin.testing import TestPluginBase
-from .._util import _disallow_empty_tables
+from .._util import _disallow_empty_tables, _safely_count_cpus
 import biom
 import numpy as np
+import unittest.mock as mock
+import psutil
 
 
 class DisallowEmptyTablesTests(TestPluginBase):
@@ -45,3 +47,37 @@ def test_decorated_lambda_with_table_param(self):
     def test_wrapped_function_has_no_table_param(self):
         with self.assertRaisesRegex(TypeError, "no parameter.*table"):
             self.function_without_table_param()
+
+
+class SafelyCountCPUSTests(TestPluginBase):
+    package = 'q2_diversity_lib.tests'
+
+    def setUp(self):
+        super().setUp()
+
+        @_safely_count_cpus
+        def function_no_args():
+            pass
+        self.function_no_args = function_no_args
+
+        @_safely_count_cpus
+        def function_w_arg(system_cpus):
+            return system_cpus
+        self.cpu_counter_w_system_cpus_arg = function_w_arg
+
+    def test_wrapped_function_does_not_take_a_cpu_count_arg(self):
+        with self.assertRaisesRegex(AttributeError, 'without \'system_cpus'):
+            self.function_no_args()
+
+    @mock.patch("q2_diversity_lib._util.psutil.Process")
+    def test_function_takes_a_cpu_count_arg(self, mock_cpu_affinity):
+        mock_cpu_affinity = psutil.Process()
+        mock_cpu_affinity.cpu_affinity = mock.MagicMock(return_value=[0, 1, 2])
+        # Is the length of our mocked cpu_affinity 3?
+        self.assertEqual(self.cpu_counter_w_system_cpus_arg(), 3)
+
+    @mock.patch('psutil.Process.cpu_affinity', side_effect=AttributeError(
+                "Here's your error"))
+    @mock.patch('psutil.cpu_count', return_value=999)
+    def test_system_has_no_cpu_affinity(self, mock_cpu_count, mock_cpu_affin):
+        self.assertEqual(self.cpu_counter_w_system_cpus_arg(), 999)