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()