Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / common / session.py: 94%

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"""Process-local ownership for mutually exclusive MegaKernel diagnostics.""" 

16 

17from __future__ import annotations 

18 

19from threading import Lock 

20from typing import Protocol 

21 

22from hyper_multicore.diagnostics.common.call import DiagnosticCall, LaunchContext 

23 

24 

25class DiagnosticSession(Protocol): 

26 """Minimal contract implemented independently by profiler and Tensor dumper.""" 

27 

28 def is_capture_enabled(self) -> bool: 

29 """Return whether this step needs diagnostic invocation preparation.""" 

30 raise NotImplementedError 

31 

32 def prepare_call(self, context: LaunchContext) -> DiagnosticCall: 

33 """Return complete launch arguments and tool-owned invocation state. 

34 

35 Args: 

36 context: Selected invocation or operand context supplied by the diagnostic framework. 

37 """ 

38 raise NotImplementedError 

39 

40 

41_ACTIVE_LOCK = Lock() 

42_ACTIVE_KIND: str | None = None 

43_ACTIVE_SESSION: DiagnosticSession | None = None 

44 

45 

46def _activate_session(kind: str, session: DiagnosticSession) -> None: 

47 """Claim the single process-local MegaKernel diagnostic session.""" 

48 # pylint: disable=global-statement 

49 global _ACTIVE_KIND, _ACTIVE_SESSION 

50 with _ACTIVE_LOCK: 

51 if _ACTIVE_SESSION is not None: 

52 raise RuntimeError( 

53 "nested MegaKernel diagnostic contexts are not supported: " 

54 f"active={_ACTIVE_KIND}, requested={kind}" 

55 ) 

56 _ACTIVE_KIND = kind 

57 _ACTIVE_SESSION = session 

58 

59 

60def _deactivate_session(session: DiagnosticSession) -> None: 

61 """Release ownership when ``session`` is still the active owner.""" 

62 # pylint: disable=global-statement 

63 global _ACTIVE_KIND, _ACTIVE_SESSION 

64 with _ACTIVE_LOCK: 

65 if _ACTIVE_SESSION is session: 

66 _ACTIVE_KIND = None 

67 _ACTIVE_SESSION = None 

68 

69 

70def _get_active_session(kind: str | None = None) -> DiagnosticSession | None: 

71 """Return the active diagnostic session when its kind matches.""" 

72 with _ACTIVE_LOCK: 

73 if kind is not None and _ACTIVE_KIND != kind: 

74 return None 

75 return _ACTIVE_SESSION 

76 

77 

78def _reset_active_session_for_test() -> None: 

79 """Clear leaked diagnostic state after a unit-test failure.""" 

80 # pylint: disable=global-statement 

81 global _ACTIVE_KIND, _ACTIVE_SESSION 

82 with _ACTIVE_LOCK: 

83 _ACTIVE_KIND = None 

84 _ACTIVE_SESSION = None