Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / shmem / _debug.py: 86%

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"""Local Runtime snapshots and best-effort SHMEM Allocation diagnostics.""" 

16 

17from __future__ import annotations 

18 

19__all__ = ["debug_state"] 

20 

21import os 

22import traceback 

23from collections.abc import Mapping 

24from typing import Any 

25 

26from ._lifecycle import _reference_count 

27from ._runtime import _load_native 

28 

29 

30_HEX_ADDRESS_FIELDS = frozenset({"allocation_base"}) 

31_LOG_LEVEL_ENV = "SHMEM_LOG_LEVEL" 

32_INTERNAL_FRAME_PREFIX = os.path.dirname(os.path.abspath(__file__)) + os.sep 

33 

34 

35def _format_debug_value(field: str, value: object) -> object: 

36 """Render known address fields in hex; everything else keeps its natural repr.""" 

37 if field in _HEX_ADDRESS_FIELDS: 

38 return hex(value) 

39 return value 

40 

41 

42def _render_debug_fields(fields: Mapping[str, object], indent: int = 0) -> list[str]: 

43 """Render diagnostic fields one per line, nesting mappings and mapping lists.""" 

44 pad = " " * indent 

45 lines: list[str] = [] 

46 for key, value in fields.items(): 

47 if isinstance(value, Mapping): 

48 lines.append(f"{pad}{key}:") 

49 lines.extend(_render_debug_fields(value, indent + 1)) 

50 elif isinstance(value, list) and value and all(isinstance(item, Mapping) for item in value): 

51 lines.append(f"{pad}{key}:") 

52 for item in value: 

53 rendered = ", ".join( 

54 f"{field}={_format_debug_value(field, item[field])}" for field in item 

55 ) 

56 lines.append(f"{pad} - {rendered}") 

57 else: 

58 lines.append(f"{pad}{key}: {_format_debug_value(key, value)}") 

59 return lines 

60 

61 

62def _log_allocation_site(tensor: Any) -> None: 

63 """Best-effort DEBUG line identifying the direct caller of :func:`empty`.""" 

64 if os.environ.get(_LOG_LEVEL_ENV) != "DEBUG": 

65 return 

66 try: 

67 frames = traceback.extract_stack() 

68 caller = next( 

69 ( 

70 frame 

71 for frame in reversed(frames) 

72 if not os.path.abspath(frame.filename).startswith(_INTERNAL_FRAME_PREFIX) 

73 ), 

74 None, 

75 ) 

76 callsite = "unknown" if caller is None else f"{caller.filename}:{caller.lineno} in {caller.name}" 

77 message = ( 

78 f"allocation_base={hex(tensor.data_ptr())} " 

79 f"bytes={tensor.numel() * tensor.element_size()} shape={tuple(tensor.shape)} dtype={tensor.dtype} " 

80 f"callsite={callsite}" 

81 ) 

82 _load_native()._log_debug(message) # pylint: disable=protected-access 

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

84 # Diagnostics must never turn a successful collective Allocation into a leaked Allocation. 

85 return 

86 

87 

88class _DebugState(dict): 

89 """Diagnostic snapshot dict rendered one field per line for human reading.""" 

90 

91 def __repr__(self) -> str: 

92 """Render diagnostic fields one per line for interactive inspection.""" 

93 return "\n".join(_render_debug_fields(self)) 

94 

95 

96def debug_state() -> dict[str, object]: 

97 """Return a local, read-only snapshot of current Runtime facts and its latest failure. 

98 

99 This query does not require an acquired Runtime reference. ``reference_count`` is the number of successful 

100 process-local ``shmem.acquire()`` calls not yet paired with ``shmem.release()``; ranks may report different values. 

101 While the Runtime is Ready, ``config`` carries the effective Runtime configuration (``heap_size_bytes``, 

102 ``timeout_seconds``, ``data_engine``, ``bootstrap_endpoint_base``). Heap usage is projected from requested bytes: 

103 ``allocated_bytes`` sums the requested sizes of all active Allocations, ``remaining_bytes`` is 

104 ``heap_size_bytes - allocated_bytes``, and ``active_allocations`` lists every active Allocation as 

105 ``{"allocation_id", "allocation_base", "allocation_bytes"}`` ordered by identity. Compare a Tensor's 

106 ``data_ptr()`` against ``allocation_base`` ranges to find its Allocation; the REPL renders addresses in hex while 

107 stored values remain integers. CANN rounds blocks to 16-byte multiples, so true device occupancy can be higher. 

108 ``leaked_allocations`` lists Allocations whose last Tensor Storage reference was dropped without 

109 ``shmem.free()``, in drop order, using the same fields as ``active_allocations``. 

110 

111 After the final clean ``shmem.release()``, the state is ``Uninitialized`` and lifecycle-specific Root, device, 

112 config, Allocation, maximum-active-allocation, leaked-allocation, and latest-failure fields are ``None`` until 

113 the next lifecycle starts. 

114 """ 

115 reference_count = _reference_count() 

116 native_state = dict(_load_native()._debug_state()) # pylint: disable=protected-access 

117 state = _DebugState() 

118 state["state"] = native_state.pop("state") 

119 state["reference_count"] = reference_count 

120 state.update(native_state) 

121 return state