• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""Repository rule to create a platform for a docker image to be used with RBE."""
2
3def _remote_platform_configure_impl(repository_ctx):
4    platform = repository_ctx.attr.platform
5    if platform == "local":
6        os = repository_ctx.os.name.lower()
7        if os.startswith("windows"):
8            platform = "windows"
9        elif os.startswith("mac os"):
10            platform = "osx"
11        else:
12            platform = "linux"
13
14    cpu = "x86_64"
15    machine_type = repository_ctx.execute(["bash", "-c", "echo $MACHTYPE"]).stdout
16    if (machine_type.startswith("ppc") or
17        machine_type.startswith("powerpc")):
18        cpu = "ppc"
19    elif machine_type.startswith("s390x"):
20        cpu = "s390x"
21    elif machine_type.startswith("aarch64"):
22        cpu = "aarch64"
23    elif machine_type.startswith("arm64"):
24        cpu = "aarch64"
25    elif machine_type.startswith("arm"):
26        cpu = "arm"
27    elif machine_type.startswith("mips64"):
28        cpu = "mips64"
29    elif machine_type.startswith("riscv64"):
30        cpu = "riscv64"
31
32    exec_properties = repository_ctx.attr.platform_exec_properties
33
34    serialized_exec_properties = "{"
35    for k, v in exec_properties.items():
36        serialized_exec_properties += "\"%s\" : \"%s\"," % (k, v)
37    serialized_exec_properties += "}"
38
39    repository_ctx.template(
40        "BUILD",
41        Label("@org_tensorflow//third_party/remote_config:BUILD.tpl"),
42        {
43            "%{platform}": platform,
44            "%{exec_properties}": serialized_exec_properties,
45            "%{cpu}": cpu,
46        },
47    )
48
49remote_platform_configure = repository_rule(
50    implementation = _remote_platform_configure_impl,
51    attrs = {
52        "platform_exec_properties": attr.string_dict(mandatory = True),
53        "platform": attr.string(default = "linux", values = ["linux", "windows", "local"]),
54    },
55)
56