Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/framework/load_library.py: 51%

69 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-03 07:57 +0000

1# Copyright 2015 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 

16"""Function for loading TensorFlow plugins.""" 

17import errno 

18import hashlib 

19import importlib 

20import os 

21import platform 

22import sys 

23 

24from tensorflow.python.client import pywrap_tf_session as py_tf 

25from tensorflow.python.eager import context 

26from tensorflow.python.framework import _pywrap_python_op_gen 

27from tensorflow.python.util import deprecation 

28from tensorflow.python.util.tf_export import tf_export 

29 

30 

31@tf_export('load_op_library') 

32def load_op_library(library_filename): 

33 """Loads a TensorFlow plugin, containing custom ops and kernels. 

34 

35 Pass "library_filename" to a platform-specific mechanism for dynamically 

36 loading a library. The rules for determining the exact location of the 

37 library are platform-specific and are not documented here. When the 

38 library is loaded, ops and kernels registered in the library via the 

39 `REGISTER_*` macros are made available in the TensorFlow process. Note 

40 that ops with the same name as an existing op are rejected and not 

41 registered with the process. 

42 

43 Args: 

44 library_filename: Path to the plugin. 

45 Relative or absolute filesystem path to a dynamic library file. 

46 

47 Returns: 

48 A python module containing the Python wrappers for Ops defined in 

49 the plugin. 

50 

51 Raises: 

52 RuntimeError: when unable to load the library or get the python wrappers. 

53 """ 

54 lib_handle = py_tf.TF_LoadLibrary(library_filename) 

55 try: 

56 wrappers = _pywrap_python_op_gen.GetPythonWrappers( 

57 py_tf.TF_GetOpList(lib_handle)) 

58 finally: 

59 # Delete the library handle to release any memory held in C 

60 # that are no longer needed. 

61 py_tf.TF_DeleteLibraryHandle(lib_handle) 

62 

63 # Get a unique name for the module. 

64 module_name = hashlib.sha1(wrappers).hexdigest() 

65 if module_name in sys.modules: 

66 return sys.modules[module_name] 

67 module_spec = importlib.machinery.ModuleSpec(module_name, None) 

68 module = importlib.util.module_from_spec(module_spec) 

69 # pylint: disable=exec-used 

70 exec(wrappers, module.__dict__) 

71 # Allow this to be recognized by AutoGraph. 

72 setattr(module, '_IS_TENSORFLOW_PLUGIN', True) 

73 sys.modules[module_name] = module 

74 return module 

75 

76 

77@deprecation.deprecated(date=None, 

78 instructions='Use `tf.load_library` instead.') 

79@tf_export(v1=['load_file_system_library']) 

80def load_file_system_library(library_filename): 

81 """Loads a TensorFlow plugin, containing file system implementation. 

82 

83 Pass `library_filename` to a platform-specific mechanism for dynamically 

84 loading a library. The rules for determining the exact location of the 

85 library are platform-specific and are not documented here. 

86 

87 Args: 

88 library_filename: Path to the plugin. 

89 Relative or absolute filesystem path to a dynamic library file. 

90 

91 Returns: 

92 None. 

93 

94 Raises: 

95 RuntimeError: when unable to load the library. 

96 """ 

97 py_tf.TF_LoadLibrary(library_filename) 

98 

99 

100def _is_shared_object(filename): 

101 """Check the file to see if it is a shared object, only using extension.""" 

102 if platform.system() == 'Linux': 

103 if filename.endswith('.so'): 

104 return True 

105 else: 

106 index = filename.rfind('.so.') 

107 if index == -1: 

108 return False 

109 else: 

110 # A shared object with the API version in filename 

111 return filename[index + 4].isdecimal() 

112 elif platform.system() == 'Darwin': 

113 return filename.endswith('.dylib') 

114 elif platform.system() == 'Windows': 

115 return filename.endswith('.dll') 

116 else: 

117 return False 

118 

119 

120@tf_export('load_library') 

121def load_library(library_location): 

122 """Loads a TensorFlow plugin. 

123 

124 "library_location" can be a path to a specific shared object, or a folder. 

125 If it is a folder, all shared objects that are named "libtfkernel*" will be 

126 loaded. When the library is loaded, kernels registered in the library via the 

127 `REGISTER_*` macros are made available in the TensorFlow process. 

128 

129 Args: 

130 library_location: Path to the plugin or the folder of plugins. 

131 Relative or absolute filesystem path to a dynamic library file or folder. 

132 

133 Returns: 

134 None 

135 

136 Raises: 

137 OSError: When the file to be loaded is not found. 

138 RuntimeError: when unable to load the library. 

139 """ 

140 if os.path.exists(library_location): 

141 if os.path.isdir(library_location): 

142 directory_contents = os.listdir(library_location) 

143 

144 kernel_libraries = [ 

145 os.path.join(library_location, f) for f in directory_contents 

146 if _is_shared_object(f)] 

147 else: 

148 kernel_libraries = [library_location] 

149 

150 for lib in kernel_libraries: 

151 py_tf.TF_LoadLibrary(lib) 

152 

153 else: 

154 raise OSError( 

155 errno.ENOENT, 

156 'The file or folder to load kernel libraries from does not exist.', 

157 library_location) 

158 

159 

160def load_pluggable_device_library(library_location): 

161 """Loads a TensorFlow PluggableDevice plugin. 

162 

163 "library_location" can be a path to a specific shared object, or a folder. 

164 If it is a folder, all shared objects will be loaded. when the library is 

165 loaded, devices/kernels registered in the library via StreamExecutor C API 

166 and Kernel/Op Registration C API are made available in TensorFlow process. 

167 

168 Args: 

169 library_location: Path to the plugin or folder of plugins. Relative or 

170 absolute filesystem path to a dynamic library file or folder. 

171 

172 Raises: 

173 OSError: When the file to be loaded is not found. 

174 RuntimeError: when unable to load the library. 

175 """ 

176 if os.path.exists(library_location): 

177 if os.path.isdir(library_location): 

178 directory_contents = os.listdir(library_location) 

179 

180 pluggable_device_libraries = [ 

181 os.path.join(library_location, f) 

182 for f in directory_contents 

183 if _is_shared_object(f) 

184 ] 

185 else: 

186 pluggable_device_libraries = [library_location] 

187 

188 for lib in pluggable_device_libraries: 

189 py_tf.TF_LoadPluggableDeviceLibrary(lib) 

190 # Reinitialized physical devices list after plugin registration. 

191 context.context().reinitialize_physical_devices() 

192 else: 

193 raise OSError( 

194 errno.ENOENT, 

195 'The file or folder to load pluggable device libraries from does not ' 

196 'exist.', library_location) 

197 

198 

199@tf_export('experimental.register_filesystem_plugin') 

200def register_filesystem_plugin(plugin_location): 

201 """Loads a TensorFlow FileSystem plugin. 

202 

203 Args: 

204 plugin_location: Path to the plugin. Relative or absolute filesystem plugin 

205 path to a dynamic library file. 

206 

207 Returns: 

208 None 

209 

210 Raises: 

211 OSError: When the file to be loaded is not found. 

212 RuntimeError: when unable to load the library. 

213 """ 

214 if os.path.exists(plugin_location): 

215 py_tf.TF_RegisterFilesystemPlugin(plugin_location) 

216 

217 else: 

218 raise OSError(errno.ENOENT, 

219 'The file to load file system plugin from does not exist.', 

220 plugin_location)