unstable to stable challenge (#313)

Author: RbRe145

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -6,14 +6,17 @@
 
 
 class UnstableToStableBackend(GraphCompilerBackend):
-    def __call__(self, model, model_path):
+    def __call__(self, model):
         # Perform unstable API check before running the model
         unstable_api = os.getenv("DISALLOWED_UNSTABLE_API", "").strip()
         self.unstable_api = unstable_api
-        self.model_path = model_path
-        self.unstable_to_stable(model)
-        self.check_unstable_api(model)
-        return self.model
+
+        def my_backend(gm, sample_inputs):
+            gm = self.unstable_to_stable(gm)
+            self.check_unstable_api(gm)
+            return gm.forward
+
+        return torch.compile(backend=my_backend)(model)
 
     """
     TODO: Implement logic to convert unstable APIs in `self.model` into their stable counterparts.
@@ -26,10 +29,11 @@ def __call__(self, model, model_path):
     **Stable API reference link:**
     """
 
-    def unstable_to_stable(self, model):
-        return
+    def unstable_to_stable(self, gm):
+        # TODO
+        return gm
 
-    def check_unstable_api(self, model):
+    def check_unstable_api(self, gm):
         """
         Check whether gm contains the API specified in the environment
         variable DISALLOWED_UNSTABLE_API. If it does, raise an exception and stop
@@ -40,20 +44,7 @@ def check_unstable_api(self, model):
         Do NOT modify, remove, or bypass this check under any circumstances.
         """
 
-        # from torch.fx import symbolic_trace
-
-        # try:
-        #     # Convert the model into a static computation graph (FX IR)
-        #     traced = symbolic_trace(self.model)
-        #     graph_text = str(traced.graph)
-        # except Exception as e:
-        #     # In case tracing fails, fallback to textual model dump
-        #     graph_text = str(*(self.model))
-
-        print(f"model path is: {self.model_path}")
-        model_file_path = self.model_path + "model.py"
-        with open(model_file_path, "r", encoding="utf-8") as f:
-            graph_text = f.read()
+        graph_text = gm.code
         # Search for the unstable API substring
         if self.unstable_api in graph_text:
             count = graph_text.count(self.unstable_api)

graph_net/torch/test_compiler.py

@@ -289,8 +289,6 @@ def test_single_model(args):
         if args.compiler == "xla":
             xla_model = get_model(args, "xla")
             compiled_model = compiler(xla_model)
-        elif args.compiler == "unstable_to_stable":
-            compiled_model = compiler(model, args.model_path)
         else:
             compiled_model = compiler(model)
 

plot_unstable_to_stable.sh

@@ -7,8 +7,14 @@ if [ -z "$DISALLOWED_UNSTABLE_API" ]; then
   exit 1
 fi
 
+if [ -z "$GRAPH_NET_WORKSPACE" ]; then
+  echo "❌ 环境变量 GRAPH_NET_WORKSPACE 未设置！"
+  echo "请使用： export GRAPH_NET_WORKSPACE=/path/to/GraphNet"
+  exit 1
+fi
+
 # === 配置区 ===
-root_dir="/root/GraphNet/todo_works/unstable_api_to_stable_api/${DISALLOWED_UNSTABLE_API}"
+root_dir="${GRAPH_NET_WORKSPACE}/todo_works/unstable_api_to_stable_api/${DISALLOWED_UNSTABLE_API}"
 file_list="${root_dir}/${DISALLOWED_UNSTABLE_API}_files.txt"
 log_file="${root_dir}/log.log"
 json_output_dir="${root_dir}/JSON_results"
@@ -68,7 +74,7 @@ else
 fi
 
 echo "📦 正在将JSON转换为结果图"
-python -m graph_net.S_analysis \
+python -m graph_net.plot_ESt \
   --benchmark-path $GRAPH_NET_BENCHMARK_PATH/JSON_results/ \
   --output-dir $GRAPH_NET_BENCHMARK_PATH \