def test_preparing_state_for_cloud_replaces_cached_inputs_with_safe(cls):
xres = Result(3, result_handler=JSONResultHandler())
state = prepare_state_for_cloud(cls(cached_inputs=dict(x=xres)))
assert isinstance(state, cls)
assert state.result is None
assert state._result == NoResult
assert state.cached_inputs == dict(x=xres)
def test_preparing_state_for_cloud_passes_if_cached_inputs_have_no_handler_for_failed(
cls, ):
xres = Result(3, result_handler=None)
state = prepare_state_for_cloud(cls(cached_inputs=dict(x=xres)))
assert isinstance(state, cls)
assert state.result == NoResult
assert state.cached_inputs == dict(x=xres)
def test_preparing_state_for_cloud_does_nothing_if_result_is_none(cls):
xres = Result(None, result_handler=JSONResultHandler())
state = prepare_state_for_cloud(cls(cached_inputs=dict(x=xres)))
assert isinstance(state, cls)
assert state.result is None
assert state._result == NoResult
assert state.cached_inputs == dict(x=xres)
assert state.serialize()["cached_inputs"]["x"]["type"] == "NoResultType"
def test_preparing_state_for_cloud_ignores_the_lack_of_result_handlers_for_cached_inputs(
cls, ):
xres = Result(3, result_handler=None)
state = prepare_state_for_cloud(cls(cached_inputs=dict(x=xres)))
assert isinstance(state, cls)
assert state.result is None
assert state._result == NoResult
assert state.cached_inputs == dict(x=xres)
assert state.serialize()["cached_inputs"]["x"]["type"] == "NoResultType"
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 = prepare_state_for_cloud(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
def test_preparing_state_for_cloud_doesnt_copy_data():
class FakeHandler(ResultHandler):
def read(self, val):
return val
def write(self, val):
return val
value = 124.090909
result = Result(value, result_handler=FakeHandler())
state = Cached(result=result)
cloud_state = prepare_state_for_cloud(state)
assert cloud_state.is_cached()
assert cloud_state.result is state.result
def call_runner_target_handlers(self, old_state: State,
new_state: State) -> State:
"""
A special state handler that the FlowRunner uses to call its flow'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.debug(msg)
if raise_on_exception:
raise exc
new_state = Failed(msg, result=exc)
flow_run_id = prefect.context.get("flow_run_id", None)
version = prefect.context.get("flow_run_version")
try:
cloud_state = prepare_state_for_cloud(new_state)
self.client.set_flow_run_state(flow_run_id=flow_run_id,
version=version,
state=cloud_state)
except Exception as exc:
self.logger.debug("Failed to set flow state with error: {}".format(
repr(exc)))
raise ENDRUN(state=new_state)
prefect.context.update(flow_run_version=version + 1)
return new_state
def test_preparing_state_for_cloud_passes_if_cached_inputs_dont_exist(cls):
state = prepare_state_for_cloud(cls())
assert isinstance(state, cls)
assert state.result == NoResult
assert state.cached_inputs is None
def test_preparing_state_for_cloud_fails_if_cached_inputs_have_no_handler(cls):
xres = Result(3, result_handler=None)
with pytest.raises(AssertionError, match="no ResultHandler"):
state = prepare_state_for_cloud(cls(cached_inputs=dict(x=xres)))
def test_preparing_state_for_cloud_replaces_cached_inputs_with_safe():
xres = Result(3, result_handler=JSONResultHandler())
state = prepare_state_for_cloud(Pending(cached_inputs=dict(x=xres)))
assert state.is_pending()
assert state.result == NoResult
assert state.cached_inputs == dict(x=xres)
