Add try_run for graph variable renamer (#491)

Author: TelGome

graph_net/sample_pass/ast_graph_variable_renamer.py

@@ -36,6 +36,7 @@ def declare_config(
         output_dir: str,
         device: str,
         resume: bool = False,
+        try_run: bool = False,
         limits_handled_models: int = None,
         data_input_predicator_filepath: str = None,
         data_input_predicator_class_name: str = None,
@@ -73,7 +74,8 @@ def resume(self, rel_model_path: str):
             )
             self._update_meta_file(temp_model_path, "weight_meta.py", rename_map)
             self._update_meta_file(temp_model_path, "input_meta.py", rename_map)
-            self._try_run(temp_model_path)
+            if self.config["try_run"]:
+                self._try_run(temp_model_path)
             shutil.copytree(temp_model_path, dst_model_path, dirs_exist_ok=True)
 
     def _get_input_and_weight_arg_names(self, graph_module, model_path):
@@ -133,7 +135,7 @@ def _update_meta_file(self, model_path, meta_filename, rename_map):
         meta_file.write_text(py_code)
 
     def _try_run(self, model_path):
-        (f"[AstGraphVariableRenamer] Try to run {model_path}")
+        print(f"[AstGraphVariableRenamer] Try to run {model_path}")
         assert self.model_runnable_predicator(
             model_path
         ), f"{model_path} is not a runnable model"

graph_net/test/ast_graph_variable_rename_test.sh

@@ -15,6 +15,7 @@ python3 -m graph_net.model_path_handler \
     "handler_config": {
         "device": "cuda",
         "resume": true,
+        "try_run": true,
         "model_path_prefix": "$GRAPH_NET_ROOT/",
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",
         "data_input_predicator_class_name": "NaiveDataInputPredicator",

graph_net/test/graph_variable_rename_test.sh

@@ -15,6 +15,7 @@ python3 -m graph_net.model_path_handler \
     "handler_config": {
         "device": "cuda",
         "resume": true,
+        "try_run": true,
         "model_path_prefix": "$GRAPH_NET_ROOT/",
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",
         "data_input_predicator_class_name": "NaiveDataInputPredicator",

graph_net/torch/graph_variable_renamer.py

@@ -39,6 +39,7 @@ def _make_model_runnable_predicator(self, config):
     def _make_config(
         self,
         resume: bool = False,
+        try_run: bool = False,
         data_input_predicator_filepath=None,
         model_runnable_predicator_filepath=None,
         output_dir="./tmp/graph_variable_renamer_dir",
@@ -58,6 +59,7 @@ def _make_config(
         return {
             "resume": resume,
             "output_dir": output_dir,
+            "try_run": try_run,
             "filter_path": filter_path,
             "filter_config": filter_config if filter_config is not None else {},
             "data_input_predicator_filepath": data_input_predicator_filepath,
@@ -94,7 +96,8 @@ def __call__(self, rel_model_path):
                 src_model_path, temp_model_path, rename_map
             )
             self._update_input_meta_py_file(src_model_path, temp_model_path, rename_map)
-            self._try_run(temp_model_path)
+            if self.config["try_run"]:
+                self._try_run(temp_model_path)
             shutil.copytree(temp_model_path, dst_model_path)
 
     def _try_run(self, model_path):