Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / debug / _module_tracker.py: 22%

1# Copyright 2026 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"""nn.Module forward/backward hook tracker for CommDebugMode.""" 

16from typing import Callable, List 

17 

18from hyper_parallel.platform import get_platform 

19from hyper_parallel.core.dtensor.dtensor import _distribute_module_named_modules 

20 

21platform = get_platform() 

22 

23 

24class ModuleTracker: 

25 """Registers hooks on *root_module* to track forward enter/exit events. 

26 

27 Args: 

28 root_module: The top-level ``nn.Module`` to instrument. 

29 on_module_event: Callback with ``(module_fqn, event_type)`` where 

30 *event_type* is ``"enter"`` or ``"exit"``. 

31 """ 

32 

33 def __init__(self, root_module, on_module_event: Callable): 

34 self._root = root_module 

35 self._callback = on_module_event 

36 self._hook_handles: List = [] 

37 self._fqn_map = {} 

38 

39 def install(self): 

40 """Register forward_pre_hook and forward_hook on all sub-modules.""" 

41 # Build fully-qualified name map. 

42 for name, mod in _distribute_module_named_modules(self._root): 

43 self._fqn_map[id(mod)] = name or "(root)" 

44 

45 for _, mod in _distribute_module_named_modules(self._root): 

46 fqn = self._fqn_map.get(id(mod), "unknown") 

47 

48 def _make_pre_hook(module_fqn): 

49 def hook(module, inputs): 

50 # pylint: disable=W0613 

51 try: 

52 self._callback(module_fqn, "enter") 

53 except Exception: # pylint: disable=W0703 

54 pass 

55 return hook 

56 

57 def _make_post_hook(module_fqn): 

58 def hook(module, inputs, output): 

59 # pylint: disable=W0613 

60 try: 

61 self._callback(module_fqn, "exit") 

62 except Exception: # pylint: disable=W0703 

63 pass 

64 return hook 

65 

66 handle_pre = mod.register_forward_pre_hook(_make_pre_hook(fqn)) 

67 handle_post = mod.register_forward_hook(_make_post_hook(fqn)) 

68 self._hook_handles.extend([handle_pre, handle_post]) 

69 

70 def uninstall(self): 

71 """Remove all registered hooks.""" 

72 for handle in self._hook_handles: 

73 handle.remove() 

74 self._hook_handles.clear() 

75 self._fqn_map.clear()