Author eryksun
Recipients davin, eryksun, jtaylor, neologix
Date 2015-02-28.06:00:51
SpamBayes Score -1.0
Marked as misclassified Yes
Message-id <1425103251.84.0.769772130826.issue23530@psf.upfronthosting.co.za>
In-reply-to
Content
> Well, we already expose CPU affinity:
> 
> >>> import os
> >>> os.sched_getaffinity(0)
> {0}

os.sched_getaffinity only exists on some POSIX systems, such as Linux. For Windows, here's a ctypes version of sched_getaffinity and sched_setaffinity:

    import sys
    from ctypes import *
    from ctypes.wintypes import *

    __all__ = ['sched_getaffinity', 'sched_setaffinity']

    kernel32 = WinDLL('kernel32')

    DWORD_PTR = WPARAM
    PDWORD_PTR = POINTER(DWORD_PTR)

    GetCurrentProcess = kernel32.GetCurrentProcess
    GetCurrentProcess.restype = HANDLE

    OpenProcess = kernel32.OpenProcess
    OpenProcess.restype = HANDLE
    OpenProcess.argtypes = (DWORD, # dwDesiredAccess,_In_
                            BOOL,  # bInheritHandle,_In_
                            DWORD) # dwProcessId, _In_

    GetProcessAffinityMask = kernel32.GetProcessAffinityMask
    GetProcessAffinityMask.argtypes = (
        HANDLE,     # hProcess, _In_
        PDWORD_PTR, # lpProcessAffinityMask, _Out_
        PDWORD_PTR) # lpSystemAffinityMask, _Out_

    SetProcessAffinityMask = kernel32.SetProcessAffinityMask
    SetProcessAffinityMask.argtypes = (
      HANDLE,    # hProcess, _In_
      DWORD_PTR) # dwProcessAffinityMask, _In_

    PROCESS_SET_INFORMATION = 0x0200
    PROCESS_QUERY_INFORMATION = 0x0400
    PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
    if sys.getwindowsversion().major < 6:
        PROCESS_QUERY_LIMITED_INFORMATION = PROCESS_QUERY_INFORMATION

    def sched_getaffinity(pid):
        if pid == 0:
            hProcess = GetCurrentProcess()
        else:
            hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION,
                                   False, pid)
        if not hProcess:
            raise WinError()
        lpProcessAffinityMask = DWORD_PTR()
        lpSystemAffinityMask = DWORD_PTR()
        if not GetProcessAffinityMask(hProcess, 
                                      byref(lpProcessAffinityMask),
                                      byref(lpSystemAffinityMask)):
            raise WinError()
        mask = lpProcessAffinityMask.value
        return {c for c in range(sizeof(DWORD_PTR) * 8) if (1 << c) & mask}

    def sched_setaffinity(pid, mask):
        if pid == 0:
            hProcess = GetCurrentProcess()
        else:
            hProcess = OpenProcess(PROCESS_SET_INFORMATION,
                                   False, pid)
        if not hProcess:
            raise WinError()
        bitmask = 0
        for cpu in mask:
            if not isinstance(cpu, int):
                raise TypeError('expected an iterator of ints, but '
                                'iterator yielded %r' % type(cpu))
            if cpu < 0:
                raise ValueError('negative CPU number')
            if cpu >= sizeof(DWORD_PTR) * 8:
                raise ValueError('CPU number too large')
            bitmask |= 1 << cpu
        if not SetProcessAffinityMask(hProcess, bitmask):
            raise WinError()
History
Date User Action Args
2015-02-28 06:00:51eryksunsetrecipients: + eryksun, neologix, jtaylor, davin
2015-02-28 06:00:51eryksunsetmessageid: <1425103251.84.0.769772130826.issue23530@psf.upfronthosting.co.za>
2015-02-28 06:00:51eryksunlinkissue23530 messages
2015-02-28 06:00:51eryksuncreate