"""Structure to create or gather pieces commonly needed to train a model.

def __init__(self,

init_op=None,

init_feed_dict=None,

init_fn=None,

ready_op=None,

ready_for_local_init_op=None,

local_init_op=None,

summary_op=None,

saver=None,

copy_from_scaffold=None):

"""Create a scaffold.

if copy_from_scaffold is not None:

if not isinstance(copy_from_scaffold, Scaffold):

raise TypeError('copy_from_scaffold is not a Scaffold instance.')

# We need _coalesce since Tensor is not converted to bool automatically,

# so the common idiom of (a or b) does not work.

coalesce = lambda a, b: a if a is not None else b

init_op = coalesce(init_op, copy_from_scaffold.init_op)

init_feed_dict = coalesce(init_feed_dict,

copy_from_scaffold.init_feed_dict)

# Use the original init_fn provided by the user to init the new Scaffold.

init_fn = coalesce(init_fn, copy_from_scaffold._user_init_fn) # pylint: disable=protected-access

ready_op = coalesce(ready_op, copy_from_scaffold.ready_op)

ready_for_local_init_op = coalesce(

ready_for_local_init_op, copy_from_scaffold.ready_for_local_init_op)

local_init_op = coalesce(local_init_op, copy_from_scaffold.local_init_op)

summary_op = coalesce(summary_op, copy_from_scaffold.summary_op)

saver = coalesce(saver, copy_from_scaffold.saver)

# NOTE(touts): modifying the init function to be passed the scaffold is a

# hack to make it easy to find the saver. Is there a better way?

self._user_init_fn = init_fn

if init_fn:

self._init_fn = lambda sess: init_fn(self, sess)

else:

self._init_fn = None

self._init_op = init_op

self._init_feed_dict = init_feed_dict

self._ready_op = ready_op

self._ready_for_local_init_op = ready_for_local_init_op

self._local_init_op = local_init_op

self._summary_op = summary_op

self._saver = saver

def finalize(self):

"""Creates operations if needed and finalizes the graph."""

if self._init_op is None:

def default_init_op():

return control_flow_ops.group(

variables.global_variables_initializer(),

resources.initialize_resources(resources.shared_resources()))

self._init_op = Scaffold.get_or_default(

'init_op',

ops.GraphKeys.INIT_OP,

default_init_op)

if self._ready_op is None:

def default_ready_op():

return array_ops.concat([

variables.report_uninitialized_variables(),

resources.report_uninitialized_resources()

], 0)

self._ready_op = Scaffold.get_or_default(

'ready_op', ops.GraphKeys.READY_OP,

default_ready_op)

if self._ready_for_local_init_op is None:

def default_ready_for_local_init_op():

return variables.report_uninitialized_variables(

variables.global_variables())

self._ready_for_local_init_op = Scaffold.get_or_default(

'ready_for_local_init_op', ops.GraphKeys.READY_FOR_LOCAL_INIT_OP,

default_ready_for_local_init_op)

if self._local_init_op is None:

self._local_init_op = Scaffold.get_or_default(

'local_init_op', ops.GraphKeys.LOCAL_INIT_OP,

Scaffold._default_local_init_op)

if self._summary_op is None:

self._summary_op = Scaffold.get_or_default('summary_op',

ops.GraphKeys.SUMMARY_OP,

summary.merge_all)

# pylint: disable=g-long-lambda

if self._saver is None:

self._saver = training_saver._get_saver_or_default() # pylint: disable=protected-access

# pylint: enable=g-long-lambda

self._saver.build()

ops.get_default_graph().finalize()

return self

@property

def init_fn(self):

return self._init_fn

@property

def init_op(self):

return self._init_op

@property

def ready_op(self):

return self._ready_op

@property

def ready_for_local_init_op(self):

return self._ready_for_local_init_op

@property

def local_init_op(self):

return self._local_init_op

@property

def summary_op(self):

return self._summary_op

@property

def saver(self):

return self._saver

@property

def init_feed_dict(self):

return self._init_feed_dict

@staticmethod

def get_or_default(arg_name, collection_key, default_constructor):

"""Get from cache or create a default operation."""

elements = ops.get_collection(collection_key)

if elements:

if len(elements) > 1:

raise RuntimeError('More than one item in the collection "%s". '

'Please indicate which one to use by passing it to '

'the tf.Scaffold constructor as: '

'tf.Scaffold(%s=item to use)', collection_key,

arg_name)

return elements[0]

op = default_constructor()

if op is not None:

ops.add_to_collection(collection_key, op)

return op

@staticmethod

def _default_local_init_op():

return control_flow_ops.group(

variables.local_variables_initializer(),

lookup_ops.tables_initializer(),

is_chief=True,

checkpoint_dir=None,

scaffold=None,

hooks=None,

chief_only_hooks=None,

save_checkpoint_secs=600,

save_summaries_steps=USE_DEFAULT,

save_summaries_secs=USE_DEFAULT,

config=None,

stop_grace_period_secs=120,

log_step_count_steps=100,

max_wait_secs=7200):

"""Creates a `MonitoredSession` for training.

if save_summaries_steps == USE_DEFAULT and save_summaries_secs == USE_DEFAULT:

save_summaries_steps = 100

save_summaries_secs = None

elif save_summaries_secs == USE_DEFAULT:

save_summaries_secs = None

elif save_summaries_steps == USE_DEFAULT:

save_summaries_steps = None

scaffold = scaffold or Scaffold()

if not is_chief:

session_creator = WorkerSessionCreator(

scaffold=scaffold,

master=master,

config=config,

max_wait_secs=max_wait_secs)

return MonitoredSession(session_creator=session_creator, hooks=hooks or [],

stop_grace_period_secs=stop_grace_period_secs)

all_hooks = []

if chief_only_hooks:

all_hooks.extend(chief_only_hooks)

session_creator = ChiefSessionCreator(

scaffold=scaffold,

checkpoint_dir=checkpoint_dir,

master=master,

config=config)

if checkpoint_dir:

if log_step_count_steps and log_step_count_steps > 0:

all_hooks.append(

basic_session_run_hooks.StepCounterHook(

output_dir=checkpoint_dir, every_n_steps=log_step_count_steps))

if (save_summaries_steps and save_summaries_steps > 0) or (

save_summaries_secs and save_summaries_secs > 0):

all_hooks.append(basic_session_run_hooks.SummarySaverHook(

scaffold=scaffold,

save_steps=save_summaries_steps,

save_secs=save_summaries_secs,

output_dir=checkpoint_dir))

if save_checkpoint_secs and save_checkpoint_secs > 0:

all_hooks.append(basic_session_run_hooks.CheckpointSaverHook(

checkpoint_dir, save_secs=save_checkpoint_secs, scaffold=scaffold))

if hooks:

all_hooks.extend(hooks)

return MonitoredSession(session_creator=session_creator, hooks=all_hooks,

"""A factory for tf.Session."""

@abc.abstractmethod

def create_session(self):

raise NotImplementedError(

"""Creates a tf.Session for a chief."""

def __init__(self,

scaffold=None,

master='',

config=None,

checkpoint_dir=None,

checkpoint_filename_with_path=None):

"""Initializes a chief session creator.

self._checkpoint_dir = checkpoint_dir

self._checkpoint_filename_with_path = checkpoint_filename_with_path

self._scaffold = scaffold or Scaffold()

self._session_manager = None

self._master = master

self._config = config

def _get_session_manager(self):

if self._session_manager:

return self._session_manager

self._session_manager = sm.SessionManager(

local_init_op=self._scaffold.local_init_op,

ready_op=self._scaffold.ready_op,

ready_for_local_init_op=self._scaffold.ready_for_local_init_op,

graph=ops.get_default_graph())

return self._session_manager

def create_session(self):

self._scaffold.finalize()

return self._get_session_manager().prepare_session(

self._master,

saver=self._scaffold.saver,

checkpoint_dir=self._checkpoint_dir,

checkpoint_filename_with_path=self._checkpoint_filename_with_path,

config=self._config,

init_op=self._scaffold.init_op,

init_feed_dict=self._scaffold.init_feed_dict,

"""Creates a tf.Session for a worker."""

def __init__(self,

scaffold=None,

master='',

config=None,

max_wait_secs=30 * 60):

"""Initializes a worker session creator.

self._scaffold = scaffold or Scaffold()

self._session_manager = None

self._master = master

self._config = config

self._max_wait_secs = max_wait_secs

def _get_session_manager(self):

if self._session_manager:

return self._session_manager

self._session_manager = sm.SessionManager(

local_init_op=self._scaffold.local_init_op,

ready_op=self._scaffold.ready_op,

ready_for_local_init_op=self._scaffold.ready_for_local_init_op,

graph=ops.get_default_graph())

return self._session_manager

def create_session(self):

self._scaffold.finalize()

return self._get_session_manager().wait_for_session(

self._master, config=self._config,

max_wait_secs=self._max_wait_secs

"""See `MonitoredSession` or `SingularMonitoredSession`."""

def __init__(self, session_creator, hooks, should_recover,

stop_grace_period_secs=120):

"""Sets up a Monitored or Hooked Session.

self._graph_was_finalized = ops.get_default_graph().finalized

self._hooks = hooks or []

for h in self._hooks:

h.begin()

# Create the session.

self._coordinated_creator = self._CoordinatedSessionCreator(

session_creator=session_creator or ChiefSessionCreator(),

hooks=self._hooks,

stop_grace_period_secs=stop_grace_period_secs)

if should_recover:

self._sess = _RecoverableSession(self._coordinated_creator)

else:

self._sess = self._coordinated_creator.create_session()

@property

def graph(self):

"""The graph that was launched in this session."""

if self._tf_sess() is None:

return None

return self._tf_sess().graph

def run(self, fetches, feed_dict=None, options=None, run_metadata=None):

"""Run ops in the monitored session.

return self._sess.run(fetches,

feed_dict=feed_dict,

options=options,

run_metadata=run_metadata)

def run_step_fn(self, step_fn):

"""Run ops using a step function.

step_fn_arguments = util.fn_args(step_fn)

if step_fn_arguments != ('step_context',) and step_fn_arguments != (

'self',

'step_context',

):

raise ValueError(

'`step_fn` may either have one `step_context` argument, or'

' `self` and `step_context` arguments if it\'s an instance'

' method. Got {} instead.'.format(step_fn_arguments))

# `self._sess` is either `_RecoverableSession` or a `_CoordinatedSession`.

# Setting `run_with_hooks` to `None` will cause `run_with_hooks` to be

# `_CoordinatedSession.run` downstream in either case. This allows

# `_PREEMPTION_ERRORS` to propage from within `step_fn` to

# `_RecoverableSession.run_step_fn`.

return self._sess.run_step_fn(step_fn, self._tf_sess(), run_with_hooks=None)

class StepContext(object):

"""Control flow instrument for the `step_fn` from `run_step_fn()`.

def __init__(self, session, run_with_hooks_fn):

"""Initializes the `step_context` argument for a `step_fn` invocation.

self._session = session

self._run_with_hooks_fn = run_with_hooks_fn

@property

def session(self):

return self._session

def run_with_hooks(self, *args, **kwargs):

"""Same as `MonitoredSession.run`. Accepts the same arguments."""

return self._run_with_hooks_fn(*args, **kwargs)

def request_stop(self):

"""Exit the training loop by causing `should_stop()` to return `True`.

raise StopIteration('step_fn has requested the iterations to stop.')

def should_stop(self):

return self._sess is None or self._sess.should_stop()

def close(self):

self._close_internal()

def __enter__(self):

return self

def __exit__(self, exception_type, exception_value, traceback):

if exception_type in [errors.OutOfRangeError, StopIteration]:

exception_type = None

self._close_internal(exception_type)

# __exit__ should return True to suppress an exception.

return exception_type is None

class _CoordinatedSessionCreator(SessionCreator):

"""Factory for the _RecoverableSession."""

def __init__(self, session_creator, hooks, stop_grace_period_secs):

self._session_creator = session_creator

self._hooks = hooks

self.coord = None

self.tf_sess = None

self._stop_grace_period_secs = stop_grace_period_secs

def create_session(self):

"""Creates a coordinated session."""

# Keep the tf_sess for unit testing.

self.tf_sess = self._session_creator.create_session()

# We don't want coordinator to suppress any exception.

self.coord = coordinator.Coordinator(clean_stop_exception_types=[])

queue_runner.start_queue_runners(sess=self.tf_sess, coord=self.coord)

# Inform the hooks that a new session has been created.

for hook in self._hooks:

hook.after_create_session(self.tf_sess, self.coord)

return _CoordinatedSession(

_HookedSession(self.tf_sess, self._hooks), self.coord,

self._stop_grace_period_secs)

def _close_internal(self, exception_type=None):

try:

if not exception_type:

for h in self._hooks:

h.end(self._coordinated_creator.tf_sess)

finally:

try:

if self._sess is None:

raise RuntimeError('Session is already closed.')

self._sess.close()

finally:

self._sess = None

self._coordinated_creator.tf_sess = None

self._coordinated_creator.coord = None

if not self._graph_was_finalized:

ops.get_default_graph()._unsafe_unfinalize() # pylint: disable=protected-access

def _is_closed(self):

"""Return True if the monitored session is closed. For tests only.

return self._coordinated_creator.tf_sess is None

def _tf_sess(self):

"""Session-like object that handles initialization, recovery and hooks.

def __init__(self, session_creator=None, hooks=None,

stop_grace_period_secs=120):

super(MonitoredSession, self).__init__(

session_creator, hooks, should_recover=True,

"""Session-like object that handles initialization, restoring, and hooks.

def __init__(self,

hooks=None,

scaffold=None,

master='',

config=None,

checkpoint_dir=None,

stop_grace_period_secs=120,

checkpoint_filename_with_path=None):

"""Creates a SingularMonitoredSession.

session_creator = ChiefSessionCreator(

scaffold=scaffold,

master=master,

config=config,

checkpoint_dir=checkpoint_dir,

checkpoint_filename_with_path=checkpoint_filename_with_path)

super(SingularMonitoredSession, self).__init__(

session_creator, hooks, should_recover=False,

stop_grace_period_secs=stop_grace_period_secs)

def raw_session(self):

"""Returns underlying `TensorFlow.Session` object."""

"""Wrapper around a `tf.Session`.

def __init__(self, sess):

"""Creates a `_WrappedSession`.

self._sess = sess

self._wrapped_is_stoppable = isinstance(self._sess, _WrappedSession)

@property

def graph(self):

return self._sess.graph

@property

def sess_str(self):

return self._sess.sess_str

def should_stop(self):

"""Return true if this session should not be used anymore.

if self._check_stop():

return True

if self._sess:

return self._wrapped_is_stoppable and self._sess.should_stop()

return True

def _check_stop(self):

"""Hook for subclasses to provide their own stop condition.

return False

def close(self):

if self._sess:

try:

self._sess.close()

except _PREEMPTION_ERRORS:

pass

finally:

self._sess = None

def run(self, *args, **kwargs):

return self._sess.run(*args, **kwargs)

def run_step_fn(self, step_fn, raw_session, run_with_hooks):

# `_RecoverableSession` sets `run_with_hooks` to `_CoordinatedSession.run`.

# It is `None` when called from `_CoordinatedSession`. In that case

# `self.run` is `_CoordinatedSession.run`.

run_with_hooks = run_with_hooks or self.run

"""A wrapped session that recreates a session upon certain kinds of errors.

def __init__(self, sess_creator):

"""Create a new `_RecoverableSession`.

self._sess_creator = sess_creator

_WrappedSession.__init__(self, self._create_session())

def _create_session(self):

while True:

try:

return self._sess_creator.create_session()

except _PREEMPTION_ERRORS as e:

logging.info('An error was raised while a session was being created. '

'This may be due to a preemption of a connected worker '

'or parameter server. A new session will be created. '

'Error: %s', e)

def _check_stop(self):

try:

if self._sess:

return self._sess._check_stop() # pylint: disable=protected-access

else:

return True

except _PREEMPTION_ERRORS as e:

logging.info('An error was raised while considering whether the '

'session is complete. This may be due to a preemption in '

'a connected worker or parameter server. The current '

'session will be closed and a new session will be '

'created. Error: %s', e)

self.close()

self._sess = self._create_session()

# Since we have just recreated the session, the overall computation should

# not stop:

return False