pressure.py

#!/usr/bin/env python3

# Output looks like this:
# $ python3 /data/pressure.py --resource all
# cpu,cpu,cpu,cpu,cpu,cpu,cpu,cpu,memory,memory,memory,memory,memory,memory,memory,memory,io,io,io,io,io,io,io,io
# some,some,some,some,full,full,full,full,some,some,some,some,full,full,full,full,some,some,some,some,full,full,full,full
# avg10,avg60,avg300,total,avg10,avg60,avg300,total,avg10,avg60,avg300,total,avg10,avg60,avg300,total,avg10,avg60,avg300,total,avg10,avg60,avg300,total
# 9.24,2.86,0.65,2300327,0.0,0.0,0.0,0,0.0,0.0,0.0,0,0.0,0.0,0.0,0,0.09,0.02,0.0,59850,0.0,0.0,0.0,50820
# 8.11,2.86,0.67,2368262,0.0,0.0,0.0,0,0.0,0.0,0.0,0,0.0,0.0,0.0,0,0.08,0.02,0.0,59850,0.0,0.0,0.0,50820
# ...

# Or:
# $ python3 /data/pressure.py --resource cpu
# some,some,some,some,full,full,full,full
# avg10,avg60,avg300,total,avg10,avg60,avg300,total
# 0.6,2.15,0.74,2850020,0.0,0.0,0.0,0
# 0.6,2.15,0.74,2850392,0.0,0.0,0.0,0

# Can be read by pandas via:
# df = pd.read_csv("pressure-all.log", header=[0,1,2])
# Or:
# df = pd.read_csv("pressure-memory.log", header=[0,1])

# Source file and its format:
# $ cat /proc/pressure/memory
# some avg10=0.00 avg60=0.00 avg300=0.00 total=0
# full avg10=0.00 avg60=0.00 avg300=0.00 total=0

import dataclasses
import time
from itertools import chain
from pathlib import Path

import fire
from parse import parse

# from rich import print


def flatten(nl):
    return list(chain(*nl))


@dataclasses.dataclass
class Info:
    avg10: float
    avg60: float
    avg300: float
    total: int

    @staticmethod
    def from_str(s: str):
        return Info(*parse("avg10={:f} avg60={:f} avg300={:f} total={:d}", s).fixed)

    def header(self):
        return [[f.name for f in dataclasses.fields(self)]]

    def row(self):
        return [getattr(self, f.name) for f in dataclasses.fields(self)]


@dataclasses.dataclass
class Resource:
    some: Info
    full: Info

    @staticmethod
    def from_str(s: str):
        return Resource(*parse("some {:Info}\nfull {:Info}", s, dict(Info=Info.from_str)).fixed)

    @staticmethod
    def from_procfs(resource):
        return Resource.from_str(Path(f"/proc/pressure/{resource}").read_text().strip())

    def header(self):
        fields = dataclasses.fields(self)
        headers = [getattr(self, f.name).header() for f in fields]
        return [
            flatten([[f.name] * len(h[0]) for f, h in zip(fields, headers)]),
            *[flatten(line) for line in zip(*headers)]
        ]

    def row(self):
        return flatten([getattr(self, f.name).row() for f in dataclasses.fields(self)])


@dataclasses.dataclass
class Pressure:
    cpu: Resource
    memory: Resource
    io: Resource
    # irq: Resource

    @staticmethod
    def from_procfs():
        fields = dataclasses.fields(Pressure)
        return Pressure(*[Resource.from_procfs(f.name) for f in fields])

    def header(self):
        fields = dataclasses.fields(self)
        headers = [getattr(self, f.name).header() for f in fields]
        return [
            flatten([[f.name] * len(h[0]) for f, h in zip(fields, headers)]),
            *[flatten(line) for line in zip(*headers)]
        ]

    def row(self):
        return flatten([getattr(self, f.name).row() for f in dataclasses.fields(self)])


def main(delay: int = 1, resource="all"):
    show_header = True
    try:
        while True:
            if resource == "all":
                p = Pressure.from_procfs()
            else:
                p = Resource.from_procfs(resource)
            if show_header:
                show_header = False
                print("\n".join([",".join(line) for line in p.header()]))
            # flush to avoid broken pipe error
            print(",".join([str(f) for f in p.row()]), flush=True)
            time.sleep(delay)
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    fire.Fire(main)