diff --git a/third_party/toolchains/BUILD b/third_party/toolchains/BUILD index 2df2c3c0050..3717ec1f8cc 100644 --- a/third_party/toolchains/BUILD +++ b/third_party/toolchains/BUILD @@ -206,6 +206,49 @@ platform( """ % container_digests["cuda10.0-cudnn7-ubuntu16.04-manylinux2010"], ) +# Built with //tensorflow/tools/ci_build/Dockerfile.rbe.cuda10.1-cudnn7-ubuntu16.04-manylinux2010. +platform( + name = "rbe_cuda10.1-cudnn7-ubuntu16.04-manylinux2010", + constraint_values = [ + "@bazel_tools//platforms:x86_64", + "@bazel_tools//platforms:linux", + ], + remote_execution_properties = """ + properties: { + name: "container-image" + value:"docker://gcr.io/tensorflow-testing/nosla-cuda10.1-cudnn7-ubuntu16.04-manylinux2010@%s" + } + properties: { + name: "Pool" + value: "default" + } + """ % container_digests["cuda10.1-cudnn7-ubuntu16.04-manylinux2010"], +) + +# The above platform with GPU support. +platform( + name = "rbe_cuda10.1-cudnn7-ubuntu16.04-manylinux2010-gpu", + constraint_values = [ + "@bazel_tools//platforms:x86_64", + "@bazel_tools//platforms:linux", + ":gpu_test", + ], + remote_execution_properties = """ + properties: { + name: "container-image" + value: "docker://gcr.io/tensorflow-testing/nosla-cuda10.1-cudnn7-ubuntu16.04-manylinux2010@%s" + } + properties: { + name: "dockerRuntime" + value: "nvidia" + } + properties: { + name: "Pool" + value: "gpu-pool" + } + """ % container_digests["cuda10.1-cudnn7-ubuntu16.04-manylinux2010"], +) + # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.rocm-ubuntu16.04 platform( name = "rbe_rocm-ubuntu16.04",