xref: /aosp_15_r20/external/tensorflow/third_party/repo.bzl (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2017 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"""Utilities for defining TensorFlow Bazel dependencies."""
16
17def tf_mirror_urls(url):
18    """A helper for generating TF-mirror versions of URLs.
19
20    Given a URL, it returns a list of the TF-mirror cache version of that URL
21    and the original URL, suitable for use in `urls` field of `tf_http_archive`.
22    """
23    if not url.startswith("https://"):
24        return [url]
25    return [
26        "https://storage.googleapis.com/mirror.tensorflow.org/%s" % url[8:],
27        url,
28    ]
29
30def _get_env_var(ctx, name):
31    if name in ctx.os.environ:
32        return ctx.os.environ[name]
33    else:
34        return None
35
36# Checks if we should use the system lib instead of the bundled one
37def _use_system_lib(ctx, name):
38    syslibenv = _get_env_var(ctx, "TF_SYSTEM_LIBS")
39    if not syslibenv:
40        return False
41    return name in [n.strip() for n in syslibenv.split(",")]
42
43def _get_link_dict(ctx, link_files, build_file):
44    link_dict = {ctx.path(v): ctx.path(Label(k)) for k, v in link_files.items()}
45    if build_file:
46        # Use BUILD.bazel because it takes precedence over BUILD.
47        link_dict[ctx.path("BUILD.bazel")] = ctx.path(Label(build_file))
48    return link_dict
49
50def _tf_http_archive_impl(ctx):
51    # Construct all paths early on to prevent rule restart. We want the
52    # attributes to be strings instead of labels because they refer to files
53    # in the TensorFlow repository, not files in repos depending on TensorFlow.
54    # See also https://github.com/bazelbuild/bazel/issues/10515.
55    link_dict = _get_link_dict(ctx, ctx.attr.link_files, ctx.attr.build_file)
56
57    # For some reason, we need to "resolve" labels once before the
58    # download_and_extract otherwise it'll invalidate and re-download the
59    # archive each time.
60    # https://github.com/bazelbuild/bazel/issues/10515
61    patch_files = ctx.attr.patch_file
62    for patch_file in patch_files:
63        if patch_file:
64            ctx.path(Label(patch_file))
65
66    if _use_system_lib(ctx, ctx.attr.name):
67        link_dict.update(_get_link_dict(
68            ctx = ctx,
69            link_files = ctx.attr.system_link_files,
70            build_file = ctx.attr.system_build_file,
71        ))
72    else:
73        ctx.download_and_extract(
74            url = ctx.attr.urls,
75            sha256 = ctx.attr.sha256,
76            type = ctx.attr.type,
77            stripPrefix = ctx.attr.strip_prefix,
78        )
79        if patch_files:
80            for patch_file in patch_files:
81                patch_file = ctx.path(Label(patch_file)) if patch_file else None
82                if patch_file:
83                    ctx.patch(patch_file, strip = 1)
84
85    for dst, src in link_dict.items():
86        ctx.delete(dst)
87        ctx.symlink(src, dst)
88
89_tf_http_archive = repository_rule(
90    implementation = _tf_http_archive_impl,
91    attrs = {
92        "sha256": attr.string(mandatory = True),
93        "urls": attr.string_list(mandatory = True),
94        "strip_prefix": attr.string(),
95        "type": attr.string(),
96        "patch_file": attr.string_list(),
97        "build_file": attr.string(),
98        "system_build_file": attr.string(),
99        "link_files": attr.string_dict(),
100        "system_link_files": attr.string_dict(),
101    },
102    environ = ["TF_SYSTEM_LIBS"],
103)
104
105def tf_http_archive(name, sha256, urls, **kwargs):
106    """Downloads and creates Bazel repos for dependencies.
107
108    This is a swappable replacement for both http_archive() and
109    new_http_archive() that offers some additional features. It also helps
110    ensure best practices are followed.
111
112    File arguments are relative to the TensorFlow repository by default. Dependent
113    repositories that use this rule should refer to files either with absolute
114    labels (e.g. '@foo//:bar') or from a label created in their repository (e.g.
115    'str(Label("//:bar"))').
116    """
117    if len(urls) < 2:
118        fail("tf_http_archive(urls) must have redundant URLs.")
119
120    if not any([mirror in urls[0] for mirror in (
121        "mirror.tensorflow.org",
122        "mirror.bazel.build",
123        "storage.googleapis.com",
124    )]):
125        fail("The first entry of tf_http_archive(urls) must be a mirror " +
126             "URL, preferrably mirror.tensorflow.org. Even if you don't have " +
127             "permission to mirror the file, please put the correctly " +
128             "formatted mirror URL there anyway, because someone will come " +
129             "along shortly thereafter and mirror the file.")
130
131    if native.existing_rule(name):
132        print("\n\033[1;33mWarning:\033[0m skipping import of repository '" +
133              name + "' because it already exists.\n")
134        return
135
136    _tf_http_archive(
137        name = name,
138        sha256 = sha256,
139        urls = urls,
140        **kwargs
141    )
142