Fix eager evaler/decoders.

PiperOrigin-RevId: 415432320

Author: Jonathan Shen

lingvo/eager_runners.py

@@ -205,14 +205,16 @@ def Start(self):
         self._task.params.eval.load_checkpoint_from)
 
     if self._eval_path:
-      self._EvalOnce(self._eval_path)
+      self._EvalOnce(path=self._eval_path)
       py_utils.UpdateProcessedCheckpoints(self._eval_dir, self._eval_path)
     elif self._task.params.eval.eval_all_checkpoints:
-      self._RunOnAllCheckpoints(self._EvalOnce, self._eval_dir)
+      self._RunOnAllCheckpoints(
+          runner_fn=self._EvalOnce, runner_dir=self._eval_dir)
     else:
-      self._RunOnLatestCheckpoints(self._EvalOnce, self._eval_dir)
+      self._RunOnLatestCheckpoints(
+          runner_fn=self._EvalOnce, runner_dir=self._eval_dir)
 
-  def _EvalOnce(self, path):
+  def _EvalOnce(self, sess=None, path=''):
     """Eval a single checkpoint."""
     with self._cluster:
       # Attempt to restore the checkpoint
@@ -337,20 +339,22 @@ def Start(self):
         self._task.params.eval.load_checkpoint_from)
 
     if self._decode_path:
-      self._DecodeOnce(self._decode_path)
+      self._DecodeOnce(path=self._decode_path)
       py_utils.UpdateProcessedCheckpoints(self._decoder_dir, self._decode_path)
     elif self._task.params.eval.decode_all_checkpoints:
-      self._RunOnAllCheckpoints(self._DecodeOnce, self._decoder_dir)
+      self._RunOnAllCheckpoints(
+          runner_fn=self._DecodeOnce, runner_dir=self._decoder_dir)
     else:
-      self._RunOnLatestCheckpoints(self._DecodeOnce, self._decoder_dir)
+      self._RunOnLatestCheckpoints(
+          runner_fn=self._DecodeOnce, runner_dir=self._decoder_dir)
 
   @classmethod
   def GetDecodeOutPath(cls, decoder_dir, checkpoint_id):
     """Gets the path to decode out file."""
     out_dir = cls._GetTtlDir(decoder_dir, duration='7d')
     return os.path.join(out_dir, 'decoder_out_%09d' % checkpoint_id)
 
-  def _DecodeOnce(self, path):
+  def _DecodeOnce(self, sess=None, path=''):
     """Decode a single checkpoint."""
     with self._cluster:
       # Attempt to restore the checkpoint

lingvo/trainer.py

@@ -415,11 +415,15 @@ def UpdateClusterParamsFromFlags(self, cluster, job_name):
     cluster.input.replicas = FLAGS.input_replicas
     cluster.input.targets = FLAGS.input_targets
 
-    cluster.evaler.name = FLAGS.evaler_job
+    if py_utils.IsEagerMode():
+      cluster.evaler.name = '/job:localhost'
+      cluster.decoder.name = '/job:localhost'
+    else:
+      cluster.evaler.name = FLAGS.evaler_job
+      cluster.decoder.name = FLAGS.decoder_job
+
     cluster.evaler.replicas = FLAGS.evaler_replicas
     cluster.evaler.gpus_per_replica = FLAGS.evaler_gpus
-
-    cluster.decoder.name = FLAGS.decoder_job
     cluster.decoder.replicas = FLAGS.decoder_replicas
     cluster.decoder.gpus_per_replica = FLAGS.decoder_gpus