affinity_helper.py

#-------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.
#--------------------------------------------------------------------------

# Get/Set cpu affinity. Currently only support part of Unix system
import logging
import os

logger = logging.getLogger(__name__)


class AffinitySetting():
    def __init__(self):
        self.pid = os.getpid()
        self.affinity = None
        self.is_os_supported = hasattr(os, 'sched_getaffinity') and hasattr(os, 'sched_setaffinity')
        if not self.is_os_supported:
            logger.warning("Current OS does not support os.get_affinity() and os.set_affinity()")

    def get_affinity(self):
        if self.is_os_supported:
            self.affinity = os.sched_getaffinity(self.pid)

    def set_affinity(self):
        if self.is_os_supported:
            current_affinity = os.sched_getaffinity(self.pid)
            if (self.affinity != current_affinity):
                logger.warning("Replacing affinity setting %s with %s", str(current_affinity), str(self.affinity))
                os.sched_setaffinity(self.pid, self.affinity)


if __name__ == '__main__':
    affi_helper = AffinitySetting()
    affi_helper.get_affinity()
    affi_helper.set_affinity()