def thread_fn():
self.assertFalse(save_context.in_save_context())
with self.assertRaisesRegex(ValueError, 'not in a SaveContext'):
save_context.get_save_options()
options = save_options.SaveOptions(save_debug_info=False)
with save_context.save_context(options):
self.assertTrue(save_context.in_save_context())
# save_debug_info has a different value in this thread.
self.assertFalse(save_context.get_save_options().save_debug_info)
entered_context_in_thread.set()
continue_thread.wait()
self.assertFalse(save_context.in_save_context())
with self.assertRaisesRegex(ValueError, 'not in a SaveContext'):
save_context.get_save_options()
def is_saving_non_distributed():
"""Returns whether we're saving a non-distributed version of the model.
It returns True iff we are in saving context and are saving a non-distributed
version of the model. That is, SaveOptions.experimental_variable_policy is
NONE.
Returns:
A boolean.
"""
if not save_context.in_save_context():
return False
options = save_context.get_save_options()
return (options is not None and options.experimental_variable_policy !=
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES)
def test_multi_thread(self):
self.assertFalse(save_context.in_save_context())
with self.assertRaisesRegex(ValueError, 'not in a SaveContext'):
save_context.get_save_options()
options = save_options.SaveOptions(save_debug_info=True)
with save_context.save_context(options):
self.assertTrue(save_context.in_save_context())
self.assertTrue(save_context.get_save_options().save_debug_info)
entered_context_in_thread = threading.Event()
continue_thread = threading.Event()
def thread_fn():
self.assertFalse(save_context.in_save_context())
with self.assertRaisesRegex(ValueError, 'not in a SaveContext'):
save_context.get_save_options()
options = save_options.SaveOptions(save_debug_info=False)
with save_context.save_context(options):
self.assertTrue(save_context.in_save_context())
# save_debug_info has a different value in this thread.
self.assertFalse(save_context.get_save_options().save_debug_info)
entered_context_in_thread.set()
continue_thread.wait()
self.assertFalse(save_context.in_save_context())
with self.assertRaisesRegex(ValueError, 'not in a SaveContext'):
save_context.get_save_options()
t = threading.Thread(target=thread_fn)
t.start()
entered_context_in_thread.wait()
# Another thread shouldn't affect this thread.
self.assertTrue(save_context.in_save_context())
self.assertTrue(save_context.get_save_options().save_debug_info)
continue_thread.set()
t.join()
# Another thread exiting SaveContext shouldn't affect this thread.
self.assertTrue(save_context.in_save_context())
self.assertTrue(save_context.get_save_options().save_debug_info)
self.assertFalse(save_context.in_save_context())
with self.assertRaisesRegex(ValueError, 'not in a SaveContext'):
save_context.get_save_options()
def _make_execution_context() -> ExecutionContext:
"""Generates an ExecutionContext based on current contextual info."""
ctx = context.context()
# Don't need to open an init_scope if the _cache_key call is in eager mode
# already.
executing_eagerly = ctx.executing_eagerly()
parent_graph = None
xla_context_id = 0
if not executing_eagerly:
# We want to force function retracing for each different
# XLAControlFlowContext, so add `xla_context_id` to the cache key.
xla_context = _enclosing_xla_context()
if xla_context is not None and xla_context.RequiresUniqueFunctionRetracing(
):
xla_context_id = id(xla_context)
with ops.init_scope():
# The graph, or whether we're executing eagerly, should be a part of the
# cache key so we don't improperly capture tensors such as variables.
executing_eagerly = ctx.executing_eagerly()
parent_graph = None if executing_eagerly else ops.get_default_graph(
)
# pylint: disable=protected-access
default_graph = ops.get_default_graph()
# TODO(b/117617952): The current distribution strategy will affect graph
# building (e.g. accessing different variables from different devices) and
# so requires retracing for each device.
strategy_stack = default_graph._distribution_strategy_stack
uses_distribution_strategy = (strategy_stack
and strategy_stack[-1].strategy.extended.
_retrace_functions_for_each_device)
if executing_eagerly:
colocation_stack = ()
if uses_distribution_strategy:
device_functions = (pydev.merge_device(ctx.device_name), )
else:
device_functions = ()
else:
colocation_stack = tuple(default_graph._colocation_stack.peek_objs())
if (uses_distribution_strategy
or func_graph_module.device_stack_has_callable(
default_graph._device_function_stack)):
# Putting the device in the cache key ensures that call-site device
# annotations are respected.
device_functions = tuple(
default_graph._device_functions_outer_to_inner)
else:
device_functions = ()
in_cross_replica_context = False
try:
in_cross_replica_context = (strategy_stack[-1].replica_context is None) # pylint: disable=protected-access
except (AttributeError, IndexError):
pass
if save_context.in_save_context():
variable_policy = (
save_context.get_save_options().experimental_variable_policy)
else:
variable_policy = None
return ExecutionContext(parent_graph, device_functions, colocation_stack,
in_cross_replica_context, variable_policy,
xla_context_id)
