System info

sys_info.py

import os
import platform
import sys

from rich.console import Console
from rich.table import Table

FILE_READ_BUFFER_SIZE = 32 * 1024


def get_cpu_limit() -> str:
    try:
        with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as quota_file:
            cpu_quota = int(quota_file.read())
            # Not useful for AWS Batch based jobs as result is -1, but works on local linux systems
            print(f"[DEBUG] {cpu_quota=}")

        with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as period_file:
            cpu_period = int(period_file.read())
            print(f"[DEBUG] {cpu_period=}")

        if cpu_quota != -1:
            # Divide quota by period, you should get num of allotted CPU to the container, rounded down if fractional.
            container_cpus = int(cpu_quota / cpu_period)
        else:
            with open("/sys/fs/cgroup/cpu/cpu.shares") as shares_file:
                cpu_shares = int(shares_file.read().rstrip())
                print(f"[DEBUG] {cpu_shares=}")
                # For AWS, gives correct value * 1024.
                container_cpus = int(cpu_shares / 1024)

        return str(container_cpus)
    except (FileNotFoundError, FileExistsError):
        return "NA"


def get_cpu_affinity() -> str:
    try:
        cpu_affinity = len(os.sched_getaffinity(0))
        return str(cpu_affinity)
    except AttributeError:
        return "NA"


def cpu_count_cores() -> str:
    """Return the number of CPU cores in the system."""
    mapping = {}
    current_info = {}
    try:
        with open("/proc/cpuinfo", "rb", buffering=FILE_READ_BUFFER_SIZE) as info_file:
            for line in info_file:
                line = line.strip().lower()
                if not line:
                    # new section
                    try:
                        mapping[current_info[b"physical id"]] = current_info[b"cpu cores"]
                    except KeyError:
                        pass
                    current_info = {}
                else:
                    # ongoing section
                    if line.startswith((b"physical id", b"cpu cores")):
                        key, value = line.split(b"\t:", 1)
                        current_info[key] = int(value)

        result = sum(mapping.values())
        return str(result)
    except (FileNotFoundError, FileExistsError):
        return "NA"


def get_cpu_model() -> str:
    try:
        with open("/proc/cpuinfo") as info_file:
            for line in info_file:
                # Ignore the blank line separating the information between
                # details about two processing units
                if line.strip():
                    if line.rstrip("\n").startswith("model name"):
                        model_name = line.rstrip("\n").split(":")[1]
                        model = model_name
                        model = model.strip()
                        break
    except (FileNotFoundError, FileExistsError):
        return "NA"

    return model


def get_simple_python_version() -> str:
    return f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"


def print_info(uname: platform.uname_result):
    info_table = Table(title="General Info", padding=(0, 2), show_edge=True, show_lines=True, show_header=False)
    info_table.add_row("System", uname.system)
    info_table.add_row("Node Name", uname.node)
    info_table.add_row("Release", uname.release)
    info_table.add_row("Machine", uname.machine)
    info_table.add_row("Python", get_simple_python_version())

    cores_table = Table(padding=(0, 2), show_edge=True, show_lines=True, show_header=False)
    cores_table.add_row("Simple", str(os.cpu_count()))
    cores_table.add_row("Physical", str(cpu_count_cores()))
    cores_table.add_row("Available", str(get_cpu_affinity()))
    cores_table.add_row("Advance", str(get_cpu_limit()))

    cpu_table = Table(title="CPU Info", padding=(0, 2), show_edge=True, show_lines=True, show_header=False)
    cpu_table.add_row("Model", str(get_cpu_model()))
    cpu_table.add_row("Cores Count", cores_table)

    console = Console()
    console.print("", info_table, cpu_table)


if __name__ == "__main__":
    print_info(platform.uname())