| Current File : //usr/bin/ptaskset |
#!/usr/bin/python3
# -*- python -*-
# -*- coding: utf-8 -*-
# Copyright (C) 2008 Red Hat Inc.
#
# Arnaldo Carvalho de Melo <acme@redhat.com>
#
# This application is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; version 2.
#
# This application is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
from __future__ import print_function
import os
import schedutils
import sys
def usage():
print('''ptaskset (python-schedutils)
usage: ptaskset [options] [mask | cpu-list] [pid | cmd [args...]]
set or get the affinity of a process
-p, --pid operate on existing given pid
-c, --cpu-list display and specify cpus in list format
-h, --help display this help''')
return
def hexbitmask(l):
hexbitmask = []
bit = 0
mask = 0
nr_entries = 1 << max(l)
for entry in range(nr_entries):
if entry in l:
mask |= (1 << bit)
bit += 1
if bit == 32:
bit = 0
hexbitmask.insert(0, mask)
mask = 0
if bit < 32 and mask != 0:
hexbitmask.insert(0, mask)
return hexbitmask
def find_last_bit(n, wordsize=32):
bit = wordsize - 1
while bit != 0:
if n & (1 << bit):
return bit
bit -= 1
return 0
def bitmasklist(line):
fields = line.strip().split(",")
bitmasklist = []
while int(fields[0], 16) == 0:
fields.pop(0)
if not fields:
return []
nr_entries = (len(fields) - 1) * 32 + find_last_bit(int(fields[0], 16)) + 1
entry = 0
for i in range(len(fields) - 1, -1, -1):
mask = int(fields[i], 16)
while mask != 0:
if mask & 1:
bitmasklist.append(entry)
mask >>= 1
entry += 1
if entry == nr_entries:
break
if entry == nr_entries:
break
return bitmasklist
def cpustring_to_list(cpustr):
"""Convert a string of numbers to an integer list.
Given a string of comma-separated numbers and number ranges,
return a simple sorted list of the integers it represents.
This function will throw exceptions for badly-formatted strings.
Returns a list of integers."""
fields = cpustr.strip().split(",")
cpu_list = []
for field in fields:
ends = field.split("-")
if len(ends) > 2:
raise "Syntax error"
if len(ends) == 2:
cpu_list += list(range(int(ends[0]), int(ends[1]) + 1))
else:
cpu_list += [int(ends[0])]
return list(set(cpu_list))
def show_settings(pid, when, cpu_list_mode):
affinity = schedutils.get_affinity(pid)
if cpu_list_mode:
mask = ",".join([str(a) for a in affinity])
else:
mask = ",".join(["%x" % a for a in hexbitmask(affinity)])
print("pid %d's %s affinity mask: %s" % (pid, when, mask))
def change_settings(pid, affinity, cpu_list_mode):
if cpu_list_mode:
try:
affinity = [int(a) for a in affinity.split(",")]
except:
affinity = cpustring_to_list(affinity)
else:
affinity = bitmasklist(affinity)
try:
schedutils.set_affinity(pid, affinity)
except SystemError as err:
print("sched_setaffinity: %s" % err[1])
print("failed to set pid %d's affinity" % pid)
def main():
args = sys.argv[1:]
if not args:
usage()
return
cpu_list_mode = False
while True:
o = args.pop(0)
if o in ("-h", "--help"):
usage()
return
elif o in ("-c", "--cpu-list"):
cpu_list_mode = True
elif o in ("-p", "--pid"):
if len(args) > 1:
affinity = args.pop(0)
pid = int(args.pop(0))
show_settings(pid, "current", cpu_list_mode)
change_settings(pid, affinity, cpu_list_mode)
show_settings(pid, "new", cpu_list_mode)
else:
pid = int(args.pop(0))
show_settings(pid, "current", cpu_list_mode)
return
else:
break
affinity = o
change_settings(0, affinity, cpu_list_mode)
os.execvp(args[0], args)
if __name__ == '__main__':
main()