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

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"""Platform monkey-patch tracer for collective communication operations.""" 

16import threading 

17from typing import Callable, Dict, Optional 

18 

19from hyper_parallel.platform import get_platform 

20 

21# Collective methods to intercept on the platform class. 

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 platform 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): 

43 self._callback = on_collective_call 

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

45 self._platform_cls: Optional[type] = None 

46 

47 def install(self): 

48 """Replace platform collective methods with tracing wrappers.""" 

49 with self._patch_lock: 

50 platform = get_platform() 

51 self._platform_cls = type(platform) 

52 cls = self._platform_cls 

53 

54 for name in _COLLECTIVE_METHODS: 

55 if name not in cls.__dict__: 

56 continue 

57 # Save the raw descriptor (staticmethod wrapper) for exact restoration. 

58 original_descriptor = cls.__dict__[name] 

59 self._originals[name] = original_descriptor 

60 

61 # Unwrap staticmethod to get the underlying function. 

62 if isinstance(original_descriptor, staticmethod): 

63 original_func = original_descriptor.__func__ 

64 else: 

65 original_func = original_descriptor 

66 

67 callback = self._callback 

68 method_name = name 

69 

70 def _make_wrapper(orig, cb, mname): 

71 def wrapper(*args, **kwargs): 

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

73 try: 

74 cb(mname, args, kwargs, result) 

75 except Exception: # pylint: disable=W0703 

76 pass # Never break production logic 

77 return result 

78 return wrapper 

79 

80 wrapper = _make_wrapper(original_func, callback, method_name) 

81 setattr(cls, name, staticmethod(wrapper)) 

82 

83 def uninstall(self): 

84 """Restore original platform collective methods.""" 

85 with self._patch_lock: 

86 if self._platform_cls is None: 

87 return 

88 cls = self._platform_cls 

89 for name, original_descriptor in self._originals.items(): 

90 # Use type.__setattr__ to precisely restore the original descriptor. 

91 type.__setattr__(cls, name, original_descriptor) 

92 self._originals.clear() 

93 self._platform_cls = None