Declare tests that uses distribute/combinations as Python3

I'll soon introduce PY3 only dependency to distribute/combinations to better
simulate multi worker.

PiperOrigin-RevId: 313669514
Change-Id: Ia986971a7fee469002614a9e272c475e9f93a250

tensorflow/core/platform/default/distribute.bzl

@@ -42,6 +42,10 @@ def distribute_py_test(
         disable_v3: whether tests for TPU version 3 should be generated.
         **kwargs: extra keyword arguments to the non-tpu test.
     """
+
+    # Default to PY3 since multi worker tests require PY3.
+    kwargs.setdefault("python_version", "PY3")
+
     _ignore = (full_precision)
     tpu_tags = tags if (tpu_tags == None) else tpu_tags
     main = main if main else "%s.py" % name

tensorflow/examples/saved_model/integration_tests/BUILD

@@ -1,4 +1,4 @@
-load("//tensorflow:tensorflow.bzl", "cuda_py_test")
+load("//tensorflow/core/platform/default:distribute.bzl", "distribute_py_test")
 
 package(
     licenses = ["notice"],  # Apache 2.0
@@ -48,7 +48,7 @@ py_library(
     ],
 )
 
-cuda_py_test(
+distribute_py_test(
     name = "saved_model_test",
     srcs = [
         "saved_model_test.py",

tensorflow/python/distribute/BUILD

@@ -416,6 +416,7 @@ py_library(
 tf_py_test(
     name = "mirrored_function_strategy_test",
     srcs = ["mirrored_function_strategy_test.py"],
+    python_version = "PY3",
     tags = ["no_pip"],
     deps = [
         ":distribute_lib",
@@ -447,6 +448,7 @@ cuda_py_test(
     srcs = [
         "multi_worker_continuous_run_test.py",
     ],
+    python_version = "PY3",
     tags = [
         "notsan",  # TODO(b/151841995)
     ],
@@ -528,6 +530,7 @@ py_library(
 cuda_py_test(
     name = "input_ops_test",
     srcs = ["input_ops_test.py"],
+    python_version = "PY3",
     deps = [
         ":input_ops",
         "//tensorflow/python:client_testlib",
@@ -764,6 +767,7 @@ py_library(
 py_test(
     name = "combinations_test",
     srcs = ["combinations_test.py"],
+    python_version = "PY3",
     tags = [
         "notap",  # TODO(b/153646955): flaky
     ],
@@ -856,6 +860,7 @@ cuda_py_test(
     name = "checkpoint_utils_test",
     size = "medium",
     srcs = ["checkpoint_utils_test.py"],
+    python_version = "PY3",
     tags = [
         "multi_and_single_gpu",
     ],
@@ -953,6 +958,7 @@ distribute_py_test(
 cuda_py_test(
     name = "cross_device_utils_test",
     srcs = ["cross_device_utils_test.py"],
+    python_version = "PY3",
     deps = [
         ":combinations",
         ":cross_device_utils",
@@ -971,6 +977,7 @@ cuda_py_test(
 cuda_py_test(
     name = "cross_device_ops_test",
     srcs = ["cross_device_ops_test.py"],
+    python_version = "PY3",
     tags = [
         "multi_and_single_gpu",
     ],
@@ -997,6 +1004,7 @@ cuda_py_test(
     name = "one_device_strategy_test",
     srcs = ["one_device_strategy_test.py"],
     grpc_enabled = True,
+    python_version = "PY3",
     deps = [
         ":combinations",
         ":strategy_combinations",
@@ -1350,6 +1358,7 @@ cuda_py_test(
     name = "warm_starting_util_test",
     size = "medium",
     srcs = ["warm_starting_util_test.py"],
+    python_version = "PY3",
     tags = [
         "multi_and_single_gpu",
     ],
@@ -1367,7 +1376,10 @@ cuda_py_test(
 cuda_py_test(
     name = "remote_mirrored_strategy_eager_test",
     srcs = ["remote_mirrored_strategy_eager_test.py"],
-    tags = ["no_oss"],  # b/154743849
+    python_version = "PY3",
+    tags = [
+        "no_oss",  # b/154743849
+    ],
     deps = [
         ":combinations",
         ":distribute_lib",
@@ -1390,6 +1402,7 @@ cuda_py_test(
 cuda_py_test(
     name = "mirrored_strategy_test",
     srcs = ["mirrored_strategy_test.py"],
+    python_version = "PY3",
     shard_count = 5,
     tags = [
         "multi_and_single_gpu",
@@ -1421,6 +1434,7 @@ cuda_py_test(
 cuda_py_test(
     name = "mirrored_variable_test",
     srcs = ["mirrored_variable_test.py"],
+    python_version = "PY3",
     tags = [
         "guitar",
         "multi_and_single_gpu",
@@ -1599,6 +1613,7 @@ distribute_py_test(
 cuda_py_test(
     name = "collective_all_reduce_strategy_test",
     srcs = ["collective_all_reduce_strategy_test.py"],
+    python_version = "PY3",
     tags = [
         "multi_and_single_gpu",
     ],
@@ -1633,6 +1648,7 @@ cuda_py_test(
 cuda_py_test(
     name = "parameter_server_strategy_test",
     srcs = ["parameter_server_strategy_test.py"],
+    python_version = "PY3",
     tags = [
         "multi_and_single_gpu",
         "no_windows",  # TODO(b/156428279): reenable this test once the image is updated.
@@ -1751,6 +1767,7 @@ py_test(
 cuda_py_test(
     name = "strategy_common_test",
     srcs = ["strategy_common_test.py"],
+    python_version = "PY3",
     tags = [
         "multi_and_single_gpu",
         # TODO(b/155301154): Enable this test on multi-gpu guitar once multi process

tensorflow/python/keras/distribute/BUILD

@@ -65,6 +65,7 @@ py_library(
 cuda_py_test(
     name = "worker_training_state_test",
     srcs = ["worker_training_state_test.py"],
+    python_version = "PY3",
     shard_count = 4,
     deps = [
         ":multi_worker_testing_utils",
@@ -312,6 +313,7 @@ py_library(
 cuda_py_test(
     name = "keras_optimizer_v2_test",
     srcs = ["keras_optimizer_v2_test.py"],
+    python_version = "PY3",
     shard_count = 4,
     tags = [
         "multi_and_single_gpu",
@@ -325,6 +327,7 @@ cuda_py_test(
 cuda_py_test(
     name = "multi_worker_test",
     srcs = ["multi_worker_test.py"],
+    python_version = "PY3",
     shard_count = 32,
     tags = [
         "multi_and_single_gpu",

tensorflow/tensorflow.bzl

@@ -2158,6 +2158,8 @@ def pywrap_tensorflow_macro(
 # 2. When --define=no_tensorflow_py_deps=false (by default), it's a normal py_test.
 def py_test(deps = [], data = [], kernels = [], **kwargs):
     # Python version placeholder
+    if kwargs.get("python_version", None) == "PY3":
+        kwargs["tags"] = kwargs.get("tags", []) + ["no_oss_py2"]
     native.py_test(
         # TODO(jlebar): Ideally we'd use tcmalloc here.,
         deps = select({