def _orchestrate_stop_initiated_pipeline(
mlmd_handle: metadata.Metadata, task_queue: tq.TaskQueue,
service_job_manager: service_jobs.ServiceJobManager,
pipeline_state: pstate.PipelineState) -> None:
"""Orchestrates stop initiated pipeline."""
stop_reason = pipeline_state.stop_initiated_reason()
assert stop_reason is not None
pipeline = pipeline_state.pipeline
has_active_executions = False
for node in pstate.get_all_pipeline_nodes(pipeline):
if service_job_manager.is_pure_service_node(pipeline_state,
node.node_info.id):
service_job_manager.stop_node_services(pipeline_state,
node.node_info.id)
elif _maybe_enqueue_cancellation_task(mlmd_handle, pipeline, node,
task_queue):
has_active_executions = True
elif service_job_manager.is_mixed_service_node(pipeline_state,
node.node_info.id):
service_job_manager.stop_node_services(pipeline_state,
node.node_info.id)
if not has_active_executions:
with pipeline_state:
# Update pipeline execution state in MLMD.
pipeline_state.update_pipeline_execution_state(stop_reason)
def _orchestrate_stop_initiated_pipeline(
mlmd_handle: metadata.Metadata, task_queue: tq.TaskQueue,
service_job_manager: service_jobs.ServiceJobManager,
pipeline_state: pstate.PipelineState) -> None:
"""Orchestrates stop initiated pipeline."""
with pipeline_state:
stop_reason = pipeline_state.stop_initiated_reason()
assert stop_reason is not None
is_active = _cancel_nodes(mlmd_handle,
task_queue,
service_job_manager,
pipeline_state,
pause=False)
if not is_active:
with pipeline_state:
# Update pipeline execution state in MLMD.
pipeline_state.set_pipeline_execution_state(
_mlmd_execution_code(stop_reason))