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

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.core.dtensor.dtensor import _distribute_module_named_modules 

19 

20 

21class ModuleTracker: 

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

23 

24 Args: 

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

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

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

28 """ 

29 

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

31 self._root = root_module 

32 self._callback = on_module_event 

33 self._hook_handles: List = [] 

34 self._fqn_map = {} 

35 

36 def install(self): 

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

38 # Build fully-qualified name map. 

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

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

41 

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

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

44 

45 def _make_pre_hook(module_fqn): 

46 def hook(module, inputs): 

47 # pylint: disable=W0613 

48 try: 

49 self._callback(module_fqn, "enter") 

50 except Exception: # pylint: disable=W0703 

51 pass 

52 return hook 

53 

54 def _make_post_hook(module_fqn): 

55 def hook(module, inputs, output): 

56 # pylint: disable=W0613 

57 try: 

58 self._callback(module_fqn, "exit") 

59 except Exception: # pylint: disable=W0703 

60 pass 

61 return hook 

62 

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

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

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

66 

67 def uninstall(self): 

68 """Remove all registered hooks.""" 

69 for handle in self._hook_handles: 

70 handle.remove() 

71 self._hook_handles.clear() 

72 self._fqn_map.clear()