tuna/sysfs.py

# -*- python -*-
# -*- coding: utf-8 -*-
"""
classes for /sys/devices/system/cpu/
so we can get topology information and do CPU hotplug operations
"""

import os


class cpu:
    """
    class to query if a cpu is online or not
    and to bring a cpu online or to take it offline
    """

    def __init__(self, basedir, name):
        self.name = name
        self.dir = os.path.join(basedir, name)
        self.online = None
        self.reload()

    def __lt__(self, other):
        # Previously it is assumed that name[:3] == "cpu" and name[3].isdigit(), so:
        return int(self.name[3:]) < int(other.name[3:])

    def readfile(self, name):
        """ read file "name", typically the per cpu "online" file """
        try:
            with open(os.path.join(self.dir, name)) as f:
                value = f.readline().strip()
                return value
        except:
            raise

    def __repr__(self):
        return f'name={self.name}, online={self.online}, physical_package_id={self.physical_package_id}'

    def reload_online(self):
        """ read the "online" file """
        self.online = True
        try:
            self.online = self.readfile("online") == "1"
        except FileNotFoundError:
            # boot CPU, usually cpu0, can't be brought offline, so
            # lacks the file and non root users can't read. In both
            # cases assume CPU is online.
            pass

    def reload(self):
        """ Load or reload the values of this class """
        self.reload_online()
        if self.online:
            try:
                self.physical_package_id = self.readfile("topology/physical_package_id")
            except:
                self.physical_package_id = "0"
        else:
            self.physical_package_id = None

    def set_online(self, online=True):
        """ Turn a cpu on or off using the online file """
        try:
            with open(os.path.join(self.dir, "online"), "w") as f:
                f.write("1" if online else "0")
        except:
            pass

        self.reload_online()
        return online == self.online


class cpus:
    """ a class that creates a dictionary of cpu information keyed by cpu  """

    def __init__(self, basedir="/sys/devices/system/cpu"):
        self.basedir = basedir
        self.cpus = {}
        self.sockets = {}
        self.reload()
        self.nr_cpus = len(self.cpus)

    def __getitem__(self, key):
        return self.cpus[key]

    def keys(self):
        """ Returns a list of keys, for example a key could be cpu9 """
        return list(self.cpus.keys())

    def __contains__(self, key):
        return key in self.cpus

    def __repr__(self):
        return f'cpus={self.cpus}, sockets={self.sockets}, nr_cpus={self.nr_cpus}'

    def reload(self):
        """ load or reload the values of this class """
        sockets_to_sort = []
        for name in os.listdir(self.basedir):
            if name[:3] != "cpu" or not name[3].isdigit():
                continue

            if name in self.cpus:
                self.cpus[name].reload(self.basedir)
            else:
                c = cpu(self.basedir, name)
                self.cpus[name] = c
                try:
                    socket = c.physical_package_id
                except:
                    socket = "0"
                if socket in self.sockets:
                    self.sockets[socket].insert(0, c)
                else:
                    self.sockets[socket] = [c, ]

                if socket not in sockets_to_sort:
                    sockets_to_sort.append(socket)

        for socket in sockets_to_sort:
            self.sockets[socket].sort(key=lambda x: int(x.name[3:]))


if __name__ == '__main__':

    cpus = cpus()
    print(f'{cpus}')

    for socket in list(cpus.sockets.keys()):
        print("Socket %s" % socket)
        for c in cpus.sockets[socket]:
            print("  %s" % c.name)
            print("    online: %s" % c.online)
            c.set_online(False)
            print("    online: %s" % c.online)
            c.set_online()
            print("    online: %s" % c.online)