Returns a session constructed using `estimator` and `serving_input_fn`.

  The Estimator API does not provide an API to construct a graph and session,
  making it necessary for this function to replicate how an estimator builds
  a graph.

  This code is based on `Estimator.export_savedmodel` (another function that
  has to replicate how an estimator builds a graph).

  Args:
    estimator: tf.Estimator to use when constructing the session.
    serving_input_fn: A function that takes no arguments and returns a
      `ServingInputReceiver`. It is used to construct the session.
    checkpoint_path: The checkpoint path to restore in the session. Must not
      be None.

def _make_estimator_serving_session(estimator, serving_input_fn,
                                    checkpoint_path):
  """Returns a session constructed using `estimator` and `serving_input_fn`.

  The Estimator API does not provide an API to construct a graph and session,
  making it necessary for this function to replicate how an estimator builds
  a graph.

  This code is based on `Estimator.export_savedmodel` (another function that
  has to replicate how an estimator builds a graph).

  Args:
    estimator: tf.Estimator to use when constructing the session.
    serving_input_fn: A function that takes no arguments and returns a
      `ServingInputReceiver`. It is used to construct the session.
    checkpoint_path: The checkpoint path to restore in the session. Must not
      be None.
  """
  with tf.Graph().as_default() as g:
    mode = tf_v1.estimator.ModeKeys.PREDICT
    tf_v1.train.create_global_step(g)
    tf_v1.set_random_seed(estimator.config.tf_random_seed)
    serving_input_receiver = serving_input_fn()

    estimator_spec = estimator.model_fn(
        features=serving_input_receiver.features,
        labels=None,
        mode=mode,
        config=estimator.config)

    # pylint: disable=protected-access
    # Note that MonitoredSession(), despite the name is not a Session, and
    # can't be used to export Modules as one can't use them with Savers.
    # As so this must use a raw tf.Session().
    session = tf_v1.Session(config=estimator._session_config)
    # pylint: enable=protected-access

    with session.as_default():
      # TODO(b/71839662): Consider if this needs to support TPUEstimatorSpec
      # which does not have a scaffold member.
      saver_for_restore = estimator_spec.scaffold.saver or tf_v1.train.Saver(
          sharded=True)
      saver_for_restore.restore(session, checkpoint_path)
    return session