def assert_device_isolation()

in optimum_benchmark/launchers/device_isolation_utils.py [0:0]
36 lines of code
13 McCabe index (conditional complexity)

def assert_device_isolation(pid: int, device_ids: str, action: str):
    log_level = os.environ.get("LOG_LEVEL", "INFO")
    log_to_file = os.environ.get("LOG_TO_FILE", "1") == "1"
    setup_logging(log_level, to_file=log_to_file, prefix="DEVICE-ISOLATION-PROCESS")

    device_isolation_pid = os.getpid()
    permitted_parent_pids = {pid, device_isolation_pid}

    while any(psutil.pid_exists(p) for p in permitted_parent_pids):
        device_pids = get_pids_running_on_system_devices(device_ids=device_ids)
        device_pids = {p for p in device_pids if psutil.pid_exists(p)}

        permitted_children_pids = set()
        for pid in permitted_parent_pids:
            permitted_children_pids |= get_children_pids(pid)

        permitted_pids = permitted_parent_pids | permitted_children_pids
        permitted_pids = {p for p in permitted_pids if psutil.pid_exists(p)}

        non_permitted_pids = device_pids - permitted_pids

        if len(non_permitted_pids) > 0:
            LOGGER.warn(
                f"Found process(es) [{non_permitted_pids}] running on device(s) [{device_ids}], "
                f"other than the isolated process [{pid}], the device isolation process [{device_isolation_pid}] "
                f"and their children [{permitted_children_pids}]."
            )

            if action == "warn":
                LOGGER.warn("Make sure no other process is running on the device(s) while benchmarking.")
            elif action == "error":
                LOGGER.error("Signaling the isolated process to error out.")
                if sys.platform == "linux":
                    os.kill(pid, signal.SIGUSR1)
                else:
                    LOGGER.error("Sending an error signal is only supported on Linux. Killing the isolated process.")
                    os.kill(pid, signal.SIGKILL)
            elif action == "kill":
                LOGGER.error("Killing the isolated process.")
                os.kill(pid, signal.SIGKILL)

            LOGGER.warn("Exiting device isolation process.")
            exit(0)

        time.sleep(1)