diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-09-29 00:59:02 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-29 01:03:16 -0700 |
commit | d78595d333c9b5c8a0705ba6852c08b107d6c462 (patch) | |
tree | 61527521ca417bedd24243d019f5f5a0711c4143 /tensorflow/tensorflow.bzl | |
parent | 2e0e934e0b3c00863918c78bf55524eea3f0c0dc (diff) |
Make cuda_py_test create a gpu and cpu target.
Currently, we run tests on machines with GPUs based on the "gpu" tag, and the
tests automatically adapt to whether a GPU is available. Creating two targets,
one tagged with "gpu" and one not, will make us run the tests in both modes.
PiperOrigin-RevId: 215045035
Diffstat (limited to 'tensorflow/tensorflow.bzl')
-rw-r--r-- | tensorflow/tensorflow.bzl | 39 |
1 files changed, 23 insertions, 16 deletions
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl index cad5de1b0c..dead44c57e 100644 --- a/tensorflow/tensorflow.bzl +++ b/tensorflow/tensorflow.bzl @@ -1798,22 +1798,29 @@ def cuda_py_test( flaky = 0, xla_enabled = False, grpc_enabled = False): - test_tags = tags + tf_cuda_tests_tags() - tf_py_test( - name = name, - size = size, - srcs = srcs, - data = data, - main = main, - args = args, - tags = test_tags, - shard_count = shard_count, - additional_deps = additional_deps, - kernels = kernels, - flaky = flaky, - xla_enabled = xla_enabled, - grpc_enabled = grpc_enabled, - ) + if main == None: + main = name + ".py" + for config in ["cpu", "gpu"]: + test_name = name + test_tags = tags + if config == "gpu": + test_name += "_gpu" + test_tags = test_tags + tf_cuda_tests_tags() + tf_py_test( + name = test_name, + size = size, + srcs = srcs, + data = data, + main = main, + args = args, + tags = test_tags, + shard_count = shard_count, + additional_deps = additional_deps, + kernels = kernels, + flaky = flaky, + xla_enabled = xla_enabled, + grpc_enabled = grpc_enabled, + ) register_extension_info( extension_name = "cuda_py_test", |