Repository: Parsl/parsl

Issue:
HTEX forcing thread affinity
I have issues around multi-threaded Parsl apps interfering with each other when placed on the same node. Instead of creating my own logic to prevent apps from attempting to use the same CPU cores, it could be nice for Parsl to enforce affinity itself (e.g., with the [`os` module](https://docs.python.org/3/library/os.html#interface-to-the-scheduler)).

**Is your feature request related to a problem? Please describe.**
It is related to my issues resulting from Tensorflow's CPU/GPU affinity (#751)

**Describe the solution you'd like**
For Parsl to force explicity HTEX workers to use certain CPUs.

**Describe alternatives you've considered**
Having users make affinity logic themselves.

**Additional context**
I'm going to work up a solution in my own parsl app and then submit it for review. 



Discussion:

Resolving patch:
diff --git a/parsl/executors/high_throughput/executor.py b/parsl/executors/high_throughput/executor.py
--- a/parsl/executors/high_throughput/executor.py
+++ b/parsl/executors/high_throughput/executor.py
@@ -129,6 +129,13 @@ class HighThroughputExecutor(StatusHandlingExecutor, RepresentationMixin):
     max_workers : int
         Caps the number of workers launched by the manager. Default: infinity
 
+    cpu_affinity: string
+        Whether or how each worker process sets thread affinity. Options are "none" to forgo
+        any CPU affinity configuration, "block" to assign adjacent cores to workers
+        (ex: assign 0-1 to worker 0, 2-3 to worker 1), and
+        "alternating" to assign cores to workers in round-robin
+        (ex: assign 0,2 to worker 0, 1,3 to worker 1).
+
     prefetch_capacity : int
         Number of tasks that could be prefetched over available worker capacity.
         When there are a few tasks (<100) or when tasks are long running, this option should
@@ -170,6 +177,7 @@ def __init__(self,
                  cores_per_worker: float = 1.0,
                  mem_per_worker: Optional[float] = None,
                  max_workers: Union[int, float] = float('inf'),
+                 cpu_affinity: str = 'none',
                  prefetch_capacity: int = 0,
                  heartbeat_threshold: int = 120,
                  heartbeat_period: int = 30,
@@ -225,6 +233,7 @@ def __init__(self,
         self.poll_period = poll_period
         self.run_dir = '.'
         self.worker_logdir_root = worker_logdir_root
+        self.cpu_affinity = cpu_affinity
 
         if not launch_cmd:
             self.launch_cmd = ("process_worker_pool.py {debug} {max_workers} "
@@ -239,7 +248,8 @@ def __init__(self,
                                "--block_id={{block_id}} "
                                "--hb_period={heartbeat_period} "
                                "{address_probe_timeout_string} "
-                               "--hb_threshold={heartbeat_threshold} ")
+                               "--hb_threshold={heartbeat_threshold} "
+                               "--cpu-affinity {cpu_affinity} ")
 
     def initialize_scaling(self):
         """ Compose the launch command and call the scale_out
@@ -270,7 +280,8 @@ def initialize_scaling(self):
                                        heartbeat_period=self.heartbeat_period,
                                        heartbeat_threshold=self.heartbeat_threshold,
                                        poll_period=self.poll_period,
-                                       logdir=worker_logdir)
+                                       logdir=worker_logdir,
+                                       cpu_affinity=self.cpu_affinity)
         self.launch_cmd = l_cmd
         logger.debug("Launch command: {}".format(self.launch_cmd))
 
diff --git a/parsl/executors/high_throughput/process_worker_pool.py b/parsl/executors/high_throughput/process_worker_pool.py
--- a/parsl/executors/high_throughput/process_worker_pool.py
+++ b/parsl/executors/high_throughput/process_worker_pool.py
@@ -64,7 +64,8 @@ def __init__(self,
                  block_id=None,
                  heartbeat_threshold=120,
                  heartbeat_period=30,
-                 poll_period=10):
+                 poll_period=10,
+                 cpu_affinity=False):
         """
         Parameters
         ----------
@@ -75,9 +76,6 @@ def __init__(self,
              Timeout in seconds for the address probe to detect viable addresses
              to the interchange. Default : 30s
 
-        worker_url : str
-             Worker url on which workers will attempt to connect back
-
         uid : str
              string unique identifier
 
@@ -116,6 +114,9 @@ def __init__(self,
 
         poll_period : int
              Timeout period used by the manager in milliseconds. Default: 10ms
+
+        cpu_affinity : str
+             Whether each worker should force its affinity to different CPUs
         """
 
         logger.info("Manager started")
@@ -186,6 +187,7 @@ def __init__(self,
         self.heartbeat_period = heartbeat_period
         self.heartbeat_threshold = heartbeat_threshold
         self.poll_period = poll_period
+        self.cpu_affinity = cpu_affinity
 
     def create_reg_message(self):
         """ Creates a registration message to identify the worker to the interchange
@@ -365,7 +367,8 @@ def worker_watchdog(self, kill_event):
                                                                      self.pending_task_queue,
                                                                      self.pending_result_queue,
                                                                      self.ready_worker_queue,
-                                                                     self._tasks_in_progress
+                                                                     self._tasks_in_progress,
+                                                                     self.cpu_affinity
                                                                  ), name="HTEX-Worker-{}".format(worker_id))
                     self.procs[worker_id] = p
                     logger.info("[WORKER_WATCHDOG_THREAD] Worker {} has been restarted".format(worker_id))
@@ -390,7 +393,8 @@ def start(self):
                                                              self.pending_task_queue,
                                                              self.pending_result_queue,
                                                              self.ready_worker_queue,
-                                                             self._tasks_in_progress
+                                                             self._tasks_in_progress,
+                                                             self.cpu_affinity
                                                          ), name="HTEX-Worker-{}".format(worker_id))
             p.start()
             self.procs[worker_id] = p
@@ -473,7 +477,7 @@ def execute_task(bufs):
         return user_ns.get(resultname)
 
 
-def worker(worker_id, pool_id, pool_size, task_queue, result_queue, worker_queue, tasks_in_progress):
+def worker(worker_id, pool_id, pool_size, task_queue, result_queue, worker_queue, tasks_in_progress, cpu_affinity):
     """
 
     Put request token into queue
@@ -496,6 +500,25 @@ def worker(worker_id, pool_id, pool_size, task_queue, result_queue, worker_queue
     if args.debug:
         logger.debug("Debug logging enabled")
 
+    # If desired, set process affinity
+    if cpu_affinity != "none":
+        # Count the number of cores per worker
+        avail_cores = sorted(os.sched_getaffinity(0))  # Get the available processors
+        cores_per_worker = len(avail_cores) // pool_size
+        assert cores_per_worker > 0, "Affinity does not work if there are more workers than cores"
+
+        # Determine this worker's cores
+        if cpu_affinity == "block":
+            my_cores = avail_cores[cores_per_worker * worker_id:cores_per_worker * (worker_id + 1)]
+        elif cpu_affinity == "alternating":
+            my_cores = avail_cores[worker_id::pool_size]
+        else:
+            raise ValueError("Affinity strategy {} is not supported".format(cpu_affinity))
+
+        # Set the affinity for this worker
+        os.sched_setaffinity(0, my_cores)
+        logger.info("Set worker CPU affinity to {}".format(my_cores))
+
     while True:
         worker_queue.put(worker_id)
 
@@ -592,6 +615,8 @@ def start_file_logger(filename, rank, name='parsl', level=logging.DEBUG, format_
                         help="Poll period used in milliseconds")
     parser.add_argument("-r", "--result_port", required=True,
                         help="REQUIRED: Result port for posting results to the interchange")
+    parser.add_argument("--cpu-affinity", type=str, choices=["none", "block", "alternating"],
+                        help="Whether/how workers should control CPU affinity.")
 
     args = parser.parse_args()
 
@@ -618,6 +643,7 @@ def start_file_logger(filename, rank, name='parsl', level=logging.DEBUG, format_
         logger.info("Prefetch capacity: {}".format(args.prefetch_capacity))
         logger.info("Heartbeat threshold: {}".format(args.hb_threshold))
         logger.info("Heartbeat period: {}".format(args.hb_period))
+        logger.info("CPU affinity: {}".format(args.cpu_affinity))
 
         manager = Manager(task_port=args.task_port,
                           result_port=args.result_port,
@@ -631,7 +657,8 @@ def start_file_logger(filename, rank, name='parsl', level=logging.DEBUG, format_
                           prefetch_capacity=int(args.prefetch_capacity),
                           heartbeat_threshold=int(args.hb_threshold),
                           heartbeat_period=int(args.hb_period),
-                          poll_period=int(args.poll))
+                          poll_period=int(args.poll),
+                          cpu_affinity=args.cpu_affinity)
         manager.start()
 
     except Exception as e:


Test patch:
diff --git a/parsl/tests/sites/test_affinity.py b/parsl/tests/sites/test_affinity.py
new file mode 100644
--- /dev/null
+++ b/parsl/tests/sites/test_affinity.py
@@ -0,0 +1,45 @@
+"""Tests related to Parsl workers being able to access their worker ID"""
+
+from parsl.providers import LocalProvider
+from parsl.channels import LocalChannel
+from parsl.config import Config
+from parsl.executors import HighThroughputExecutor
+from parsl import python_app
+import pytest
+import os
+
+local_config = Config(
+    executors=[
+        HighThroughputExecutor(
+            label="htex_Local",
+            worker_debug=True,
+            max_workers=2,
+            cpu_affinity='block',
+            provider=LocalProvider(
+                channel=LocalChannel(),
+                init_blocks=1,
+                max_blocks=1,
+            ),
+        )
+    ],
+    strategy=None,
+)
+
+
+@python_app
+def get_worker_info():
+    from time import sleep
+    import os
+    rank = int(os.environ['PARSL_WORKER_RANK'])
+    aff = os.sched_getaffinity(0)
+    sleep(1.0)
+    return rank, aff
+
+
+@pytest.mark.local
+@pytest.mark.skipif('sched_getaffinity' not in dir(os), reason='System does not support sched_setaffinity')
+@pytest.mark.skipif(os.cpu_count() == 1, reason='Must have a more than one CPU')
+def test_htex():
+    worker_info = [get_worker_info() for _ in range(4)]
+    worker_affinity = dict([r.result() for r in worker_info])
+    assert worker_affinity[0] != worker_affinity[1]