def check_flow_is_pending_or_running(self, state: State) -> State:
"""
Checks if the flow is in either a Pending state or Running state. Either are valid
starting points (because we allow simultaneous runs of the same flow run).
Args:
- state (State): the current state of this flow
Returns:
- State: the state of the flow after running the check
Raises:
- ENDRUN: if the flow is not pending or running
"""
# the flow run is already finished
if state.is_finished() is True:
self.logger.info("Flow run has already finished.")
raise ENDRUN(state)
# the flow run must be either pending or running (possibly redundant with above)
elif not (state.is_pending() or state.is_running()):
self.logger.info("Flow is not ready to run.")
raise ENDRUN(state)
return state
def set_task_to_running(self, state: State, inputs: Dict[str,
Result]) -> State:
"""
Sets the task to running
Args:
- state (State): the current state of this task
- inputs (Dict[str, Result]): a dictionary of inputs whose keys correspond
to the task's `run()` arguments.
Returns:
- State: the state of the task after running the check
Raises:
- ENDRUN: if the task is not ready to run
"""
if not state.is_pending():
self.logger.debug(
"Task '{name}': can't set state to Running because it "
"isn't Pending; ending run.".format(name=prefect.context.get(
"task_full_name", self.task.name)))
raise ENDRUN(state)
new_state = Running(message="Starting task run.")
return new_state
def get_task_inputs(
self, state: State, upstream_states: Dict[Edge, State]
) -> Dict[str, Result]:
"""
Given the task's current state and upstream states, generates the inputs for this task.
Upstream state result values are used. If the current state has `cached_inputs`, they
will override any upstream values.
Args:
- state (State): the task's current state.
- upstream_states (Dict[Edge, State]): the upstream state_handlers
Returns:
- Dict[str, Result]: the task inputs
"""
task_inputs = {} # type: Dict[str, Result]
results = {} # type: Dict[str, Result]
for edge, upstream_state in upstream_states.items():
# construct task inputs
if edge.key is not None:
task_inputs[edge.key] = upstream_state._result # type: ignore
if state.is_pending() and state.cached_inputs:
task_inputs.update(
{
k: r
for k, r in state.cached_inputs.items()
if task_inputs.get(k, NoResult) == NoResult
}
)
return task_inputs
def check_task_is_ready(self, state: State) -> State:
"""
Checks to make sure the task is ready to run (Pending or Mapped).
Args:
- state (State): the current state of this task
Returns:
- State: the state of the task after running the check
Raises:
- ENDRUN: if the task is not ready to run
"""
# the task is ready
if state.is_pending():
return state
# the task is mapped, in which case we still proceed so that the children tasks
# are generated (note that if the children tasks)
elif state.is_mapped():
self.logger.debug(
"Task '{name}': task is mapped, but run will proceed so children are generated.".format(
name=prefect.context.get("task_full_name", self.task.name)
)
)
return state
# this task is already running
elif state.is_running():
self.logger.debug(
"Task '{name}': task is already running.".format(
name=prefect.context.get("task_full_name", self.task.name)
)
)
raise ENDRUN(state)
elif state.is_cached():
return state
# this task is already finished
elif state.is_finished():
self.logger.debug(
"Task '{name}': task is already finished.".format(
name=prefect.context.get("task_full_name", self.task.name)
)
)
raise ENDRUN(state)
# this task is not pending
else:
self.logger.debug(
"Task '{name}' is not ready to run or state was unrecognized ({state}).".format(
name=prefect.context.get("task_full_name", self.task.name),
state=state,
)
)
raise ENDRUN(state)
def checkpoint_handler(task_runner: DSTaskRunner, old_state: State, new_state: State) -> State:
"""
A handler designed to implement result caching by filename. If the result handler's ``read``
method can be successfully run, this handler loads the result of that method as the task result
and sets the task state to ``Success``. Similarly, on successful
completion of the task, if the task was actually run and not loaded from cache, this handler
will apply the result handler's ``write`` method to the task.
Parameters
----------
task_runner : instance of DSTaskRunner
The task runner associated with the flow the handler is used in.
old_state : instance of prefect.engine.state.State
The current state of the task.
new_state : instance of prefect.engine.state.State
The expected new state of the task.
Returns
-------
new_state : instance of prefect.engine.state.State
The actual new state of the task.
"""
if "PREFECT__FLOWS__CHECKPOINTING" in os.environ and os.environ["PREFECT__FLOWS__CHECKPOINTING"] == "true":
raise AttributeError("Cannot use standard prefect checkpointing with this handler")
if task_runner.result_handler is not None and old_state.is_pending() and new_state.is_running():
if not hasattr(task_runner, "upstream_states"):
raise TypeError(
"upstream_states not found in task runner. Make sure to use "
"prefect_ds.task_runner.DSTaskRunner."
)
input_mapping = _create_input_mapping(task_runner.upstream_states)
try:
data = task_runner.task.result_handler.read(input_mapping=input_mapping)
except FileNotFoundError:
return new_state
except TypeError: # unexpected argument input_mapping
raise TypeError(
"Result handler could not accept input_mapping argument. "
"Please ensure that you are using a handler from prefect_ds."
)
result = Result(value=data, result_handler=task_runner.task.result_handler)
state = Success(result=result, message="Task loaded from disk.")
return state
if task_runner.result_handler is not None and old_state.is_running() and new_state.is_successful():
input_mapping = _create_input_mapping(task_runner.upstream_states)
task_runner.task.result_handler.write(new_state.result, input_mapping=input_mapping)
return new_state
def set_flow_to_running(self, state: State) -> State:
"""
Puts Pending flows in a Running state; leaves Running flows Running.
Args:
- state (State): the current state of this flow
Returns:
- State: the state of the flow after running the check
Raises:
- ENDRUN: if the flow is not pending or running
"""
if state.is_pending():
return Running(message="Running flow.")
elif state.is_running():
return state
else:
raise ENDRUN(state)
def set_task_to_running(self, state: State) -> State:
"""
Sets the task to running
Args:
- state (State): the current state of this task
Returns:
- State: the state of the task after running the check
Raises:
- ENDRUN: if the task is not ready to run
"""
if not state.is_pending():
self.logger.debug(
"Task '{name}': can't set state to Running because it "
"isn't Pending; ending run.".format(name=prefect.context.get(
"task_full_name", self.task.name)))
raise ENDRUN(state)
return Running(message="Starting task run.")
def get_task_inputs(
self, state: State, upstream_states: Dict[Edge, State]
) -> Dict[str, Result]:
"""
Given the task's current state and upstream states, generates the inputs for this task.
Upstream state result values are used. If the current state has `cached_inputs`, they
will override any upstream values which are `NoResult`.
Args:
- state (State): the task's current state.
- upstream_states (Dict[Edge, State]): the upstream state_handlers
Returns:
- Dict[str, Result]: the task inputs
"""
task_inputs = {} # type: Dict[str, Result]
handlers = {} # type: Dict[str, ResultHandler]
for edge, upstream_state in upstream_states.items():
# construct task inputs
if edge.key is not None:
handlers[edge.key] = handler = getattr(
edge.upstream_task, "result_handler", None
)
task_inputs[ # type: ignore
edge.key
] = upstream_state._result.to_result( # type: ignore
handler
) # type: ignore
if state.is_pending() and state.cached_inputs is not None: # type: ignore
task_inputs.update(
{
k: r.to_result(handlers.get(k))
for k, r in state.cached_inputs.items() # type: ignore
if task_inputs.get(k, NoResult) == NoResult
}
)
return task_inputs
