use tempfile, fix sys problem, remove unsless configs

Author: roll-away

graph_net/test/graph_decompose_and_look_for_fully_fusable_subgraph_test.sh

@@ -13,21 +13,16 @@ decorator_config_json_str=$(cat <<EOF
         "name": "$MODEL_NAME",
         "custom_extractor_path": "$GRAPH_NET_ROOT/torch/fully_fusable_subgraph_extractor.py",
         "custom_extractor_config": {
-            "output_dir": "/tmp/naive_decompose_workspace",
             "split_positions": [],
             "group_head_and_tail": true,
-            "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
-            "filter_config": {},
-            "post_extract_process_path":"$GRAPH_NET_ROOT/torch/post_extract_process_count_kernels.py",
-            "post_extract_process_class_name": "GraphFullyFusable",
             "max_step": 5,
             "min_step": 2,
-            "max_nodes": 32
+            "max_nodes": 6
         }
     }
 }
 EOF
 )
 DECORATOR_CONFIG=$(echo $decorator_config_json_str | base64 -w 0)
 
-python3 -m graph_net.torch.run_model --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --decorator-config=$DECORATOR_CONFIG
+python3 -m graph_net.torch.run_model --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --decorator-config=$DECORATOR_CONFIG

graph_net/torch/fully_fusable_subgraph_extractor.py

@@ -1,7 +1,8 @@
 import os
 import torch
-import sys
 import graph_net
+import tempfile
+from graph_net.torch import constraint_util
 
 
 class GraphExtractor:
@@ -25,11 +26,6 @@ def make_config(
         split_positions=(),
         group_head_and_tail=False,
         chain_style=False,
-        output_dir="./tmp/naive_decomposer_dir",
-        filter_path=None,
-        filter_config=None,
-        post_extract_process_path=None,
-        post_extract_process_class_name=None,
         max_step=8,
         min_step=2,
         max_nodes=32,
@@ -42,40 +38,38 @@ def make_config(
             "split_positions": split_positions,
             "group_head_and_tail": group_head_and_tail,
             "chain_style": chain_style,
-            "output_dir": output_dir,
-            "filter_path": filter_path,
-            "filter_config": filter_config if filter_config is not None else {},
-            "post_extract_process_path": post_extract_process_path,
-            "post_extract_process_class_name": post_extract_process_class_name,
             "max_step": max_step,
             "min_step": min_step,
             "max_nodes": max_nodes,
         }
 
     def _get_sub_ranges(self):
-        kMinLenOps = self.config["min_step"]
-        num_ops = self.config["max_nodes"]
-        for length in reversed(range(kMinLenOps, num_ops)):
-            for start_pos in range(num_ops - length):
-                end_pos = start_pos + length
+        for step in reversed(
+            range(self.config["min_step"], self.config["max_step"] + 1)
+        ):
+            for start_pos in range(self.config["max_nodes"] - step):
+                end_pos = start_pos + step
                 yield start_pos, end_pos
 
     def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
-        import json
-        import base64
-
+        temp_dir_obj = tempfile.TemporaryDirectory(prefix="_check_fusable_")
+        temp_output_dir = temp_dir_obj.name
+        found_fusable_subgraph = False
+        print(f"Using temp output dir: {temp_output_dir}")
         for start_pos, end_pos in self._get_sub_ranges():
             self.config["split_positions"] = [start_pos, end_pos]
             print("current split_positions:", self.config["split_positions"])
             graph_net_root = os.path.dirname(graph_net.__file__)
-            model_path = f"{graph_net_root}/../samples//timm/{self.name}"
+            model_path = os.path.join(
+                graph_net_root, "..", "samples", "timm", self.name
+            )
             check_fusable_config = {
                 "decorator_path": f"{graph_net_root}/torch/extractor.py",
                 "decorator_config": {
                     "name": f"{self.name}",
                     "custom_extractor_path": f"{graph_net_root}/torch/naive_graph_decomposer.py",
                     "custom_extractor_config": {
-                        "output_dir": "/tmp/naive_decompose_workspace",
+                        "output_dir": temp_output_dir,
                         "split_positions": self.config["split_positions"],
                         "group_head_and_tail": False,
                         "filter_path": f"{graph_net_root}/torch/naive_subgraph_filter.py",
@@ -85,15 +79,20 @@ def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
                     },
                 },
             }
-            json_string = json.dumps(check_fusable_config)
-            json_bytes = json_string.encode("utf-8")
-            b64_encoded_bytes = base64.b64encode(json_bytes)
-            checker_config = b64_encoded_bytes.decode("utf-8")
-            cmd = f"{sys.executable} -m graph_net.torch.run_model --model-path {model_path} --decorator-config '{checker_config}'"
-            res_code = os.system(cmd)
-            if res_code == 0:
-                print("find the biggest fully fusable subgraph")
+            success = constraint_util.RunModelPredicator(check_fusable_config)(
+                model_path
+            )
+            if success:
+                found_fusable_subgraph = True
+                temp_dir_obj.cleanup = lambda: None
+                print(
+                    f"SUCCESS in finding the biggest fully fusable subgraph saved in: {temp_output_dir}."
+                )
                 break
             else:
+                print("Failed attempt. clean up the workspace and continue the search.")
+                temp_dir_obj.cleanup()
                 continue
+        if not found_fusable_subgraph:
+            print("No fusable subgraph found")
         return gm

graph_net/torch/naive_graph_decomposer.py

@@ -1,7 +1,5 @@
 import os
 import torch
-import sys
-import shutil
 from graph_net.torch.decompose_util import convert_to_submodules_graph
 from graph_net.torch.extractor import GraphExtractor as BuiltinGraphExtractor
 import graph_net.imp_util as imp_util
@@ -104,15 +102,7 @@ def _post_extract_process(self):
         model_path = os.path.join(
             self.parent_graph_extractor.config["output_dir"], self.model_name
         )
-        fully_fusable = self.post_extract_process(model_path)
-        if fully_fusable:
-            print(f"{model_path} is the biggest fully fusable subgraph!")
-            sys.exit(0)
-        else:
-            # remove if not fully fusable
-            shutil.rmtree(model_path)
-            print(f"remove: {model_path}")
-            sys.exit(1)
+        return self.post_extract_process(model_path)
 
     def make_filter(self, config):
         if config["filter_path"] is None:

graph_net/torch/post_extract_process_count_kernels.py

@@ -1,6 +1,7 @@
 from graph_net.torch import utils
 import importlib.util
 import torch
+import sys
 from typing import Type
 from torch.profiler import profile, record_function, ProfilerActivity
 
@@ -12,7 +13,7 @@ def __init__(self, config):
     def __call__(self, model_path=None):
         torch._dynamo.reset()
         if model_path is None:
-            return False
+            sys.exit(1)
         # model
         model_class = load_class_from_file(
             f"{model_path}/model.py", class_name="GraphModule"
@@ -30,20 +31,20 @@ def __call__(self, model_path=None):
             model(**state_dict)
         except Exception as e:
             print(f"failed in running model:{e}")
-            return False
+            sys.exit(1)
         # try to compile the model
         try:
             compiled_model = torch.compile(model)
         except Exception as e:
             print(f"failed in compiling model:{e}")
-            return False
+            sys.exit(1)
         compiled_num_of_kernels = count_kernels(compiled_model, state_dict)
         if compiled_num_of_kernels == 1:
             print(model_path, "can be fully integrated!!!!!!!!!!!")
-            return True
+            sys.exit(0)
         else:
             print(f"{model_path} can not be fully integrated, to be removed...")
-            return False
+            sys.exit(1)
 
 
 def load_class_from_file(file_path: str, class_name: str) -> Type[torch.nn.Module]: