136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
from data_stat import Cpu, Format
|
|
|
|
import argparse
|
|
import glob
|
|
import os
|
|
import subprocess
|
|
import random
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('cpu', choices=[x.name.lower() for x in Cpu])
|
|
parser.add_argument('output_dir')
|
|
parser.add_argument('matrix_dir')
|
|
parser.add_argument('format', type=str,
|
|
choices=[fmt.name.lower() for fmt in Format])
|
|
parser.add_argument('base_iterations', type=int)
|
|
parser.add_argument('min_time_s', 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', action='store_true')
|
|
args = parser.parse_args()
|
|
args.cpu = Cpu[args.cpu.upper()]
|
|
args.format = Format[args.format.upper()]
|
|
|
|
srun_args = {
|
|
Cpu.ALTRA: [
|
|
'--account', 'oasis',
|
|
'--partition', 'oasis',
|
|
'--qos', 'oasis-exempt',
|
|
#'--cpus-per-task 160',
|
|
'--cpus-per-task', '160',
|
|
#'--mem 28114',
|
|
'--mem', '16G',
|
|
'--ntasks-per-node', '1',
|
|
'--time', '1-00:00:00'
|
|
#'--exclusive',
|
|
#'--output', '/dev/null',
|
|
#'--error', '/dev/null'
|
|
],
|
|
Cpu.EPYC_7313P: [
|
|
'--account', 'nexus',
|
|
'--partition', 'tron',
|
|
'--qos', 'high',
|
|
'--cpus-per-task', '16',
|
|
'--ntasks-per-node', '1',
|
|
'--time', '1-00:00:00',
|
|
'--prefer', 'EPYC-7313P'
|
|
],
|
|
Cpu.XEON_4216: [
|
|
'--account', 'nexus',
|
|
'--partition', 'tron',
|
|
'--qos', 'tron-exempt',
|
|
'--cpus-per-task', '32',
|
|
'--ntasks-per-node', '1',
|
|
'--time', '1-00:00:00',
|
|
'--prefer', 'Xeon,4216'
|
|
]
|
|
}
|
|
python = {
|
|
Cpu.ALTRA: 'python3',
|
|
Cpu.EPYC_7313P: 'python3.11',
|
|
Cpu.XEON_4216: 'python3.11'
|
|
}
|
|
|
|
def run(run_args, matrix_file: str, srun_args_list: list = None) -> list:
|
|
run_args_list = [
|
|
args.cpu.name.lower(),
|
|
matrix_file,
|
|
args.format.name.lower(),
|
|
str(args.base_iterations),
|
|
str(args.min_time_s),
|
|
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]
|
|
|
|
if srun_args_list is None:
|
|
command = [python[args.cpu], 'run.py'] + run_args_list
|
|
else:
|
|
command = ['srun'] + srun_args_list + [python[args.cpu], 'run.py'] + run_args_list
|
|
|
|
print(command)
|
|
return command
|
|
|
|
processes = list()
|
|
|
|
for i, matrix in enumerate(glob.glob(f'{args.matrix_dir.rstrip("/")}/*.mtx')):
|
|
#if args.distribute:
|
|
# if args.cpu == Cpu.ALTRA:
|
|
# i = i % 40
|
|
# srun_args_temp = srun_args[args.cpu] + ['--nodelist', f'oasis{i:02}']
|
|
# elif args.cpu == Cpu.EPYC_7313P:
|
|
# srun_args_temp = srun_args[args.cpu]
|
|
#else:
|
|
srun_args_temp = srun_args[args.cpu]
|
|
|
|
output_filename = '_'.join([
|
|
args.cpu.name.lower(),
|
|
str(args.min_time_s),
|
|
str(args.baseline_time_s),
|
|
str(args.baseline_delay_s),
|
|
os.path.splitext(os.path.basename(matrix))[0]])
|
|
|
|
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(json_filepath)
|
|
print(raw_filepath)
|
|
|
|
if args.distribute:
|
|
processes.append(subprocess.Popen(
|
|
run(args, matrix, srun_args_temp),
|
|
stdout=json_file,
|
|
stderr=raw_file))
|
|
else:
|
|
subprocess.run(
|
|
run(args, matrix),
|
|
stdout=json_file,
|
|
stderr=raw_file)
|
|
|
|
if args.distribute:
|
|
# Wait on every 10 jobs to avoid socket timeout.
|
|
if i % 10 == 9:
|
|
print("Waiting on 10 jobs")
|
|
for process in processes:
|
|
process.wait()
|
|
|
|
processes = list()
|
|
|
|
if args.distribute:
|
|
for process in processes:
|
|
process.wait()
|