Python 多线程/多处理和限制 CPU 核心亲和力

Question

在 Python 中，您可以使用 multiprocessing.Pool、multiprocessing.ThreadPool、concurrent.futures.ProcessPoolExecutor 和 concurrent.futures.ThreadPoolExecutor 创建新线程和进程来运行给定任务。

默认情况下，这些线程/进程以与其父进程相同的 CPU 核心亲和性运行，即所有可用的核心/线程。

在 Linux/Unix 系统上，可以使用 os.sched_setaffinity(pid, mask) 更改 CPU 内核关联。问题在于这仅限于某些 Linux/Unix 系统。

psutil python 库公开了使用psutil.Process().cpu_affinity(CPUS) 设置CPU 核心亲和性的能力，其中CPUS 是一个整数列表，标识进程应该使用哪些CPU 核心/线程，从0 开始.

问题在于，通常操作系统 CPU 调度程序可以处理挑选和选择应该为给定进程使用哪个内核/线程，而不是让最终用户决定要使用哪些 CPU 内核/线程。

我的问题是，是否可以创建线程/进程池并将每个实例限制为使用 X 个 CPU 内核/线程，但不限制它们的确切内核亲和性？

例如，如果我有 16 核的 PC，并且想要创建 4 个进程，我可以创建一个 multiprocessing.Pool(processes=4) 对象。现在，如果我希望这 4 个孩子中的每一个被限制为每个仅使用 2 个 CPU 内核，我将不得不使用 psutil 抢先选择 2 个 CPU 内核并将它们分配给那个进程，从可用的 CPU 内核中移除这 2 个 CPU 内核CPU 内核列表，然后对所有 4 个进程重复该过程。

但这并不理想，如果我给一个进程分配系统中最弱的两个内核会怎样？或者如果这 2 个内核在物理上相距更远（例如现代多芯片 AMD Ryzen CPU 或双 CPU 插槽系统）。

我想让操作系统自动为每个进程安排 2 个内核，并根据需要调整它们，而不是手动设置和取消设置每个进程的 CPU 内核。

有没有一种方法可以在 Python 中完成？

Answer 1

前段时间我也有类似的需求，所以我写了一个 CPUResourceManager 类来跟踪我分配给进程的内核。在这里你会调用 get_processors 方法来获取您将要使用的内核列表。您设置核心亲和力使用 PSUTIL 就像您已经在做的那样。当您的过程完成后，使用free_processors 方法返回内核。

from typing import NamedTuple
from enum import Enum


class CPUResponse(NamedTuple):
    """
    This is the response when the CPUResourceManager is asked for some cores.
    """
    success: bool     # whether or not there are enough cores available
    processors: list  # the list of processors to be used by the process


class ProcState(Enum):
    """
    Represents the state of a processor core.  This is only 
    represents what we are having the cores do.  Not what other unrelated
    processes on the machine are doing
    """
    idle = 0
    busy = 1


class CPUResourceManager:
    def __init__(self, cpu_count=max(cpu_count() - 2, 1)) -> None:

        self.cpu_count = cpu_count

        self.processors = {i: ProcState.idle for i in range(self.cpu_count)}

    def cpu_avalaible_count(self):
        available = [
            p for p, state in self.processors.items() if state == ProcState.idle
        ]
        return len(available)

    def get_processors(self, count=1):
        """Get some available cores"""
        available = [
            p for p, state in self.processors.items() if state == ProcState.idle
        ]

        if len(available) >= count:
            cpus = available[:count]
            for p in cpus:
                self.processors[p] = ProcState.busy
            return CPUResponse(True, available[:count])
        else:
            return CPUResponse(False, [])

    def free_processors(self, processors: list):
        """return the cores when you are done"""
        for p in processors:
            if p in self.processors:
                self.processors[p] = ProcState.idle
            else:
                # manager was likely resized and this processor
                # should no longer be considered available
                pass