• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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."""
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import errno
22import hashlib
23import imp
24import os
25import platform
26import sys
27
28from tensorflow.python import _pywrap_python_op_gen
29from tensorflow.python.client import pywrap_tf_session as py_tf
30from tensorflow.python.eager import context
31from tensorflow.python.util import deprecation
32from tensorflow.python.util.tf_export import tf_export
33
34
35@tf_export('load_op_library')
36def load_op_library(library_filename):
37  """Loads a TensorFlow plugin, containing custom ops and kernels.
38
39  Pass "library_filename" to a platform-specific mechanism for dynamically
40  loading a library. The rules for determining the exact location of the
41  library are platform-specific and are not documented here. When the
42  library is loaded, ops and kernels registered in the library via the
43  `REGISTER_*` macros are made available in the TensorFlow process. Note
44  that ops with the same name as an existing op are rejected and not
45  registered with the process.
46
47  Args:
48    library_filename: Path to the plugin.
49      Relative or absolute filesystem path to a dynamic library file.
50
51  Returns:
52    A python module containing the Python wrappers for Ops defined in
53    the plugin.
54
55  Raises:
56    RuntimeError: when unable to load the library or get the python wrappers.
57  """
58  lib_handle = py_tf.TF_LoadLibrary(library_filename)
59  try:
60    wrappers = _pywrap_python_op_gen.GetPythonWrappers(
61        py_tf.TF_GetOpList(lib_handle))
62  finally:
63    # Delete the library handle to release any memory held in C
64    # that are no longer needed.
65    py_tf.TF_DeleteLibraryHandle(lib_handle)
66
67  # Get a unique name for the module.
68  module_name = hashlib.md5(wrappers).hexdigest()
69  if module_name in sys.modules:
70    return sys.modules[module_name]
71  module = imp.new_module(module_name)
72  # pylint: disable=exec-used
73  exec(wrappers, module.__dict__)
74  # Allow this to be recognized by AutoGraph.
75  setattr(module, '_IS_TENSORFLOW_PLUGIN', True)
76  sys.modules[module_name] = module
77  return module
78
79
80@deprecation.deprecated(date=None,
81                        instructions='Use `tf.load_library` instead.')
82@tf_export(v1=['load_file_system_library'])
83def load_file_system_library(library_filename):
84  """Loads a TensorFlow plugin, containing file system implementation.
85
86  Pass `library_filename` to a platform-specific mechanism for dynamically
87  loading a library. The rules for determining the exact location of the
88  library are platform-specific and are not documented here.
89
90  Args:
91    library_filename: Path to the plugin.
92      Relative or absolute filesystem path to a dynamic library file.
93
94  Returns:
95    None.
96
97  Raises:
98    RuntimeError: when unable to load the library.
99  """
100  py_tf.TF_LoadLibrary(library_filename)
101
102
103def _is_shared_object(filename):
104  """Check the file to see if it is a shared object, only using extension."""
105  if platform.system() == 'Linux':
106    if filename.endswith('.so'):
107      return True
108    else:
109      index = filename.rfind('.so.')
110      if index == -1:
111        return False
112      else:
113        # A shared object with the API version in filename
114        return filename[index + 4].isdecimal()
115  elif platform.system() == 'Darwin':
116    return filename.endswith('.dylib')
117  elif platform.system() == 'Windows':
118    return filename.endswith('.dll')
119  else:
120    return False
121
122
123@tf_export('load_library')
124def load_library(library_location):
125  """Loads a TensorFlow plugin.
126
127  "library_location" can be a path to a specific shared object, or a folder.
128  If it is a folder, all shared objects that are named "libtfkernel*" will be
129  loaded. When the library is loaded, kernels registered in the library via the
130  `REGISTER_*` macros are made available in the TensorFlow process.
131
132  Args:
133    library_location: Path to the plugin or the folder of plugins.
134      Relative or absolute filesystem path to a dynamic library file or folder.
135
136  Returns:
137    None
138
139  Raises:
140    OSError: When the file to be loaded is not found.
141    RuntimeError: when unable to load the library.
142  """
143  if os.path.exists(library_location):
144    if os.path.isdir(library_location):
145      directory_contents = os.listdir(library_location)
146
147      kernel_libraries = [
148          os.path.join(library_location, f) for f in directory_contents
149          if _is_shared_object(f)]
150    else:
151      kernel_libraries = [library_location]
152
153    for lib in kernel_libraries:
154      py_tf.TF_LoadLibrary(lib)
155
156  else:
157    raise OSError(
158        errno.ENOENT,
159        'The file or folder to load kernel libraries from does not exist.',
160        library_location)
161
162
163def load_pluggable_device_library(library_location):
164  """Loads a TensorFlow PluggableDevice plugin.
165
166  "library_location" can be a path to a specific shared object, or a folder.
167  If it is a folder, all shared objects will be loaded. when the library is
168  loaded, devices/kernels registered in the library via StreamExecutor C API
169  and Kernel/Op Registration C API are made available in TensorFlow process.
170
171  Args:
172    library_location: Path to the plugin or folder of plugins. Relative or
173      absolute filesystem path to a dynamic library file or folder.
174
175  Raises:
176    OSError: When the file to be loaded is not found.
177    RuntimeError: when unable to load the library.
178  """
179  if os.path.exists(library_location):
180    if os.path.isdir(library_location):
181      directory_contents = os.listdir(library_location)
182
183      pluggable_device_libraries = [
184          os.path.join(library_location, f)
185          for f in directory_contents
186          if _is_shared_object(f)
187      ]
188    else:
189      pluggable_device_libraries = [library_location]
190
191    for lib in pluggable_device_libraries:
192      py_tf.TF_LoadPluggableDeviceLibrary(lib)
193    # Reinitialized physical devices list after plugin registration.
194    context.context().reinitialize_physical_devices()
195  else:
196    raise OSError(
197        errno.ENOENT,
198        'The file or folder to load pluggable device libraries from does not '
199        'exist.', library_location)
200
201
202@tf_export('experimental.register_filesystem_plugin')
203def register_filesystem_plugin(plugin_location):
204  """Loads a TensorFlow FileSystem plugin.
205
206  Args:
207    plugin_location: Path to the plugin. Relative or absolute filesystem plugin
208      path to a dynamic library file.
209
210  Returns:
211    None
212
213  Raises:
214    OSError: When the file to be loaded is not found.
215    RuntimeError: when unable to load the library.
216  """
217  if os.path.exists(plugin_location):
218    py_tf.TF_RegisterFilesystemPlugin(plugin_location)
219
220  else:
221    raise OSError(errno.ENOENT,
222                  'The file to load file system plugin from does not exist.',
223                  plugin_location)
224