wehub-resource-sync
239 lines
7.5 KiB
Python
239 lines
7.5 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from contextlib import contextmanager
|
|
|
|
from paddle.utils.environments import (
|
|
BooleanEnvironmentVariable,
|
|
EnvironmentVariable,
|
|
EnvironmentVariableGuard,
|
|
IntegerEnvironmentVariable,
|
|
StringEnvironmentVariable,
|
|
)
|
|
|
|
|
|
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 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
|
|
items = self.split_by_unbracketed_commas(env_var)
|
|
ret = {}
|
|
for item in items:
|
|
ret.update(self.parse_parameterized_key(item))
|
|
return ret
|
|
|
|
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()), (
|
|
"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 ""))
|
|
|
|
return ",".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)
|
|
ENV_SOT_WITH_CONTROL_FLOW = BooleanEnvironmentVariable(
|
|
"SOT_WITH_CONTROL_FLOW", True
|
|
)
|
|
ENV_SOT_EXPORT = StringEnvironmentVariable("SOT_EXPORT", "")
|
|
ENV_SOT_ALLOW_DYNAMIC_SHAPE = BooleanEnvironmentVariable(
|
|
"SOT_ALLOW_DYNAMIC_SHAPE",
|
|
# Enable SOT dynamic shape as default in PIR mode
|
|
True,
|
|
)
|
|
ENV_SOT_ENABLE_FASTER_GUARD = BooleanEnvironmentVariable(
|
|
"SOT_ENABLE_FASTER_GUARD",
|
|
False,
|
|
)
|
|
ENV_SOT_ENABLE_STRICT_GUARD_CHECK = BooleanEnvironmentVariable(
|
|
"SOT_ENABLE_STRICT_GUARD_CHECK",
|
|
False,
|
|
)
|
|
ENV_SOT_ENABLE_GUARD_TREE = BooleanEnvironmentVariable(
|
|
"SOT_ENABLE_GUARD_TREE",
|
|
False,
|
|
)
|
|
ENV_ENABLE_SOT_STEP_PROFILER = BooleanEnvironmentVariable(
|
|
"ENABLE_SOT_STEP_PROFILER", False
|
|
)
|
|
ENV_SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE = BooleanEnvironmentVariable(
|
|
"SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False
|
|
)
|
|
ENV_SOT_COLLECT_INFO = PEP508LikeEnvironmentVariable("SOT_COLLECT_INFO", {})
|
|
ENV_SOT_SERIALIZE_INFO = BooleanEnvironmentVariable("SOT_SERIALIZE_INFO", False)
|
|
ENV_SOT_CE_DEBUG_MODE = BooleanEnvironmentVariable("SOT_CE_DEBUG_MODE", False)
|
|
ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(
|
|
"SOT_FORCE_FALLBACK_SIR_IDS", ""
|
|
)
|
|
ENV_SOT_TRACE_NUMPY = BooleanEnvironmentVariable("ENV_SOT_TRACE_NUMPY", True)
|
|
ENV_SOT_UNSAFE_CACHE_FASTPATH = BooleanEnvironmentVariable(
|
|
"SOT_UNSAFE_CACHE_FASTPATH", False
|
|
)
|
|
ENV_SOT_ENABLE_0_SIZE_FALLBACK = BooleanEnvironmentVariable(
|
|
"SOT_ENABLE_0_SIZE_FALLBACK", True
|
|
)
|
|
ENV_SOT_SPECIALIZED_DIM_NUMBERS = StringEnvironmentVariable(
|
|
"SOT_SPECIALIZED_DIM_NUMBERS", "0"
|
|
)
|
|
ENV_SOT_ENABLE_COMPILE_TIME_LIMIT = BooleanEnvironmentVariable(
|
|
"SOT_ENABLE_COMPILE_TIME_LIMIT", True
|
|
)
|
|
|
|
|
|
def update_ce_flags():
|
|
if not ENV_SOT_CE_DEBUG_MODE.get():
|
|
return
|
|
# Enable information collection flags to facilitate debugging and analysis
|
|
|
|
collected_info_item: dict[str, list[str]] = ENV_SOT_COLLECT_INFO.get()
|
|
collected_info_item.setdefault("breakgraph_reason", [])
|
|
collected_info_item.setdefault("subgraph_info", [])
|
|
|
|
ENV_SOT_COLLECT_INFO.set(collected_info_item)
|
|
ENV_SOT_SERIALIZE_INFO.set(True)
|
|
|
|
|
|
update_ce_flags()
|
|
|
|
|
|
@contextmanager
|
|
def strict_mode_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_STRICT_MODE, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def min_graph_size_guard(value: int):
|
|
with EnvironmentVariableGuard(ENV_MIN_GRAPH_SIZE, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def with_control_flow_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_SOT_WITH_CONTROL_FLOW, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def export_guard(value: str):
|
|
with EnvironmentVariableGuard(ENV_SOT_EXPORT, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def allow_dynamic_shape_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_SOT_ALLOW_DYNAMIC_SHAPE, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def faster_guard_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_SOT_ENABLE_FASTER_GUARD, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def guard_tree_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_SOT_ENABLE_GUARD_TREE, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def sot_step_profiler_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_ENABLE_SOT_STEP_PROFILER, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def specialized_dim_numbers_guard(value: str):
|
|
with EnvironmentVariableGuard(ENV_SOT_SPECIALIZED_DIM_NUMBERS, value):
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def enable_0_size_fallback_guard(value: bool):
|
|
with EnvironmentVariableGuard(ENV_SOT_ENABLE_0_SIZE_FALLBACK, value):
|
|
yield
|