codekingpro/portable-devtools

1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation.  All rights reserved.
3# Licensed under the MIT License.
4# --------------------------------------------------------------------------
5
6# Get/Set cpu affinity. Currently only support part of Unix system
7import logging
8import os
9
10logger = logging.getLogger(__name__)
11
12
13class AffinitySetting:
14    def __init__(self):
15        self.pid = os.getpid()
16        self.affinity = None
17        self.is_os_supported = hasattr(os, "sched_getaffinity") and hasattr(os, "sched_setaffinity")
18        if not self.is_os_supported:
19            logger.warning("Current OS does not support os.get_affinity() and os.set_affinity()")
20
21    def get_affinity(self):
22        if self.is_os_supported:
23            self.affinity = os.sched_getaffinity(self.pid)
24
25    def set_affinity(self):
26        if self.is_os_supported:
27            current_affinity = os.sched_getaffinity(self.pid)
28            if self.affinity != current_affinity:
29                logger.warning(
30                    "Replacing affinity setting %s with %s",
31                    str(current_affinity),
32                    str(self.affinity),
33                )
34                os.sched_setaffinity(self.pid, self.affinity)
35
36
37if __name__ == "__main__":
38    affi_helper = AffinitySetting()
39    affi_helper.get_affinity()
40    affi_helper.set_affinity()
41