def get_task_run_state(
self,
state: State,
inputs: Dict[str, Result],
timeout_handler: Optional[Callable],
) -> State:
"""
Runs the task and traps any signals or errors it raises.
Also checkpoints the result of a successful task, if `task.checkpoint` is `True`.
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.
- timeout_handler (Callable, optional): function for timing out
task execution, with call signature `handler(fn, *args, **kwargs)`. Defaults to
`prefect.utilities.executors.main_thread_timeout`
Returns:
- State: the state of the task after running the check
Raises:
- signals.PAUSE: if the task raises PAUSE
- ENDRUN: if the task is not ready to run
"""
if not state.is_running():
self.logger.debug(
"Task '{name}': can't run task because it's not in a "
"Running state; ending run.".format(name=prefect.context.get(
"task_full_name", self.task.name)))
raise ENDRUN(state)
try:
self.logger.debug(
"Task '{name}': Calling task.run() method...".format(
name=prefect.context.get("task_full_name",
self.task.name)))
timeout_handler = timeout_handler or main_thread_timeout
raw_inputs = {k: r.value for k, r in inputs.items()}
result = timeout_handler(self.task.run,
timeout=self.task.timeout,
**raw_inputs)
# inform user of timeout
except TimeoutError as exc:
if prefect.context.get("raise_on_exception"):
raise exc
state = TimedOut("Task timed out during execution.",
result=exc,
cached_inputs=inputs)
return state
result = Result(value=result, result_handler=self.result_handler)
state = Success(result=result, message="Task run succeeded.")
if state.is_successful() and self.task.checkpoint is True:
state._result.store_safe_value()
return state
dict(state=Looped(), assert_true={"is_finished", "is_looped"}),
dict(state=Mapped(),
assert_true={"is_finished", "is_mapped", "is_successful"}),
dict(state=Paused(), assert_true={"is_pending", "is_scheduled"}),
dict(state=Pending(), assert_true={"is_pending"}),
dict(state=Queued(), assert_true={"is_meta_state", "is_queued"}),
dict(state=Resume(), assert_true={"is_pending", "is_scheduled"}),
dict(state=Retrying(),
assert_true={"is_pending", "is_scheduled", "is_retrying"}),
dict(state=Running(), assert_true={"is_running"}),
dict(state=Scheduled(), assert_true={"is_pending", "is_scheduled"}),
dict(state=Skipped(),
assert_true={"is_finished", "is_successful", "is_skipped"}),
dict(state=Submitted(), assert_true={"is_meta_state", "is_submitted"}),
dict(state=Success(), assert_true={"is_finished", "is_successful"}),
dict(state=TimedOut(), assert_true={"is_finished", "is_failed"}),
dict(state=TriggerFailed(), assert_true={"is_finished", "is_failed"}),
],
)
def test_state_is_methods(state_check):
"""
Iterates over all of the "is_*()" methods of the state, asserting that each one is
False, unless the name of that method is provided as `assert_true`.
For example, if `state_check == (Pending(), {'is_pending'})`, then this method will
assert that `state.is_running()` is False, `state.is_successful()` is False, etc. but
`state.is_pending()` is True.
"""
state = state_check["state"]
assert_true = state_check["assert_true"]
def get_task_run_state(self, state: State, inputs: Dict[str,
Result]) -> State:
"""
Runs the task and traps any signals or errors it raises.
Also checkpoints the result of a successful task, if `task.checkpoint` is `True`.
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
Raises:
- signals.PAUSE: if the task raises PAUSE
- ENDRUN: if the task is not ready to run
"""
if not state.is_running():
self.logger.debug(
"Task '{name}': can't run task because it's not in a "
"Running state; ending run.".format(name=prefect.context.get(
"task_full_name", self.task.name)))
raise ENDRUN(state)
value = None
raw_inputs = {k: r.value for k, r in inputs.items()}
try:
self.logger.debug(
"Task '{name}': Calling task.run() method...".format(
name=prefect.context.get("task_full_name",
self.task.name)))
timeout_handler = prefect.utilities.executors.timeout_handler
if getattr(self.task, "log_stdout", False):
with redirect_stdout(
prefect.utilities.logging.RedirectToLog(
self.logger) # type: ignore
):
value = timeout_handler(self.task.run,
timeout=self.task.timeout,
**raw_inputs)
else:
value = timeout_handler(self.task.run,
timeout=self.task.timeout,
**raw_inputs)
# inform user of timeout
except TimeoutError as exc:
if prefect.context.get("raise_on_exception"):
raise exc
state = TimedOut("Task timed out during execution.", result=exc)
return state
except signals.LOOP as exc:
new_state = exc.state
assert isinstance(new_state, Looped)
new_state.result = self.result.from_value(value=new_state.result)
new_state.message = exc.state.message or "Task is looping ({})".format(
new_state.loop_count)
return new_state
# 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 value is not None):
try:
formatting_kwargs = {
**prefect.context.get("parameters", {}).copy(),
**raw_inputs,
**prefect.context,
}
result = self.result.write(value, **formatting_kwargs)
except NotImplementedError:
result = self.result.from_value(value=value)
else:
result = self.result.from_value(value=value)
state = Success(result=result, message="Task run succeeded.")
return state
def get_task_run_state(self, state: State, inputs: Dict[str,
Result]) -> State:
"""
Runs the task and traps any signals or errors it raises.
Also checkpoints the result of a successful task, if `task.checkpoint` is `True`.
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
Raises:
- signals.PAUSE: if the task raises PAUSE
- ENDRUN: if the task is not ready to run
"""
task_name = prefect.context.get("task_full_name", self.task.name)
if not state.is_running():
self.logger.debug(
f"Task {task_name!r}: Can't run task because it's not in a Running "
"state; ending run.")
raise ENDRUN(state)
value = None
raw_inputs = {k: r.value for k, r in inputs.items()}
new_state = None
try:
self.logger.debug(
f"Task {task_name!r}: Calling task.run() method...")
# Create a stdout redirect if the task has log_stdout enabled
log_context = (
redirect_stdout(
prefect.utilities.logging.RedirectToLog(self.logger))
if getattr(self.task, "log_stdout", False) else nullcontext()
) # type: AbstractContextManager
with log_context:
value = prefect.utilities.executors.run_task_with_timeout(
task=self.task,
args=(),
kwargs=raw_inputs,
logger=self.logger,
)
except TaskTimeoutSignal as exc: # Convert timeouts to a `TimedOut` state
if prefect.context.get("raise_on_exception"):
raise exc
state = TimedOut("Task timed out during execution.", result=exc)
return state
except signals.LOOP as exc: # Convert loop signals to a `Looped` state
new_state = exc.state
assert isinstance(new_state, Looped)
value = new_state.result
new_state.message = exc.state.message or "Task is looping ({})".format(
new_state.loop_count)
except signals.SUCCESS as exc:
# Success signals can be treated like a normal result
new_state = exc.state
assert isinstance(new_state, Success)
value = new_state.result
except Exception as exc: # Handle exceptions in the task
if prefect.context.get("raise_on_exception"):
raise
self.logger.error(
f"Task {task_name!r}: Exception encountered during task execution!",
exc_info=True,
)
state = Failed(f"Error during execution of task: {exc!r}",
result=exc)
return state
# 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 value is not None):
try:
formatting_kwargs = {
**prefect.context.get("parameters", {}).copy(),
**prefect.context,
**raw_inputs,
}
result = self.result.write(value, **formatting_kwargs)
except ResultNotImplementedError:
result = self.result.from_value(value=value)
else:
result = self.result.from_value(value=value)
if new_state is not None:
new_state.result = result
return new_state
state = Success(result=result, message="Task run succeeded.")
return state
def get_task_run_state(
self,
state: State,
inputs: Dict[str, Result],
timeout_handler: Optional[Callable] = None,
) -> State:
"""
Runs the task and traps any signals or errors it raises.
Also checkpoints the result of a successful task, if `task.checkpoint` is `True`.
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.
- timeout_handler (Callable, optional): function for timing out
task execution, with call signature `handler(fn, *args, **kwargs)`. Defaults to
`prefect.utilities.executors.timeout_handler`
Returns:
- State: the state of the task after running the check
Raises:
- signals.PAUSE: if the task raises PAUSE
- ENDRUN: if the task is not ready to run
"""
if not state.is_running():
self.logger.debug(
"Task '{name}': can't run task because it's not in a "
"Running state; ending run.".format(name=prefect.context.get(
"task_full_name", self.task.name)))
raise ENDRUN(state)
try:
self.logger.debug(
"Task '{name}': Calling task.run() method...".format(
name=prefect.context.get("task_full_name",
self.task.name)))
timeout_handler = (timeout_handler
or prefect.utilities.executors.timeout_handler)
raw_inputs = {k: r.value for k, r in inputs.items()}
if getattr(self.task, "log_stdout", False):
with redirect_stdout(
prefect.utilities.logging.RedirectToLog(
self.logger)): # type: ignore
result = timeout_handler(self.task.run,
timeout=self.task.timeout,
**raw_inputs)
else:
result = timeout_handler(self.task.run,
timeout=self.task.timeout,
**raw_inputs)
except KeyboardInterrupt:
self.logger.debug("Interrupt signal raised, cancelling task run.")
state = Cancelled(
message="Interrupt signal raised, cancelling task run.")
return state
# inform user of timeout
except TimeoutError as exc:
if prefect.context.get("raise_on_exception"):
raise exc
state = TimedOut("Task timed out during execution.",
result=exc,
cached_inputs=inputs)
return state
except signals.LOOP as exc:
new_state = exc.state
assert isinstance(new_state, Looped)
new_state.result = Result(value=new_state.result,
result_handler=self.result_handler)
new_state.cached_inputs = inputs
new_state.message = exc.state.message or "Task is looping ({})".format(
new_state.loop_count)
return new_state
result = Result(value=result, result_handler=self.result_handler)
state = Success(result=result,
message="Task run succeeded.",
cached_inputs=inputs)
## checkpoint tasks if a result_handler is present, except for when the user has opted out by disabling checkpointing
if (state.is_successful()
and prefect.context.get("checkpointing") is True
and self.task.checkpoint is not False
and self.result_handler is not None):
state._result.store_safe_value()
return state
def test_state_type_methods_with_timedout_state(self):
state = TimedOut(message="")
assert not state.is_cached()
assert not state.is_pending()
assert not state.is_retrying()
assert not state.is_running()
assert state.is_finished()
assert not state.is_skipped()
assert not state.is_scheduled()
assert not state.is_successful()
assert state.is_failed()
assert not state.is_mapped()
assert not state.is_meta_state()