ampere_research/pytorch/power.py

#! /bin/python3

# Why does this program continuously print instead of accumulating the results
# and printing them at the end? Because of indefinite program runtime!

# For the sake of simplicity, this script will only return the power info
# for the first socket of Ampere Altra nodes.

import argparse, subprocess
import sys
parser = argparse.ArgumentParser()
parser.add_argument('program', nargs='*')
parser.add_argument('-s', '--seconds', type=int)
args = parser.parse_args()

arch = subprocess.run(['uname', '-m'], stdout=subprocess.PIPE, text=True).stdout.strip()
if arch == "aarch64":
    import time
    upper_bound = args.seconds if args.seconds is not None else -1
    i = 0
    while i != upper_bound:
        proc = subprocess.Popen(['sensors'], stdout=subprocess.PIPE)
        proc = subprocess.Popen(
                ['awk', '/CPU power:/ {print $3; exit}'],
                stdin=proc.stdout,
                stdout=subprocess.PIPE,
                text=True)
        power = proc.communicate()[0].strip().split('\n')[0]
        print(power)

        i += 1
        time.sleep(1)
elif arch == "x86_64":
    if args.seconds is None:
        proc = subprocess.Popen(
                ['turbostat', '-s', 'PkgWatt'] + args.program,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE)
        proc = subprocess.Popen(
                ['sed', '-n', '/PkgWatt/{n;p}'],
                stdin=proc.stderr,
                stdout=subprocess.PIPE,
                text=True)
    else:
        proc = subprocess.Popen(
                ['turbostat', '-s', 'PkgWatt', '-n', str(args.seconds), '-i', '1'],
                stdout=subprocess.PIPE)
        proc = subprocess.Popen(
                ['sed', '-n', '/PkgWatt/{n;p}'],
                stdin=proc.stdout,
                stdout=subprocess.PIPE,
                text=True)
    power = proc.communicate()[0].strip().split('\n')
    print('\n'.join(power))