限制代码使用的CPU数量

限制代码使用的CPU数量

最近在实验室的服务器上跑代码被其他成员抱怨我的代码cpu占用较高，他们的代码跑起来很卡，用htop观察了一下确实占用很高，于是开始百度如何限制cpu使用，发现基本都是如下这样的python代码：

cpu_num = 4
os.environ ['OMP_NUM_THREADS']        = str(cpu_num) 
os.environ ['MKL_NUM_THREADS']        = str(cpu_num) 
os.environ ['NUMEXPR_NUM_THREADS']    = str(cpu_num) 
os.environ ['OPENBLAS_NUM_THREADS']   = str(cpu_num) 
os.environ ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num) 
torch.set_num_threads(cpu_num)

我实际测试了一下，发现并没有真的限制到cpu的使用。
然后问了下GPT4，除了提到以上方法，还提到可以使用task命令设置进程的cpu亲和性，比如taskset -c 0-3 python train.py可以限制当前任务只使用前四个cpu核心，实际测试了一下是可行的，但是现在我想在代码中进行限制，该怎么操作呢？

使用os.sched_setaffinity设置代码的cpu亲和性

我们可以使用python的os.sched_setaffinity函数来限制cpu的使用数量，代码如下：

import os
from multiprocessing import cpu_count 
cpu_num = cpu_count()
cpu_use = 4
cur_pid = os.getpid()
os.sched_setaffinity(cur_pid, list(range(cpu_num))[:cpu_use])
print(f"set the max number of cpu used to {cpu_use}")

首先使用os.getpid()函数获取到当前任务的pid，然后根据其pid使用os.sched_setaffinity为该任务设置cpu使用数量