def check_for_retry(self, state: State, inputs: Dict[str, Result]) -> State:
"""
Checks to see if a FAILED task should be retried.
Args:
- state (State): the current state of this task
- inputs (Dict[str, Result], optional): a dictionary of inputs whose keys correspond
to the task's `run()` arguments.
Returns:
- State: the state of the task after running the check
"""
if state.is_failed():
run_count = prefect.context.get("task_run_count", 1)
if run_count <= self.task.max_retries:
start_time = pendulum.now("utc") + self.task.retry_delay
msg = "Retrying Task (after attempt {n} of {m})".format(
n=run_count, m=self.task.max_retries + 1
)
retry_state = Retrying(
start_time=start_time,
cached_inputs=inputs,
message=msg,
run_count=run_count,
)
return retry_state
return state
def prepare_state_for_cloud(state: State) -> State:
"""
Prepares a Prefect State for being sent to Cloud; this ensures that any data attributes
are properly handled prior to being shipped off to a database.
Args:
- state (State): the Prefect State to prepare
Returns:
- State: a sanitized copy of the original state
"""
if state.is_cached():
state._result.store_safe_value()
if (isinstance(state, Failed) and state.cached_inputs is not None
and all(r.result_handler is not None
for k, r in state.cached_inputs.items())): # type: ignore
for res in state.cached_inputs.values():
res.store_safe_value()
elif (hasattr(state, "cached_inputs")
and state.cached_inputs is not None # type: ignore
and not state.is_failed()):
for res in state.cached_inputs.values(): # type: ignore
res.store_safe_value()
return state
def cache_result(self, state: State, inputs: Dict[str, Result]) -> State:
"""
Caches the result of a successful task, if appropriate. Alternatively,
if the task is failed, caches the inputs.
Tasks are cached if:
- task.cache_for is not None
- the task state is Successful
- the task state is not Skipped (which is a subclass of Successful)
Args:
- state (State): the current state of this task
- inputs (Dict[str, Result], optional): a dictionary of inputs whose keys correspond
to the task's `run()` arguments.
Returns:
- State: the state of the task after running the check
"""
if state.is_failed():
state.cached_inputs = inputs # type: ignore
if (state.is_successful() and not state.is_skipped()
and self.task.cache_for is not None):
expiration = pendulum.now("utc") + self.task.cache_for
cached_state = Cached(
result=state._result,
cached_inputs=inputs,
cached_result_expiration=expiration,
cached_parameters=prefect.context.get("parameters"),
message=state.message,
)
return cached_state
return state
def check_for_retry(self, state: State, inputs: Dict[str, Result]) -> State:
"""
Checks to see if a FAILED task should be retried.
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
"""
if state.is_failed():
run_count = prefect.context.get("task_run_count", 1)
loop_result = None
state_context = None
if prefect.context.get("task_loop_count") is not None:
loop_result = self.result.from_value(
value=prefect.context.get("task_loop_result")
)
# checkpoint tasks if a result is present, except for when the user has opted
# out by disabling checkpointing
if (
prefect.context.get("checkpointing") is True
and self.task.checkpoint is not False
and loop_result.value is not None
):
try:
raw_inputs = {k: r.value for k, r in inputs.items()}
formatting_kwargs = {
**prefect.context.get("parameters", {}).copy(),
**raw_inputs,
**prefect.context,
}
loop_result = self.result.write(
loop_result.value, **formatting_kwargs
)
except NotImplementedError:
pass
state_context = {"_loop_count": prefect.context["task_loop_count"]}
if run_count <= self.task.max_retries:
start_time = pendulum.now("utc") + self.task.retry_delay
msg = "Retrying Task (after attempt {n} of {m})".format(
n=run_count, m=self.task.max_retries + 1
)
retry_state = Retrying(
start_time=start_time,
context=state_context,
message=msg,
run_count=run_count,
result=loop_result,
)
return retry_state
return state
def check_for_retry(self, state: State, inputs: Dict[str, Result]) -> State:
"""
Checks to see if a FAILED task should be retried.
Args:
- state (State): the current state of this task
- inputs (Dict[str, Result], optional): a dictionary of inputs whose keys correspond
to the task's `run()` arguments.
Returns:
- State: the state of the task after running the check
"""
if state.is_failed():
run_count = prefect.context.get("task_run_count", 1)
if prefect.context.get("task_loop_count") is not None:
loop_result = self.result.from_value(
value=prefect.context.get("task_loop_result")
)
## checkpoint tasks if a result is present, except for when the user has opted out by disabling checkpointing
if (
prefect.context.get("checkpointing") is True
and self.task.checkpoint is not False
and loop_result.value is not None
):
try:
value = prefect.context.get("task_loop_result")
loop_result = self.result.write(
value, filename="output", **prefect.context
)
except NotImplementedError:
pass
loop_context = {
"_loop_count": PrefectResult(
location=json.dumps(prefect.context["task_loop_count"]),
),
"_loop_result": loop_result,
}
inputs.update(loop_context)
if run_count <= self.task.max_retries:
start_time = pendulum.now("utc") + self.task.retry_delay
msg = "Retrying Task (after attempt {n} of {m})".format(
n=run_count, m=self.task.max_retries + 1
)
retry_state = Retrying(
start_time=start_time,
cached_inputs=inputs,
message=msg,
run_count=run_count,
)
return retry_state
return state
def check_for_retry(self, state: State, inputs: Dict[str, Result]) -> State:
"""
Checks to see if a FAILED task should be retried.
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
"""
if state.is_failed():
run_count = prefect.context.get("task_run_count", 1)
loop_result = None
state_context = None
if prefect.context.get("task_loop_count") is not None:
loop_result = self.result.from_value(
value=prefect.context.get("task_loop_result")
)
# checkpoint tasks if a result is present, except for when the user has opted
# out by disabling checkpointing
if (
prefect.context.get("checkpointing") is True
and self.task.checkpoint is not False
and loop_result.value is not None
):
try:
raw_inputs = {k: r.value for k, r in inputs.items()}
formatting_kwargs = {
**prefect.context.get("parameters", {}).copy(),
**prefect.context,
**raw_inputs,
}
loop_result = self.result.write(
loop_result.value, **formatting_kwargs
)
