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 )