[SOT] Add PEP508LikeEnvironmentVariable (#71430)

DrRyanHuang · web-flow · commit 512ed6bbc23e · 2025-03-06T10:18:00.000+08:00
diff --git a/python/paddle/jit/sot/symbolic/compile_cache.py b/python/paddle/jit/sot/symbolic/compile_cache.py
@@ -41,8 +41,7 @@
 from .interpreter import compile_sir
 
 if TYPE_CHECKING:
-    from paddle.framework import Program
-    from paddle.static import InputSpec
+    from paddle.static import InputSpec, Program
 
     from .symbolic_context import SymbolicTraceContext
 
@@ -226,8 +225,6 @@ def collect_subgraph_info(self, program: Program):
             SubGraphInfo,
             program,
             self.graph_size(),
-            self.SIR.name,
-            self.is_first_call,
         )
 
     def __call__(self, *args, **kwargs):
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -23,6 +23,7 @@
     ENV_SOT_LOG_LEVEL,
     ENV_SOT_WITH_CONTROL_FLOW,
     ENV_STRICT_MODE,
+    PEP508LikeEnvironmentVariable,
     allow_dynamic_shape_guard,
     export_guard,
     faster_guard_guard,
diff --git a/python/paddle/jit/sot/utils/envs.py b/python/paddle/jit/sot/utils/envs.py
@@ -14,17 +14,104 @@
 
 from __future__ import annotations
 
+import os
 from contextlib import contextmanager
+from typing import Dict, List
 
 import paddle
 from paddle.utils.environments import (
     BooleanEnvironmentVariable,
+    EnvironmentVariable,
     EnvironmentVariableGuard,
     IntegerEnvironmentVariable,
     StringEnvironmentVariable,
-    StringListEnvironmentVariable,
 )
 
+
+class PEP508LikeEnvironmentVariable(EnvironmentVariable[Dict[str, List[str]]]):
+    """
+    Environment variable parser following PEP 508 extras specification syntax.
+    https://peps.python.org/pep-0508/
+
+    Processes strings using PEP 508-style bracket notation for optional components:
+    "feat1[opt1,opt2], feat2[opt3,opt4]" -> {'feat1': ['opt1', 'opt2'], 'feat2': ['opt3', 'opt4']}
+    """
+
+    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]]:
+        env_var = os.getenv(self.name)
+        if env_var is None or env_var == "":
+            return self.default
+        items = self.split_by_unbracketed_commas(env_var)
+        ret = {}
+        for item in items:
+            ret.update(self.parse_parameterized_key(item))
+        return ret
+
+    def set(self, value: dict[str, list[str]]) -> None:
+        assert isinstance(value, dict), "The input must be a dict"
+        assert all(
+            isinstance(x, str) for x in value.keys()
+        ), "Keys must be a string"
+        assert all(
+            isinstance(x, list) for x in value.values()
+        ), "Values must be a list"
+
+        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)
+
+    @staticmethod
+    def split_by_unbracketed_commas(input_str: str) -> list[str]:
+        """Split string by commas that are not enclosed in square brackets"""
+        # "feat1[opt1,opt2], feat2[opt3], feat3" -> ["feat1[opt1,opt2]", "feat2[opt3]", "feat3"]
+        bracket_depth = 0
+        split_parts = []
+        _start = 0
+
+        for _current, char in enumerate(input_str):
+            if char == "[":
+                bracket_depth += 1
+            elif char == "]":
+                bracket_depth = max(
+                    0, bracket_depth - 1
+                )  # Prevent negative depth
+
+            if char == "," and bracket_depth == 0:
+                split_parts.append(input_str[_start:_current].strip())
+                _start = _current + 1  # Skip comma
+
+        # Add remaining content after last comma
+        if remaining := input_str[_start:].strip():
+            split_parts.append(remaining)
+
+        return split_parts
+
+    @staticmethod
+    def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
+        """Parse key with parameters in brackets into a dictionary."""
+
+        start_bracket = input_str.find("[")
+        end_bracket = input_str.rfind("]")
+
+        if start_bracket == -1 or end_bracket == -1:
+            return {input_str: []}
+
+        parameter_key = input_str[:start_bracket].strip()
+
+        # Extract and clean parameters
+        parameters_str = input_str[start_bracket + 1 : end_bracket]
+        parameter_values = [
+            v.strip() for v in parameters_str.split(",") if v.strip()
+        ]
+
+        return {parameter_key: parameter_values}
+
+
 ENV_MIN_GRAPH_SIZE = IntegerEnvironmentVariable("MIN_GRAPH_SIZE", 10)
 ENV_SOT_LOG_LEVEL = IntegerEnvironmentVariable("SOT_LOG_LEVEL", 0)
 ENV_STRICT_MODE = BooleanEnvironmentVariable("STRICT_MODE", False)
@@ -52,7 +139,7 @@
 ENV_SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE = BooleanEnvironmentVariable(
     "SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False
 )
-ENV_SOT_COLLECT_INFO = StringListEnvironmentVariable("SOT_COLLECT_INFO", [])
+ENV_SOT_COLLECT_INFO = PEP508LikeEnvironmentVariable("SOT_COLLECT_INFO", {})
 ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(
     "SOT_FORCE_FALLBACK_SIR_IDS", ""
 )
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()
+        return cls.SHORT_NAME in ENV_SOT_COLLECT_INFO.get_with_cache()
 
     def clear_step_info(self):
         self._step_info.clear()
@@ -307,17 +307,12 @@ def collect_break_graph_reason(reason: BreakGraphReasonBase):
 
 class SubGraphInfo(InfoBase):
     SHORT_NAME = "subgraph_info"
-    TYPE = InfoType.E2E_INFO
+    TYPE = InfoType.STEP_INFO
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, graph, op_num):
         super().__init__()
-        self.clear()
-
-        self.graph, self.op_num, *_ = args
-
-    def clear(self):
-        self.graph = None
-        self.op_num = 0
+        self.graph = graph
+        self.op_num = op_num
 
     def __str__(self):
         return f"OpNum: {self.op_num}\n{self.graph}"
diff --git a/test/sot/test_envs.py b/test/sot/test_envs.py
@@ -17,6 +17,7 @@
 import os
 import unittest
 
+from paddle.jit.sot.utils import PEP508LikeEnvironmentVariable
 from paddle.utils.environments import (
     BooleanEnvironmentVariable,
     EnvironmentVariableGuard,
@@ -179,5 +180,79 @@ def test_int_env_guard(self):
             self.assertEqual(env_int.get(), 1000)
 
 
+class TestPEP508LikeEnvironmentVariable(unittest.TestCase):
+    def test_PEP508_like_env_get(self):
+        env_name = "___TEST_ENV_PEP508LikeDict_GET"
+        env_PEP508_like = PEP508LikeEnvironmentVariable(env_name, {})
+
+        if env_name in os.environ:
+            del os.environ[env_name]
+
+        self.assertEqual(env_PEP508_like.get(), {})
+
+        os.environ[env_name] = "subgraph_info"
+        self.assertEqual(env_PEP508_like.get(), {"subgraph_info": []})
+
+        os.environ[env_name] = ""
+        self.assertEqual(env_PEP508_like.get(), {})
+
+        os.environ[env_name] = "subgraph_info[ops, data, relationships]"
+        self.assertEqual(
+            env_PEP508_like.get(),
+            {"subgraph_info": ["ops", "data", "relationships"]},
+        )
+
+        os.environ[env_name] = (
+            "subgraph_info[ops, data, relationships], breakgraph_reason[inlinecall, others]"
+        )
+        self.assertEqual(
+            env_PEP508_like.get(),
+            {
+                "subgraph_info": ["ops", "data", "relationships"],
+                "breakgraph_reason": ["inlinecall", "others"],
+            },
+        )
+
+        os.environ[env_name] = "subgraph_info[ops, [[[[[[[[]], relationships], "
+        self.assertEqual(
+            env_PEP508_like.get(),
+            {
+                "subgraph_info": ["ops", "[[[[[[[[]]", "relationships"],
+            },
+        )
+
+        os.environ[env_name] = "subgraph_info[ops, , , relationships], "
+        self.assertEqual(
+            env_PEP508_like.get(),
+            {
+                "subgraph_info": ["ops", "relationships"],
+            },
+        )
+
+    def test_PEP508_like_env_set(self):
+        env_name = "___TEST_ENV_PEP508LikeDict_GET"
+        env_PEP508_like = PEP508LikeEnvironmentVariable(env_name, {})
+
+        env_PEP508_like.set(
+            {
+                "subgraph_info": ["ops", "data", "relationships"],
+                "breakgraph_reason": ["inlinecall", "others"],
+            }
+        )
+        self.assertEqual(
+            os.environ[env_name],
+            "subgraph_info[ops,data,relationships],breakgraph_reason[inlinecall,others]",
+        )
+
+        env_PEP508_like.set({"subgraph_info": []})
+        self.assertEqual(os.environ[env_name], "subgraph_info")
+
+        env_PEP508_like.set({})
+        self.assertEqual(
+            os.environ[env_name],
+            "",
+        )
+
+
 if __name__ == "__main__":
     unittest.main()
