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

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"""Tracer for the collective communication operations of core.dtensor.""" 

16import threading 

17from typing import Any, Callable, Dict 

18 

19from hyper_parallel.core.dtensor import _utils 

20 

21# Collective functions in hyper_parallel.core.dtensor._utils to intercept. 

22_COLLECTIVE_METHODS = ( 

23 "differentiable_all_gather_concat", 

24 "differentiable_all_to_all", 

25 "differentiable_all_reduce", 

26 "differentiable_reduce_scatter", 

27 "differentiable_all_to_all_single", 

28 "differentiable_all_to_all_single_async", 

29) 

30 

31 

32class CollectiveTracer: 

33 """Intercepts core.dtensor collective operations via monkey-patching. 

34 

35 Args: 

36 on_collective_call: Callback invoked after each collective with 

37 ``(method_name, args, kwargs, result)``. 

38 """ 

39 

40 _patch_lock = threading.Lock() 

41 

42 def __init__(self, on_collective_call: Callable) -> None: 

43 """Record the callback and start with nothing patched. 

44 

45 Args: 

46 on_collective_call: Callback invoked after each collective with 

47 ``(method_name, args, kwargs, result)``. 

48 """ 

49 self._callback = on_collective_call 

50 self._originals: Dict[str, object] = {} 

51 self._wrappers: Dict[str, object] = {} 

52 

53 def install(self) -> None: 

54 """Replace the ``_utils`` collective functions with tracing wrappers. 

55 

56 Names in :data:`_COLLECTIVE_METHODS` that ``_utils`` does not define are 

57 skipped, so a missing collective degrades to no tracing rather than an 

58 error. Installations are idempotent per name: a second tracer replaces 

59 the first tracer's wrapper and saves that wrapper as its own original. 

60 """ 

61 with self._patch_lock: 

62 for name in _COLLECTIVE_METHODS: 

63 if not hasattr(_utils, name): 

64 continue 

65 # Save the raw function for exact restoration. 

66 original_func = getattr(_utils, name) 

67 self._originals[name] = original_func 

68 

69 callback = self._callback 

70 method_name = name 

71 

72 def _make_wrapper(orig: Callable, cb: Callable, mname: str) -> Callable: 

73 def wrapper(*args: Any, **kwargs: Any) -> Any: 

74 """Call the original collective, then notify the callback. 

75 

76 A callback failure must never break the collective, so it 

77 is swallowed and the original result returned untouched. 

78 """ 

79 result = orig(*args, **kwargs) 

80 try: 

81 cb(mname, args, kwargs, result) 

82 except Exception: # pylint: disable=W0703 

83 pass # Never break production logic 

84 return result 

85 return wrapper 

86 

87 wrapper = _make_wrapper(original_func, callback, method_name) 

88 self._wrappers[name] = wrapper 

89 setattr(_utils, name, wrapper) 

90 

91 def uninstall(self) -> None: 

92 """Restore the original ``_utils`` collective functions. 

93 

94 Only restores the attributes this tracer installed. If an inner tracer 

95 patched the same name afterwards, restoring our saved original would 

96 silently strip that tracer's wrapper, so it is left in place. 

97 """ 

98 with self._patch_lock: 

99 for name, original_func in self._originals.items(): 

100 if getattr(_utils, name, None) is self._wrappers.get(name): 

101 setattr(_utils, name, original_func) 

102 self._originals.clear() 

103 self._wrappers.clear()