Skip to content

Track scheduling dependencies for workflow invocations #22128

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

Draft
mvdbeek wants to merge 1 commit into
base: dev
Choose a base branch
from
Draft

Track scheduling dependencies for workflow invocations #22128

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
20 changes: 15 additions & 5 deletions lib/galaxy/model/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2324,17 +2324,27 @@ def update_hdca_update_time_for_job(self, update_time, sa_session):

def set_final_state(self, final_state):
self.set_state(final_state)
# TODO: migrate to where-in subqueries?
sa_session = required_object_session(self)
update_time = now()
self.update_hdca_update_time_for_job(update_time=update_time, sa_session=sa_session)
params = {"job_id": self.id, "update_time": update_time}
# Update workflow_invocation_step for direct job_id
statement = text("""
UPDATE workflow_invocation_step
SET update_time = :update_time
WHERE job_id = :job_id;
""")
sa_session = required_object_session(self)
update_time = now()
self.update_hdca_update_time_for_job(update_time=update_time, sa_session=sa_session)
params = {"job_id": self.id, "update_time": update_time}
sa_session.execute(statement, params)
# Also update via implicit_collection_jobs link
statement_icj = text("""
UPDATE workflow_invocation_step
SET update_time = :update_time
WHERE implicit_collection_jobs_id IN (
SELECT implicit_collection_jobs_id FROM implicit_collection_jobs_job_association
WHERE job_id = :job_id
);
""")
sa_session.execute(statement_icj, params)

def get_destination_configuration(self, dest_params, config, key, default=None):
"""Get a destination parameter that can be defaulted back
Expand Down
28 changes: 25 additions & 3 deletions lib/galaxy/workflow/modules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
import re
from collections import defaultdict
from collections.abc import Iterable
from dataclasses import dataclass
from enum import Enum
from typing import (
Any,
cast,
Expand Down Expand Up @@ -176,7 +178,7 @@ def to_cwl(
if step:
if not value.dataset.in_ready_state():
why = f"dataset [{value.id}] is needed for valueFrom expression and is non-ready"
raise DelayedWorkflowEvaluation(why=why)
raise DelayedWorkflowEvaluation(why=why, dependency=SchedulingDependency(DependencyType.HDA, value.id))
if not value.is_ok:
raise FailWorkflowEvaluation(
why=InvocationFailureDatasetFailed(
Expand Down Expand Up @@ -905,6 +907,9 @@ def execute(
subworkflow_invoker.invoke()
subworkflow = subworkflow_invoker.workflow
subworkflow_progress = subworkflow_invoker.progress
# Propagate scheduling dependencies from subworkflow to parent
if subworkflow_progress.scheduling_dependencies:
progress.scheduling_dependencies.update(subworkflow_progress.scheduling_dependencies)
outputs = {}
for workflow_output in subworkflow.workflow_outputs:
workflow_output_label = (
Expand Down Expand Up @@ -1935,7 +1940,10 @@ def recover_mapping(self, invocation_step, progress):
)
)
delayed_why = "workflow paused at this step waiting for review"
raise DelayedWorkflowEvaluation(why=delayed_why)
dependency = None
if invocation_step:
dependency = SchedulingDependency(DependencyType.WORKFLOW_INVOCATION_STEP, invocation_step.id)
raise DelayedWorkflowEvaluation(why=delayed_why, dependency=dependency)

def do_invocation_step_action(self, step, action):
"""Update or set the workflow invocation state action - generic
Expand Down Expand Up @@ -2782,9 +2790,23 @@ def from_workflow_step(self, trans, step: WorkflowStep, **kwargs) -> WorkflowMod
module_factory = WorkflowModuleFactory(module_types)


class DependencyType(str, Enum):
JOB = "job"
HDA = "hda"
HDCA = "hdca"
WORKFLOW_INVOCATION_STEP = "workflow_invocation_step"


@dataclass(frozen=True)
class SchedulingDependency:
dependency_type: DependencyType
id: int


class DelayedWorkflowEvaluation(Exception):
def __init__(self, why=None):
def __init__(self, why=None, dependency: Optional[SchedulingDependency] = None):
self.why = why
self.dependency = dependency


class CancelWorkflowEvaluation(Exception):
Expand Down
43 changes: 34 additions & 9 deletions lib/galaxy/workflow/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,11 @@ def schedule(
workflow: "Workflow",
workflow_run_config: WorkflowRunConfig,
workflow_invocation: WorkflowInvocation,
) -> tuple[WorkflowOutputsType, WorkflowInvocation]:
return __invoke(trans, workflow, workflow_run_config, workflow_invocation)
) -> set[modules.SchedulingDependency]:
_outputs, _workflow_invocation, scheduling_dependencies = __invoke(
trans, workflow, workflow_run_config, workflow_invocation
)
return scheduling_dependencies


def __invoke(
Expand All @@ -78,7 +81,7 @@ def __invoke(
workflow_run_config: WorkflowRunConfig,
workflow_invocation: Optional[WorkflowInvocation] = None,
populate_state: bool = False,
) -> tuple[WorkflowOutputsType, WorkflowInvocation]:
) -> tuple[WorkflowOutputsType, WorkflowInvocation, set[modules.SchedulingDependency]]:
"""Run the supplied workflow in the supplied target_history."""
if populate_state:
modules.populate_module_and_state(
Expand Down Expand Up @@ -118,11 +121,13 @@ def __invoke(
workflow_invocation.fail()
workflow_invocation.add_message(failure)

scheduling_dependencies = invoker.progress.scheduling_dependencies

# Be sure to update state of workflow_invocation.
trans.sa_session.add(workflow_invocation)
trans.sa_session.commit()

return outputs, workflow_invocation
return outputs, workflow_invocation, scheduling_dependencies


def queue_invoke(
Expand Down Expand Up @@ -259,6 +264,8 @@ def invoke(self) -> dict[int, Any]:
except modules.DelayedWorkflowEvaluation as de:
step_delayed = delayed_steps = True
self.progress.mark_step_outputs_delayed(step, why=de.why)
if de.dependency:
self.progress.scheduling_dependencies.add(de.dependency)
except Exception as e:
log_function = log.error
failure_details = []
Expand Down Expand Up @@ -349,7 +356,10 @@ def __check_implicitly_dependent_step(self, output_id: int, step_id: int):
delayed_why = (
f"depends on step [{output_id}] but one or more jobs created from that step have not finished yet"
)
raise modules.DelayedWorkflowEvaluation(why=delayed_why)
raise modules.DelayedWorkflowEvaluation(
why=delayed_why,
dependency=modules.SchedulingDependency(modules.DependencyType.JOB, job.id),
)

if job.state != job.states.OK:
raise modules.FailWorkflowEvaluation(
Expand Down Expand Up @@ -403,6 +413,7 @@ def __init__(
when_values=None,
) -> None:
self.outputs: dict[int, Any] = {}
self.scheduling_dependencies: set[modules.SchedulingDependency] = set()
self.module_injector = module_injector
self.workflow_invocation = workflow_invocation
self.inputs_by_step_id = inputs_by_step_id
Expand Down Expand Up @@ -550,7 +561,10 @@ def replacement_for_connection(self, connection: "WorkflowStepConnection", is_da
)

delayed_why = f"dependent collection [{replacement.id}] not yet populated with datasets"
raise modules.DelayedWorkflowEvaluation(why=delayed_why)
raise modules.DelayedWorkflowEvaluation(
why=delayed_why,
dependency=modules.SchedulingDependency(modules.DependencyType.HDCA, replacement.id),
)

if isinstance(replacement, model.DatasetCollection):
raise NotImplementedError
Expand All @@ -559,7 +573,9 @@ def replacement_for_connection(self, connection: "WorkflowStepConnection", is_da
):
if isinstance(replacement, model.HistoryDatasetAssociation):
if replacement.is_pending:
raise modules.DelayedWorkflowEvaluation()
raise modules.DelayedWorkflowEvaluation(
dependency=modules.SchedulingDependency(modules.DependencyType.HDA, replacement.id)
)
if not replacement.is_ok:
raise modules.FailWorkflowEvaluation(
why=InvocationFailureDatasetFailed(
Expand All @@ -571,11 +587,15 @@ def replacement_for_connection(self, connection: "WorkflowStepConnection", is_da
)
else:
if not replacement.collection.populated:
raise modules.DelayedWorkflowEvaluation()
raise modules.DelayedWorkflowEvaluation(
dependency=modules.SchedulingDependency(modules.DependencyType.HDCA, replacement.id)
)
pending = False
pending_dataset_instance = None
for dataset_instance in replacement.dataset_instances:
if dataset_instance.is_pending:
pending = True
pending_dataset_instance = dataset_instance
elif not dataset_instance.is_ok:
raise modules.FailWorkflowEvaluation(
why=InvocationFailureDatasetFailed(
Expand All @@ -586,7 +606,10 @@ def replacement_for_connection(self, connection: "WorkflowStepConnection", is_da
)
)
if pending:
raise modules.DelayedWorkflowEvaluation()
assert pending_dataset_instance is not None
raise modules.DelayedWorkflowEvaluation(
dependency=modules.SchedulingDependency(modules.DependencyType.HDA, pending_dataset_instance.id)
)

return replacement

Expand Down Expand Up @@ -831,6 +854,8 @@ def _recover_mapping(self, step_invocation: WorkflowInvocationStep) -> None:
step_invocation.workflow_step.module.recover_mapping(step_invocation, self)
except modules.DelayedWorkflowEvaluation as de:
self.mark_step_outputs_delayed(step_invocation.workflow_step, de.why)
if de.dependency:
self.scheduling_dependencies.add(de.dependency)


__all__ = ("queue_invoke", "WorkflowRunConfig")
11 changes: 7 additions & 4 deletions lib/galaxy/workflow/schedulers/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@
ABCMeta,
abstractmethod,
)
from typing import TYPE_CHECKING

if TYPE_CHECKING:
from galaxy.workflow.modules import SchedulingDependency


class WorkflowSchedulingPlugin(metaclass=ABCMeta):
Expand Down Expand Up @@ -35,8 +39,7 @@ def shutdown(self):

class ActiveWorkflowSchedulingPlugin(WorkflowSchedulingPlugin, metaclass=ABCMeta):
@abstractmethod
def schedule(self, workflow_invocation):
"""Optionally return one or more commands to instrument job. These
commands will be executed on the compute server prior to the job
running.
def schedule(self, workflow_invocation) -> "set[SchedulingDependency]":
"""Schedule the workflow invocation and return any scheduling
dependencies that should be tracked for the next iteration.
"""
5 changes: 3 additions & 2 deletions lib/galaxy/workflow/schedulers/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

if TYPE_CHECKING:
from galaxy.model import WorkflowInvocation
from galaxy.workflow.modules import SchedulingDependency


log = logging.getLogger(__name__)
Expand All @@ -31,14 +32,14 @@ def startup(self, app):
def shutdown(self):
pass

def schedule(self, workflow_invocation: "WorkflowInvocation") -> None:
def schedule(self, workflow_invocation: "WorkflowInvocation") -> "set[SchedulingDependency]":
workflow = workflow_invocation.workflow
history = workflow_invocation.history
request_context = context.WorkRequestContext(
app=self.app, history=history, user=history.user
) # trans-like object not tied to a web-thread.
workflow_run_config = run_request.workflow_request_to_run_config(workflow_invocation)
run.schedule(
return run.schedule(
trans=request_context,
workflow=workflow,
workflow_run_config=workflow_run_config,
Expand Down
Loading
Loading