github.com/alwaysproblem/mlserving-tutorial@v0.0.0-20221124033215-121cfddbfbf4/TFserving/CustomOp/custom-op/tf/tf_configure.bzl (about)

     1  """Setup TensorFlow as external dependency"""
     2  
     3  _TF_HEADER_DIR = "TF_HEADER_DIR"
     4  _TF_SHARED_LIBRARY_DIR = "TF_SHARED_LIBRARY_DIR"
     5  _TF_SHARED_LIBRARY_NAME = "TF_SHARED_LIBRARY_NAME"
     6  
     7  def _tpl(repository_ctx, tpl, substitutions = {}, out = None):
     8      if not out:
     9          out = tpl
    10      repository_ctx.template(
    11          out,
    12          Label("//tf:%s.tpl" % tpl),
    13          substitutions,
    14      )
    15  
    16  def _fail(msg):
    17      """Output failure message when auto configuration fails."""
    18      red = "\033[0;31m"
    19      no_color = "\033[0m"
    20      fail("%sPython Configuration Error:%s %s\n" % (red, no_color, msg))
    21  
    22  def _is_windows(repository_ctx):
    23      """Returns true if the host operating system is windows."""
    24      os_name = repository_ctx.os.name.lower()
    25      if os_name.find("windows") != -1:
    26          return True
    27      return False
    28  
    29  def _execute(
    30          repository_ctx,
    31          cmdline,
    32          error_msg = None,
    33          error_details = None,
    34          empty_stdout_fine = False):
    35      """Executes an arbitrary shell command.
    36      Args:
    37        repository_ctx: the repository_ctx object
    38        cmdline: list of strings, the command to execute
    39        error_msg: string, a summary of the error if the command fails
    40        error_details: string, details about the error or steps to fix it
    41        empty_stdout_fine: bool, if True, an empty stdout result is fine, otherwise
    42          it's an error
    43      Return:
    44        the result of repository_ctx.execute(cmdline)
    45      """
    46      result = repository_ctx.execute(cmdline)
    47      if result.stderr or not (empty_stdout_fine or result.stdout):
    48          _fail("\n".join([
    49              error_msg.strip() if error_msg else "Repository command failed",
    50              result.stderr.strip(),
    51              error_details if error_details else "",
    52          ]))
    53      return result
    54  
    55  def _read_dir(repository_ctx, src_dir):
    56      """Returns a string with all files in a directory.
    57      Finds all files inside a directory, traversing subfolders and following
    58      symlinks. The returned string contains the full path of all files
    59      separated by line breaks.
    60      """
    61      if _is_windows(repository_ctx):
    62          src_dir = src_dir.replace("/", "\\")
    63          find_result = _execute(
    64              repository_ctx,
    65              ["cmd.exe", "/c", "dir", src_dir, "/b", "/s", "/a-d"],
    66              empty_stdout_fine = True,
    67          )
    68  
    69          # src_files will be used in genrule.outs where the paths must
    70          # use forward slashes.
    71          result = find_result.stdout.replace("\\", "/")
    72      else:
    73          find_result = _execute(
    74              repository_ctx,
    75              ["find", src_dir, "-follow", "-type", "f"],
    76              empty_stdout_fine = True,
    77          )
    78          result = find_result.stdout
    79      return result
    80  
    81  def _genrule(genrule_name, command, outs):
    82      """Returns a string with a genrule.
    83  
    84      Genrule executes the given command and produces the given outputs.
    85  
    86      Args:
    87          genrule_name: A unique name for genrule target.
    88          command: The command to run.
    89          outs: A list of files generated by this rule.
    90  
    91      Returns:
    92          A genrule target.
    93      """
    94      return (
    95          "genrule(\n" +
    96          '    name = "' +
    97          genrule_name + '",\n' +
    98          "    outs = [\n" +
    99          outs +
   100          "\n    ],\n" +
   101          '    cmd = """\n' +
   102          command +
   103          '\n   """,\n' +
   104          ")\n"
   105      )
   106  
   107  def _norm_path(path):
   108      """Returns a path with '/' and remove the trailing slash."""
   109      path = path.replace("\\", "/")
   110      if path[-1] == "/":
   111          path = path[:-1]
   112      return path
   113  
   114  def _symlink_genrule_for_dir(
   115          repository_ctx,
   116          src_dir,
   117          dest_dir,
   118          genrule_name,
   119          src_files = [],
   120          dest_files = [],
   121          tf_pip_dir_rename_pair = []):
   122      """Returns a genrule to symlink(or copy if on Windows) a set of files.
   123  
   124      If src_dir is passed, files will be read from the given directory; otherwise
   125      we assume files are in src_files and dest_files.
   126  
   127      Args:
   128          repository_ctx: the repository_ctx object.
   129          src_dir: source directory.
   130          dest_dir: directory to create symlink in.
   131          genrule_name: genrule name.
   132          src_files: list of source files instead of src_dir.
   133          dest_files: list of corresonding destination files.
   134          tf_pip_dir_rename_pair: list of the pair of tf pip parent directory to
   135            replace. For example, in TF pip package, the source code is under
   136            "tensorflow_core", and we might want to replace it with
   137            "tensorflow" to match the header includes.
   138  
   139      Returns:
   140          genrule target that creates the symlinks.
   141      """
   142      # Check that tf_pip_dir_rename_pair has the right length
   143      tf_pip_dir_rename_pair_len = len(tf_pip_dir_rename_pair)
   144      if tf_pip_dir_rename_pair_len != 0 and tf_pip_dir_rename_pair_len !=2:
   145        _fail("The size of argument tf_pip_dir_rename_pair should be either 0 or 2, but %d is given." % tf_pip_dir_rename_pair_len)
   146  
   147      if src_dir != None:
   148          src_dir = _norm_path(src_dir)
   149          dest_dir = _norm_path(dest_dir)
   150          files = "\n".join(sorted(_read_dir(repository_ctx, src_dir).splitlines()))
   151  
   152          # Create a list with the src_dir stripped to use for outputs.
   153          if tf_pip_dir_rename_pair_len:
   154            dest_files = files.replace(src_dir, "").replace(tf_pip_dir_rename_pair[0], tf_pip_dir_rename_pair[1]).splitlines()
   155          else:
   156            dest_files = files.replace(src_dir, "").splitlines()
   157          src_files = files.splitlines()
   158      command = []
   159      outs = []
   160  
   161      for i in range(len(dest_files)):
   162          if dest_files[i] != "":
   163              # If we have only one file to link we do not want to use the dest_dir, as
   164              # $(@D) will include the full path to the file.
   165              dest = "$(@D)/" + dest_dir + dest_files[i] if len(dest_files) != 1 else "$(@D)/" + dest_files[i]
   166  
   167              # Copy the headers to create a sandboxable setup.
   168              cmd = "cp -f"
   169              command.append(cmd + ' "%s" "%s"' % (src_files[i], dest))
   170              outs.append('        "' + dest_dir + dest_files[i] + '",')
   171      dest_dir = "abc"
   172      genrule = _genrule(
   173          genrule_name,
   174          " && ".join(command),
   175          "\n".join(outs),
   176      )
   177      return genrule
   178  
   179  def _tf_pip_impl(repository_ctx):
   180      tf_header_dir = repository_ctx.os.environ[_TF_HEADER_DIR]
   181      tf_header_rule = _symlink_genrule_for_dir(
   182          repository_ctx,
   183          tf_header_dir,
   184          "include",
   185          "tf_header_include",
   186          tf_pip_dir_rename_pair = ["tensorflow_core", "tensorflow"]
   187      )
   188  
   189      tf_shared_library_dir = repository_ctx.os.environ[_TF_SHARED_LIBRARY_DIR]
   190      tf_shared_library_name = repository_ctx.os.environ[_TF_SHARED_LIBRARY_NAME]
   191      tf_shared_library_path = "%s/%s" % (tf_shared_library_dir, tf_shared_library_name)
   192  
   193      tf_shared_library_rule = _symlink_genrule_for_dir(
   194          repository_ctx,
   195          None,
   196          "",
   197          "libtensorflow_framework.so",
   198          [tf_shared_library_path],
   199          ["_pywrap_tensorflow_internal.lib"  if _is_windows(repository_ctx) else "libtensorflow_framework.so"],
   200      )
   201  
   202      _tpl(repository_ctx, "BUILD", {
   203          "%{TF_HEADER_GENRULE}": tf_header_rule,
   204          "%{TF_SHARED_LIBRARY_GENRULE}": tf_shared_library_rule,
   205      })
   206  
   207  tf_configure = repository_rule(
   208      implementation = _tf_pip_impl,
   209      environ = [
   210          _TF_HEADER_DIR,
   211          _TF_SHARED_LIBRARY_DIR,
   212          _TF_SHARED_LIBRARY_NAME,
   213      ],
   214  )