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 AnnotatedArgsContext, EagerOpExecutionContext 

23 

24 

25@dataclass 

26class _ExecuteContextBinding: 

27 ctx: EagerOpExecutionContext 

28 active: bool = True 

29 

30 

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

32 "ge_custom_op_execute_context", default=None 

33) 

34 

35 

36def get_execute_ctx() -> EagerOpExecutionContext: 

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

38 

39 binding = _CURRENT_EXECUTE_CONTEXT.get() 

40 if binding is None or not binding.active: 

41 raise RuntimeError( 

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

43 ) 

44 return binding.ctx 

45 

46 

47@contextmanager 

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

49 binding = _ExecuteContextBinding(ctx=ctx) 

50 token = _CURRENT_EXECUTE_CONTEXT.set(binding) 

51 try: 

52 yield 

53 finally: 

54 binding.active = False 

55 _CURRENT_EXECUTE_CONTEXT.reset(token) 

56 

57 

58@dataclass 

59class _DeclareLaunchArgsContextBinding: 

60 ctx: AnnotatedArgsContext 

61 active: bool = True 

62 

63 

64_CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT: ContextVar[ 

65 Optional[_DeclareLaunchArgsContextBinding] 

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

67 

68 

69def get_declare_launch_args_ctx() -> AnnotatedArgsContext: 

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

71 

72 binding = _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.get() 

73 if binding is None or not binding.active: 

74 raise RuntimeError( 

75 "get_declare_launch_args_ctx() is only available inside declare_launch_args" 

76 ) 

77 return binding.ctx 

78 

79 

80@contextmanager 

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

82 binding = _DeclareLaunchArgsContextBinding(ctx=ctx) 

83 token = _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.set(binding) 

84 try: 

85 yield 

86 finally: 

87 binding.active = False 

88 _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.reset(token)