smt_manager.py

#!/usr/bin/env python3

# smt-manager.py - a script for managing logical cores
# Python implementation of the original Perl script by Steven Barrett
# https://github.com/damentz/smt-manager

import argparse
import os
import subprocess
import sys

# This is the top folder where CPUs can be enumerated and more detail retrieved
SYS_CPU = "/sys/devices/system/cpu"
DEBUG = False


def get_cpu_indexes() -> list[int]:
    """
    Get a list of CPU indexes by reading the system CPU directory

    Returns:
        list[int]: A sorted list of CPU indexes
    """
    try:
        cpu_dirs = [
            d for d in os.listdir(SYS_CPU) if d.startswith("cpu") and d[3:].isdigit()
        ]
        cpu_indexes = [int(d[3:]) for d in cpu_dirs]
        return sorted(cpu_indexes)
    except OSError as e:
        sys.exit(f"Cannot open folder: {SYS_CPU}. Error: {e}")


def get_cpu_settings() -> dict[int, dict[str, str]]:
    """
    Get settings for all CPUs including core type and power state

    Returns:
        dict[int, dict[str, str]]: A dictionary mapping CPU indexes to their settings,
            where each setting is a dictionary with 'core_type' and 'power' keys
    """
    cpu_indexes = get_cpu_indexes()
    cpus = {}

    for cpu in cpu_indexes:
        siblings_file = f"{SYS_CPU}/cpu{cpu}/topology/thread_siblings_list"
        power_file = f"{SYS_CPU}/cpu{cpu}/online"

        cpu_settings = {"core_type": "unknown", "power": "offline"}

        # Populate core topology, primary / logical
        try:
            with open(siblings_file, "r") as f:
                siblings_line = f.readline().strip()
                # Handle both comma-separated and hyphen-separated formats
                if "," in siblings_line:
                    siblings = [int(s) for s in siblings_line.split(",")]
                elif "-" in siblings_line:
                    start, end = map(int, siblings_line.split("-"))
                    siblings = list(range(start, end + 1))
                else:
                    siblings = [int(siblings_line)]

                if cpu == siblings[0]:
                    cpu_settings["core_type"] = "primary"
                else:
                    cpu_settings["core_type"] = "logical"
        except (OSError, IOError):
            if DEBUG:
                print(f"[ERROR] Could not open: {siblings_file}")

        # Populate core status, online / offline
        try:
            # CPU0 is always online and doesn't have an 'online' file
            if cpu == 0:
                cpu_settings["power"] = "online"
            else:
                with open(power_file, "r") as f:
                    cpu_power = f.readline().strip()
                    if cpu_power == "1":
                        cpu_settings["power"] = "online"
        except (OSError, IOError):
            if DEBUG:
                print(f"[ERROR] Could not open: {power_file}, assuming online")
            cpu_settings["power"] = "online"

        cpus[cpu] = cpu_settings

    return cpus


def set_logical_cpus(power_state: str) -> bool:
    """
    Set all logical CPUs to the specified power state (online/offline)

    Args:
        power_state (str): The desired power state ('online' or 'offline')

    Returns:
        bool: True if any CPU state was changed, False otherwise
    """
    cpus = get_cpu_settings()
    state_changed = False
    changed_cpus = []

    for cpu in sorted(cpus.keys()):
        # Skip CPU0 as it can't be disabled
        if cpu == 0:
            continue

        if (
            cpus[cpu]["core_type"] == "logical" or cpus[cpu]["core_type"] == "unknown"
        ) and cpus[cpu]["power"] != power_state:
            power_file = f"{SYS_CPU}/cpu{cpu}/online"

            try:
                with open(power_file, "w") as f:
                    state_changed = True
                    print(f"Setting CPU {cpu} to {power_state} ... ", end="")

                    if power_state == "online":
                        f.write("1")
                    elif power_state == "offline":
                        f.write("0")

                    changed_cpus.append(cpu)
                    print("done!")
            except (OSError, IOError):
                print(
                    f"[ERROR] failed to open file for writing: {power_file}. Are you root?"
                )

    if state_changed:
        # Rebalance the interrupts after power state changes
        try:
            subprocess.run(["irqbalance", "--oneshot"], check=True)
        except (subprocess.SubprocessError, FileNotFoundError):
            print(
                "[ERROR] Failed to balance interrupts with 'irqbalance --oneshot', "
                "you may experience strange behavior.",
                file=sys.stderr,
            )

        print()

    return state_changed


def pretty_print_topology() -> None:
    """
    Print the current CPU topology in a readable format using only standard library
    """
    cpus = get_cpu_settings()

    # Get the maximum width needed for each column
    cpu_width = max(len(str(cpu)) for cpu in cpus.keys())
    type_width = max(len(cpus[cpu]["core_type"]) for cpu in cpus.keys())
    power_width = max(len(cpus[cpu]["power"]) for cpu in cpus.keys())

    # Add header width to the calculation
    cpu_width = max(cpu_width, len("CPU"))
    type_width = max(type_width, len("Core Type"))
    power_width = max(power_width, len("Power State"))

    # Calculate total table width for the title
    total_width = cpu_width + type_width + power_width + 8  # 8 for borders and padding

    # Print table title
    print("CPU Topology".center(total_width))
    print("-" * total_width)

    # Print header
    print(
        f" {'CPU':<{cpu_width}} | {'Core Type':<{type_width}} | {'Power State':<{power_width}} "
    )
    print(f" {'-' * cpu_width} | {'-' * type_width} | {'-' * power_width} ")

    # Print rows
    for cpu in sorted(cpus.keys()):
        print(
            f" {cpu:<{cpu_width}} | {cpus[cpu]['core_type']:<{type_width}} | {cpus[cpu]['power']:<{power_width}} "
        )

    print()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="View current status of CPU topology or set logical cores to offline or online.",
        epilog="This script provides details about whether each CPU is physical or logical. "
        "When provided an optional parameter, the logical CPUs can be enabled or disabled.",
    )

    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--online", action="store_true", help="Enables all logical CPU cores"
    )
    group.add_argument(
        "--offline", action="store_true", help="Disables all logical CPU cores"
    )
    parser.add_argument("--debug", action="store_true", help="Enable debug output")

    args = parser.parse_args()

    global DEBUG
    DEBUG = args.debug

    power_state = None
    if args.online:
        power_state = "online"
    elif args.offline:
        power_state = "offline"

    pretty_print_topology()

    if power_state and set_logical_cpus(power_state):
        # If there was a change, print the new state
        pretty_print_topology()


if __name__ == "__main__":
    # Check if running as root, which is required for changing CPU states
    if os.geteuid() != 0 and (
        len(sys.argv) > 1 and ("--online" in sys.argv or "--offline" in sys.argv)
    ):
        print("You need to have root privileges to change CPU states.")
        print("Please run the script with sudo or as root.")
        sys.exit(1)

    main()