如何用Python实时更改Linux上的CPU亲和性？

Question

我知道可以用 os.sched_setaffinity 来设置 CPU 亲和性，但似乎无法实时更改它。下面是我的代码：

首先，我有一个 C++ 程序

// test.cpp
#include <iostream>
#include <thread>
#include <vector>

void workload() {
    unsigned long long int sum = 0;
    for (long long int i = 0; i < 50000000000; ++i) {
        sum += i;
    }
    std::cout << "Sum: " << sum << std::endl;
}

int main() {
    unsigned int num_threads = std::thread::hardware_concurrency();
    std::cout << "Creating " << num_threads << " threads." << std::endl;

    std::vector<std::thread> threads;
    for (unsigned int i = 0; i < num_threads; ++i) {
        threads.push_back(std::thread(workload));
    }

    for (auto& thread : threads) {
        thread.join();
    }

    return 0;
}

然后，我编译它

g++ test.cpp -O0

这样我就会在同一个目录下得到一个 a.out 文件。

接着，在同一个目录下，我还有一个 Python 文件

# test.py
from subprocess import Popen
import os
import time
a = set(range(8, 16))
b = set(range(4, 12))

if __name__ == "__main__":
    proc = Popen("./a.out", shell=True)
    pid = proc.pid
    print("pid", pid)
    tic = time.time()
    while True:
        if time.time() - tic < 10:
            os.sched_setaffinity(pid, a)
            print("a", os.sched_getaffinity(pid))
        else:
            os.sched_setaffinity(pid, b)
            print("b", os.sched_getaffinity(pid))
        res = proc.poll()
        if res is None:
            time.sleep(1)
        else:
            break

a.out 会运行很长时间，我希望 test.py 的表现是：在前 10 秒内，CPU 8~15 会很忙，而 0~7 会闲置；而在 10 秒后，我希望看到 CPU 4~11 忙碌，其他的闲置。但是，当我用 htop 观察时，发现前 10 秒的确符合我的预期，但在 10 秒后，我每秒都能看到 b {4, 5, 6, 7, 8, 9, 10, 11}，就好像我成功设置了亲和性；然而在 htop 上，我仍然看到 CPU 8~15 忙碌，而 0~7 闲置，直到程序正常停止，这意味着我没有成功设置亲和性。

我想问一下，为什么会这样？我看过 手册，但没有找到相关的说明。而且似乎 Python 的 os.sched_setaffinity 不会返回任何结果，所以我看不到效果。

我使用的是 AMD 的 CPU，但我觉得这并不重要。

Answer 1

问题的根本在于，在Linux系统上，sched_setaffinity()是针对线程的，而不是针对整个进程的。主线程和进程有相同的ID，但后续创建的线程会有不同的ID，它们会从父线程那里继承CPU亲和性。

显然，Python能够快速设置主线程的CPU亲和性，因此子线程会从主线程继承最初的CPU设置。但是，如果之后再改变主线程的CPU亲和性，已经在运行的线程不会受到影响（不过理论上，如果后续从主线程创建新线程，它们会受到影响）。

为了绕过这个问题，你可以引入一个辅助函数，像这样：

def ChangeProcessAffinity(pid, cpus):
  for tid in map(int, os.listdir(f'/proc/{pid}/task/')):
    try:
      os.sched_setaffinity(tid, cpus)
    except e:
      pass  # maybe log the error instead

...然后在你的循环中调用这个函数。需要注意的是，这里有两个问题：

1. 使用procfs的方法是特定于Linux的，
2. 这个解决方案本身是有竞争条件的：在列出线程后，线程可能会开始或停止，这也是你需要异常处理的原因。

我觉得没有一种原子性的方法可以同时设置一个进程中所有线程的CPU亲和性；你已经发现的taskset命令行工具大致上做的就是我上面提到的事情。

如果你想要一个更稳健的解决方案，我认为最好是在C++的二进制文件内部实现这个功能（如果你能控制它的话），因为它知道线程何时被启动或停止。