ampere_research/pytorch/batch.py
2024-12-02 23:32:33 -05:00

74 lines
2.4 KiB
Python
Executable File

#! /bin/python3
import argparse
import glob
import os
import subprocess
parser = argparse.ArgumentParser()
parser.add_argument('arch')
parser.add_argument('output_dir')
parser.add_argument('matrix_dir')
parser.add_argument('iterations', type=int)
parser.add_argument('baseline_time_s', type=int)
parser.add_argument('baseline_delay_s', type=int)
parser.add_argument('--perf', action='store_const', const='--perf')
parser.add_argument('--power', action='store_const', const='--power')
parser.add_argument('--distribute', type=bool)
args = parser.parse_args()
srun_args_altra = [
'--account', 'oasis',
'--partition', 'oasis',
'--qos', 'oasis-exempt',
#'--cpus-per-task 160',
'--cpus-per-task', '160',
#'--mem 28114',
'--mem', '16G',
'--ntasks-per-node', '1'#,
#'--exclusive',
#'--output', '/dev/null',
#'--error', '/dev/null'
]
def srun(srun_args_list: list, run_args, matrix_file: str) -> list:
run_args_list = [
args.arch,
matrix_file,
str(args.iterations),
str(args.baseline_time_s),
str(args.baseline_delay_s)]
if args.perf is not None:
run_args_list += [args.perf]
if args.power is not None:
run_args_list += [args.power]
return ['srun'] + srun_args_list + ['run.py'] + run_args_list
for i, matrix in enumerate(glob.glob(f'{args.matrix_dir.rstrip("/")}/*.mtx')):
if args.arch == 'altra':
if args.distribute == True:
i = i % 40
srun_args_altra += [f'--nodelist oasis{i:02}']
output_filename = '_'.join([
args.arch,
str(args.iterations),
os.path.splitext(os.path.basename(matrix))[0],
str(args.baseline_time_s),
str(args.baseline_delay_s)])
json_filepath = f'{args.output_dir.rstrip("/")}/{output_filename}.json'
raw_filepath = f'{args.output_dir.rstrip("/")}/{output_filename}.output'
with open(json_filepath, 'w') as json_file, open(raw_filepath, 'w') as raw_file:
print(srun(srun_args_altra, args, matrix))
proc = subprocess.run(
srun(srun_args_altra, args, matrix),
stdout=json_file,
stderr=raw_file,
text=True)
#output = proc.communicate()
#print(output[0])
#print(output[1])
break;