def _apply_operation_on_fine_grained_view(self, operation_def,
fine_grained_views,
next_hashed_path):
"""Applies a shardable operation on a fine grained view.
This also updates `cache_output_nodes` when necessary.
Args:
operation_def: A shardable `OperationDef`.
fine_grained_views: A tuple of `_OptimizationView.fine_grained_view`s.
next_hashed_path: The hashed path for the currently processed
operation_def.
Returns:
The resulting list of `_OptimizationView.fine_grained_view`s.
"""
result_fine_grained_view = collections.OrderedDict()
cache_entry_key = analyzer_cache.make_cache_entry_key(
tf.compat.as_bytes(operation_def.label) + b'-' + next_hashed_path)
for (dataset_idx, dataset_key) in enumerate(self._sorted_dataset_keys):
# We use an index for the label in order to make beam labels more stable.
infix = 'AnalysisIndex{}'.format(dataset_idx)
if (operation_def.cache_coder and self._cache_dict.get(
dataset_key, {}).get(cache_entry_key) is not None):
self._dataset_has_cache_misses[dataset_key] |= False
decode_cache = analyzer_nodes.DecodeCache(
dataset_key,
cache_entry_key,
coder=operation_def.cache_coder,
label='DecodeCache[{}][{}]'.format(operation_def.label,
infix))
(op_output, ) = nodes.OperationNode(decode_cache,
tuple()).outputs
else:
value_nodes = tuple(v[dataset_key] for v in fine_grained_views)
(op_output, ) = nodes.OperationNode(
operation_def._replace(
label='{}[{}]'.format(operation_def.label, infix)),
value_nodes).outputs
if operation_def.cache_coder:
self._dataset_has_cache_misses[dataset_key] = True
encode_cache = nodes.apply_operation(
analyzer_nodes.EncodeCache,
op_output,
coder=operation_def.cache_coder,
label='EncodeCache[{}][{}]'.format(
operation_def.label, infix))
self.cache_output_nodes[(dataset_key,
cache_entry_key)] = encode_cache
result_fine_grained_view[dataset_key] = op_output
return result_fine_grained_view
def _apply_operation_on_fine_grained_view(self, operation_def,
fine_grained_view):
"""Applies a shardable operation on a fine grained view.
This also updates `cache_output_nodes` when necessary.
Args:
operation_def: A shardable `OperationDef`.
fine_grained_view: A `_OptimizationView.fine_grained_view`.
Returns:
The resulting list of `_OptimizationView.fine_grained_view`s.
"""
result_fine_grained_view = collections.OrderedDict()
# TODO(b/37788560): Use a better cache key than label. A good alternative is
# to reuse graph_tools logic to compose names that include properties and
# fingerprint it.
cache_entry_key = analyzer_cache.make_cache_entry_key(
operation_def.label)
for dataset_key in self._dataset_keys:
# TODO(b/37788560): Add instrumentation.
if self._cache_dict.get(dataset_key,
{}).get(cache_entry_key) is not None:
(op_output, ) = nodes.OperationNode(
analyzer_nodes.DecodeCache(
dataset_key,
cache_entry_key,
coder=operation_def.cache_coder), tuple()).outputs
else:
value_node = fine_grained_view[dataset_key]
(op_output, ) = nodes.OperationNode(
operation_def._replace(label='{}[{}]'.format(
operation_def.label, dataset_key)),
(value_node, )).outputs
if operation_def.cache_coder:
encoded_cache = nodes.apply_operation(
analyzer_nodes.EncodeCache,
op_output,
coder=operation_def.cache_coder,
label='EncodeCache[{}][{}]'.format(
operation_def.label, dataset_key))
self.cache_output_nodes[(dataset_key,
cache_entry_key)] = encoded_cache
result_fine_grained_view[dataset_key] = op_output
return result_fine_grained_view
def _apply_operation_on_fine_grained_view(self, operation_def,
fine_grained_view,
next_hashed_path):
"""Applies a shardable operation on a fine grained view.
This also updates `cache_output_nodes` when necessary.
Args:
operation_def: A shardable `OperationDef`.
fine_grained_view: A `_OptimizationView.fine_grained_view`.
next_hashed_path: The hashed path for the currently processed
operation_def.
Returns:
The resulting list of `_OptimizationView.fine_grained_view`s.
"""
result_fine_grained_view = collections.OrderedDict()
cache_entry_key = analyzer_cache.make_cache_entry_key(
tf.compat.as_bytes(operation_def.label) + b'-' + next_hashed_path)
for dataset_key in self._dataset_keys:
if (operation_def.cache_coder and self._cache_dict.get(
dataset_key, {}).get(cache_entry_key) is not None):
(op_output,) = nodes.OperationNode(
analyzer_nodes.DecodeCache(
dataset_key,
cache_entry_key,
operation_def.label,
coder=operation_def.cache_coder), tuple()).outputs
else:
value_node = fine_grained_view[dataset_key]
(op_output,) = nodes.OperationNode(
operation_def._replace(
label='{}[{}]'.format(operation_def.label, dataset_key)),
(value_node,)).outputs
if operation_def.cache_coder:
encoded_cache = nodes.apply_operation(
analyzer_nodes.EncodeCache,
op_output,
coder=operation_def.cache_coder,
label='EncodeCache[{}][{}]'.format(operation_def.label,
dataset_key))
self.cache_output_nodes[(dataset_key,
cache_entry_key)] = encoded_cache
result_fine_grained_view[dataset_key] = op_output
return result_fine_grained_view