Win7 Zugehörigkeit Prozessorkerne

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# $Id: cpuaffinity.py 770 2011-10-19 21:00:15Z Stef $
# Copyright (C) 2011 Stefan Schwendeler ( kungpfui@gmail.com )

""" Set CPU affinity and process priority of a new process. """  

__version__ = u'$Revision: 770 $'  
__author__ = u'Stefan Schwendeler'  

import os, sys
import optparse
import multiprocessing
import time
from ctypes import *


check_interval	= 5.0
validation_nb	= 12


# some constants of windows os
priorities = {
	'low':			0x00000040,  
	'below_normal':	0x00004000,  
	'normal':		0x00000020,  
	'above_normal':	0x00008000,  
	'high':			0x00000080,  
	'realtime':		0x00000100,  
}


BYTE      = c_ubyte
WORD      = c_ushort
DWORD     = c_ulong
DWORD_PTR = c_ulong if sizeof(c_void_p) == 4 else c_ulonglong
LPBYTE    = POINTER(c_ubyte)
LPTSTR    = POINTER(c_char)
HANDLE    = c_void_p
PVOID     = c_void_p
LPVOID    = c_void_p

class STARTUPINFO(Structure):
    _fields_ = [("cb",            DWORD),  
                ("lpReserved",    LPTSTR),  
                ("lpDesktop",     LPTSTR),  
                ("lpTitle",       LPTSTR),  
                ("dwX",           DWORD),  
                ("dwY",           DWORD),  
                ("dwXSize",       DWORD),  
                ("dwYSize",       DWORD),  
                ("dwXCountChars", DWORD),  
                ("dwYCountChars", DWORD),  
                ("dwFillAttribute",DWORD),  
                ("dwFlags",       DWORD),  
                ("wShowWindow",   WORD),  
                ("cbReserved2",   WORD),  
                ("lpReserved2",   LPBYTE),  
                ("hStdInput",     HANDLE),  
                ("hStdOutput",    HANDLE),  
                ("hStdError",     HANDLE),]  

class PROCESS_INFORMATION(Structure):
    _fields_ = [("hProcess",    HANDLE),  
                ("hThread",     HANDLE),  
                ("dwProcessId", DWORD),  
                ("dwThreadId",  DWORD),]  


DETACHED_PROCESS = 0x00000008


def main():
	""" the main routine :-)"""  
	usage = "usage: %prog [options] <program path> [<arg 1> <arg n>]"  
	parser = optparse.OptionParser(usage=usage)
	parser.disable_interspersed_args()
	parser.add_option("-c", "--cpuid", default="0", dest="cpuid", help="list of comma separated cpu IDs. E.g. 0,3 or 4,5")  
	parser.add_option("-p", "--priority", default="normal", type="choice", choices=priorities.keys(), dest="priority", help="process priority class. [low, normal, high, realtime] - Default: %default")  
	options, args = parser.parse_args()

	if len(args) < 1: parser.error("incorrect number of arguments")  

	# check cpu IDs
	try:
		cpu_ids = [ int(id) for id in options.cpuid.split(',') ]  
	except:
		parser.error("incorrect cpu id parameter")  

	cpu_count = multiprocessing.cpu_count()
	if not all([ id < cpu_count for id in cpu_ids ]):
		parser.error("invalid cpu id value")  

	# check program path
	if not os.path.exists(args) or not os.path.isfile(args):
		parser.error("incorrect program path")  

	cpu_affinity_mask = 0
	for id in cpu_ids:
		cpu_affinity_mask += (1 << id);

	# re-join the commands to a string
	command = ''  
	for arg in args:

		if ' ' in arg: arg = '"'+arg+'"'  
		if command != '': command += ' '  
		command += arg

	startup_dir = os.path.dirname(args) or None

	priority_class = priorities[options.priority]

	# start the process
	si = STARTUPINFO();
	si.cb = sizeof(si)
	pi = PROCESS_INFORMATION();
	if not windll.kernel32.CreateProcessA(None, command, None, None, False, DETACHED_PROCESS | priority_class, None, startup_dir, byref(si), byref(pi)):
		print("Can't create process.");  
		sys.exit(1)

	# wait some seconds
	time.sleep(check_interval)

	dwProcessMask = DWORD_PTR(0)
	dwSystemMask = DWORD_PTR(0)
	if windll.kernel32.GetProcessAffinityMask ( pi.hProcess,  byref(dwProcessMask), byref(dwSystemMask) ) \
		and (dwSystemMask.value & cpu_affinity_mask):

		validations = 0

		#~ print dwSystemMask.value, dwProcessMask.value, cpu_affinity_mask
		while validations < validation_nb:
			if not windll.kernel32.SetProcessAffinityMask( pi.hProcess, DWORD_PTR(dwSystemMask.value & cpu_affinity_mask) ):
				error_code = windll.kernel32.GetLastError()
				print("Error [%d]: Can't set process affinity mask.", error_code)  
				break

			time.sleep(check_interval)
			if not windll.kernel32.GetProcessAffinityMask ( pi.hProcess, byref(dwProcessMask), byref(dwSystemMask) ):
				print("Error Can't set process affinity mask.")  
				break
			#~ print dwSystemMask.value, dwProcessMask.value, cpu_affinity_mask
			if ( dwSystemMask.value & cpu_affinity_mask ) != dwProcessMask.value:
				validations = 0
			else:
				validations += 1

	else:
		print("Error: Can't set process affinity mask. CPU ID does not exist.")  

	windll.kernel32.CloseHandle(pi.hThread)
	windll.kernel32.CloseHandle(pi.hProcess)

	sys.exit(0)



if __name__ == "__main__":  
	main()