Skip to content

feat: Orchestrator - Conditional execution #301

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Ark-kun wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: Orchestrator - Conditional execution #301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions cloud_pipelines_backend/api_server_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1826,6 +1826,53 @@ def _recursively_create_all_executions_and_artifacts(
for link in child_execution_node.output_artifact_links
}

# Handling conditional execution (`child_task_spec.is_enabled`)
if child_task_spec.is_enabled is not None:
if not isinstance(
child_component_spec.implementation,
structures.ContainerImplementation,
):
# We do not support `is_enabled` on graph component tasks since it's unclear
# how it should interact with the `is_enabled` settings of child component tasks`.
raise ApiServiceError(
f"TaskSpec.is_enabled is only supported for container component tasks. {child_task_id=}, {child_task_spec=}"
)
is_enabled_argument = child_task_spec.is_enabled
if isinstance(is_enabled_argument, (bool, str)):
# Do not create an ArtifactNode for a constant value
is_enabled_artifact_node = None
pass
elif isinstance(is_enabled_argument, structures.GraphInputArgument):
is_enabled_artifact_node = input_artifact_nodes.get(
is_enabled_argument.graph_input.input_name
)
if is_enabled_artifact_node is None:
# Warning: unconnected upstream for is_enabled
pass
elif isinstance(is_enabled_argument, structures.TaskOutputArgument):
task_output_source = is_enabled_argument.task_output
is_enabled_artifact_node = task_output_artifact_nodes[
task_output_source.task_id
][task_output_source.output_name]
else:
raise ApiServiceError(
f"Unsupported TaskSpec.is_enabled value: {child_task_id=}, {is_enabled_argument=}"
)
if is_enabled_artifact_node:
if not isinstance(is_enabled_artifact_node, bts.ArtifactNode):
raise ApiServiceError(
f"Unsupported TaskSpec.is_enabled value: {child_task_id=}, {is_enabled_artifact_node=}"
)
child_task_input_artifact_nodes[
bts.EXECUTION_NODE_TASK_IS_ENABLED_SPECIAL_INPUT_NAME
] = is_enabled_artifact_node
input_artifact_link = bts.InputArtifactLink(
execution=child_execution_node,
input_name=bts.EXECUTION_NODE_TASK_IS_ENABLED_SPECIAL_INPUT_NAME,
artifact=is_enabled_artifact_node,
)
session.add(input_artifact_link)

# Processing root graph output artifacts
for output_name, output_source in (graph_spec.output_values or {}).items():
if not isinstance(output_source, structures.TaskOutputArgument):
Expand Down
4 changes: 4 additions & 0 deletions cloud_pipelines_backend/backend_types_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,10 @@ class ExecutionNode(_TableBase):

_status_changed: bool = dataclasses.field(default=False, init=False, repr=False)

# The name of a special input that links to the the `TaskSpec.is_enabled` artifact.
EXECUTION_NODE_TASK_IS_ENABLED_SPECIAL_INPUT_NAME = (
"tangleml.com/orchestration/conditional_execution/is_enabled"
)

EXECUTION_NODE_EXTRA_DATA_STATUS_HISTORY_KEY = "container_execution_status_history"
EXECUTION_NODE_EXTRA_DATA_SYSTEM_ERROR_EXCEPTION_MESSAGE_KEY = (
Expand Down
59 changes: 59 additions & 0 deletions cloud_pipelines_backend/orchestrator_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,65 @@ def internal_process_one_queued_execution(
session.commit()
return

# region Conditional execution
if task_spec.is_enabled is not None:
is_enabled: bool | None = None
if isinstance(task_spec.is_enabled, bool):
is_enabled = task_spec.is_enabled
else:
is_enabled_str: str | None = None
if isinstance(task_spec.is_enabled, str):
is_enabled_str = task_spec.is_enabled
elif isinstance(
task_spec.is_enabled,
(structures.TaskOutputArgument, structures.GraphInputArgument),
):
if (
bts.EXECUTION_NODE_TASK_IS_ENABLED_SPECIAL_INPUT_NAME
not in input_artifact_data
):
# TODO: Change this to a warning. Missing input artifact should lead to is_enabled = True
raise OrchestratorError(
f"Conditional execution check: is_enabled input does not have an artifact (this should not happen). {task_spec.is_enabled=}"
)
# Getting and removing the is_enabled artifact from input_artifact_data so that it does not affect caching.
is_enabled_artifact = input_artifact_data.pop(
bts.EXECUTION_NODE_TASK_IS_ENABLED_SPECIAL_INPUT_NAME
)
if not (is_enabled_artifact and is_enabled_artifact.value):
raise OrchestratorError(
f"Conditional execution check: is_enabled artifact does not have value {is_enabled_artifact=}, {task_spec.is_enabled=}."
)
is_enabled_str = is_enabled_artifact.value
# DynamicData or secrets is not supported here.
else:
raise OrchestratorError(
f"Conditional execution check: Unexpected {task_spec.is_enabled=}."
)
if is_enabled_str.lower() == "true":
is_enabled = True
elif is_enabled_str.lower() == "false":
is_enabled = False
else:
raise OrchestratorError(
f"Conditional execution check: Unexpected {is_enabled_str=}, {task_spec.is_enabled=}."
)
if not is_enabled:
_logger.info(
f"Conditionally skipping execution {execution.id} and all downstream executions."
)
execution.container_execution_status = (
bts.ContainerExecutionStatus.SKIPPED
)
_mark_all_downstream_executions_as_skipped(
session=session,
execution=execution,
)
session.commit()
# Do not process the ExecutionNode any further.
return
# endregion

# We need to extract dynamic data arguments so that we can use them in cache key calculation.
# We read secrets and dynamic data from execution_node.extra_data rather than from task_spec.arguments,
# because some secrets might have been passed from upstream graph inputs.
Expand Down
Loading
Loading