Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / hook_base.py: 98%

1# Copyright 2025 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""hook base module""" 

16from __future__ import annotations 

17from typing import TYPE_CHECKING 

18 

19from abc import ABCMeta 

20from abc import abstractmethod 

21import inspect 

22 

23if TYPE_CHECKING: 

24 from typing import Dict, Callable 

25 

26 

27class HookMetaclass(ABCMeta): 

28 """meta class""" 

29 

30 def __init__(cls, name, bases, namespace): 

31 super().__init__(name, bases, namespace) 

32 hook_registry = getattr(cls, "hook_registry", {}) 

33 # For each abstract method in base classes, compare signatures 

34 for base in bases: 

35 for meth in getattr(base, "__abstractmethods__", set()): 

36 base_method = getattr(base, meth) 

37 sub_method = getattr(cls, meth, None) 

38 base_sig = inspect.signature(base_method) 

39 sub_sig = inspect.signature(sub_method) 

40 base_params = list(base_sig.parameters.values()) 

41 sub_params = list(sub_sig.parameters.values()) 

42 if getattr(sub_method, "__isabstractmethod__", False): 

43 raise TypeError( 

44 f"{meth}() abstract method needs " 

45 f"to be implemented in subclass" 

46 ) 

47 if sub_method not in hook_registry.values(): 

48 raise TypeError( 

49 f"{meth}() needs to be decorated with " 

50 f"@hook_runner(model_name), " 

51 f"with model_name as a string" 

52 ) 

53 if len(base_params) != len(sub_params): 

54 raise TypeError( 

55 f"Mismatch signature of {meth}() " 

56 f"in class {cls.__name__}: " 

57 f"Expected {[p.name for p in base_params]} " 

58 f"instead of {[p.name for p in sub_params]})" 

59 ) 

60 

61 

62def hook_runner(model_name: str) -> Callable: 

63 """decorator""" 

64 if not model_name or not isinstance(model_name, str): 

65 raise TypeError( 

66 f"@hook_runner decorator " f"has invalid model_name ({model_name})" 

67 ) 

68 if ( 

69 model_name in MemEvalHook.hook_registry 

70 ): # pylint: disable=protected-access 

71 raise TypeError( 

72 f"@hook_runner model_name " f"'{model_name}' is already defined" 

73 ) 

74 

75 def wrapper(func): 

76 """register""" 

77 MemEvalHook.hook_registry[model_name] = ( 

78 func # pylint: disable=protected-access 

79 ) 

80 return func 

81 

82 return wrapper 

83 

84 

85class MemEvalHook(metaclass=HookMetaclass): 

86 """abstract base hook class""" 

87 

88 hook_registry = {} 

89 

90 def get_hooks(self) -> Dict: 

91 """gather all the child class hooks""" 

92 target_cls = [self.__class__] + list(self.__class__.__bases__) 

93 cls_names = [c.__name__ for c in target_cls] 

94 res = { 

95 k: v 

96 for k, v in MemEvalHook.hook_registry.items() 

97 if v.__qualname__.split(".")[-2] in cls_names 

98 } 

99 # print(self.__class__.__name__) 

100 # print(self.__class__.__subclasses__()) 

101 # print(MemEvalHook.__subclasses__()) 

102 return res 

103 

104 @staticmethod 

105 @abstractmethod 

106 def run_hooks(e): 

107 """must implement""" 

108 pass # pylint: disable=unnecessary-pass 

109 

110 

111# Testing 

112 

113# class A(MemEvalHook) : 

114# @hook_runner("a") 

115# def run_hooks(e) : 

116# print("a stuff") 

117 

118# @hook_runner("a2") 

119# def run_hooks(e) : 

120# print("a2 stuff") 

121 

122# class B(MemEvalHook) : 

123# @hook_runner("b") 

124# def run_hooks(e) : 

125# print("b stuff") 

126 

127# class C(A,B) : pass 

128 

129# k = C() 

130# h = k.get_hooks() 

131# print(h) 

132# for r in h.values() : r(None) 

133# k = A() 

134# h = k.get_hooks() 

135# print(h) 

136# for r in h.values() : r(None) 

137 

138# class D(A) : pass 

139# k = D() 

140# h = k.get_hooks() 

141# print(h) 

142# for r in h.values() : r(None) 

143 

144# class E() : 

145# def a() : print(1) 

146# def a() : print(2) 

147# E.a()