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('-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}'], 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":
    pass