From 57e3474884d0ab1e36827c9db454e110268ad0d2 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Wed, 13 Aug 2025 11:35:32 +0800 Subject: [PATCH 1/9] Update script_wrapper.py --- .../pytorch/hook_module/script_wrapper.py | 53 ++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index c6d611d5c..5ffa43ca9 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -13,11 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools +import importlib import types import torch from msprobe.pytorch.hook_module.api_register import get_api_register from msprobe.pytorch.common.utils import torch_version_above_or_equal_2 + if torch_version_above_or_equal_2: from torch._dynamo.convert_frame import convert_frame as _orig_convert_frame, Hooks @@ -70,7 +73,55 @@ def wrap_compile_script_func(): _cf_mod.convert_frame = _patched_convert_frame +def patch_dynamo__compile() -> bool: + cf = importlib.import_module("torch._dynamo.convert_frame") + if not hasattr(cf, "_compile"): + raise RuntimeError("未找到 torch._dynamo.convert_frame._compile") + + original = cf._compile + if getattr(original, "__msprobe_patched__", False): + return False + + @functools.wraps(original) + def wrapped(*args, **kwargs): + try: + from msprobe.pytorch.dump.module_dump.module_dump import get_api_register + reg = get_api_register() + reg.restore_all_api() + print("415") + except Exception as e: + print(f"[msprobe] 前置 restore_all_api 异常: {e}") + + try: + return original(*args, **kwargs) + finally: + try: + print("213321415") + from msprobe.pytorch.dump.module_dump.module_dump import get_api_register + reg = get_api_register() + reg.register_all_api() # ✅ 改成注册 + except Exception as e: + print(f"[msprobe] 后置 register_all_api 异常: {e}") + + wrapped.__msprobe_patched__ = True + wrapped.__msprobe_original__ = original + cf._compile = wrapped + return True + +def unpatch_dynamo__compile() -> bool: + # 预留取消patch接口 + cf = importlib.import_module("torch._dynamo.convert_frame") + current = getattr(cf, "_compile", None) + if current is None: + return False + original = getattr(current, "__msprobe_original__", None) + if original is None: + return False + cf._compile = original + return True + + def wrap_script_func(): wrap_jit_script_func() if torch_version_above_or_equal_2: - wrap_compile_script_func() \ No newline at end of file + patch_dynamo__compile() \ No newline at end of file -- Gitee From 302cf630278905bc7f3cc2718ccbe0d026c9ff85 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Thu, 14 Aug 2025 16:31:19 +0800 Subject: [PATCH 2/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index 5ffa43ca9..566d1bb8c 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -17,6 +17,7 @@ import functools import importlib import types import torch +from msprobe.core.common.log import logger from msprobe.pytorch.hook_module.api_register import get_api_register from msprobe.pytorch.common.utils import torch_version_above_or_equal_2 @@ -85,23 +86,19 @@ def patch_dynamo__compile() -> bool: @functools.wraps(original) def wrapped(*args, **kwargs): try: - from msprobe.pytorch.dump.module_dump.module_dump import get_api_register reg = get_api_register() reg.restore_all_api() - print("415") except Exception as e: - print(f"[msprobe] 前置 restore_all_api 异常: {e}") + logger.warning(f"[msprobe] 前置 restore_all_api 异常: {e}") try: return original(*args, **kwargs) finally: try: - print("213321415") - from msprobe.pytorch.dump.module_dump.module_dump import get_api_register reg = get_api_register() - reg.register_all_api() # ✅ 改成注册 + reg.register_all_api() # 改成注册hook except Exception as e: - print(f"[msprobe] 后置 register_all_api 异常: {e}") + logger.warning(f"[msprobe] 后置 register_all_api 异常: {e}") wrapped.__msprobe_patched__ = True wrapped.__msprobe_original__ = original -- Gitee From b07ddbfa56d9a5163509aaf4a097daf90e88cee5 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Thu, 14 Aug 2025 16:38:17 +0800 Subject: [PATCH 3/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index 566d1bb8c..deff6aa53 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -74,7 +74,7 @@ def wrap_compile_script_func(): _cf_mod.convert_frame = _patched_convert_frame -def patch_dynamo__compile() -> bool: +def patch_dynamo__compile(): cf = importlib.import_module("torch._dynamo.convert_frame") if not hasattr(cf, "_compile"): raise RuntimeError("未找到 torch._dynamo.convert_frame._compile") @@ -96,14 +96,14 @@ def patch_dynamo__compile() -> bool: finally: try: reg = get_api_register() - reg.register_all_api() # 改成注册hook + reg.register_all_api() # 改成注册hook except Exception as e: logger.warning(f"[msprobe] 后置 register_all_api 异常: {e}") wrapped.__msprobe_patched__ = True wrapped.__msprobe_original__ = original cf._compile = wrapped - return True + def unpatch_dynamo__compile() -> bool: # 预留取消patch接口 @@ -121,4 +121,4 @@ def unpatch_dynamo__compile() -> bool: def wrap_script_func(): wrap_jit_script_func() if torch_version_above_or_equal_2: - patch_dynamo__compile() \ No newline at end of file + patch_dynamo__compile() -- Gitee From 1249b534bb71c670fa7120e422f17bb562e223d2 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Fri, 15 Aug 2025 09:15:08 +0800 Subject: [PATCH 4/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index deff6aa53..a7afaf218 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -81,7 +81,7 @@ def patch_dynamo__compile(): original = cf._compile if getattr(original, "__msprobe_patched__", False): - return False + return @functools.wraps(original) def wrapped(*args, **kwargs): -- Gitee From 63e78a1a906a6997bbaa6b23c649bfee78049318 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Mon, 18 Aug 2025 15:49:55 +0800 Subject: [PATCH 5/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index a7afaf218..494c75a10 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -93,6 +93,8 @@ def patch_dynamo__compile(): try: return original(*args, **kwargs) + except Exception: + pass finally: try: reg = get_api_register() -- Gitee From fb18556b3deab1bf13005a8d4b473b7e4e8db522 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Tue, 19 Aug 2025 18:47:17 +0800 Subject: [PATCH 6/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index 494c75a10..b6a5e060c 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -91,17 +91,19 @@ def patch_dynamo__compile(): except Exception as e: logger.warning(f"[msprobe] 前置 restore_all_api 异常: {e}") + result = None try: - return original(*args, **kwargs) + result = original(*args, **kwargs) except Exception: - pass + logger.warning("[msprobe] _compile execution failed (returning None)") + result = None finally: try: reg = get_api_register() reg.register_all_api() # 改成注册hook except Exception as e: logger.warning(f"[msprobe] 后置 register_all_api 异常: {e}") - + return result wrapped.__msprobe_patched__ = True wrapped.__msprobe_original__ = original cf._compile = wrapped -- Gitee From 7d77c679a069ba43fe44a99ec372b66805640956 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Tue, 19 Aug 2025 20:31:02 +0800 Subject: [PATCH 7/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index b6a5e060c..0e4de41de 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -74,10 +74,10 @@ def wrap_compile_script_func(): _cf_mod.convert_frame = _patched_convert_frame -def patch_dynamo__compile(): +def patch_dynamo_compile(): cf = importlib.import_module("torch._dynamo.convert_frame") if not hasattr(cf, "_compile"): - raise RuntimeError("未找到 torch._dynamo.convert_frame._compile") + logger.warning("No found torch._dynamo.convert_frame._compile") original = cf._compile if getattr(original, "__msprobe_patched__", False): @@ -89,7 +89,8 @@ def patch_dynamo__compile(): reg = get_api_register() reg.restore_all_api() except Exception as e: - logger.warning(f"[msprobe] 前置 restore_all_api 异常: {e}") + logger.warning(f"[msprobe] Pre restore_all_api failed: {e}") + return result = None try: @@ -102,14 +103,14 @@ def patch_dynamo__compile(): reg = get_api_register() reg.register_all_api() # 改成注册hook except Exception as e: - logger.warning(f"[msprobe] 后置 register_all_api 异常: {e}") + logger.warning(f"[msprobe] Post register_all_api failed: {e}") return result wrapped.__msprobe_patched__ = True wrapped.__msprobe_original__ = original cf._compile = wrapped -def unpatch_dynamo__compile() -> bool: +def unpatch_dynamo_compile() -> bool: # 预留取消patch接口 cf = importlib.import_module("torch._dynamo.convert_frame") current = getattr(cf, "_compile", None) @@ -125,4 +126,4 @@ def unpatch_dynamo__compile() -> bool: def wrap_script_func(): wrap_jit_script_func() if torch_version_above_or_equal_2: - patch_dynamo__compile() + patch_dynamo_compile() -- Gitee From 900b2648e742321b845c9505388e2723f0a5bcd9 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Wed, 20 Aug 2025 10:46:10 +0800 Subject: [PATCH 8/9] Update script_wrapper.py --- .../msprobe/pytorch/hook_module/script_wrapper.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index 0e4de41de..6cca90dad 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -85,14 +85,14 @@ def patch_dynamo_compile(): @functools.wraps(original) def wrapped(*args, **kwargs): + result = None try: reg = get_api_register() reg.restore_all_api() except Exception as e: logger.warning(f"[msprobe] Pre restore_all_api failed: {e}") - return + return result - result = None try: result = original(*args, **kwargs) except Exception: @@ -104,6 +104,7 @@ def patch_dynamo_compile(): reg.register_all_api() # 改成注册hook except Exception as e: logger.warning(f"[msprobe] Post register_all_api failed: {e}") + return result return result wrapped.__msprobe_patched__ = True wrapped.__msprobe_original__ = original -- Gitee From 614056f30f73ef8614a9fd090c6963fec92b1363 Mon Sep 17 00:00:00 2001 From: mnhdxnh Date: Wed, 20 Aug 2025 11:11:09 +0800 Subject: [PATCH 9/9] Update script_wrapper.py --- .../accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py | 1 - 1 file changed, 1 deletion(-) diff --git a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py index 6cca90dad..2bca42656 100644 --- a/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py +++ b/debug/accuracy_tools/msprobe/pytorch/hook_module/script_wrapper.py @@ -104,7 +104,6 @@ def patch_dynamo_compile(): reg.register_all_api() # 改成注册hook except Exception as e: logger.warning(f"[msprobe] Post register_all_api failed: {e}") - return result return result wrapped.__msprobe_patched__ = True wrapped.__msprobe_original__ = original -- Gitee