Merge pull request #31 from nomad010/port-process-affinity-fix-to-master

Port process affinity fix to master

processfamily/__init__.py

@@ -381,9 +381,10 @@ def _handle_response_line(self, line):
 #being killed
 _global_process_job_handle = None
 
-CPU_AFFINITY_STRATEGY_NONE = 0
-CPU_AFFINITY_STRATEGY_CHILDREN_ONLY = 1
-CPU_AFFINITY_STRATEGY_PARENT_INCLUDED = 2
+CPU_AFFINITY_STRATEGY_INHERIT = 0  # By default, process affinity masks are inherited by child processes. This retains that behavior
+CPU_AFFINITY_STRATEGY_CHILDREN_ONLY = 1  # Only the child processes that are spawned will get their processes affinity tied to a core
+CPU_AFFINITY_STRATEGY_PARENT_INCLUDED = 2  # Like the above, except that the parent process also gets it affinity tied to a core
+CPU_AFFINITY_STRATEGY_NONE = 3  # This overrides any default defined by the OS and allows spawned processes to float regardless of the affinity of the parent
 
 
 class NoCommsStrategy(ChildCommsStrategy):
@@ -679,6 +680,10 @@ def set_child_affinity_mask(self, pid, child_index):
         i = child_index+1 if self.CPU_AFFINITY_STRATEGY == CPU_AFFINITY_STRATEGY_PARENT_INCLUDED else child_index
         set_process_affinity({i%self.cpu_count}, pid=pid)
 
+    def allow_child_to_float(self, pid):
+        """ Allows the process given by pid to not be tied to any of the cores. """
+        set_process_affinity(list(range(self.cpu_count)), pid=pid)
+
     def start(self, timeout=30):
         if self.child_processes:
             raise Exception("Invalid state: start() can only be called once")
@@ -697,11 +702,15 @@ def start(self, timeout=30):
             logger.debug("Commandline for %s: %s", self.get_child_name(i), json.dumps(cmd))
             p = self.get_Popen_class()(cmd, **self.get_Popen_kwargs(i, close_fds=self.CLOSE_FDS))
 
-            if self.CPU_AFFINITY_STRATEGY and p.poll() is None:
+            if p.poll() is None:
                 try:
-                    self.set_child_affinity_mask(p.pid, i)
+                    if self.CPU_AFFINITY_STRATEGY in [CPU_AFFINITY_STRATEGY_CHILDREN_ONLY, CPU_AFFINITY_STRATEGY_PARENT_INCLUDED]:
+                        self.set_child_affinity_mask(p.pid, i)
+                    elif self.CPU_AFFINITY_STRATEGY == CPU_AFFINITY_STRATEGY_NONE:
+                        self.allow_child_to_float(p.pid)
                 except Exception as e:
                     logger.error("Unable to set affinity for %s process %d: %s", self.get_child_name(i), p.pid, e)
+
             self.child_processes.append(self.CHILD_COMMS_STRATEGY(p, self.ECHO_STD_ERR, i, self))
 
         if sys.platform.startswith('win') and self.WIN_PASS_HANDLES_OVER_COMMANDLINE:

processfamily/test/ParentProcess.py

@@ -64,8 +64,10 @@ def __init__(self, number_of_child_processes=None, run_as_script=True):
                 self.CLOSE_FDS = False
             elif command == 'use_job_object_off':
                 self.WIN_USE_JOB_OBJECT = False
+            elif command == 'cpu_affinity_inherit':
+                self.CPU_AFFINITY_STRATEGY = processfamily.CPU_AFFINITY_STRATEGY_INHERIT
             elif command == 'cpu_affinity_off':
-                self.CPU_AFFINITY_STRATEGY = None
+                self.CPU_AFFINITY_STRATEGY = processfamily.CPU_AFFINITY_STRATEGY_NONE
             elif command == 'use_cat' or command == 'use_cat_comms_none':
                 self.WIN_PASS_HANDLES_OVER_COMMANDLINE = False
                 self.CHILD_COMMS_STRATEGY = processfamily.CHILD_COMMS_STRATEGY_PIPES_CLOSE if command == 'use_cat' else processfamily.CHILD_COMMS_STRATEGY_NONE

processfamily/test_processfamily.py

@@ -31,6 +31,7 @@
 from pytest_lazyfixture import lazy_fixture
 
 from processfamily.futurecompat import get_env_dict, list_to_native_str
+from processfamily.processes import get_process_affinity, set_process_affinity, cpu_count
 
 if sys.platform.startswith('win'):
     from processfamily._winprocess_ctypes import CAN_USE_EXTENDED_STARTUPINFO, CREATE_BREAKAWAY_FROM_JOB
@@ -76,6 +77,14 @@ def get_pid_files():
     return glob.glob(os.path.join(pid_dir, "*.pid"))
 
 
+def get_pids():
+    pids = []
+    for filename in get_pid_files():
+        with open(filename) as pidfile:
+            pids.append(int(pidfile.read()))
+    return pids
+
+
 def kill_parent():
     for pid_file in get_pid_files():
         if os.path.basename(pid_file).startswith('c'):
@@ -627,9 +636,40 @@ def test_use_job_object_off(self, fws):
                      'use_job_object_off')
         check_stop()
 
-    def test_cpu_affinity_off(self, fws):
+    def test_cpu_affinity_inherit(self, fws):
+        fws.start_up(test_command='cpu_affinity_inherit')
+        check_stop()
+
+    def test_affinity_inherited_by_children(self, fws):
+        tied_to_cores = [0]
+        set_process_affinity(tied_to_cores)
+        fws.start_up(test_command='cpu_affinity_inherit')
+        children_pids = []
+        for pid in get_pids():
+            children_pids.append(list(get_process_affinity(pid)))
+        check_stop()
+        for child_pids in children_pids:
+            assert child_pids == tied_to_cores
+        set_process_affinity(range(cpu_count()))
+
+    def test_no_affinity_children_float(self, fws):
+        tied_to_cores = {0}
+        all_cores = set(range(cpu_count()))
+        set_process_affinity(tied_to_cores)
         fws.start_up(test_command='cpu_affinity_off')
+        children_affinity = []
+        for pid in get_pids():
+            children_affinity.append(set(get_process_affinity(pid)))
         check_stop()
+        parent_pid_checked = False
+        for affinity in children_affinity:
+            if affinity == tied_to_cores:
+                assert not parent_pid_checked, "More than 1 process detected with parent affinity"
+                parent_pid_checked = True
+            else:
+                assert affinity == all_cores
+        set_process_affinity(all_cores)
+
 
     def test_handles_over_commandline_off_file_open_by_parent(self, fws):
         if not sys.platform.startswith('win') or not CAN_USE_EXTENDED_STARTUPINFO:
@@ -658,4 +698,4 @@ def test_service_stop_child_freeze_on_start(self, windows_service):
         assert_middle_child_port_unbound()
         win32serviceutil.StopService(Config.svc_name)
         # This still needs time to wait for the child to stop for 10 seconds:
-        windows_service.wait_for_parent_to_stop(11)
+        windows_service.wait_for_parent_to_stop(11)

setup.py

@@ -14,7 +14,7 @@
 
 setup(
     name='processfamily',
-    version='0.8',
+    version='0.9',
     packages = find_packages(),
     license='Apache License, Version 2.0',
     description='A library for launching, maintaining, and terminating a family of long-lived python child processes on Windows and *nix.',