Enable try-run the unittest when main is generated.

Author: Xreki

graph_net/test/agent_unittest_generator_test.sh

@@ -2,19 +2,18 @@
 
 ROOT_DIR="$(cd "$(dirname "$0")/../.." && pwd)"
 GRAPH_NET_ROOT=$(python -c "import graph_net, os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
-
-MODEL_PATH_PREFIX="$ROOT_DIR"
 OUTPUT_DIR="/tmp/agent_unittests"
 
 HANDLER_CONFIG=$(base64 -w 0 <<EOF
 {
     "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_passes/agent_unittest_generator.py",
     "handler_class_name": "AgentUnittestGeneratorPass",
     "handler_config": {
-        "model_path_prefix": "$MODEL_PATH_PREFIX",
+        "model_path_prefix": "${GRAPH_NET_ROOT}",
         "output_dir": "$OUTPUT_DIR",
         "device": "auto",
         "generate_main": true,
+        "try_run": true,
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/graph_net/torch/constraint_util.py",                                                                                     
         "data_input_predicator_class_name": "NaiveDataInputPredicator"
     }
@@ -24,14 +23,13 @@ EOF
 
 run_case() {
   local rel_sample_path="$1"
-  local name="$2"
-  echo "[AgentTest] running $name sample at $rel_sample_path"
+  echo "[AgentTest] running $rel_sample_path"
   python -m graph_net.model_path_handler \
     --model-path "$rel_sample_path" \
     --handler-config "$HANDLER_CONFIG"
 }
 
-run_case "samples/torchvision/resnet18" "CV (torchvision/resnet18)"
-run_case "samples/transformers-auto-model/albert-base-v2" "NLP (transformers-auto-model/albert-base-v2)"
+run_case "samples/torchvision/resnet18"
+run_case "samples/transformers-auto-model/albert-base-v2"
 
 echo "[AgentTest] done. Generated *_test.py files should now exist beside the samples."

graph_net/torch/sample_passes/agent_unittest_generator.py

@@ -1,8 +1,11 @@
 import re
+import os
+import sys
 import ast
 import inspect
 import jinja2
 import textwrap
+import tempfile
 from pathlib import Path
 from typing import Literal
 from collections import namedtuple
@@ -127,18 +130,21 @@ def __init__(
         output_dir: str,
         device: Literal["auto", "cpu", "cuda"] = "auto",
         generate_main: bool = False,
+        try_run: bool = False,
         data_input_predicator_filepath: str = None,
         data_input_predicator_class_name: str = None,
     ):
         self.model_path = Path(model_path).resolve()
         self.output_dir = Path(output_dir)
         self.device = self._choose_device(device)
         self.generate_main = generate_main
+        self.try_run = try_run and generate_main
         self.data_input_predicator = self._make_data_input_predicator(
             data_input_predicator_filepath, data_input_predicator_class_name
         )
 
     def generate(self):
+        print(f"[AgentUnittestGenerator] Generate unittest for {self.model_path}")
         model_name = "".join(
             word.capitalize() for word in re.split(r"[_.-]", self.model_path.name)
         )
@@ -152,7 +158,6 @@ def generate(self):
             input_tensor_metas,
             weight_tensor_metas,
         ) = self._get_input_and_weight_tensor_metas(input_arg_names, weight_arg_names)
-        print(f"{input_arg_names=}")
         graph_module_desc = GraphModuleDescriptor(
             device=self.device,
             generate_main=self.generate_main,
@@ -166,7 +171,8 @@ def generate(self):
             ),
         )
         unittest = self._render_template(graph_module_desc)
-        self._write_to_file(unittest)
+        if self._try_to_run_unittest(unittest):
+            self._write_to_file(unittest, self.output_dir)
 
     def _choose_device(self, device) -> str:
         if device in ["cpu", "cuda"]:
@@ -180,15 +186,25 @@ def _make_data_input_predicator(
             module = imp_util.load_module(data_input_predicator_filepath)
             cls = getattr(module, data_input_predicator_class_name)
             return cls(config={})
-        return lambda *args, **kwargs: False
+        return lambda *args, **kwargs: True
 
-    def _write_to_file(self, unittest):
-        output_path = Path(self.output_dir) / f"{self.model_path.name}_test.py"
+    def _write_to_file(self, unittest, output_dir):
+        output_path = Path(output_dir) / f"{self.model_path.name}_test.py"
         output_path.parent.mkdir(parents=True, exist_ok=True)
         output_path.write_text(unittest, encoding="utf-8")
         print(
             f"[AgentUnittestGenerator] Generate unittest: {output_path} (device={self.device})"
         )
+        return output_path
+
+    def _try_to_run_unittest(self, unittest):
+        if not self.try_run:
+            return True
+
+        with tempfile.TemporaryDirectory(prefix="unittest_") as temp_dir:
+            output_path = self._write_to_file(unittest, temp_dir)
+            cmd = f"{sys.executable} {output_path}"
+            return os.system(cmd) == 0
 
     def _get_input_and_weight_arg_names(self, graph_module):
         input_arg_names = []
@@ -287,28 +303,28 @@ class AgentUnittestGeneratorPass(SamplePass):
 
     def __init__(self, config=None):
         super().__init__(config)
-        print(f"[AgentUnittestGeneratorPass] {self.config=}")
 
     def declare_config(
         self,
         model_path_prefix: str,
         output_dir: str,
         device: str = "auto",
         generate_main: bool = False,
+        try_run: bool = False,
         data_input_predicator_filepath: str = None,
         data_input_predicator_class_name: str = None,
     ):
         pass
 
     def __call__(self, rel_model_path: str):
-        print(f"[AgentUnittestGeneratorPass] {rel_model_path=}")
         model_path_prefix = Path(self.config["model_path_prefix"])
         output_dir = Path(self.config["output_dir"])
         generator = AgentUnittestGenerator(
             model_path=str(model_path_prefix / rel_model_path),
             output_dir=str(output_dir / rel_model_path),
             device=self.config["device"],
             generate_main=self.config["generate_main"],
+            try_run=self.config["try_run"],
             data_input_predicator_filepath=self.config[
                 "data_input_predicator_filepath"
             ],