Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / module.py: 92%
219 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
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"""Execution-resource ownership shared by Torch multicore modules."""
17from __future__ import annotations
19import itertools
20import weakref
21from collections.abc import Callable, Iterable
22from typing import Any, Optional
24import torch
25import torch.distributed as dist
27from .. import _automatic
30class _ExecutionResourceGroup:
31 """Process-owned runtime resources shared by compatible modules."""
33 def __init__(
34 self,
35 identifier: int,
36 member_token: int,
37 specification: Any,
38 compatibility_key: Any,
39 scope_key: Any,
40 ) -> None:
41 """Initialize one unbound process-owned resource group."""
42 self.identifier = identifier
43 self.members = {member_token}
44 self.specification = specification
45 self.compatibility_key = compatibility_key
46 self.scope_key = scope_key
47 self.resources = None
48 self.binding = None
49 self.shared = False
50 self.binding_id = None
51 self.closing = False
52 self.closed = False
53 self.module_type = None
56class _MulticoreResourceManager:
57 """Track lazy resources while their owning modules are alive."""
59 def __init__(self) -> None:
60 """Initialize empty process resource state."""
61 self._next_group_id = itertools.count()
62 self._next_member_token = itertools.count()
63 self._next_binding_id = itertools.count()
64 self._groups: dict[int, _ExecutionResourceGroup] = {}
65 self.runtime_release: Optional[Callable[[], None]] = None
66 self._runtime_released = False
68 def recycle(self, target: _ExecutionResourceGroup, binding: Any, signature: Any) -> None:
69 """Reconcile orphan slots on first binding, never on a local GC decision.
71 Args:
72 target: New owner requesting resources.
73 binding: Local device and dtype binding.
74 signature: Rank-independent configuration to compare across peers.
75 """
76 bound = self._bound_groups(self._groups.values())
77 manifest = tuple((group.binding_id, group.resources.lifecycle_signature()) for group in bound)
78 states = []
79 for group in bound:
80 idle = not group.members and not group.closing and group.resources.can_close()
81 compatible = (
82 (group.compatibility_key, group.binding, group.scope_key, group.shared, group.module_type)
83 == (target.compatibility_key, binding, target.scope_key, target.shared, target.module_type)
84 )
85 states.append((idle, compatible))
86 peers = self._exchange((manifest, signature, states))
87 if any(peer[:2] != (manifest, signature) for peer in peers):
88 raise RuntimeError("multicore resource binding order or specifications differ across ranks")
89 reusable = None
90 for index, group in enumerate(bound):
91 if not all(peer[2][index][0] for peer in peers):
92 continue
93 if reusable is None and all(peer[2][index][1] for peer in peers):
94 reusable = group
95 else:
96 self._close_group_collectively(group)
97 if reusable is not None:
98 target.resources, reusable.resources = reusable.resources, None
99 target.binding, target.binding_id = reusable.binding, reusable.binding_id
100 reusable.closed = True
101 self._groups.pop(reusable.identifier)
103 def create_group(
104 self,
105 specification: Any,
106 compatibility_key: Any,
107 scope_key: Any,
108 ) -> tuple[_ExecutionResourceGroup, int]:
109 """Create one unbound resource group and its first membership token.
111 Args:
112 specification: Resource specification supplied by the owning module.
113 compatibility_key: Key used to determine whether resources can be shared.
114 scope_key: Execution scope that owns the resource group.
115 """
116 group_id = next(self._next_group_id)
117 member_token = next(self._next_member_token)
118 group = _ExecutionResourceGroup(
119 group_id,
120 member_token,
121 specification,
122 compatibility_key,
123 scope_key,
124 )
125 self._groups[group_id] = group
126 return group, member_token
128 def add_member(self, group: _ExecutionResourceGroup) -> int:
129 """Add a new membership token to an existing group.
131 Args:
132 group: Existing resource group receiving the member.
133 """
134 member_token = next(self._next_member_token)
135 group.members.add(member_token)
136 return member_token
138 def retire(self, group_id: int, member_token: int) -> None:
139 """Retire membership; native resources wait for coordinated shutdown.
141 Args:
142 group_id: Identifier of the resource group.
143 member_token: Membership token to retire.
144 """
145 group = self._groups.get(group_id)
146 if group is None:
147 return
148 group.members.discard(member_token)
149 if group.members or group.resources is not None:
150 return
151 group.closed = True
152 self._groups.pop(group_id, None)
154 @staticmethod
155 def _exchange(value: Any) -> list[Any]:
156 """Exchange lifecycle metadata on the registered communication root."""
157 return _automatic.exchange(value)
159 def close_groups(self, groups: Iterable[_ExecutionResourceGroup] | None = None) -> None:
160 """Close all automatic owners, or the last explicit member's group.
162 Args:
163 groups: Selected groups, or all registered groups for automatic cleanup.
164 """
165 groups = list(self._groups.values() if groups is None else groups)
166 bound = self._bound_groups(groups)
167 self._validate_collective_close(bound)
168 for group in bound:
169 self._close_group_collectively(group)
170 for group in groups:
171 if group.resources is None:
172 group.closed = True
173 group.members.clear()
174 self._groups.pop(group.identifier, None)
175 if self.runtime_release is not None and not any(group.resources is not None
176 for group in self._groups.values()):
177 _automatic.collective_call(
178 lambda: self._release_runtime(self.runtime_release),
179 stage="multicore Runtime release",
180 )
181 self.runtime_release = None
182 self._runtime_released = False
184 def _release_runtime(self, release: Callable[[], None]) -> None:
185 """Keep successful local release idempotent until every peer acknowledges it."""
186 if not self._runtime_released:
187 release()
188 self._runtime_released = True
190 @staticmethod
191 def _bound_groups(groups: Iterable[_ExecutionResourceGroup]) -> list[_ExecutionResourceGroup]:
192 """Order native owners by collective binding rather than local construction."""
193 return sorted((group for group in groups if group.resources is not None), key=lambda group: group.binding_id)
195 def _validate_collective_close(self, bound: list[_ExecutionResourceGroup]) -> None:
196 """Require matching idle resources on every rank before release."""
197 manifest = tuple((group.binding_id, group.resources.lifecycle_signature()) for group in bound)
198 ready = all(group.resources.can_close() for group in bound)
199 peers = self._exchange((manifest, ready))
200 if any(peer[0] != manifest for peer in peers):
201 raise RuntimeError("multicore resource manifests differ across ranks")
202 if not all(peer[1] for peer in peers):
203 raise RuntimeError("multicore cleanup requires idle workspaces and no pending backward graphs")
205 def _close_group_collectively(self, group: _ExecutionResourceGroup) -> None:
206 """Close group resources collectively while retaining ownership on failure."""
207 group.closing = True
208 _automatic.collective_call(group.resources.close, stage="multicore resource close")
209 # Ownership transfers only after every peer succeeds; retries require the same manifest.
210 group.resources = None
211 group.binding = None
212 group.closed = True
213 group.members.clear()
214 self._groups.pop(group.identifier, None)
216 def active_specifications(self, scope_key: Any) -> tuple[Any, ...]:
217 """Return one specification per live or native-bound resource group.
219 Args:
220 scope_key: Execution scope whose resource groups are selected.
221 """
222 specifications = []
223 for group in self._groups.values():
224 same_scope = group.scope_key == scope_key
225 is_active = bool(group.members) or group.resources is not None
226 if same_scope and is_active:
227 specifications.append(group.specification)
228 return tuple(specifications)
231_RESOURCE_MANAGER = _MulticoreResourceManager()
234class MulticoreModule(torch.nn.Module):
235 """Base class for Torch modules that lazily own multicore resources."""
237 def __init__(
238 self,
239 *,
240 resource_specification: Any,
241 resource_compatibility_key: Any,
242 resource_scope_key: Any,
243 ) -> None:
244 """Initialize an unbound, process-owned execution-resource group."""
245 super().__init__()
246 group, member_token = _RESOURCE_MANAGER.create_group(
247 resource_specification,
248 resource_compatibility_key,
249 resource_scope_key,
250 )
251 self._resource_group = group
252 self._resource_member_token = member_token
253 self._resource_closed = False
254 self._resource_finalizer = weakref.finalize(
255 self,
256 _RESOURCE_MANAGER.retire,
257 group.identifier,
258 member_token,
259 )
261 @classmethod
262 def share_execution_resources(cls, modules: Iterable[MulticoreModule]) -> None:
263 """Share one synchronous workspace among compatible serial modules.
265 Sharing must be configured before any member executes. Modules retain
266 independent parameters and optimizer state.
268 Args:
269 modules: Compatible modules that execute serially.
270 """
271 members = tuple(modules)
272 first_group = cls._validate_shared_members(members)
273 if all(
274 member._resource_group is first_group # pylint: disable=protected-access
275 for member in members
276 ):
277 return
278 cls._validate_unbound_groups(members, first_group)
279 first_group.shared = True
280 for member in members[1:]:
281 member._move_to_resource_group(first_group) # pylint: disable=protected-access
283 @classmethod
284 def _validate_shared_members(
285 cls,
286 members: tuple[MulticoreModule, ...],
287 ) -> _ExecutionResourceGroup:
288 """Validate the requested members and return the leading group."""
289 if not members:
290 raise ValueError("shared multicore execution requires at least one module.")
291 if len({id(member) for member in members}) != len(members):
292 raise ValueError("shared multicore execution cannot contain duplicate modules.")
293 if any(not isinstance(member, cls) for member in members):
294 actual = [type(member).__name__ for member in members]
295 raise TypeError(
296 f"all shared modules must be {cls.__name__} instances, got {actual}."
297 )
298 if any(member._resource_closed or member._resource_group.closed # pylint: disable=protected-access
299 for member in members):
300 raise RuntimeError("closed multicore modules cannot share execution resources.")
301 if any(member._resource_group.closing # pylint: disable=protected-access
302 for member in members):
303 raise RuntimeError("closing multicore modules cannot share execution resources.")
304 concrete_types = {type(member) for member in members}
305 if len(concrete_types) != 1:
306 raise TypeError("shared multicore execution requires one concrete module type.")
307 return members[0]._resource_group # pylint: disable=protected-access
309 @staticmethod
310 def _validate_unbound_groups(
311 members: tuple[MulticoreModule, ...],
312 first_group: _ExecutionResourceGroup,
313 ) -> None:
314 """Require compatible, single-member groups before merging them."""
315 for member in members:
316 group = member._resource_group # pylint: disable=protected-access
317 if group.compatibility_key != first_group.compatibility_key:
318 raise ValueError(
319 "shared multicore modules must have identical configuration and scope."
320 )
321 if group.resources is not None or group.binding is not None or len(group.members) != 1:
322 raise RuntimeError(
323 "execution resources must be shared before first use or previous grouping."
324 )
326 def _move_to_resource_group(self, target: _ExecutionResourceGroup) -> None:
327 """Move this module to ``target`` without invoking native cleanup."""
328 old_group = self._resource_group
329 old_token = self._resource_member_token
330 self._resource_finalizer.detach()
331 _RESOURCE_MANAGER.retire(
332 old_group.identifier,
333 old_token,
334 )
335 member_token = _RESOURCE_MANAGER.add_member(target)
336 self._resource_group = target
337 self._resource_member_token = member_token
338 self._resource_finalizer = weakref.finalize(
339 self,
340 _RESOURCE_MANAGER.retire,
341 target.identifier,
342 member_token,
343 )
345 def _get_execution_resources(self, tensor: Any) -> Any:
346 """Create or return resources bound to ``tensor`` device and dtype."""
347 if self._resource_closed or self._resource_group.closed:
348 raise RuntimeError("cannot execute a closed multicore module.")
349 if self._resource_group.closing:
350 raise RuntimeError("cannot execute a closing multicore module; finish cleanup before reuse.")
351 group = self._resource_group
352 binding = self._execution_binding(tensor)
353 if group.resources is None:
354 group.module_type = type(self)
355 _automatic.register(_RESOURCE_MANAGER.close_groups, self._root_group())
356 _RESOURCE_MANAGER.recycle(group, binding, self._resource_signature(tensor))
357 if group.resources is None:
358 active_specs = _RESOURCE_MANAGER.active_specifications(group.scope_key)
359 group.resources = self._create_execution_resources(
360 tensor,
361 shared=group.shared,
362 active_specifications=active_specs,
363 )
364 group.binding = binding
365 group.binding_id = next(_RESOURCE_MANAGER._next_binding_id) # pylint: disable=protected-access
366 if _RESOURCE_MANAGER.runtime_release is None:
367 self._register_runtime_release(_RESOURCE_MANAGER, group.resources)
368 elif binding != group.binding:
369 raise ValueError(
370 "multicore resources are bound to the first input device/dtype "
371 f"{group.binding}, got {binding}."
372 )
373 return group.resources
375 def _resource_signature(self, tensor: Any) -> Any:
376 """Identify a collective first binding without rank-local device indices."""
377 return type(self).__module__, type(self).__qualname__, str(tensor.dtype)
379 def _register_runtime_release(
380 self,
381 manager: _MulticoreResourceManager,
382 resources: Any,
383 ) -> None:
384 """Optionally register a final release for resources with a shared Runtime."""
386 def _root_group(self) -> Any:
387 """Return the communication dependency used by these resources."""
388 return dist.group.WORLD
390 @staticmethod
391 def _execution_binding(tensor: Any) -> tuple[Any, Any]:
392 """Return the tensor binding used for resource compatibility."""
393 try:
394 return tensor.device, tensor.dtype
395 except AttributeError as error:
396 raise TypeError("multicore modules require a Torch tensor input.") from error
398 def _create_execution_resources(
399 self,
400 tensor: Any,
401 *,
402 shared: bool,
403 active_specifications: tuple[Any, ...],
404 ) -> Any:
405 """Create resources for the first input binding."""
406 raise NotImplementedError
408 def close(self) -> None:
409 """Release this member, closing shared resources only after the last member.
411 Call in the same order on all WORLD ranks after the last backward.
412 Recoverable failures retain ownership for retry; partially closed modules
413 cannot execute again. Native communication failures may require restart.
414 """
415 if self._resource_closed:
416 return
417 group = self._resource_group
418 if group.resources is not None and group.members == {self._resource_member_token}:
419 _RESOURCE_MANAGER.close_groups((group,))
420 elif group.closed:
421 # A previous attempt may have freed buffers but failed the final runtime release.
422 _RESOURCE_MANAGER.close_groups(())
423 else:
424 _RESOURCE_MANAGER.retire(group.identifier, self._resource_member_token)
425 self._resource_closed = True
426 self._resource_finalizer.detach()