|
|
import json |
|
|
from tqdm import tqdm |
|
|
from args_config import load_config |
|
|
from dataloaders.TritonBench import TritonBench |
|
|
|
|
|
|
|
|
def main(): |
|
|
args = load_config("configs/parallel_scaling_config.yaml") |
|
|
file_list = args.file_list |
|
|
output_file = args.output_file |
|
|
assert output_file.endswith(".jsonl"), f"expect output file to be a jsonl file, but got {file} instead" |
|
|
|
|
|
|
|
|
dataset = TritonBench(statis_path=args.statis_path, |
|
|
py_folder=args.py_folder, |
|
|
instruction_path=args.instruction_path, |
|
|
py_interpreter=args.py_interpreter, |
|
|
golden_metrics=args.golden_metrics, |
|
|
perf_ref_folder=args.perf_ref_folder, |
|
|
perf_G_path=args.perf_G_path) |
|
|
|
|
|
|
|
|
result_dict = {} |
|
|
call_num = [0 for _ in range(args.data_len)] |
|
|
exe_num = [0 for _ in range(args.data_len)] |
|
|
for file in file_list: |
|
|
assert file.endswith(".jsonl"), f"expect jsonl file, but got {file} instead" |
|
|
with open(file, "r") as f: |
|
|
lines = f.readlines() |
|
|
assert args.data_len == len(lines), f"expect {args.data_len} entries, but got {len(lines)} entries instead for file {file}" |
|
|
with tqdm(total=args.data_len) as pbar: |
|
|
for i, line in enumerate(lines): |
|
|
content = json.loads(line) |
|
|
filename = content["filename"] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
tmp_dir = "par_scaling_tmp" |
|
|
exe_dir = "par_scaling_exe" |
|
|
pass_call, pass_exe, call_stdout, call_stderr, exe_stdout, exe_stderr = dataset.test_opt_correctness(content["predict"], |
|
|
filename=filename, |
|
|
tmp_dir=tmp_dir, |
|
|
save_scripts=True, |
|
|
exe_dir=exe_dir) |
|
|
ms = None |
|
|
if pass_call: |
|
|
call_num[i] = 1 |
|
|
if pass_exe: |
|
|
exe_num[i] = 1 |
|
|
|
|
|
|
|
|
result_dir = os.path.join(exe_dir, "result") |
|
|
script_dir = os.path.join(exe_dir, "script") |
|
|
log_dir = os.path.join(exe_dir, "log") |
|
|
path_gen = os.path.join(result_dir, filename[:-3] + ".json") |
|
|
path_ref = os.path.join(dataset.perf_ref_folder, filename[:-3] + ".json") |
|
|
|
|
|
dataset.write_perf_file(input_folder_path=exe_dir, results_path=result_dir, tmp_dir=script_dir) |
|
|
dataset.run_perf_scripts(script_dir=script_dir, log_dir=log_dir) |
|
|
try: |
|
|
_, _, ms = dataset.calculate(path_gen=path_gen, path_ref=path_ref) |
|
|
except: |
|
|
pass |
|
|
|
|
|
if not filename in result_dict: |
|
|
result_dict[filename] = { |
|
|
"predict": content["predict"], |
|
|
"pass_call": pass_call, |
|
|
"pass_exe": pass_exe, |
|
|
"instruction": content["instruction"], |
|
|
"label": content["label"] if "label" in content else None, |
|
|
"latency": ms |
|
|
} |
|
|
else: |
|
|
if (not result_dict[filename]["pass_exe"]) and pass_exe: |
|
|
result_dict[filename]["predict"] = content["predict"] |
|
|
result_dict[filename]["pass_call"] = pass_call |
|
|
result_dict[filename]["pass_exe"] = pass_exe |
|
|
result_dict[filename]["latency"] = ms |
|
|
elif (not result_dict[filename]["pass_call"]) and pass_call: |
|
|
result_dict[filename]["predict"] = content["predict"] |
|
|
result_dict[filename]["pass_call"] = pass_call |
|
|
result_dict[filename]["pass_exe"] = pass_exe |
|
|
result_dict[filename]["latency"] = ms |
|
|
elif result_dict[filename]["pass_exe"] and ms: |
|
|
if result_dict[filename]["latency"] is None or result_dict[filename]["latency"] > ms: |
|
|
result_dict[filename]["predict"] = content["predict"] |
|
|
result_dict[filename]["latency"] = ms |
|
|
|
|
|
os.system(f'rm -rf {tmp_dir}') |
|
|
os.system(f'rm -rf {exe_dir}') |
|
|
pbar.update(1) |
|
|
|
|
|
call_acc = sum(call_num) / 184.0 |
|
|
exe_acc = sum(exe_num) / 184.0 |
|
|
|
|
|
print(f"call acc: {call_acc}") |
|
|
print(f"exe acc: {exe_acc}") |
|
|
|
|
|
with open(output_file, "w") as f: |
|
|
for filename, result in result_dict.items(): |
|
|
output = { |
|
|
"instruction": result["instruction"], |
|
|
"label": result["label"], |
|
|
"filename": filename, |
|
|
"predict": result["predict"] |
|
|
} |
|
|
f.write(json.dumps(output) + "\n") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
