Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/core/distributed_runtime/preemption/gen_check_preemption_op.py: 32%

1"""Python wrappers around TensorFlow ops. 

2 

3This file is MACHINE GENERATED! Do not edit. 

4""" 

5 

6import collections 

7 

8from tensorflow.python import pywrap_tfe as pywrap_tfe 

9from tensorflow.python.eager import context as _context 

10from tensorflow.python.eager import core as _core 

11from tensorflow.python.eager import execute as _execute 

12from tensorflow.python.framework import dtypes as _dtypes 

13from tensorflow.security.fuzzing.py import annotation_types as _atypes 

14 

15from tensorflow.python.framework import op_def_registry as _op_def_registry 

16from tensorflow.python.framework import ops as _ops 

17from tensorflow.python.framework import op_def_library as _op_def_library 

18from tensorflow.python.util.deprecation import deprecated_endpoints 

19from tensorflow.python.util import dispatch as _dispatch 

20from tensorflow.python.util.tf_export import tf_export 

21 

22from typing import TypeVar 

23 

24@_dispatch.add_fallback_dispatch_list 

25@_dispatch.add_type_based_api_dispatcher 

26@tf_export('check_preemption') 

27def check_preemption(preemption_key="TF_DEFAULT_PREEMPTION_NOTICE_KEY", name=None): 

28 r"""Check if a preemption notice has been received in coordination service. 

29 

30 Args: 

31 preemption_key: An optional `string`. Defaults to `"TF_DEFAULT_PREEMPTION_NOTICE_KEY"`. 

32 Key for preemption check in coordination service. 

33 name: A name for the operation (optional). 

34 

35 Returns: 

36 The created Operation. 

37 """ 

38 _ctx = _context._context or _context.context() 

39 tld = _ctx._thread_local_data 

40 if tld.is_eager: 

41 try: 

42 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

43 _ctx, "CheckPreemption", name, "preemption_key", preemption_key) 

44 return _result 

45 except _core._NotOkStatusException as e: 

46 _ops.raise_from_not_ok_status(e, name) 

47 except _core._FallbackException: 

48 pass 

49 try: 

50 _result = _dispatcher_for_check_preemption( 

51 (preemption_key, name,), None) 

52 if _result is not NotImplemented: 

53 return _result 

54 return check_preemption_eager_fallback( 

55 preemption_key=preemption_key, name=name, ctx=_ctx) 

56 except _core._SymbolicException: 

57 pass # Add nodes to the TensorFlow graph. 

58 except (TypeError, ValueError): 

59 _result = _dispatch.dispatch( 

60 check_preemption, (), dict(preemption_key=preemption_key, 

61 name=name) 

62 ) 

63 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

64 return _result 

65 raise 

66 else: 

67 _result = _dispatcher_for_check_preemption( 

68 (preemption_key, name,), None) 

69 if _result is not NotImplemented: 

70 return _result 

71 # Add nodes to the TensorFlow graph. 

72 if preemption_key is None: 

73 preemption_key = "TF_DEFAULT_PREEMPTION_NOTICE_KEY" 

74 preemption_key = _execute.make_str(preemption_key, "preemption_key") 

75 try: 

76 _, _, _op, _outputs = _op_def_library._apply_op_helper( 

77 "CheckPreemption", preemption_key=preemption_key, name=name) 

78 except (TypeError, ValueError): 

79 _result = _dispatch.dispatch( 

80 check_preemption, (), dict(preemption_key=preemption_key, name=name) 

81 ) 

82 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

83 return _result 

84 raise 

85 return _op 

86CheckPreemption = tf_export("raw_ops.CheckPreemption")(_ops.to_raw_op(check_preemption)) 

87_dispatcher_for_check_preemption = check_preemption._tf_type_based_dispatcher.Dispatch 

88 

89 

90def check_preemption_eager_fallback(preemption_key, name, ctx): 

91 if preemption_key is None: 

92 preemption_key = "TF_DEFAULT_PREEMPTION_NOTICE_KEY" 

93 preemption_key = _execute.make_str(preemption_key, "preemption_key") 

94 _inputs_flat = [] 

95 _attrs = ("preemption_key", preemption_key) 

96 _result = _execute.execute(b"CheckPreemption", 0, inputs=_inputs_flat, 

97 attrs=_attrs, ctx=ctx, name=name) 

98 _result = None 

99 return _result 

100