You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
40 lines
1.4 KiB
40 lines
1.4 KiB
# -------------------------------------------------------------------------
|
|
# 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()
|