Adds tracing spans inside of function invoked on EagerTFExecutor CreateCall.
@@ -226,6 +226,8 @@ def embed_tensorflow_computation(comp, type_spec=None, device=None):
# resource on iteration, as well as usages that expect to reinitialize a
# table with new data. Revisit the semantics implied by this cleanup
# pattern.
+ with tracing.span(
+ 'EagerTFExecutor.create_call', 'hashtable_cleanup', span=True):
eager_cleanup_resources = []
for op in wrapped_fn.graph.get_operations():
if op.type == 'HashTableV2':
@@ -237,6 +239,8 @@ def embed_tensorflow_computation(comp, type_spec=None, device=None):
param_elements = []
if arg is not None:
+ with tracing.span(
+ 'EagerTFExecutor.create_call', 'arg_ingestion', span=True):
arg_parts = structure.flatten(arg)
if len(arg_parts) != len(param_fns):
raise RuntimeError('Expected {} arguments, found {}.'.format(
@@ -251,6 +255,8 @@ def embed_tensorflow_computation(comp, type_spec=None, device=None):
# manually.
# TODO(b/144127474): Remove this manual cleanup once tf.wrap_function(...)
# is fixed.
+ with tracing.span(
+ 'EagerTFExecutor.create_call', 'variable_cleanup', span=True):
resources = []
for op in wrapped_fn.graph.get_operations():
if op.type == 'VarHandleOp':
@@ -259,6 +265,8 @@ def embed_tensorflow_computation(comp, type_spec=None, device=None):
for resource in wrapped_fn.prune(feeds={}, fetches=resources)():
tf.raw_ops.DestroyResourceOp(resource=resource)
+ with tracing.span(
+ 'EagerTFExecutor.create_call', 'result_packing', span=True):
result_elements = []
for result_part, result_fn in zip(result_parts, result_fns):
result_elements.append(result_fn(result_part))
