Guest User

pin-cpu-cores.py

a guest
Dec 13th, 2020
120
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #! /bin/python3
  2.  
  3. import subprocess
  4. import os
  5. import operator
  6.  
  7. #########################################################
  8.  
  9. import re
  10.  
  11. def natural_sort_key(s, _nsre=re.compile('([0-9]+)')):
  12.     return [int(text) if text.isdigit() else text.lower()
  13.             for text in _nsre.split(s)]
  14.  
  15. #########################################################
  16.  
  17. all_cpus = [i for i in range(23)]
  18. vm_cpus = [0, 12, 1, 13, 2, 14, 3, 15, 4, 16, 5, 17, 6, 7, 8]
  19. vm_cpus_string = ",".join(str(c) for c in vm_cpus)
  20.  
  21. host_cpus = [c for c in all_cpus if c not in vm_cpus]
  22. host_cpus_string = ",".join(str(c) for c in host_cpus)
  23.  
  24. class QemuThread:
  25.     def __init__(self, pid):
  26.         self.pid = pid
  27.         self.proc_path = f"/proc/{pid}/"
  28.         with open(f"{self.proc_path}/comm", 'r') as file:
  29.             self.name = file.read().replace('\n', '')
  30.     def __repr__(self):
  31.        return f"{self.name}={self.pid}"
  32.  
  33. def pin(thread, cpu, pid):
  34.     thread = thread.replace(" ", "-").replace("/", "-")
  35.     # create cpuset
  36.     process = subprocess.Popen(f"cset set -c {cpu} -s {thread} -m 0".split(), stdout=subprocess.PIPE, universal_newlines=True)
  37.     cout, cerr = process.communicate()
  38.     if cerr:
  39.         print(f"Creating cpuset failed: {cerr}")
  40.         exit()
  41.     # attach thread to cpuset
  42.     process = subprocess.Popen(f"cset proc -m -p {pid} -t {thread}".split(), stdout=subprocess.PIPE, universal_newlines=True)
  43.     cout, cerr = process.communicate()
  44.     if cerr:
  45.         print(f"Creating cpuset failed: {cerr}")
  46.         exit()
  47.     print(f"Successfully attached {thread} ({pid}) to {cpu}")
  48.  
  49.  
  50.  
  51. # Find qemu pid
  52. process = subprocess.Popen("pgrep qemu-system-x86".split(), stdout=subprocess.PIPE, universal_newlines=True)
  53. pid, error = process.communicate()
  54. pid = pid.rstrip("\n")
  55. if not pid:
  56.        print("Qemu process is not found")
  57.        exit()
  58. print(f"=== Detected qemu pid: {pid}")
  59.  
  60.  
  61. # Find all the thread pids of qemu process
  62. tasks_dir = f"/proc/{pid}/task/"
  63. print(f"=== Searching \"{tasks_dir}\"")
  64. threads = [QemuThread(thread_pid.name) for thread_pid in os.scandir(tasks_dir)]
  65. print(f"=== Found threads: {threads}")
  66.  
  67. # Find cpu threads
  68. cpu_threads = list(filter(lambda x: x.name.startswith("CPU"), threads))
  69. cpu_threads.sort(key=lambda t: natural_sort_key(t.name))
  70. print(f"=== CPU threads: {cpu_threads}")
  71.  
  72. # Find io threads
  73. io_threads = list(filter(lambda x: x.name.startswith("IO"), threads))
  74. io_threads.sort(key=lambda t: t.name)
  75. print(f"=== IO threads: {io_threads}")
  76.  
  77. # Find emulator threads:
  78. emulator_threads = list(filter(lambda x: not x.name.startswith("IO") and not x.name.startswith("CPU"), threads))
  79. emulator_threads.sort(key=lambda t: t.name)
  80. print(f"=== Emulator threads: {emulator_threads}")
  81.  
  82. if (len(cpu_threads) + len(io_threads)) >= len(vm_cpus):
  83.     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)")
  84.     exit()
  85.  
  86. # Isolate cpus
  87. process = subprocess.Popen(f"cset set -c {host_cpus_string} -s system -m 1".split(), stdout=subprocess.PIPE, universal_newlines=True)
  88. cout, cerr = process.communicate()
  89. process = subprocess.Popen(f"cset proc -m -f root -t system".split(), stdout=subprocess.PIPE, universal_newlines=True)
  90. cout, cerr = process.communicate()
  91. print("Moved all tasks to 'system' cpuset")
  92.  
  93. # Pin threads
  94. cpu_id = 0
  95. cpu_used_for_non_emulator_threads = []
  96. for t in cpu_threads:
  97.     pin(t.name, vm_cpus[cpu_id], t.pid)
  98.     cpu_used_for_non_emulator_threads.append(vm_cpus[cpu_id])
  99.     cpu_id = cpu_id + 1
  100.  
  101. for t in io_threads:
  102.     pin(t.name, vm_cpus[cpu_id], t.pid)
  103.     # cpu_used_for_non_emulator_threads.append(vm_cpus[cpu_id])
  104.     cpu_id = cpu_id + 1
  105.  
  106. pin("emulator", ','.join(str(c) for c in vm_cpus[cpu_id:]), ",".join(str(x.pid) for x in emulator_threads))
  107.  
  108. # IRQ affinity
  109. irq_cpu_affinity_list = ",".join(str(x) for x in cpu_used_for_non_emulator_threads)
  110. process = subprocess.Popen(['/bin/bash'], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
  111. 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'))
  112. print(f"Set IRQ affinity to {irq_cpu_affinity_list}")
RAW Paste Data