Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / _automatic.py: 91%

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"""Version-isolated adapters for cleanup before Torch communication teardown.""" 

16 

17from __future__ import annotations 

18 

19import atexit 

20import functools 

21import logging 

22from collections.abc import Callable 

23from typing import Any, Optional 

24 

25import torch.distributed as dist 

26import torch.distributed.distributed_c10d as c10d 

27 

28from . import _collective 

29 

30_LOGGER = logging.getLogger(__name__) 

31 

32 

33def exchange(value: Any) -> list[Any]: 

34 """Exchange resource metadata on the registered cleanup Root.""" 

35 return _collective.exchange( 

36 value, 

37 group=_CLEANUP.root_group, 

38 failure_prefix="Multicore metadata exchange failed while converging rank results", 

39 ) 

40 

41 

42def collective_call( 

43 callback: Callable[[], None], 

44 *, 

45 stage: str, 

46) -> None: 

47 """Converge resource callback failures on the registered cleanup Root.""" 

48 _collective.collective_call( 

49 callback, 

50 stage=stage, 

51 group=_CLEANUP.root_group, 

52 ) 

53 

54 

55class _Cleanup: 

56 """Own a single exit hook and dependencies for the current runtime generation.""" 

57 

58 def __init__(self) -> None: 

59 """Start without native owners or patched framework entry points.""" 

60 self.callback: Optional[Callable[[], None]] = None 

61 self.dependencies: set[Any] = set() 

62 self.installed = False 

63 self.root_group: Any = None 

64 

65 def close(self) -> None: 

66 """Keep both the callback and communication available if release fails.""" 

67 if self.callback is not None: 

68 self.callback() 

69 self.callback = None 

70 self.dependencies.clear() 

71 

72 def at_exit(self) -> None: 

73 """Best-effort cleanup before torch_npu closes its device runtime.""" 

74 try: 

75 self.close() 

76 # Exit must continue through the framework's own device cleanup on failure. 

77 except Exception: # pylint: disable=broad-exception-caught 

78 _LOGGER.exception("Automatic multicore cleanup failed during process exit") 

79 

80 def install(self) -> None: 

81 """Intercept the backend boundary, including previously imported destroy aliases.""" 

82 if self.installed: 

83 return 

84 if hasattr(c10d.ProcessGroup, "shutdown"): 

85 owner, name = c10d.ProcessGroup, "shutdown" 

86 elif hasattr(c10d, "_shutdown_backend"): 

87 owner, name = c10d, "_shutdown_backend" 

88 else: 

89 raise RuntimeError("unsupported Torch process-group lifecycle; no automatic cleanup boundary") 

90 original = getattr(owner, name) 

91 

92 @functools.wraps(original) 

93 def shutdown(group: Any, *args: Any, **kwargs: Any) -> Any: 

94 """Run cleanup before a dependent communicator is destroyed. 

95 

96 Args: 

97 group: Communication dependency being destroyed. 

98 *args: Original framework arguments. 

99 **kwargs: Original framework keyword arguments. 

100 """ 

101 if group in self.dependencies: 

102 self.close() 

103 return original(group, *args, **kwargs) 

104 

105 setattr(owner, name, shutdown) 

106 # A real NPU input requires torch_npu initialization before this first binding. 

107 atexit.register(self.at_exit) 

108 self.installed = True 

109 

110 

111_CLEANUP = _Cleanup() 

112 

113 

114def register(callback: Callable[[], None], root_group: Any) -> None: 

115 """Register native owners without changing the user's process-group API. 

116 

117 Args: 

118 callback: Collective release operation; failures must preserve ownership. 

119 root_group: SHMEM Root dependency and lifecycle metadata group. 

120 """ 

121 _CLEANUP.install() 

122 _CLEANUP.root_group = root_group 

123 _CLEANUP.callback = callback 

124 _CLEANUP.dependencies.update(group for group in (root_group, dist.group.WORLD) if group is not None)