Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/eager/cancellation.py: 58%

1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 

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"""Cancellation support for eager execution.""" 

16 

17from tensorflow.python import pywrap_tfe 

18 

19 

20class CancellationManager(object): 

21 """A mechanism for cancelling blocking computation.""" 

22 

23 __slots__ = ["_impl"] 

24 

25 def __init__(self): 

26 self._impl = pywrap_tfe.TFE_NewCancellationManager() 

27 

28 @property 

29 def is_cancelled(self): 

30 """Returns `True` if `CancellationManager.start_cancel` has been called.""" 

31 return pywrap_tfe.TFE_CancellationManagerIsCancelled(self._impl) 

32 

33 def start_cancel(self): 

34 """Cancels blocking operations that have been registered with this object.""" 

35 pywrap_tfe.TFE_CancellationManagerStartCancel(self._impl) 

36 

37 def get_cancelable_function(self, concrete_function): 

38 def cancellable(*args, **kwargs): 

39 with CancellationManagerContext(self): 

40 return concrete_function(*args, **kwargs) 

41 return cancellable 

42 

43_active_context = None 

44 

45 

46def context(): 

47 return _active_context 

48 

49 

50class CancellationManagerContext: 

51 """A Python context for wrapping a cancellable ConcreteFunction.""" 

52 

53 def __init__(self, cancellation_manager): 

54 self._cancellation_manager = cancellation_manager 

55 

56 def __enter__(self): 

57 global _active_context 

58 _active_context = self._cancellation_manager 

59 

60 def __exit__(self, exc_type, exc_value, exc_tb): 

61 global _active_context 

62 _active_context = None