def components(self, components: List[base_node.BaseNode]):
deduped_components = set(components)
producer_map = {}
instances_per_component_type = collections.defaultdict(set)
# Fills in producer map.
for component in deduped_components:
# Guarantees every component of a component type has unique component_id.
if component.id in instances_per_component_type[component.type]:
raise RuntimeError('Duplicated component_id %s for component type %s' %
(component.id, component.type))
instances_per_component_type[component.type].add(component.id)
for key, output_channel in component.outputs.items():
assert not producer_map.get(
output_channel), '{} produced more than once'.format(output_channel)
producer_map[output_channel] = component
output_channel.producer_component_id = component.id
output_channel.output_key = key
# Connects nodes based on producer map.
for component in deduped_components:
for i in component.inputs.values():
if producer_map.get(i):
component.add_upstream_node(producer_map[i])
producer_map[i].add_downstream_node(component)
layers = topsort.topsorted_layers(
list(deduped_components),
get_node_id_fn=lambda c: c.id,
get_parent_nodes=lambda c: c.upstream_nodes,
get_child_nodes=lambda c: c.downstream_nodes)
self._components = []
for layer in layers:
for component in layer:
self._components.append(component)
def test_topsorted_layers_empty(self):
layers = topsort.topsorted_layers(
nodes=[],
get_node_id_fn=lambda n: n.name,
get_parent_nodes=lambda n: [],
get_child_nodes=lambda n: [])
self.assertEqual([], layers)
def components(self, components: List[base_node.BaseNode]):
deduped_components = set(components)
node_by_id = {}
# TODO(b/202822834): Use better distinction for bound channels.
# bound_channels stores the exhaustive list of all nodes' output channels,
# which is implicitly *bound* to the single pipeline run, as opposed to
# manually constructed channels to fetch artifacts beyond the current
# pipeline run.
bound_channels = set()
# Fills in producer map.
for component in deduped_components:
# Checks every node has an unique id.
if component.id in node_by_id:
raise RuntimeError(
f'Duplicated node_id {component.id} for component type'
f'{component.type}.')
node_by_id[component.id] = component
for key, output_channel in component.outputs.items():
if (output_channel.producer_component_id is not None and
output_channel.producer_component_id != component.id
and output_channel.output_key != key):
raise AssertionError(
f'{output_channel} is produced more than once: '
f'{output_channel.producer_id}[{output_channel.output_key}], '
f'{component.id}[{key}]')
output_channel.producer_component_id = component.id
output_channel.output_key = key
bound_channels.add(output_channel)
# Connects nodes based on producer map.
for component in deduped_components:
channels = list(component.inputs.values())
for exec_property in component.exec_properties.values():
if isinstance(exec_property, ph.ChannelWrappedPlaceholder):
channels.append(exec_property.channel)
for base_channel in channels:
for ch in channel_utils.get_individual_channels(base_channel):
if ch not in bound_channels:
continue
upstream_node = node_by_id.get(ch.producer_component_id)
if upstream_node:
component.add_upstream_node(upstream_node)
upstream_node.add_downstream_node(component)
layers = topsort.topsorted_layers(
list(deduped_components),
get_node_id_fn=lambda c: c.id,
get_parent_nodes=lambda c: c.upstream_nodes,
get_child_nodes=lambda c: c.downstream_nodes)
self._components = []
for layer in layers:
for component in layer:
self._components.append(component)
if self.beam_pipeline_args:
for component in self._components:
add_beam_pipeline_args_to_component(component,
self.beam_pipeline_args)
def test_topsorted_layers_error_if_cycle(self):
nodes = [
Node('A', [], ['B', 'E']),
Node('B', ['A', 'D'], ['C']),
Node('C', ['B'], ['D']),
Node('D', ['C'], ['B']),
Node('E', ['A'], [])
]
node_map = {node.name: node for node in nodes}
with self.assertRaisesRegex(topsort.InvalidDAGError, 'Cycle detected.'):
topsort.topsorted_layers(
nodes,
get_node_id_fn=lambda n: n.name,
get_parent_nodes=(
lambda n: [node_map[name] for name in n.upstream_nodes]),
get_child_nodes=(
lambda n: [node_map[name] for name in n.downstream_nodes]))
def _topsorted_layers(
pipeline: pipeline_pb2.Pipeline
) -> List[List[pipeline_pb2.PipelineNode]]:
"""Returns pipeline nodes in topologically sorted layers."""
node_by_id = _node_by_id(pipeline)
return topsort.topsorted_layers(
[node.pipeline_node for node in pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [node_by_id[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node: [node_by_id[n] for n in node.downstream_nodes]))
def generate(self) -> List[task_lib.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
result = []
for nodes in layers:
# Boolean that's set if there's at least one successfully executed node
# in the current layer.
executed_nodes = False
for node in nodes:
if node.node_info.id in self._ignore_node_ids:
logging.info(
'Ignoring node for task generation: %s',
task_lib.NodeUid.from_pipeline_node(
self._pipeline, node))
continue
# If a task for the node is already tracked by the task queue, it need
# not be considered for generation again.
if self._is_task_id_tracked_fn(
task_lib.exec_node_task_id_from_pipeline_node(
self._pipeline, node)):
continue
executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
if (executions
and task_gen_utils.is_latest_execution_successful(
executions)):
executed_nodes = True
continue
# If all upstream nodes are executed but current node is not executed,
# the node is deemed ready for execution.
if self._upstream_nodes_executed(node):
task = self._generate_task(node)
if task:
result.append(task)
# If there are no executed nodes in the current layer, downstream nodes
# need not be checked.
if not executed_nodes:
break
return result
def test_topsorted_layers_ignore_duplicate_parent_node(self):
nodes = [
Node('A', [], ['B']),
Node('B', ['A', 'A'], []), # Duplicate parent node 'A'
]
node_map = {node.name: node for node in nodes}
layers = topsort.topsorted_layers(
nodes,
get_node_id_fn=lambda n: n.name,
get_parent_nodes=(
lambda n: [node_map[name] for name in n.upstream_nodes]),
get_child_nodes=(
lambda n: [node_map[name] for name in n.downstream_nodes]))
self.assertEqual([['A'], ['B']],
[[node.name for node in layer] for layer in layers])
def test_topsorted_layers_ignore_unknown_parent_node(self):
nodes = [
Node('A', [], ['B']),
Node('B', ['A'], ['C']),
Node('C', ['B'], []),
]
node_map = {node.name: node for node in nodes}
# Exclude node A. Node B now has a parent node 'A' that should be ignored.
layers = topsort.topsorted_layers(
[node_map['B'], node_map['C']],
get_node_id_fn=lambda n: n.name,
get_parent_nodes=(
lambda n: [node_map[name] for name in n.upstream_nodes]),
get_child_nodes=(
lambda n: [node_map[name] for name in n.downstream_nodes]))
self.assertEqual([['B'], ['C']],
[[node.name for node in layer] for layer in layers])
def test_topsorted_layers_DAG(self):
nodes = [
Node('A', [], ['B', 'C', 'D']),
Node('B', ['A'], []),
Node('C', ['A'], ['D']),
Node('D', ['A', 'C', 'F'], ['E']),
Node('E', ['D'], []),
Node('F', [], ['D'])
]
node_map = {node.name: node for node in nodes}
layers = topsort.topsorted_layers(
nodes,
get_node_id_fn=lambda n: n.name,
get_parent_nodes=(
lambda n: [node_map[name] for name in n.upstream_nodes]),
get_child_nodes=(
lambda n: [node_map[name] for name in n.downstream_nodes]))
self.assertEqual([['A', 'F'], ['B', 'C'], ['D'], ['E']],
[[node.name for node in layer] for layer in layers])
def generate(self) -> List[task_pb2.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
tasks = []
with self._mlmd_connection as m:
# TODO(goutham): Cache executions and/or use TaskQueue so that we don't
# have to make MLMD queries for upstream nodes in each iteration.
for nodes in layers:
# Boolean that's set if there's at least one successfully executed node
# in the current layer.
executed_nodes = False
for node in nodes:
executions = task_gen_utils.get_executions(m, node)
if (executions
and task_gen_utils.is_latest_execution_successful(
executions)):
executed_nodes = True
continue
# If all upstream nodes are executed but current node is not executed,
# the node is deemed ready for execution.
if self._upstream_nodes_executed(m, node):
task = self._generate_task(m, node)
if task:
tasks.append(task)
# If there are no executed nodes in the current layer, downstream nodes
# need not be checked.
if not executed_nodes:
break
return tasks
def generate(self) -> List[task_lib.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
result = []
for layer_num, nodes in enumerate(layers):
# Boolean that's set if there's at least one successfully executed node
# in the current layer.
completed_node_ids = set()
for node in nodes:
node_uid = task_lib.NodeUid.from_pipeline_node(
self._pipeline, node)
node_id = node.node_info.id
if self._service_job_manager.is_pure_service_node(
self._pipeline_state, node.node_info.id):
if not self._upstream_nodes_executed(node):
continue
service_status = self._service_job_manager.ensure_node_services(
self._pipeline_state, node_id)
if service_status == service_jobs.ServiceStatus.SUCCESS:
logging.info('Service node completed successfully: %s',
node_uid)
completed_node_ids.add(node_id)
elif service_status == service_jobs.ServiceStatus.FAILED:
logging.error('Failed service node: %s', node_uid)
return [
task_lib.FinalizePipelineTask(
pipeline_uid=self._pipeline_state.pipeline_uid,
status=status_lib.Status(
code=status_lib.Code.ABORTED,
message=
(f'Aborting pipeline execution due to service '
f'node failure; failed node uid: {node_uid}'
)))
]
else:
logging.info('Pure service node in progress: %s',
node_uid)
continue
# If a task for the node is already tracked by the task queue, it need
# not be considered for generation again.
if self._is_task_id_tracked_fn(
task_lib.exec_node_task_id_from_pipeline_node(
self._pipeline, node)):
continue
executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
if (executions
and task_gen_utils.is_latest_execution_successful(
executions)):
completed_node_ids.add(node_id)
continue
# If all upstream nodes are executed but current node is not executed,
# the node is deemed ready for execution.
if self._upstream_nodes_executed(node):
task = self._generate_task(node)
if task_lib.is_finalize_pipeline_task(task):
return [task]
else:
result.append(task)
# If there are no completed nodes in the current layer, downstream nodes
# need not be checked.
if not completed_node_ids:
break
# If all nodes in the final layer are completed successfully , the
# pipeline can be finalized.
# TODO(goutham): If there are conditional eval nodes, not all nodes may be
# executed in the final layer. Handle this case when conditionals are
# supported.
if layer_num == len(layers) - 1 and completed_node_ids == set(
node.node_info.id for node in nodes):
return [
task_lib.FinalizePipelineTask(
pipeline_uid=self._pipeline_state.pipeline_uid,
status=status_lib.Status(code=status_lib.Code.OK))
]
return result
def generate(self) -> List[task_lib.Task]:
"""Generates tasks for executing the next executable nodes in the pipeline.
The returned tasks must have `exec_task` populated. List may be empty if
no nodes are ready for execution.
Returns:
A `list` of tasks to execute.
"""
layers = topsort.topsorted_layers(
[node.pipeline_node for node in self._pipeline.nodes],
get_node_id_fn=lambda node: node.node_info.id,
get_parent_nodes=(
lambda node: [self._node_map[n] for n in node.upstream_nodes]),
get_child_nodes=(
lambda node:
[self._node_map[n] for n in node.downstream_nodes]))
result = []
successful_node_ids = set()
for layer_num, layer_nodes in enumerate(layers):
for node in layer_nodes:
node_uid = task_lib.NodeUid.from_pipeline_node(
self._pipeline, node)
node_id = node.node_info.id
if self._in_successful_nodes_cache(node_uid):
successful_node_ids.add(node_id)
continue
if not self._upstream_nodes_successful(node,
successful_node_ids):
continue
# If this is a pure service node, there is no ExecNodeTask to generate
# but we ensure node services and check service status.
service_status = self._ensure_node_services_if_pure(node_id)
if service_status is not None:
if service_status == service_jobs.ServiceStatus.FAILED:
return [
self._abort_task(
f'service job failed; node uid: {node_uid}')
]
if service_status == service_jobs.ServiceStatus.SUCCESS:
logging.info('Service node successful: %s', node_uid)
successful_node_ids.add(node_id)
continue
# If a task for the node is already tracked by the task queue, it need
# not be considered for generation again but we ensure node services
# in case of a mixed service node.
if self._is_task_id_tracked_fn(
task_lib.exec_node_task_id_from_pipeline_node(
self._pipeline, node)):
service_status = self._ensure_node_services_if_mixed(
node_id)
if service_status == service_jobs.ServiceStatus.FAILED:
return [
self._abort_task(
f'associated service job failed; node uid: {node_uid}'
)
]
continue
node_executions = task_gen_utils.get_executions(
self._mlmd_handle, node)
latest_execution = task_gen_utils.get_latest_execution(
node_executions)
# If the latest execution is successful, we're done.
if latest_execution and execution_lib.is_execution_successful(
latest_execution):
logging.info('Node successful: %s', node_uid)
successful_node_ids.add(node_id)
continue
# If the latest execution failed, the pipeline should be aborted.
if latest_execution and not execution_lib.is_execution_active(
latest_execution):
error_msg_value = latest_execution.custom_properties.get(
constants.EXECUTION_ERROR_MSG_KEY)
error_msg = data_types_utils.get_metadata_value(
error_msg_value) if error_msg_value else ''
return [
self._abort_task(
f'node failed; node uid: {node_uid}; error: {error_msg}'
)
]
# Finally, we are ready to generate an ExecNodeTask for the node.
task = self._maybe_generate_task(node, node_executions,
successful_node_ids)
if task:
if task_lib.is_finalize_pipeline_task(task):
return [task]
else:
result.append(task)
layer_node_ids = set(node.node_info.id for node in layer_nodes)
successful_layer_node_ids = layer_node_ids & successful_node_ids
self._update_successful_nodes_cache(successful_layer_node_ids)
# If all nodes in the final layer are completed successfully , the
# pipeline can be finalized.
# TODO(goutham): If there are conditional eval nodes, not all nodes may be
# executed in the final layer. Handle this case when conditionals are
# supported.
if (layer_num == len(layers) - 1
and successful_layer_node_ids == layer_node_ids):
return [
task_lib.FinalizePipelineTask(
pipeline_uid=self._pipeline_uid,
status=status_lib.Status(code=status_lib.Code.OK))
]
return result
