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 )