Fix chief-0 timeout in tf_keras premade models distributed test by specifying steps_per_epoch when training on Dataset.

PiperOrigin-RevId: 937158025

Author: tensorflower-gardener

tf_keras/distribute/keras_premade_models_test.py

@@ -85,85 +85,87 @@ def dataset_fn(input_context):
 
 
 class KerasPremadeModelsTest(tf.test.TestCase, parameterized.TestCase):
-    @tf.__internal__.distribute.combinations.generate(
+  @tf.__internal__.distribute.combinations.generate(
         strategy_combinations_eager_data_fn()
     )
-    def test_linear_model(self, distribution, use_dataset_creator, data_fn):
-        if (not use_dataset_creator) and isinstance(
+  def test_linear_model(self, distribution, use_dataset_creator, data_fn):
+    if (not use_dataset_creator) and isinstance(
             distribution, tf.distribute.experimental.ParameterServerStrategy
         ):
-            self.skipTest(
+      self.skipTest(
                 "Parameter Server strategy requires dataset creator to be used "
                 "in model.fit."
             )
-        if (
+    if (
             not tf.__internal__.tf2.enabled()
             and use_dataset_creator
             and isinstance(
                 distribution, tf.distribute.experimental.ParameterServerStrategy
             )
         ):
-            self.skipTest(
+      self.skipTest(
                 "Parameter Server strategy with dataset creator needs to be "
                 "run when eager execution is enabled."
             )
-        with distribution.scope():
-            model = linear.LinearModel()
-            opt = gradient_descent.SGD(learning_rate=0.1)
-            model.compile(opt, "mse")
-            if use_dataset_creator:
-                x = dataset_creator.DatasetCreator(dataset_fn)
-                hist = model.fit(x, epochs=3, steps_per_epoch=INPUT_SIZE)
-            else:
-                if data_fn == "numpy":
-                    inputs, output = get_numpy()
-                    hist = model.fit(inputs, output, epochs=3)
-                else:
-                    hist = model.fit(get_dataset(), epochs=3)
-                self.assertLess(hist.history["loss"][2], 0.2)
-
-    @tf.__internal__.distribute.combinations.generate(
+    with distribution.scope():
+      model = linear.LinearModel()
+      opt = gradient_descent.SGD(learning_rate=0.1)
+      model.compile(opt, "mse")
+      if use_dataset_creator:
+        x = dataset_creator.DatasetCreator(dataset_fn)
+        hist = model.fit(x, epochs=3, steps_per_epoch=INPUT_SIZE)
+      else:
+        if data_fn == "numpy":
+          inputs, output = get_numpy()
+          hist = model.fit(inputs, output, epochs=3)
+        else:
+          hist = model.fit(get_dataset(), epochs=3, steps_per_epoch=INPUT_SIZE)
+        self.assertLess(hist.history["loss"][2], 0.2)
+
+  @tf.__internal__.distribute.combinations.generate(
         strategy_combinations_eager_data_fn()
     )
-    def test_wide_deep_model(self, distribution, use_dataset_creator, data_fn):
-        if (not use_dataset_creator) and isinstance(
+  def test_wide_deep_model(self, distribution, use_dataset_creator, data_fn):
+    if (not use_dataset_creator) and isinstance(
             distribution, tf.distribute.experimental.ParameterServerStrategy
         ):
-            self.skipTest(
+      self.skipTest(
                 "Parameter Server strategy requires dataset creator to be used "
                 "in model.fit."
             )
-        if (
+    if (
             not tf.__internal__.tf2.enabled()
             and use_dataset_creator
             and isinstance(
                 distribution, tf.distribute.experimental.ParameterServerStrategy
             )
         ):
-            self.skipTest(
+      self.skipTest(
                 "Parameter Server strategy with dataset creator needs to be "
                 "run when eager execution is enabled."
             )
-        with distribution.scope():
-            linear_model = linear.LinearModel(units=1)
-            dnn_model = sequential.Sequential([core.Dense(units=1)])
-            wide_deep_model = wide_deep.WideDeepModel(linear_model, dnn_model)
-            linear_opt = gradient_descent.SGD(learning_rate=0.05)
-            dnn_opt = adagrad.Adagrad(learning_rate=0.1)
-            wide_deep_model.compile(optimizer=[linear_opt, dnn_opt], loss="mse")
-
-            if use_dataset_creator:
-                x = dataset_creator.DatasetCreator(dataset_fn)
-                hist = wide_deep_model.fit(
+    with distribution.scope():
+      linear_model = linear.LinearModel(units=1)
+      dnn_model = sequential.Sequential([core.Dense(units=1)])
+      wide_deep_model = wide_deep.WideDeepModel(linear_model, dnn_model)
+      linear_opt = gradient_descent.SGD(learning_rate=0.05)
+      dnn_opt = adagrad.Adagrad(learning_rate=0.1)
+      wide_deep_model.compile(optimizer=[linear_opt, dnn_opt], loss="mse")
+
+      if use_dataset_creator:
+        x = dataset_creator.DatasetCreator(dataset_fn)
+        hist = wide_deep_model.fit(
                     x, epochs=3, steps_per_epoch=INPUT_SIZE
                 )
-            else:
-                if data_fn == "numpy":
-                    inputs, output = get_numpy()
-                    hist = wide_deep_model.fit(inputs, output, epochs=3)
-                else:
-                    hist = wide_deep_model.fit(get_dataset(), epochs=3)
-            self.assertLess(hist.history["loss"][2], 0.2)
+      else:
+        if data_fn == "numpy":
+          inputs, output = get_numpy()
+          hist = wide_deep_model.fit(inputs, output, epochs=3)
+        else:
+          hist = wide_deep_model.fit(
+              get_dataset(), epochs=3, steps_per_epoch=INPUT_SIZE
+          )
+      self.assertLess(hist.history["loss"][2], 0.2)
 
 
 if __name__ == "__main__":