77 lines
2.8 KiB
Python
Executable File
77 lines
2.8 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
#============================================================================
|
|
# - really simple script, which just prints out the numactl cmd to
|
|
# prefix before your actual command. it determines this based on free
|
|
# memory size attached to every node.
|
|
# - when you run this on a machine without `numactl`, the output is empty,
|
|
# so `$(numa_prefix) <cmd> <args>` turns in to `<cmd> <args>`.
|
|
# - when the machine has `numactl` installed, regardless of the socket-count
|
|
# on the machine, the resulting command is:
|
|
# `numactl -m <socket> -C <core-id list> -- <cmd> <args>`
|
|
# - example output from `numactl -H` on a 2 socket machine:
|
|
# available: 2 nodes (0-1)
|
|
# node 0 cpus: 0 2 4 6 8 10 12 14 16 18 20 22
|
|
# node 0 size: 131026 MB
|
|
# node 0 free: 7934 MB
|
|
# node 1 cpus: 1 3 5 7 9 11 13 15 17 19 21 23
|
|
# node 1 size: 65536 MB
|
|
# node 1 free: 429 MB
|
|
# node distances:
|
|
# node 0 1
|
|
# 0: 10 20
|
|
# 1: 20 10
|
|
#============================================================================
|
|
|
|
import subprocess
|
|
import re
|
|
import sys
|
|
|
|
which_proc = subprocess.Popen(["which", "numactl"], stdout=subprocess.PIPE)
|
|
out, err = which_proc.communicate()
|
|
|
|
if out != "":
|
|
numactl_proc = subprocess.Popen(["numactl", "-H"], stdout=subprocess.PIPE)
|
|
out, err = numactl_proc.communicate()
|
|
|
|
lines = out.split("\n")
|
|
line_idx = 0
|
|
|
|
head_line = lines[line_idx]
|
|
line_idx += 1
|
|
node_match = re.match(r"^ *available: +(\d+) nodes", head_line)
|
|
if node_match:
|
|
avail_nodes = node_match.group(1)
|
|
best_node_id = ""
|
|
best_cpus = ""
|
|
best_free_size = 0
|
|
|
|
# loop through available nodes, selecting the node with the most free mem
|
|
for i in avail_nodes:
|
|
cpu_line = lines[line_idx]
|
|
mem_size_line = lines[line_idx + 1]
|
|
mem_free_line = lines[line_idx + 2]
|
|
line_idx += 3
|
|
|
|
cpu_match = re.match(r"^ *node (\d+) cpus: (\d.*\d)$", cpu_line)
|
|
if cpu_match:
|
|
node_id = cpu_match.group(1)
|
|
cpus = cpu_match.group(2).replace(" ", ",")
|
|
|
|
mem_free_match = re.match(r"^ *node " + node_id + " free: (\d+) \S+$", mem_free_line)
|
|
if mem_free_match:
|
|
free_size = mem_free_match.group(1)
|
|
if int(free_size) > int(best_free_size):
|
|
best_node_id = node_id
|
|
best_cpus = cpus
|
|
best_free_size = free_size
|
|
else:
|
|
sys.exit("[ERROR] Malformed mem free line: " + mem_free_line)
|
|
|
|
else:
|
|
sys.exit("[ERROR] Malformed cpus line: " + cpu_line)
|
|
|
|
sys.stdout.write("numactl -m " + best_node_id + " -C " + best_cpus + " --")
|
|
else:
|
|
sys.exit("[ERROR] Malformed head line: " + head_line)
|