Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/context.py: 100%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ----------------------------------------------------------------------------------------------------------- 

4# Copyright (c) 2026 Huawei Technologies Co., Ltd. 

5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 

6# CANN Open Software License Agreement Version 2.0 (the "License"). 

7# Please refer to the License for details. You may not use this file except in compliance with the License. 

8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 

9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. 

10# See LICENSE in the root of the software repository for the full text of the License. 

11# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Execution context access for schema-bound Python custom ops.""" 

14 

15from __future__ import annotations 

16 

17from contextlib import contextmanager 

18from contextvars import ContextVar 

19from dataclasses import dataclass 

20from typing import Iterator, Optional 

21 

22from ._native import ( 

23 AnnotatedArgsContext, 

24 CompilePlatformInfo, 

25 EagerOpExecutionContext, 

26 OpCompileContext, 

27) 

28 

29 

30@dataclass 

31class _ExecuteContextBinding: 

32 ctx: EagerOpExecutionContext 

33 active: bool = True 

34 

35 

36_CURRENT_EXECUTE_CONTEXT: ContextVar[Optional[_ExecuteContextBinding]] = ContextVar( 

37 "ge_custom_op_execute_context", default=None 

38) 

39 

40 

41def get_execute_ctx() -> EagerOpExecutionContext: 

42 """Return the borrowed context of the active schema-bound execute callback.""" 

43 

44 binding = _CURRENT_EXECUTE_CONTEXT.get() 

45 if binding is None or not binding.active: 

46 raise RuntimeError( 

47 "get_execute_ctx() is only available inside schema-bound execute" 

48 ) 

49 return binding.ctx 

50 

51 

52@contextmanager 

53def _execute_ctx_scope(ctx: EagerOpExecutionContext) -> Iterator[None]: 

54 binding = _ExecuteContextBinding(ctx=ctx) 

55 token = _CURRENT_EXECUTE_CONTEXT.set(binding) 

56 try: 

57 yield 

58 finally: 

59 binding.active = False 

60 _CURRENT_EXECUTE_CONTEXT.reset(token) 

61 

62 

63@dataclass 

64class _DeclareLaunchArgsContextBinding: 

65 ctx: AnnotatedArgsContext 

66 active: bool = True 

67 

68 

69_CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT: ContextVar[ 

70 Optional[_DeclareLaunchArgsContextBinding] 

71] = ContextVar("ge_custom_op_declare_launch_args_context", default=None) 

72 

73 

74def get_declare_launch_args_ctx() -> AnnotatedArgsContext: 

75 """Return the borrowed context of the active declare_launch_args callback.""" 

76 

77 binding = _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.get() 

78 if binding is None or not binding.active: 

79 raise RuntimeError( 

80 "get_declare_launch_args_ctx() is only available inside declare_launch_args" 

81 ) 

82 return binding.ctx 

83 

84 

85@contextmanager 

86def _declare_launch_args_ctx_scope(ctx: AnnotatedArgsContext) -> Iterator[None]: 

87 binding = _DeclareLaunchArgsContextBinding(ctx=ctx) 

88 token = _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.set(binding) 

89 try: 

90 yield 

91 finally: 

92 binding.active = False 

93 _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.reset(token) 

94 

95 

96@dataclass 

97class _CompileContextBinding: 

98 ctx: OpCompileContext 

99 active: bool = True 

100 

101 

102_CURRENT_COMPILE_CONTEXT: ContextVar[Optional[_CompileContextBinding]] = ContextVar( 

103 "ge_custom_op_compile_context", default=None 

104) 

105 

106 

107def get_compile_ctx() -> OpCompileContext: 

108 """Return the borrowed context of the active schema-bound compile callback.""" 

109 

110 binding = _CURRENT_COMPILE_CONTEXT.get() 

111 if binding is None or not binding.active: 

112 raise RuntimeError( 

113 "get_compile_ctx() is only available inside schema-bound compile" 

114 ) 

115 return binding.ctx 

116 

117 

118def get_compile_platform_info() -> CompilePlatformInfo: 

119 """Return the platform information view of the active compile callback.""" 

120 

121 return get_compile_ctx()._get_platform_info() 

122 

123 

124@contextmanager 

125def _compile_ctx_scope(ctx: OpCompileContext) -> Iterator[None]: 

126 binding = _CompileContextBinding(ctx=ctx) 

127 token = _CURRENT_COMPILE_CONTEXT.set(binding) 

128 try: 

129 yield 

130 finally: 

131 binding.active = False 

132 _CURRENT_COMPILE_CONTEXT.reset(token)