Example #1
0
    def call_runner_target_handlers(self, old_state: State,
                                    new_state: State) -> State:
        """
        A special state handler that the TaskRunner uses to call its task's state handlers.
        This method is called as part of the base Runner's `handle_state_change()` method.

        Args:
            - old_state (State): the old (previous) state
            - new_state (State): the new (current) state

        Returns:
            - State: the new state
        """
        raise_on_exception = prefect.context.get("raise_on_exception", False)

        try:
            new_state = super().call_runner_target_handlers(
                old_state=old_state, new_state=new_state)
        except Exception as exc:
            msg = "Exception raised while calling state handlers: {}".format(
                repr(exc))
            self.logger.exception(msg)
            if raise_on_exception:
                raise exc
            new_state = Failed(msg, result=exc)

        task_run_id = prefect.context.get("task_run_id")
        version = prefect.context.get("task_run_version")

        try:
            cloud_state = new_state
            state = self.client.set_task_run_state(
                task_run_id=task_run_id,
                version=version,
                state=cloud_state,
                cache_for=self.task.cache_for,
            )
        except Exception as exc:
            self.logger.exception(
                "Failed to set task state with error: {}".format(repr(exc)))
            raise ENDRUN(state=ClientFailed(state=new_state))

        if state.is_queued():
            state.state = old_state  # type: ignore
            raise ENDRUN(state=state)

        if version is not None:
            prefect.context.update(task_run_version=version +
                                   1)  # type: ignore

        return new_state
Example #2
0
    def call_runner_target_handlers(self, old_state: State,
                                    new_state: State) -> State:
        """
        A special state handler that the TaskRunner uses to call its task's state handlers.
        This method is called as part of the base Runner's `handle_state_change()` method.

        Args:
            - old_state (State): the old (previous) state
            - new_state (State): the new (current) state

        Returns:
            - State: the new state
        """
        raise_on_exception = prefect.context.get("raise_on_exception", False)

        try:
            new_state = super().call_runner_target_handlers(
                old_state=old_state, new_state=new_state)

        # PrefectStateSignals are trapped and turned into States
        except prefect.engine.signals.PrefectStateSignal as exc:
            self.logger.info("{name} signal raised: {rep}".format(
                name=type(exc).__name__, rep=repr(exc)))
            if raise_on_exception:
                raise exc
            new_state = exc.state

        except Exception as exc:
            msg = "Exception raised while calling state handlers: {}".format(
                repr(exc))
            self.logger.exception(msg)
            if raise_on_exception:
                raise exc
            new_state = Failed(msg, result=exc)

        task_run_id = prefect.context.get("task_run_id")
        version = prefect.context.get("task_run_version")

        try:
            cloud_state = new_state
            state = self.client.set_task_run_state(
                task_run_id=task_run_id,
                version=version if cloud_state.is_running() else None,
                state=cloud_state,
                cache_for=self.task.cache_for,
            )
        except VersionLockError as exc:
            state = self.client.get_task_run_state(task_run_id=task_run_id)

            if state.is_running():
                self.logger.debug(
                    "Version lock encountered and task {} is already in a running state."
                    .format(self.task.name))
                raise ENDRUN(state=state) from exc

            self.logger.debug(
                "Version lock encountered for task {}, proceeding with state {}..."
                .format(self.task.name,
                        type(state).__name__))

            try:
                new_state = state.load_result(self.result)
            except Exception as exc_inner:
                self.logger.debug(
                    "Error encountered attempting to load result for state of {} task..."
                    .format(self.task.name))
                self.logger.error(repr(exc_inner))
                raise ENDRUN(state=state) from exc_inner
        except Exception as exc:
            self.logger.exception(
                "Failed to set task state with error: {}".format(repr(exc)))
            raise ENDRUN(state=ClientFailed(state=new_state)) from exc

        if state.is_queued():
            state.state = old_state  # type: ignore
            raise ENDRUN(state=state)

        prefect.context.update(task_run_version=(version or 0) + 1)

        return new_state
Example #3
0
        assert issubclass(Skipped, Success)

    def test_timedout_is_failed(self):
        assert issubclass(TimedOut, Failed)

    def test_trigger_failed_is_failed(self):
        assert issubclass(TriggerFailed, Failed)


@pytest.mark.parametrize(
    "state_check",
    [
        dict(state=Cancelled(), assert_true={"is_finished"}),
        dict(state=Cached(),
             assert_true={"is_cached", "is_finished", "is_successful"}),
        dict(state=ClientFailed(), assert_true={"is_meta_state"}),
        dict(state=Failed(), assert_true={"is_finished", "is_failed"}),
        dict(state=Finished(), assert_true={"is_finished"}),
        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"}),