[SOT] Optimize env and subgraph information printing (#71460)

DrRyanHuang · web-flow · commit 4cf32bbf5ac8 · 2025-03-08T11:51:43.000+08:00
diff --git a/python/paddle/jit/sot/profiler/profiler.py b/python/paddle/jit/sot/profiler/profiler.py
@@ -39,7 +39,7 @@ def disable(self):
 def EventGuard(event_name, event_level=1):
     need_pop = False
     try:
-        if ENV_SOT_EVENT_LEVEL.get_with_cache() >= event_level:
+        if ENV_SOT_EVENT_LEVEL.get() >= event_level:
             core.nvprof_nvtx_push(event_name)
             need_pop = True
         yield
@@ -60,7 +60,7 @@ def call_with_event(*args, **kwargs):
     def do_nothing(func):
         return func
 
-    if ENV_SOT_EVENT_LEVEL.get_with_cache() >= event_level:
+    if ENV_SOT_EVENT_LEVEL.get() >= event_level:
         return event_wrapper
     else:
         return do_nothing
diff --git a/python/paddle/jit/sot/symbolic/compile_cache.py b/python/paddle/jit/sot/symbolic/compile_cache.py
@@ -225,6 +225,7 @@ def collect_subgraph_info(self, program: Program):
             SubGraphInfo,
             program,
             self.graph_size(),
+            self.SIR.name,
         )
 
     def __call__(self, *args, **kwargs):
diff --git a/python/paddle/jit/sot/utils/envs.py b/python/paddle/jit/sot/utils/envs.py
@@ -41,7 +41,7 @@ def __init__(self, name: str, default: dict[str, list[str]]):
         super().__init__(name, default)
         assert isinstance(default, dict), "default must be a dict"
 
-    def get(self) -> dict[str, list[str]]:
+    def parse_from_string(self) -> dict[str, list[str]]:
         env_var = os.getenv(self.name)
         if env_var is None or env_var == "":
             return self.default
@@ -51,7 +51,7 @@ def get(self) -> dict[str, list[str]]:
             ret.update(self.parse_parameterized_key(item))
         return ret
 
-    def set(self, value: dict[str, list[str]]) -> None:
+    def convert_to_string(self, value: dict[str, list[str]]) -> str:
         assert isinstance(value, dict), "The input must be a dict"
         assert all(
             isinstance(x, str) for x in value.keys()
@@ -63,7 +63,8 @@ def set(self, value: dict[str, list[str]]) -> None:
         env_list = []
         for k, v in value.items():
             env_list.append(f"{k}" + (f"[{','.join(v)}]" if len(v) else ""))
-        os.environ[self.name] = ",".join(env_list)
+
+        return ",".join(env_list)
 
     @staticmethod
     def split_by_unbracketed_commas(input_str: str) -> list[str]:
diff --git a/python/paddle/jit/sot/utils/info_collector.py b/python/paddle/jit/sot/utils/info_collector.py
@@ -70,7 +70,7 @@ def register(self, info: InfoBase) -> None:
         info_dict[info_class_name].append(info)
 
     def need_collect(self, cls: type[InfoBase]) -> bool:
-        return cls.SHORT_NAME in ENV_SOT_COLLECT_INFO.get_with_cache()
+        return cls.SHORT_NAME in ENV_SOT_COLLECT_INFO.get()
 
     def clear_step_info(self):
         self._step_info.clear()
@@ -309,19 +309,32 @@ class SubGraphInfo(InfoBase):
     SHORT_NAME = "subgraph_info"
     TYPE = InfoType.STEP_INFO
 
-    def __init__(self, graph, op_num):
+    def __init__(self, graph, op_num, sir_name):
         super().__init__()
         self.graph = graph
         self.op_num = op_num
+        self.sir_name = sir_name
 
     def __str__(self):
-        return f"OpNum: {self.op_num}\n{self.graph}"
+        return f"[SIR Name]: {self.sir_name}   [OpNum]: {self.op_num}\n{self.graph}"
 
     @classmethod
     def summary(cls, history: list[Self]) -> str:
-        return "\n".join(
-            [
-                f"SubGraphIdx: {idx} {info}"
-                for idx, info in enumerate(map(str, history))
-            ]
-        )
+
+        num_of_subgraph = len(history)
+        sum_of_op_num = sum(item.op_num for item in history)
+
+        need_details = "details" in ENV_SOT_COLLECT_INFO.get()[cls.SHORT_NAME]
+
+        details = ""
+        if need_details:
+            details = "\n".join(
+                [
+                    f"[SubGraphIdx]: {idx}   {info}"
+                    for idx, info in enumerate(map(str, history))
+                ]
+            )
+
+        summary = f"[Number of subgraph]: {num_of_subgraph} [Sum of opnum]: {sum_of_op_num}"
+
+        return f"{summary}\n{details}"
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -125,25 +125,25 @@ def next(self):
 
 
 def log(level, *args):
-    cur_level = ENV_SOT_LOG_LEVEL.get_with_cache()
+    cur_level = ENV_SOT_LOG_LEVEL.get()
     if level <= cur_level:
         print(*args, end="", flush=True)
 
 
 def log_do(level, fn):
-    cur_level = ENV_SOT_LOG_LEVEL.get_with_cache()
+    cur_level = ENV_SOT_LOG_LEVEL.get()
     if level <= cur_level:
         fn()
 
 
 def log_format(level, str, *args):
-    cur_level = ENV_SOT_LOG_LEVEL.get_with_cache()
+    cur_level = ENV_SOT_LOG_LEVEL.get()
     if level <= cur_level:
         print(str.format(*args), end="", flush=True)
 
 
 def log_enabled(level):
-    return level <= ENV_SOT_LOG_LEVEL.get_with_cache()
+    return level <= ENV_SOT_LOG_LEVEL.get()
 
 
 def no_eval_frame(func):
diff --git a/python/paddle/utils/environments.py b/python/paddle/utils/environments.py
@@ -39,17 +39,27 @@ class EnvironmentVariable(Generic[T]):
     def __init__(self, name: str, default: T):
         self.name = name
         self.default = default
+        self._last_env_value: str | None = None
         self._cached_value: T | None = None
 
-    def get_with_cache(self) -> T:
-        if self._cached_value is None:
-            self._cached_value = self.get()
+    def get(self) -> T:
+        _current_env_value = os.getenv(self.name)
+        if (
+            self._cached_value is None
+            or self._last_env_value != _current_env_value
+        ):
+            self._cached_value = self.parse_from_string()
+            self._last_env_value = _current_env_value
         return self._cached_value
 
-    def get(self) -> T:
+    def set(self, value: T) -> None:
+        os.environ[self.name] = self.convert_to_string(value)
+        self._cached_value = value
+
+    def parse_from_string(self) -> T:
         raise NotImplementedError
 
-    def set(self, value: T) -> None:
+    def convert_to_string(self, value: T) -> str:
         raise NotImplementedError
 
     def delete(self) -> None:
@@ -64,27 +74,27 @@ def __init__(self, name: str, default: str):
         super().__init__(name, default)
         assert isinstance(default, str), "default must be a string"
 
-    def get(self) -> str:
+    def parse_from_string(self) -> str:
         return os.getenv(self.name, self.default)
 
-    def set(self, value: str) -> None:
+    def convert_to_string(self, value: str) -> str:
         assert isinstance(value, str), "value must be a string"
-        os.environ[self.name] = value
+        return value
 
 
 class BooleanEnvironmentVariable(EnvironmentVariable[bool]):
     def __init__(self, name: str, default: bool):
         super().__init__(name, default)
         assert isinstance(default, bool), "default must be a boolean"
 
-    def get(self) -> bool:
+    def parse_from_string(self) -> bool:
         default = str(self.default)
         env_str = os.getenv(self.name, default)
         return strtobool(env_str)
 
-    def set(self, value: bool) -> None:
+    def convert_to_string(self, value: bool) -> str:
         assert isinstance(value, bool), "value must be a boolean"
-        os.environ[self.name] = str(value).lower()
+        return str(value).lower()
 
     def __bool__(self) -> bool:
         raise ValueError(
@@ -100,33 +110,33 @@ def __init__(self, name: str, default: int):
             default, bool
         ), "default must be an integer"
 
-    def get(self) -> int:
+    def parse_from_string(self) -> int:
         try:
             return int(os.getenv(self.name, str(self.default)))
         except ValueError:
             return self.default
 
-    def set(self, value: int) -> None:
+    def convert_to_string(self, value: int) -> str:
         assert isinstance(value, int) and not isinstance(
             value, bool
         ), "value must be an integer"
-        os.environ[self.name] = str(value)
+        return str(value)
 
 
 class StringListEnvironmentVariable(EnvironmentVariable[List[str]]):
     def __init__(self, name: str, default: list[str]):
         super().__init__(name, default)
         assert isinstance(default, list), "default must be a list"
 
-    def get(self) -> list[str]:
+    def parse_from_string(self) -> list[str]:
         return os.getenv(self.name, ",".join(self.default)).split(",")
 
-    def set(self, value: list[str]) -> None:
+    def convert_to_string(self, value: list[str]) -> str:
         assert isinstance(value, list), "value must be a list"
         assert all(
             isinstance(x, str) for x in value
         ), "value must be a list of strings"
-        os.environ[self.name] = ",".join(value)
+        return ",".join(value)
 
 
 class EnvironmentVariableGuard(Generic[T]):
diff --git a/test/sot/test_binary_operator_tracker.py b/test/sot/test_binary_operator_tracker.py
@@ -17,9 +17,6 @@
 
 from __future__ import annotations
 
-import os
-
-os.environ['MIN_GRAPH_SIZE'] = '-1'
 import operator
 import unittest
 
@@ -38,6 +35,9 @@
     BinaryOperatorTracker,
 )
 from paddle.jit.sot.opcode_translator.executor.variables import ConstantVariable
+from paddle.jit.sot.utils import ENV_MIN_GRAPH_SIZE
+
+ENV_MIN_GRAPH_SIZE.set(-1)
 
 
 def register_dispatch(fn, parameters, handler):

Original file line number	Diff line number	Diff line change
`@@ -225,6 +225,7 @@ def collect_subgraph_info(self, program: Program):`
`225`	`225`	`SubGraphInfo,`
`226`	`226`	`program,`
`227`	`227`	`self.graph_size(),`
	`228`	`+ self.SIR.name,`
`228`	`229`	`)`
`229`	`230`
`230`	`231`	`def __call__(self, args, *kwargs):`