ampere_research/pytorch/batch.py

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()