pin-cpu-cores.py

#! /bin/python3

import subprocess
import os
import operator

#########################################################

import re

def natural_sort_key(s, _nsre=re.compile('([0-9]+)')):
    return [int(text) if text.isdigit() else text.lower()
            for text in _nsre.split(s)]

#########################################################

all_cpus = [i for i in range(23)]
vm_cpus = [0, 12, 1, 13, 2, 14, 3, 15, 4, 16, 5, 17, 6, 7, 8]
vm_cpus_string = ",".join(str(c) for c in vm_cpus)

host_cpus = [c for c in all_cpus if c not in vm_cpus]
host_cpus_string = ",".join(str(c) for c in host_cpus)

class QemuThread:
    def __init__(self, pid):
        self.pid = pid
        self.proc_path = f"/proc/{pid}/"
        with open(f"{self.proc_path}/comm", 'r') as file:
            self.name = file.read().replace('\n', '')
    def __repr__(self):
       return f"{self.name}={self.pid}"

def pin(thread, cpu, pid):
    thread = thread.replace(" ", "-").replace("/", "-")
    # create cpuset
    process = subprocess.Popen(f"cset set -c {cpu} -s {thread} -m 0".split(), stdout=subprocess.PIPE, universal_newlines=True)
    cout, cerr = process.communicate()
    if cerr:
        print(f"Creating cpuset failed: {cerr}")
        exit()
    # attach thread to cpuset
    process = subprocess.Popen(f"cset proc -m -p {pid} -t {thread}".split(), stdout=subprocess.PIPE, universal_newlines=True)
    cout, cerr = process.communicate()
    if cerr:
        print(f"Creating cpuset failed: {cerr}")
        exit()
    print(f"Successfully attached {thread} ({pid}) to {cpu}")



# Find qemu pid
process = subprocess.Popen("pgrep qemu-system-x86".split(), stdout=subprocess.PIPE, universal_newlines=True)
pid, error = process.communicate()
pid = pid.rstrip("\n")
if not pid:
       print("Qemu process is not found")
       exit()
print(f"=== Detected qemu pid: {pid}")


# Find all the thread pids of qemu process
tasks_dir = f"/proc/{pid}/task/"
print(f"=== Searching \"{tasks_dir}\"")
threads = [QemuThread(thread_pid.name) for thread_pid in os.scandir(tasks_dir)]
print(f"=== Found threads: {threads}")

# Find cpu threads
cpu_threads = list(filter(lambda x: x.name.startswith("CPU"), threads))
cpu_threads.sort(key=lambda t: natural_sort_key(t.name))
print(f"=== CPU threads: {cpu_threads}")

# Find io threads
io_threads = list(filter(lambda x: x.name.startswith("IO"), threads))
io_threads.sort(key=lambda t: t.name)
print(f"=== IO threads: {io_threads}")

# Find emulator threads:
emulator_threads = list(filter(lambda x: not x.name.startswith("IO") and not x.name.startswith("CPU"), threads))
emulator_threads.sort(key=lambda t: t.name)
print(f"=== Emulator threads: {emulator_threads}")

if (len(cpu_threads) + len(io_threads)) >= len(vm_cpus):
    print(f"Not enough isolated cpus ({len(vm_cpus)}) to pin cpu threads ({len(cpu_threads)}), io threads ({len(io_threads)}) and emulator threads (requires at least 1 cpu)")
    exit()

# Isolate cpus
process = subprocess.Popen(f"cset set -c {host_cpus_string} -s system -m 1".split(), stdout=subprocess.PIPE, universal_newlines=True)
cout, cerr = process.communicate()
process = subprocess.Popen(f"cset proc -m -f root -t system".split(), stdout=subprocess.PIPE, universal_newlines=True)
cout, cerr = process.communicate()
print("Moved all tasks to 'system' cpuset")

# Pin threads
cpu_id = 0
cpu_used_for_non_emulator_threads = []
for t in cpu_threads:
    pin(t.name, vm_cpus[cpu_id], t.pid)
    cpu_used_for_non_emulator_threads.append(vm_cpus[cpu_id])
    cpu_id = cpu_id + 1

for t in io_threads:
    pin(t.name, vm_cpus[cpu_id], t.pid)
    # cpu_used_for_non_emulator_threads.append(vm_cpus[cpu_id])
    cpu_id = cpu_id + 1

pin("emulator", ','.join(str(c) for c in vm_cpus[cpu_id:]), ",".join(str(x.pid) for x in emulator_threads))

# IRQ affinity
irq_cpu_affinity_list = ",".join(str(x) for x in cpu_used_for_non_emulator_threads)
process = subprocess.Popen(['/bin/bash'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
cout, cerr = process.communicate(bytes(f"grep vfio /proc/interrupts | cut -d ':' -f 1 | while read -r i; do echo {irq_cpu_affinity_list} > /proc/irq/$i/smp_affinity_list; done", 'utf-8'))
print(f"Set IRQ affinity to {irq_cpu_affinity_list}")