已合并
patch: fixed pickle serialization for precompile cache #40165
Jingming Liu创建于 7月6日
patch: fixed pickle serialization for precompile cache #40165
已合并
共 3 个文件变更+120-0
| @@ -588,5 +588,47 @@ class TestTorchNpuBootstrap(TestCase): | |||
| 588 | """ | 588 | """ |
| 589 | ) | 589 | ) |
| 590 | 590 | ||
| 591 | + def test_13_patch_dataclass_before_dynamo_source_import(self): | ||
| 592 | + self._run_python( | ||
| 593 | + """ | ||
| 594 | + import torch | ||
| 595 | + import torch_npu | ||
| 596 | + | ||
| 597 | + patched_decorator = torch._guards.dataclass_with_cached_hash | ||
| 598 | + assert patched_decorator.__module__ == "torch_npu.utils._guards" | ||
| 599 | + | ||
| 600 | + from torch._dynamo import source as source_module | ||
| 601 | + from torch._dynamo.source import DefaultsSource, LocalSource | ||
| 602 | + | ||
| 603 | + assert source_module.dataclass_with_cached_hash is patched_decorator | ||
| 604 | + | ||
| 605 | + source = DefaultsSource(LocalSource("fn"), 0) | ||
| 606 | + reduce_callable, reduce_args = source.__reduce__() | ||
| 607 | + assert reduce_callable is DefaultsSource | ||
| 608 | + assert reduce_args == (LocalSource("fn"), 0, False) | ||
| 609 | + """ | ||
| 610 | + ) | ||
| 611 | + | ||
| 612 | + def test_14_patch_dataclass_after_dynamo_source_import(self): | ||
| 613 | + self._run_python( | ||
| 614 | + """ | ||
| 615 | + import torch | ||
| 616 | + from torch._dynamo import source as source_module | ||
| 617 | + from torch._dynamo.source import DefaultsSource, LocalSource | ||
| 618 | + | ||
| 619 | + source = DefaultsSource(LocalSource("fn"), 0) | ||
| 620 | + assert len(source.__reduce__()[1]) == 5 | ||
| 621 | + | ||
| 622 | + import torch_npu | ||
| 623 | + | ||
| 624 | + patched_decorator = torch._guards.dataclass_with_cached_hash | ||
| 625 | + assert source_module.dataclass_with_cached_hash is patched_decorator | ||
| 626 | + | ||
| 627 | + reduce_callable, reduce_args = source.__reduce__() | ||
| 628 | + assert reduce_callable is DefaultsSource | ||
| 629 | + assert reduce_args == (LocalSource("fn"), 0, False) | ||
| 630 | + """ | ||
| 631 | + ) | ||
| 632 | + | ||
| 591 | if __name__ == "__main__": | 633 | if __name__ == "__main__": |
| 592 | run_tests() | 634 | run_tests() |
| @@ -13,3 +13,11 @@ def apply_npugraph_tree_patch(): | |||
| 13 | from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods | 13 | from torch_npu.utils._graph_tree import _apply_npugraph_tree_methods |
| 14 | 14 | ||
| 15 | _apply_npugraph_tree_methods() | 15 | _apply_npugraph_tree_methods() |
| 16 | + | ||
| 17 | + | ||
| 18 | +# Source subclasses bind and execute this decorator while torch._dynamo.source is imported. | ||
| 19 | + | ||
| 20 | +def apply_dataclass_with_cached_hash_patch(): | ||
| 21 | + from torch_npu.utils._guards import patch_dataclass_with_cached_hash | ||
| 22 | + | ||
| 23 | + patch_dataclass_with_cached_hash() | ||
 | |||
| @@ -0,0 +1,70 @@ | |||
| 1 | +import dataclasses | ||
| 2 | +import sys | ||
| 3 | +from typing import Any, Callable, overload, TypeVar | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +T = TypeVar("T") | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +def _reduce_without_cached_hash(self): | ||
| 12 | + fields = dataclasses.fields(self) | ||
| 13 | + field_values = tuple(getattr(self, field.name) for field in fields if field.init) | ||
| 14 | + return (self.__class__, field_values) | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +def _patch_existing_source_classes(): | ||
| 18 | + source_class = torch._guards.Source | ||
| 19 | + pending = [source_class] | ||
| 20 | + while pending: | ||
| 21 | + cls = pending.pop() | ||
| 22 | + reduce_method = cls.__dict__.get("__reduce__") | ||
| 23 | + if ( | ||
| 24 | + reduce_method is not None | ||
| 25 | + and reduce_method.__module__ == "torch._guards" | ||
| 26 | + ): | ||
| 27 | + cls.__reduce__ = _reduce_without_cached_hash | ||
| 28 | + pending.extend(cls.__subclasses__()) | ||
| 29 | + | ||
| 30 | + # torch._dynamo.source binds the decorator during import. Keep that binding | ||
| 31 | + # consistent when the module was imported before torch_npu. | ||
| 32 | + source_module = sys.modules.get("torch._dynamo.source") | ||
| 33 | + if source_module is not None: | ||
| 34 | + source_module.dataclass_with_cached_hash = ( | ||
| 35 | + torch._guards.dataclass_with_cached_hash | ||
| 36 | + ) | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +def patch_dataclass_with_cached_hash(): | ||
| 40 | + | ||
| 41 | + def dataclass_with_cached_hash(cls: type[T], **kwargs: Any) -> type[T]: ... | ||
| 42 | + | ||
| 43 | + | ||
| 44 | + def dataclass_with_cached_hash( | ||
| 45 | + cls: None = None, **kwargs: Any | ||
| 46 | + ) -> Callable[[type[T]], type[T]]: ... | ||
| 47 | + | ||
| 48 | + def dataclass_with_cached_hash( | ||
| 49 | + cls: type[T] | None = None, **kwargs: Any | ||
| 50 | + ) -> type[T] | Callable[[type[T]], type[T]]: | ||
| 51 | + def wrap(cls_inner: type[T]) -> type[T]: | ||
| 52 | + new_cls = dataclasses.dataclass(cls_inner, **kwargs) | ||
| 53 | + old_hash = cls_inner.__hash__ | ||
| 54 | + | ||
| 55 | + def __hash__(self) -> int: | ||
| 56 | + if not hasattr(self, "_hash"): | ||
| 57 | + object.__setattr__(self, "_hash", old_hash(self)) | ||
| 58 | + return self._hash | ||
| 59 | + | ||
| 60 | + new_cls.__hash__ = __hash__ | ||
| 61 | + new_cls.__reduce__ = _reduce_without_cached_hash | ||
| 62 | + return new_cls # type: ignore[return-value] | ||
| 63 | + | ||
| 64 | + if cls is None: | ||
| 65 | + return wrap | ||
| 66 | + | ||
| 67 | + return wrap(cls) | ||
| 68 | + | ||
| 69 | + torch._guards.dataclass_with_cached_hash = dataclass_with_cached_hash | ||
| 70 | + _patch_existing_source_classes() | ||
🟡 Medium Priority
在
dynamo_patches.py中，新增的第 20 行定义了与第 12 行同名的函数apply_npugraph_tree_patch。Python 中后定义的同名函数会覆盖先前的定义，导致原有的 graph tree patch 函数（第 12-15 行）在模块命名空间中不可访问。虽然两个函数都通过
@PatchManager.register_patch("dynamo")装饰器在定义时注册到了_patch_groups["dynamo"]列表中（_add_patch使用身份比较，两个不同函数对象都会被添加），两者在补丁执行时都会运行，但存在以下风险：apply_npugraph_tree_patch，但其实际功能是 patchdataclass_with_cached_hash用于 pickle 序列化修复，与 "npugraph tree" 完全无关——名称极具误导性。from torch_npu._init.patches.dynamo_patches import apply_npugraph_tree_patch导入，增加了维护风险。建议：将新增的补丁函数重命名为能准确描述其功能的名称，例如
apply_dataclass_pickle_patch或apply_guards_patch。同时建议将注释中的 "npugraph tree" 相关内容更新为对 guards patch 的描述。