-
Notifications
You must be signed in to change notification settings - Fork 60
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
2561c33
commit 4deb5df
Showing
1 changed file
with
34 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,34 @@ | ||
#!/usr/bin/env python3 | ||
# -*- coding: utf-8 -*- | ||
import codecs | ||
import re | ||
import sys | ||
|
||
import torch_neuronx | ||
from torch_neuronx.parallel_compile.neuron_parallel_compile import LOGGER as torch_neuronx_logger | ||
from torch_neuronx.parallel_compile.neuron_parallel_compile import main | ||
|
||
|
||
def get_hlos_from_run_log(trial_run_log): | ||
# New graphs are detected by specific message matching key | ||
hlo_key = "Extracting graphs" | ||
new_hlo_list = [] | ||
with codecs.open(trial_run_log, "r", encoding="utf-8", errors="ignore") as f: | ||
for line in f.readlines(): | ||
# Move temporary MODULE_* files into workdir before checking if there are any | ||
# new graphs. In try_compilations, compile only new graphs (those without | ||
# corresponding neffs). | ||
if hlo_key in line: | ||
model_path = line.split("Extracting graphs (")[1].split(")")[0] | ||
new_hlo_list.append(model_path) | ||
|
||
format_str = "\n\t" | ||
torch_neuronx_logger.info(f"New graph list from script: {format_str.join(new_hlo_list)}") | ||
return new_hlo_list | ||
|
||
|
||
torch_neuronx.parallel_compile.neuron_parallel_compile.get_hlos_from_run_log = get_hlos_from_run_log | ||
|
||
if __name__ == "__main__": | ||
sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0]) | ||
sys.exit(main()) |