easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/compat/embedding_ops.py

@@ -0,0 +1,162 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Add embedding column for EmbeddingVariable which is only available on pai."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+
+
+def _prune_invalid_ids(sparse_ids, sparse_weights):
+  """Prune invalid IDs (< 0) from the input ids and weights."""
+  is_id_valid = math_ops.greater_equal(sparse_ids.values, 0)
+  if sparse_weights is not None:
+    is_id_valid = math_ops.logical_and(
+        is_id_valid,
+        array_ops.ones_like(sparse_weights.values, dtype=dtypes.bool))
+  sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_id_valid)
+  if sparse_weights is not None:
+    sparse_weights = sparse_ops.sparse_retain(sparse_weights, is_id_valid)
+  return sparse_ids, sparse_weights
+
+
+def _prune_invalid_weights(sparse_ids, sparse_weights):
+  """Prune invalid weights (< 0) from the input ids and weights."""
+  if sparse_weights is not None:
+    is_weights_valid = math_ops.greater(sparse_weights.values, 0)
+    sparse_ids = sparse_ops.sparse_retain(sparse_ids, is_weights_valid)
+    sparse_weights = sparse_ops.sparse_retain(sparse_weights, is_weights_valid)
+  return sparse_ids, sparse_weights
+
+
+def safe_embedding_lookup_sparse(embedding_weights,
+                                 sparse_ids,
+                                 sparse_weights=None,
+                                 combiner='mean',
+                                 default_id=None,
+                                 name=None,
+                                 partition_strategy='div',
+                                 max_norm=None):
+  """Lookup embedding results, accounting for invalid IDs and empty features.
+
+  Fixed so that could be used with Pai EmbeddingVariables.
+
+  The partitioned embedding in `embedding_weights` must all be the same shape
+  except for the first dimension. The first dimension is allowed to vary as the
+  vocabulary size is not necessarily a multiple of `P`.  `embedding_weights`
+  may be a `PartitionedVariable` as returned by using `tf.get_variable()` with a
+  partitioner.
+
+  Invalid IDs (< 0) are pruned from input IDs and weights, as well as any IDs
+  with non-positive weight. For an entry with no features, the embedding vector
+  for `default_id` is returned, or the 0-vector if `default_id` is not supplied.
+
+  The ids and weights may be multi-dimensional. Embeddings are always aggregated
+  along the last dimension.
+
+  Args:
+    embedding_weights:  A list of `P` float `Tensor`s or values representing
+        partitioned embedding `Tensor`s.  Alternatively, a `PartitionedVariable`
+        created by partitioning along dimension 0.  The total unpartitioned
+        shape should be `[e_0, e_1, ..., e_m]`, where `e_0` represents the
+        vocab size and `e_1, ..., e_m` are the embedding dimensions.
+    sparse_ids: `SparseTensor` of shape `[d_0, d_1, ..., d_n]` containing the
+        ids. `d_0` is typically batch size.
+    sparse_weights: `SparseTensor` of same shape as `sparse_ids`, containing
+        float weights corresponding to `sparse_ids`, or `None` if all weights
+        are be assumed to be 1.0.
+    combiner: A string specifying how to combine embedding results for each
+        entry. Currently "mean", "sqrtn" and "sum" are supported, with "mean"
+        the default.
+    default_id: The id to use for an entry with no features.
+    name: A name for this operation (optional).
+    partition_strategy: A string specifying the partitioning strategy.
+        Currently `"div"` and `"mod"` are supported. Default is `"div"`.
+    max_norm: If not `None`, all embeddings are l2-normalized to max_norm before
+        combining.
+
+
+  Returns:
+    Dense `Tensor` of shape `[d_0, d_1, ..., d_{n-1}, e_1, ..., e_m]`.
+
+  Raises:
+    ValueError: if `embedding_weights` is empty.
+  """
+  if embedding_weights is None:
+    raise ValueError('Missing embedding_weights %s.' % embedding_weights)
+
+  embed_tensors = [ops.convert_to_tensor(embedding_weights)]
+  with ops.name_scope(name, 'embedding_lookup',
+                      embed_tensors + [sparse_ids, sparse_weights]) as scope:
+    # Reshape higher-rank sparse ids and weights to linear segment ids.
+    original_shape = sparse_ids.dense_shape
+    original_rank_dim = sparse_ids.dense_shape.get_shape()[0]
+    original_rank = (
+        array_ops.size(original_shape)
+        if original_rank_dim.value is None else original_rank_dim.value)
+    sparse_ids = sparse_ops.sparse_reshape(sparse_ids, [
+        math_ops.reduce_prod(
+            array_ops.slice(original_shape, [0], [original_rank - 1])),
+        array_ops.gather(original_shape, original_rank - 1)
+    ])
+    if sparse_weights is not None:
+      sparse_weights = sparse_tensor.SparseTensor(sparse_ids.indices,
+                                                  sparse_weights.values,
+                                                  sparse_ids.dense_shape)
+
+    # Prune invalid ids and weights.
+    sparse_ids, sparse_weights = _prune_invalid_ids(sparse_ids, sparse_weights)
+    if combiner != 'sum':
+      sparse_ids, sparse_weights = _prune_invalid_weights(
+          sparse_ids, sparse_weights)
+
+    # Fill in dummy values for empty features, if necessary.
+    sparse_ids, is_row_empty = sparse_ops.sparse_fill_empty_rows(
+        sparse_ids, default_id or 0)
+    if sparse_weights is not None:
+      sparse_weights, _ = sparse_ops.sparse_fill_empty_rows(sparse_weights, 1.0)
+
+    indices = sparse_ids.indices
+    values = sparse_ids.values
+    if values.dtype != dtypes.int64:
+      values = math_ops.to_int64(values)
+    sparse_ids = sparse_tensor.SparseTensor(
+        indices=indices, values=values, dense_shape=sparse_ids.dense_shape)
+
+    result = embedding_ops.embedding_lookup_sparse(
+        embedding_weights,
+        sparse_ids,
+        sparse_weights,
+        combiner=combiner,
+        partition_strategy=partition_strategy,
+        name=None if default_id is None else scope,
+        max_norm=max_norm)
+
+    if default_id is None:
+      # Broadcast is_row_empty to the same shape as embedding_lookup_result,
+      # for use in Select.
+      is_row_empty = array_ops.tile(
+          array_ops.reshape(is_row_empty, [-1, 1]),
+          array_ops.stack([1, array_ops.shape(result)[1]]))
+
+      result = array_ops.where(
+          is_row_empty, array_ops.zeros_like(result), result, name=scope)
+
+    # Reshape back from linear ids back into higher-dimensional dense result.
+    final_result = array_ops.reshape(
+        result,
+        array_ops.concat([
+            array_ops.slice(
+                math_ops.cast(original_shape, dtypes.int32), [0],
+                [original_rank - 1]),
+            array_ops.slice(array_ops.shape(result), [1], [-1])
+        ], 0))
+    final_result.set_shape(
+        tensor_shape.unknown_shape(
+            (original_rank_dim - 1).value).concatenate(result.get_shape()[1:]))
+    return final_result

easy_rec/python/compat/embedding_parallel_saver.py

@@ -0,0 +1,316 @@
+# -*- encoding:utf-8 -*-
+
+import logging
+import os
+
+import numpy as np
+from tensorflow.core.protobuf import saver_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+# from tensorflow.python.ops import math_ops
+# from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.platform import gfile
+from tensorflow.python.training import saver
+
+from easy_rec.python.utils import constant
+
+try:
+  import horovod.tensorflow as hvd
+  from sparse_operation_kit.experiment import raw_ops as dynamic_variable_ops
+  from easy_rec.python.compat import dynamic_variable
+except Exception:
+  dynamic_variable_ops = None
+  dynamic_variable = None
+
+try:
+  from tensorflow.python.framework.load_library import load_op_library
+  import easy_rec
+  load_embed_lib_path = os.path.join(easy_rec.ops_dir, 'libload_embed.so')
+  load_embed_lib = load_op_library(load_embed_lib_path)
+except Exception as ex:
+  logging.warning('load libload_embed.so failed: %s' % str(ex))
+  load_embed_lib = None
+
+
+def _get_embed_part_id(embed_file):
+  embed_file = embed_file.split('/')[-1]
+  embed_file = embed_file.split('.')[0]
+  embed_id = embed_file.split('-')[-1]
+  return int(embed_id)
+
+
+class EmbeddingParallelSaver(saver.Saver):
+
+  def __init__(self,
+               var_list=None,
+               reshape=False,
+               sharded=False,
+               max_to_keep=5,
+               keep_checkpoint_every_n_hours=10000.0,
+               name=None,
+               restore_sequentially=False,
+               saver_def=None,
+               builder=None,
+               defer_build=False,
+               allow_empty=False,
+               write_version=saver_pb2.SaverDef.V2,
+               pad_step_number=False,
+               save_relative_paths=False,
+               filename=None):
+    self._kv_vars = []
+    self._embed_vars = []
+    tf_vars = []
+    embed_para_vars = ops.get_collection(constant.EmbeddingParallel)
+    for var in var_list:
+      if dynamic_variable is not None and isinstance(
+          var, dynamic_variable.DynamicVariable):
+        self._kv_vars.append(var)
+      elif var.name in embed_para_vars:
+        logging.info('save shard embedding %s part_id=%d part_shape=%s' %
+                     (var.name, hvd.rank(), var.get_shape()))
+        self._embed_vars.append(var)
+      else:
+        tf_vars.append(var)
+    super(EmbeddingParallelSaver, self).__init__(
+        tf_vars,
+        reshape=reshape,
+        sharded=sharded,
+        max_to_keep=max_to_keep,
+        keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,
+        name=name,
+        restore_sequentially=restore_sequentially,
+        saver_def=saver_def,
+        builder=builder,
+        defer_build=defer_build,
+        allow_empty=allow_empty,
+        write_version=write_version,
+        pad_step_number=pad_step_number,
+        save_relative_paths=save_relative_paths,
+        filename=filename)
+    self._is_build = False
+
+  def _has_embed_vars(self):
+    return (len(self._kv_vars) + len(self._embed_vars)) > 0
+
+  def _save_dense_embedding(self, embed_var):
+    logging.info('task[%d] save_dense_embed: %s' % (hvd.rank(), embed_var.name))
+
+    def _save_embed(embed, filename, var_name):
+      task_id = hvd.rank()
+      filename = filename.decode('utf-8')
+      var_name = var_name.decode('utf-8').replace('/', '__')
+      embed_dir = filename + '-embedding/'
+      logging.info('task[%d] save_dense_embed: %s to %s' %
+                   (task_id, var_name, embed_dir))
+      if not gfile.Exists(embed_dir):
+        gfile.MakeDirs(embed_dir)
+      embed_file = filename + '-embedding/embed-' + var_name + '-part-%d.bin' % task_id
+      with gfile.GFile(embed_file, 'wb') as fout:
+        fout.write(embed.tobytes())
+
+      if task_id == 0:
+        # clear old embedding tables
+        embed_pattern = filename + '-embedding/embed-' + var_name + '-part-*.bin'
+        embed_files = gfile.Glob(embed_pattern)
+        for embed_file in embed_files:
+          embed_id = _get_embed_part_id(embed_file)
+          if embed_id >= hvd.size():
+            gfile.DeleteRecursively(embed_file)
+      return np.asarray([embed_file], order='C', dtype=np.object)
+
+    file_name = ops.get_default_graph().get_tensor_by_name(
+        self.saver_def.filename_tensor_name)
+    save_paths = script_ops.py_func(_save_embed,
+                                    [embed_var, file_name, embed_var.name],
+                                    dtypes.string)
+    return save_paths
+
+  def _load_dense_embedding(self, embed_var):
+    file_name = ops.get_default_graph().get_tensor_by_name(
+        self.saver_def.filename_tensor_name)
+    embed_dim = embed_var.get_shape()[-1]
+    embed_part_size = embed_var.get_shape()[0]
+
+    def _load_embed(embed, embed_dim, embed_part_size, part_id, part_num,
+                    filename, var_name):
+      filename = filename.decode('utf-8')
+      var_name = var_name.decode('utf-8').replace('/', '__')
+      embed_pattern = filename + '-embedding/embed-' + var_name + '-part-*.bin'
+      embed_files = gfile.Glob(embed_pattern)
+
+      embed_files.sort(key=_get_embed_part_id)
+
+      logging.info('task[%d] embed_files=%s embed_dim=%d embed_part_size=%d' %
+                   (part_id, ','.join(embed_files), embed_dim, embed_part_size))
+
+      part_embed_vals = np.zeros([embed_part_size, embed_dim], dtype=np.float32)
+      part_update_cnt = 0
+      for embed_file in embed_files:
+        part_id_o = _get_embed_part_id(embed_file)
+        with gfile.GFile(embed_file, 'rb') as fin:
+          embed_val = np.frombuffer(fin.read(), np.float32)
+          embed_val = embed_val.reshape([-1, embed_dim])
+          embed_ids_o = np.arange(len(embed_val))
+          embed_ids_o = part_id_o + embed_ids_o * len(embed_files)
+          sel_ids = np.where(
+              np.logical_and((embed_ids_o % part_num) == part_id,
+                             embed_ids_o < embed_part_size * part_num))[0]
+          part_update_cnt += len(sel_ids)
+          embed_ids = embed_ids_o[sel_ids]
+          embed_ids_n = np.array(embed_ids / part_num, dtype=np.int64)
+          part_embed_vals[embed_ids_n] = embed_val[sel_ids]
+      logging.info('task[%d] load_part_cnt=%d' % (part_id, part_update_cnt))
+      return part_embed_vals
+
+    with ops.control_dependencies([embed_var._initializer_op]):
+      if load_embed_lib is not None:
+        embed_val = load_embed_lib.load_embed(
+            task_index=hvd.rank(),
+            task_num=hvd.size(),
+            embed_dim=embed_dim,
+            embed_part_size=embed_part_size,
+            var_name='embed-' + embed_var.name.replace('/', '__'),
+            ckpt_path=file_name)
+      else:
+        embed_val = script_ops.py_func(_load_embed, [
+            embed_var, embed_dim, embed_part_size,
+            hvd.rank(),
+            hvd.size(), file_name, embed_var.name
+        ], dtypes.float32)
+      embed_val.set_shape(embed_var.get_shape())
+      return state_ops.assign(embed_var, embed_val)
+
+  def _save_kv_embedding(self, sok_var):
+    indices, values = dynamic_variable_ops.dummy_var_export(
+        sok_var.handle, key_type=sok_var.key_type, dtype=sok_var.handle_dtype)
+    file_name = ops.get_default_graph().get_tensor_by_name(
+        self.saver_def.filename_tensor_name)
+
+    def _save_key_vals(indices, values, filename, var_name):
+      var_name = var_name.decode('utf-8').replace('/', '__')
+      filename = filename.decode('utf-8')
+      sok_dir = filename + '-embedding/'
+      if not gfile.Exists(sok_dir):
+        gfile.MakeDirs(sok_dir)
+      task_id = hvd.rank()
+      key_file = filename + '-embedding/embed-' + var_name + '-part-%d.key' % task_id
+      with gfile.GFile(key_file, 'wb') as fout:
+        fout.write(indices.tobytes())
+      val_file = filename + '-embedding/embed-' + var_name + '-part-%d.val' % task_id
+      with gfile.GFile(val_file, 'wb') as fout:
+        fout.write(values.tobytes())
+
+      if task_id == 0:
+        key_file_pattern = filename + '-embedding/embed-' + var_name + '-part-*.key'
+        key_files = gfile.Glob(key_file_pattern)
+        for key_file in key_files:
+          embed_id = _get_embed_part_id(key_file)
+          if embed_id >= hvd.size():
+            gfile.DeleteRecursively(key_file)
+            val_file = key_file[:-4] + '.val'
+            if gfile.Exists(val_file):
+              gfile.DeleteRecursively(val_file)
+
+      return np.asarray([key_file, val_file], order='C', dtype=np.object)
+
+    save_paths = script_ops.py_func(_save_key_vals,
+                                    [indices, values, file_name, sok_var.name],
+                                    dtypes.string)
+    return save_paths
+
+  def _load_kv_embedding(self, sok_var):
+
+    def _load_key_vals(filename, var_name):
+      var_name = var_name.decode('utf-8').replace('/', '__')
+      filename = filename.decode('utf-8')
+      key_file_pattern = filename + '-embedding/embed-' + var_name + '-part-*.key'
+      logging.info('key_file_pattern=%s filename=%s var_name=%s var=%s' %
+                   (key_file_pattern, filename, var_name, str(sok_var)))
+      key_files = gfile.Glob(key_file_pattern)
+      logging.info('key_file_pattern=%s file_num=%d' %
+                   (key_file_pattern, len(key_files)))
+      all_keys = []
+      all_vals = []
+      for key_file in key_files:
+        with gfile.GFile(key_file, 'rb') as fin:
+          tmp_keys = np.frombuffer(fin.read(), dtype=np.int64)
+          tmp_ids = tmp_keys % hvd.size()
+          tmp_ids = np.where(tmp_ids == hvd.rank())[0]
+          if len(tmp_ids) == 0:
+            break
+          all_keys.append(tmp_keys.take(tmp_ids, axis=0))
+          logging.info('part_keys.shape=%s %s %s' % (str(
+              tmp_keys.shape), str(tmp_ids.shape), str(all_keys[-1].shape)))
+
+        val_file = key_file[:-4] + 'vals'
+        with gfile.GFile(val_file, 'rb') as fin:
+          tmp_vals = np.frombuffer(
+              fin.read(), dtype=np.float32).reshape([-1, sok_var._dimension])
+          all_vals.append(tmp_vals.take(tmp_ids, axis=0))
+          logging.info('part_vals.shape=%s %s %s' % (str(
+              tmp_vals.shape), str(tmp_ids.shape), str(all_vals[-1].shape)))
+
+      all_keys = np.concatenate(all_keys, axis=0)
+      all_vals = np.concatenate(all_vals, axis=0)
+
+      shuffle_ids = np.array(range(len(all_keys)))
+      np.random.shuffle(shuffle_ids)
+      all_keys = all_keys.take(shuffle_ids, axis=0)
+      all_vals = all_vals.take(shuffle_ids, axis=0)
+      return all_keys, all_vals
+
+    file_name = ops.get_default_graph().get_tensor_by_name(
+        self.saver_def.filename_tensor_name)
+    if load_embed_lib is not None:
+      keys, vals = load_embed_lib.load_kv_embed(
+          task_index=hvd.rank(),
+          task_num=hvd.size(),
+          embed_dim=sok_var._dimension,
+          var_name='embed-' + sok_var.name.replace('/', '__'),
+          ckpt_path=file_name)
+    else:
+      logging.warning('libload_embed.so not loaded, will use python script_ops')
+      keys, vals = script_ops.py_func(_load_key_vals, [file_name, sok_var.name],
+                                      (dtypes.int64, dtypes.float32))
+    with ops.control_dependencies([sok_var._initializer_op]):
+      return dynamic_variable_ops.dummy_var_assign(sok_var.handle, keys, vals)
+
+  def build(self):
+    if self._is_built:
+      return
+    super(EmbeddingParallelSaver, self).build()
+    if self.saver_def.restore_op_name and self._has_embed_vars():
+      # load data from the model
+      restore_ops = []
+      for sok_var in self._kv_vars:
+        restore_ops.append(self._load_kv_embedding(sok_var))
+      for embed_var in self._embed_vars:
+        restore_ops.append(self._load_dense_embedding(embed_var))
+      old_restore_op = ops.get_default_graph().get_operation_by_name(
+          self.saver_def.restore_op_name)
+      restore_ops.append(old_restore_op)
+      restore_op_n = control_flow_ops.group(restore_ops)
+      self.saver_def.restore_op_name = restore_op_n.name
+
+    if self.saver_def.save_tensor_name and self._has_embed_vars():
+      file_name = ops.get_default_graph().get_tensor_by_name(
+          self.saver_def.filename_tensor_name)
+      save_part_ops = []
+      for sok_var in self._kv_vars:
+        save_part_op = self._save_kv_embedding(sok_var)
+        save_part_ops.append(save_part_op)
+      for embed_var in self._embed_vars:
+        save_part_op = self._save_dense_embedding(embed_var)
+        save_part_ops.append(save_part_op)
+      old_save_op = ops.get_default_graph().get_tensor_by_name(
+          self.saver_def.save_tensor_name)
+      # only the first worker needs to save non embedding variables
+      if hvd.rank() == 0:
+        save_part_ops.append(old_save_op)
+      with ops.control_dependencies(save_part_ops):
+        save_op_n = array_ops.identity(file_name)
+      self.saver_def.save_tensor_name = save_op_n.name

easy_rec/python/compat/estimator_train.py

@@ -0,0 +1,116 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+import os
+
+import tensorflow as tf
+from tensorflow.python.estimator import run_config as run_config_lib
+from tensorflow.python.util import compat
+from tensorflow_estimator.python.estimator.training import _assert_eval_spec
+from tensorflow_estimator.python.estimator.training import _TrainingExecutor
+
+from easy_rec.python.compat.exporter import FinalExporter
+from easy_rec.python.utils import estimator_utils
+
+from tensorflow_estimator.python.estimator.training import _ContinuousEvalListener  # NOQA
+
+from tensorflow.python.distribute import estimator_training as distribute_coordinator_training  # NOQA
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+gfile = tf.gfile
+
+
+class TrainDoneListener(_ContinuousEvalListener):
+  """Interface for listeners that take action before or after evaluation."""
+
+  def __init__(self, estimator):
+    self._model_dir = estimator.model_dir
+    self._train_done_file = os.path.join(self._model_dir,
+                                         'ESTIMATOR_TRAIN_DONE')
+
+  @property
+  def train_done_file(self):
+    return self._train_done_file
+
+  def after_eval(self, eval_result):
+    """Called after the evaluation is executed.
+
+    Args:
+      eval_result: An `_EvalResult` instance.
+
+    Returns:
+      False if you want to early stop continuous evaluation; `True` otherwise.
+    """
+    last_ckpt_path = eval_result.checkpoint_path
+    if last_ckpt_path is not None:
+      model_dir = os.path.dirname(last_ckpt_path).rstrip('/') + '/'
+      latest_ckpt_path = estimator_utils.latest_checkpoint(model_dir)
+      if latest_ckpt_path != last_ckpt_path:
+        logging.info(
+            'TrainDoneListener: latest_ckpt_path[%s] != last_ckpt_path[%s]' %
+            (latest_ckpt_path, last_ckpt_path))
+        # there are more checkpoints wait to be evaluated
+        return True
+    return not gfile.Exists(self._train_done_file)
+
+
+def train_and_evaluate(estimator, train_spec, eval_spec):
+  _assert_eval_spec(eval_spec)  # fail fast if eval_spec is invalid.
+
+  train_done_listener = TrainDoneListener(estimator)
+  executor = _TrainingExecutor(
+      estimator=estimator,
+      train_spec=train_spec,
+      eval_spec=eval_spec,
+      continuous_eval_listener=train_done_listener)
+  config = estimator.config
+
+  # If `distribute_coordinator_mode` is set and running in distributed
+  # environment, we run `train_and_evaluate` via distribute coordinator.
+  if distribute_coordinator_training.should_run_distribute_coordinator(config):
+    logging.info('Running `train_and_evaluate` with Distribute Coordinator.')
+    distribute_coordinator_training.train_and_evaluate(estimator, train_spec,
+                                                       eval_spec,
+                                                       _TrainingExecutor)
+    return
+
+  if (config.task_type == run_config_lib.TaskType.EVALUATOR and
+      config.task_id > 0):
+    raise ValueError(
+        'For distributed training, there can only be one `evaluator` task '
+        '(with task id 0).  Given task id {}'.format(config.task_id))
+
+  result = executor.run()
+
+  # fix for the bug evaluator fails to export in case num_epoch is reached
+  # before num_steps is reached or num_steps is set to infinite
+  if estimator_utils.is_evaluator():
+    export_dir_base = os.path.join(
+        compat.as_str_any(estimator.model_dir), compat.as_str_any('export'))
+    for exporter in eval_spec.exporters:
+      if isinstance(exporter, FinalExporter):
+        export_path = os.path.join(
+            compat.as_str_any(export_dir_base),
+            compat.as_str_any(exporter.name))
+        # avoid duplicate export
+        if gfile.IsDirectory(export_path + '/'):
+          continue
+        exporter.export(
+            estimator=estimator,
+            export_path=export_path,
+            checkpoint_path=estimator_utils.latest_checkpoint(
+                estimator.model_dir),
+            eval_result=None,
+            is_the_final_export=True)
+
+  if estimator_utils.is_chief():
+    with gfile.GFile(train_done_listener.train_done_file, 'w') as fout:
+      fout.write('Train Done.')
+
+  return result
+
+
+def estimator_train_done(estimator):
+  train_done_file = os.path.join(estimator.model_dir, 'ESTIMATOR_TRAIN_DONE')
+  return gfile.Exists(train_done_file)