56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
"""Repository rule to create a platform for a docker image to be used with RBE."""
|
|
|
|
def _remote_platform_configure_impl(repository_ctx):
|
|
platform = repository_ctx.attr.platform
|
|
if platform == "local":
|
|
os = repository_ctx.os.name.lower()
|
|
if os.startswith("windows"):
|
|
platform = "windows"
|
|
elif os.startswith("mac os"):
|
|
platform = "osx"
|
|
else:
|
|
platform = "linux"
|
|
|
|
cpu = "x86_64"
|
|
machine_type = repository_ctx.execute(["bash", "-c", "echo $MACHTYPE"]).stdout
|
|
if (machine_type.startswith("ppc") or
|
|
machine_type.startswith("powerpc")):
|
|
cpu = "ppc"
|
|
elif machine_type.startswith("s390x"):
|
|
cpu = "s390x"
|
|
elif machine_type.startswith("aarch64"):
|
|
cpu = "aarch64"
|
|
elif machine_type.startswith("arm64"):
|
|
cpu = "aarch64"
|
|
elif machine_type.startswith("arm"):
|
|
cpu = "arm"
|
|
elif machine_type.startswith("mips64"):
|
|
cpu = "mips64"
|
|
elif machine_type.startswith("riscv64"):
|
|
cpu = "riscv64"
|
|
|
|
exec_properties = repository_ctx.attr.platform_exec_properties
|
|
|
|
serialized_exec_properties = "{"
|
|
for k, v in exec_properties.items():
|
|
serialized_exec_properties += "\"%s\" : \"%s\"," % (k, v)
|
|
serialized_exec_properties += "}"
|
|
|
|
repository_ctx.template(
|
|
"BUILD",
|
|
Label("@org_tensorflow//third_party/remote_config:BUILD.tpl"),
|
|
{
|
|
"%{platform}": platform,
|
|
"%{exec_properties}": serialized_exec_properties,
|
|
"%{cpu}": cpu,
|
|
},
|
|
)
|
|
|
|
remote_platform_configure = repository_rule(
|
|
implementation = _remote_platform_configure_impl,
|
|
attrs = {
|
|
"platform_exec_properties": attr.string_dict(mandatory = True),
|
|
"platform": attr.string(default = "linux", values = ["linux", "windows", "local"]),
|
|
},
|
|
)
|