fix eval hooks with distribute strategy

PiperOrigin-RevId: 314226004

Author: w-xinyi

tensorflow_estimator/python/estimator/distribute_strategy_estimator_integration_test.py

@@ -23,8 +23,13 @@
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
+from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.distribute import combinations
 from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import training_util
+from tensorflow_estimator.python.estimator import estimator as estimator_lib
+from tensorflow_estimator.python.estimator import model_fn as model_fn_lib
 from tensorflow_estimator.python.estimator import run_config
 from tensorflow_estimator.python.estimator import training
 from tensorflow_estimator.python.estimator.canned import dnn_linear_combined
@@ -50,6 +55,61 @@ def input_fn():
 
     return input_fn
 
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'],
+          distribution=[
+              strategy_combinations.one_device_strategy,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.mirrored_strategy_with_two_gpus
+          ],
+          use_train_and_evaluate=[True, False]))
+  def test_estimator_with_strategy_hooks(self, distribution,
+                                         use_train_and_evaluate):
+    config = run_config.RunConfig(eval_distribute=distribution)
+
+    def _input_map_fn(tensor):
+      return {'feature': tensor}, tensor
+
+    def input_fn():
+      return dataset_ops.Dataset.from_tensors(
+          [1.]).repeat(10).batch(5).map(_input_map_fn)
+
+    def model_fn(features, labels, mode):
+      del features, labels
+      global_step = training_util.get_global_step()
+      if mode == model_fn_lib.ModeKeys.TRAIN:
+        train_hook1 = basic_session_run_hooks.StepCounterHook(
+            every_n_steps=1, output_dir=self.get_temp_dir())
+        train_hook2 = tf.compat.v1.test.mock.MagicMock(
+            wraps=tf.compat.v1.train.SessionRunHook(),
+            spec=tf.compat.v1.train.SessionRunHook)
+        return model_fn_lib.EstimatorSpec(
+            mode,
+            loss=tf.constant(1.),
+            train_op=global_step.assign_add(1),
+            training_hooks=[train_hook1, train_hook2])
+      if mode == model_fn_lib.ModeKeys.EVAL:
+        eval_hook1 = basic_session_run_hooks.StepCounterHook(
+            every_n_steps=1, output_dir=self.get_temp_dir())
+        eval_hook2 = tf.compat.v1.test.mock.MagicMock(
+            wraps=tf.compat.v1.train.SessionRunHook(),
+            spec=tf.compat.v1.train.SessionRunHook)
+        return model_fn_lib.EstimatorSpec(
+            mode=mode,
+            loss=tf.constant(1.),
+            evaluation_hooks=[eval_hook1, eval_hook2])
+    num_steps = 10
+    estimator = estimator_lib.EstimatorV2(
+        model_fn=model_fn, model_dir=self.get_temp_dir(), config=config)
+    if use_train_and_evaluate:
+      training.train_and_evaluate(
+          estimator, training.TrainSpec(input_fn, max_steps=num_steps),
+          training.EvalSpec(input_fn))
+    else:
+      estimator.train(input_fn, steps=num_steps)
+      estimator.evaluate(input_fn, steps=num_steps)
+
   @combinations.generate(
       combinations.combine(
           mode=['graph'],

tensorflow_estimator/python/estimator/estimator.py

@@ -1619,8 +1619,16 @@ def step_fn(ctx, inputs):
 
     scaffold = _combine_distributed_scaffold(grouped_estimator_spec.scaffold,
                                              self._eval_distribution)
-    evaluation_hooks = self._eval_distribution.experimental_local_results(
-        grouped_estimator_spec.evaluation_hooks)[0]
+
+    def get_hooks_from_the_first_device(per_device_hooks):
+      return [
+          self._eval_distribution.experimental_local_results(per_device_hook)[0]
+          for per_device_hook in per_device_hooks
+      ]
+
+    evaluation_hooks = get_hooks_from_the_first_device(
+        grouped_estimator_spec.evaluation_hooks)
+
     return (scaffold, evaluation_hooks, input_hooks, update_op, eval_dict)
 
   def _evaluate_run(self, checkpoint_path, scaffold, update_op, eval_dict,