#956: cpu_affinity([]) can now be used as an alias to set affinity ag…

…ainst all eligible CPUs.

HISTORY.rst

@@ -9,6 +9,8 @@
 
 - 357_: added psutil.Process.cpu_num() (what CPU a process is on).
 - 941_: added psutil.cpu_freq() (CPU frequency).
+- 956_: cpu_affinity([]) can now be used as an alias to set affinity against
+  all eligible CPUs.
 
 **Bug fixes**
 

README.rst

@@ -248,9 +248,9 @@ Process management
     12.1
     >>> p.cpu_affinity()
     [0, 1, 2, 3]
-    >>> p.cpu_affinity([0])  # set
+    >>> p.cpu_affinity([0, 1])  # set
     >>> p.cpu_num()
-    2
+    1
     >>>
     >>> p.memory_info()
     pmem(rss=10915840, vms=67608576, shared=3313664, text=2310144, lib=0, data=7262208, dirty=0)

appveyor.yml

@@ -1,3 +1,5 @@
+# Build: 0 (bump this up by 1 to force an appveyor run)
+
 os: Visual Studio 2015
 
 environment:
@@ -124,4 +126,3 @@ only_commits:
     psutil/tests/test_windows.py
     scripts/*
     setup.py
-

docs/index.rst

@@ -1214,32 +1214,35 @@ Process class
 
     Get or set process current
     `CPU affinity <http://www.linuxjournal.com/article/6799?page=0,0>`__.
-    CPU affinity consists in telling the OS to run a certain process on a
-    limited set of CPUs only.
+    CPU affinity consists in telling the OS to run a process on a limited set
+    of CPUs only.
     On Linux this is done via the ``taskset`` command.
-    The number of eligible CPUs can be obtained with
-    ``list(range(psutil.cpu_count()))``.
-    ``ValueError`` will be raised on set in case an invalid CPU number is
-    specified.
+    If no argument is passed it returns the current CPU affinity as a list
+    of integers.
+    If passed it must be a list of integers specifying the new CPUs affinity.
+    If an empty list is passed all eligible CPUs are assumed (and set);
+    on Linux this may not necessarily mean all available CPUs as in
+    ``list(range(psutil.cpu_count()))``).
 
       >>> import psutil
       >>> psutil.cpu_count()
       4
       >>> p = psutil.Process()
-      >>> p.cpu_affinity()  # get
+      >>> # get
+      >>> p.cpu_affinity()
       [0, 1, 2, 3]
-      >>> p.cpu_affinity([0])  # set; from now on, process will run on CPU #0 only
+      >>> # set; from now on, process will run on CPU #0 and #1 only
+      >>> p.cpu_affinity([0, 1])
       >>> p.cpu_affinity()
-      [0]
-      >>>
-      >>> # reset affinity against all CPUs
-      >>> all_cpus = list(range(psutil.cpu_count()))
-      >>> p.cpu_affinity(all_cpus)
-      >>>
+      [0, 1]
+      >>> # reset affinity against all eligible CPUs
+      >>> p.cpu_affinity([])
 
     Availability: Linux, Windows, FreeBSD
 
     .. versionchanged:: 2.2.0 added support for FreeBSD
+    .. versionchanged:: 5.1.0 an empty list can be passed to set affinity
+      against all eligible CPUs.
 
   .. method:: cpu_num()
 

psutil/__init__.py

@@ -828,6 +828,8 @@ def cpu_affinity(self, cpus=None):
             """Get or set process CPU affinity.
             If specified 'cpus' must be a list of CPUs for which you
             want to set the affinity (e.g. [0, 1]).
+            If an empty list is passed, all egible CPUs are assumed
+            (and set).
             (Windows, Linux and BSD only).
             """
             # Automatically remove duplicates both on get and
@@ -836,6 +838,11 @@ def cpu_affinity(self, cpus=None):
             if cpus is None:
                 return list(set(self._proc.cpu_affinity_get()))
             else:
+                if not cpus:
+                    if hasattr(self._proc, "_get_eligible_cpus"):
+                        cpus = self._proc._get_eligible_cpus()
+                    else:
+                        cpus = tuple(range(len(cpu_times(percpu=True))))
                 self._proc.cpu_affinity_set(list(set(cpus)))
 
     # Linux, FreeBSD, SunOS

psutil/_pslinux.py

@@ -1521,18 +1521,33 @@ def nice_set(self, value):
     def cpu_affinity_get(self):
         return cext.proc_cpu_affinity_get(self.pid)
 
+    def _get_eligible_cpus(
+            self, _re=re.compile("Cpus_allowed_list:\t(\d+)-(\d+)")):
+        # See: https://github.com/giampaolo/psutil/issues/956
+        data = self._read_status_file()
+        match = _re.findall(data)
+        if match:
+            return tuple(range(int(match[0][0]), int(match[0][1]) + 1))
+        else:
+            return tuple(range(len(per_cpu_times())))
+
     @wrap_exceptions
     def cpu_affinity_set(self, cpus):
         try:
             cext.proc_cpu_affinity_set(self.pid, cpus)
         except (OSError, ValueError) as err:
             if isinstance(err, ValueError) or err.errno == errno.EINVAL:
-                allcpus = tuple(range(len(per_cpu_times())))
+                eligible_cpus = self._get_eligible_cpus()
+                all_cpus = tuple(range(len(per_cpu_times())))
                 for cpu in cpus:
-                    if cpu not in allcpus:
+                    if cpu not in all_cpus:
                         raise ValueError(
                             "invalid CPU number %r; choose between %s" % (
-                                cpu, allcpus))
+                                cpu, eligible_cpus))
+                    if cpu not in eligible_cpus:
+                        raise ValueError(
+                            "CPU number %r is not eligible; choose "
+                            "between %s" % (cpu, eligible_cpus))
             raise
 
     # only starting from kernel 2.6.13

psutil/tests/test_linux.py

@@ -1277,6 +1277,11 @@ def test_cpu_affinity(self):
         self.assertEqual(
             self.proc.cpu_affinity(), list(range(min_, max_ + 1)))
 
+    def test_cpu_affinity_eligible_cpus(self):
+        with mock.patch("psutil._pslinux.per_cpu_times") as m:
+            self.proc._proc._get_eligible_cpus()
+        assert not m.called
+
 
 if __name__ == '__main__':
     run_test_module_by_name(__file__)

psutil/tests/test_process.py

@@ -849,10 +849,13 @@ def test_cwd_2(self):
     def test_cpu_affinity(self):
         p = psutil.Process()
         initial = p.cpu_affinity()
+        assert initial, initial
         self.addCleanup(p.cpu_affinity, initial)
+
         if hasattr(os, "sched_getaffinity"):
             self.assertEqual(initial, list(os.sched_getaffinity(p.pid)))
         self.assertEqual(len(initial), len(set(initial)))
+
         all_cpus = list(range(len(psutil.cpu_percent(percpu=True))))
         # setting on travis doesn't seem to work (always return all
         # CPUs on get):
@@ -867,9 +870,14 @@ def test_cpu_affinity(self):
             if hasattr(p, "num_cpu"):
                 self.assertEqual(p.cpu_affinity()[0], p.num_cpu())
 
-        #
-        p.cpu_affinity(all_cpus)
-        self.assertEqual(p.cpu_affinity(), all_cpus)
+        # [] is an alias for "all eligible CPUs"; on Linux this may
+        # not be equal to all available CPUs, see:
+        # https://github.com/giampaolo/psutil/issues/956
+        p.cpu_affinity([])
+        if LINUX:
+            self.assertEqual(p.cpu_affinity(), p._proc._get_eligible_cpus())
+        else:
+            self.assertEqual(p.cpu_affinity(), all_cpus)
         if hasattr(os, "sched_getaffinity"):
             self.assertEqual(p.cpu_affinity(),
                              list(os.sched_getaffinity(p.pid)))