"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Wednesday, 18th Dec., 2024

    OBJECTIVE:

        To set the CPU affinity of your process. This makes the process prefer a specific set of CPU cores over others.
        Apparently, this makes cache access faster and provides other optimizations. Here's what ChatGPT had to say:

        "Setting CPU affinity in Python allows controlling which CPU cores a process or thread runs on,
        optimizing performance by improving cache locality, reducing context switching, and ensuring efficient resource
        allocation. It's beneficial in high-performance, real-time, and NUMA systems but may add overhead in general
        tasks." - ChatGPT

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# System-level activities:
import psutil


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


def set_affinity(requested_cpus: list):

    """
    Sets the affinity of the current process to certain CPUs so that performance is boosted. The main factors that
    contribute to gains are cache-locality, reduced context switching, and effective resource management.
    :param requested_cpus: The array of integers of which CPU cores are preferred.
    :return: None.
    """

    # Get the number of available CPUs:
    num_cpus = psutil.cpu_count()

    # Wrap around logic for when a core has been request that doesn't exist on this machine.
    # This is useful in cases like developing on a local machine with just 4 cores, but your server has dozens of cores.
    valid_cpus = [cpu % num_cpus for cpu in requested_cpus]

    # Set the CPU affinity:
    psutil.Process().cpu_affinity(valid_cpus)


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    pass