def test_starting_at_arbitrary_loop_index_from_cloud_context(client):
@prefect.task
def looper(x):
if prefect.context.get("task_loop_count", 1) < 20:
raise LOOP(result=prefect.context.get("task_loop_result", 0) + x)
return prefect.context.get("task_loop_result", 0) + x
@prefect.task
def downstream(l):
return l**2
with prefect.Flow(name="looping", result_handler=JSONResultHandler()) as f:
inter = looper(10)
final = downstream(inter)
client.get_flow_run_info = MagicMock(return_value=MagicMock(
context={"task_loop_count": 20}))
client.set_flow_run_state = MagicMock()
flow_state = CloudFlowRunner(flow=f).run(return_tasks=[inter, final])
assert flow_state.is_successful()
assert flow_state.result[inter].result == 10
assert flow_state.result[final].result == 100
def test_simple_map(monkeypatch):
flow_run_id = str(uuid.uuid4())
task_run_id_1 = str(uuid.uuid4())
with prefect.Flow(name="test", result_handler=JSONResultHandler()) as flow:
t1 = plus_one.map([0, 1, 2])
client = MockedCloudClient(
flow_runs=[FlowRun(id=flow_run_id)],
task_runs=[
TaskRun(id=task_run_id_1, task_slug=flow.slugs[t1], flow_run_id=flow_run_id)
]
+ [
TaskRun(
id=str(uuid.uuid4()), task_slug=flow.slugs[t], flow_run_id=flow_run_id
)
for t in flow.tasks
if t is not t1
],
monkeypatch=monkeypatch,
)
with prefect.context(flow_run_id=flow_run_id):
state = CloudFlowRunner(flow=flow).run(
return_tasks=flow.tasks, executor=LocalExecutor()
)
assert state.is_successful()
assert client.flow_runs[flow_run_id].state.is_successful()
assert client.task_runs[task_run_id_1].state.is_mapped()
# there should be a total of 4 task runs corresponding to the mapped task
assert (
len([tr for tr in client.task_runs.values() if tr.task_slug == flow.slugs[t1]])
== 4
)
def test_roundtrip(self):
schema = ResultHandlerSchema()
handler = schema.load(schema.dump(JSONResultHandler()))
assert isinstance(handler, JSONResultHandler)
assert handler.write(3) == "3"
def test_create_task_with_and_without_result_handler(self):
t1 = Task()
assert t1.result_handler is None
t2 = Task(result_handler=JSONResultHandler())
assert isinstance(t2.result_handler, ResultHandler)
assert isinstance(t2.result_handler, JSONResultHandler)
def test_basic_safe_result_repr():
r = SafeResult(2, result_handler=JSONResultHandler())
assert repr(r) == "<SafeResult: 2>"
def test_result_inits_with_handled_and_result_handler(self):
handler = JSONResultHandler()
r = Result(value=3, result_handler=handler)
assert r.value == 3
assert r.safe_value is NoResult
assert r.result_handler == handler
def test_store_safe_value_for_safe_results(self):
r = SafeResult(value=4, result_handler=JSONResultHandler())
output = r.store_safe_value()
assert output is None
assert isinstance(r.safe_value, SafeResult)
assert r.value == 4
def test_safe_results_with_different_handlers_are_not_same(self):
r = SafeResult("3", result_handler=JSONResultHandler())
s = SafeResult("3", result_handler=LocalResultHandler())
assert r != s
def test_basic_safe_result_serializes():
res = SafeResult("3", result_handler=JSONResultHandler())
serialized = SafeResultSchema().dump(res)
assert serialized["__version__"] == prefect.__version__
assert serialized["value"] == "3"
assert serialized["result_handler"]["type"] == "JSONResultHandler"
def test_result_allows_none_value():
schema = SafeResultSchema()
r = SafeResult(value=None, result_handler=JSONResultHandler())
handled = schema.load(schema.dump(r))
assert handled.value is None
assert handled.result_handler == JSONResultHandler()
def test_safe_result_disallows_none_result_handler_at_deserialization_time():
schema = SafeResultSchema()
r = SafeResult(value=None, result_handler=None)
serialized = schema.dump(r)
with pytest.raises(marshmallow.exceptions.ValidationError):
schema.load(serialized)
@pytest.mark.parametrize(
"obj", [SafeResult(value=19, result_handler=JSONResultHandler()), NoResult]
)
def test_state_result_schema_chooses_schema(obj):
schema = StateResultSchema()
assert type(schema.load(schema.dump(obj))) == type(obj)
def test_value_raises_error_on_dump_if_not_valid_json():
r = SafeResult(
value={"x": {"y": {"z": lambda: 1}}}, result_handler=JSONResultHandler()
)
with pytest.raises(marshmallow.exceptions.ValidationError):
StateResultSchema().dump(r)
def test_json_handler_is_pickleable(self):
handler = JSONResultHandler()
new = cloudpickle.loads(cloudpickle.dumps(handler))
assert isinstance(new, JSONResultHandler)
def test_json_handler_raises_normally(self):
handler = JSONResultHandler()
with pytest.raises(TypeError):
handler.write(type(None))
def test_json_handler_writes_and_reads(self, res):
handler = JSONResultHandler()
final = handler.read(handler.write(res))
assert final == res
def test_json_handler_writes(self, res):
handler = JSONResultHandler()
blob = handler.write(res)
assert isinstance(blob, str)
def test_json_handler_initializes_with_no_args(self):
handler = JSONResultHandler()
def test_results_are_same_if_handled(self):
r = Result("3", result_handler=JSONResultHandler())
s = Result("3", result_handler=JSONResultHandler())
r.store_safe_value()
s.store_safe_value()
assert s == r
def test_safe_results_are_same(self):
r = SafeResult("3", result_handler=JSONResultHandler())
s = SafeResult("3", result_handler=JSONResultHandler())
assert r == s
def test_roundtrip_for_results():
r = SafeResult(value={"y": 4}, result_handler=JSONResultHandler())
s = SafeResultSchema().load(SafeResultSchema().dump(r))
assert r == s
def test_safe_results_to_results_remain_the_same(self):
r = SafeResult("3", result_handler=JSONResultHandler())
s = SafeResult("3", result_handler=JSONResultHandler())
assert r.to_result() == s.to_result()
def test_safe_result_serializes_result_handlers():
r = SafeResult(value=3, result_handler=JSONResultHandler())
handled = SafeResultSchema().dump(r)
assert handled["result_handler"]["type"] == "JSONResultHandler"
def test_results_are_pickleable_with_their_safe_values():
res = Result(3, result_handler=JSONResultHandler())
res.store_safe_value()
assert cloudpickle.loads(cloudpickle.dumps(res)) == res
def test_result_allows_none_value():
schema = SafeResultSchema()
r = SafeResult(value=None, result_handler=JSONResultHandler())
handled = schema.load(schema.dump(r))
assert handled.value is None
assert handled.result_handler == JSONResultHandler()
def test_safe_result_inits_with_both_args(self):
res = SafeResult(value="3", result_handler=JSONResultHandler())
assert res.value == "3"
assert res.result_handler == JSONResultHandler()
assert res.safe_value is res
def test_value_raises_error_on_dump_if_not_valid_json():
r = SafeResult(
value={"x": {"y": {"z": lambda: 1}}}, result_handler=JSONResultHandler()
)
with pytest.raises(marshmallow.exceptions.ValidationError):
StateResultSchema().dump(r)
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 == NoResult
assert state.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)
def test_serialize(self):
serialized = ResultHandlerSchema().dump(JSONResultHandler())
assert isinstance(serialized, dict)
def test_to_result_returns_self_for_no_results(self):
assert NoResult.to_result() is NoResult
def test_to_result_returns_hydrated_result_for_safe(self):
s = SafeResult("3", result_handler=JSONResultHandler())
res = s.to_result()
assert isinstance(res, Result)
assert res.value == 3
assert res.safe_value is s
assert res.result_handler is s.result_handler
@pytest.mark.parametrize(
"obj",
[
Result(3),
Result(object, result_handler=LocalResultHandler()),
NoResult,
SafeResult("3", result_handler=JSONResultHandler()),
],
)
def test_everything_is_pickleable_after_init(obj):
assert cloudpickle.loads(cloudpickle.dumps(obj)) == obj
def test_results_are_pickleable_with_their_safe_values():
res = Result(3, result_handler=JSONResultHandler())
res.store_safe_value()
assert cloudpickle.loads(cloudpickle.dumps(res)) == res
def test_cloud_task_runners_submitted_to_remote_machines_respect_original_config(
monkeypatch, ):
"""
This test is meant to simulate the behavior of running a Cloud Flow against an external
cluster which has _not_ been configured for Prefect. The idea is that the configuration
settings which were present on the original machine are respected in the remote job, reflected
here by having the CloudHandler called during logging and the special values present in context.
"""
class CustomFlowRunner(CloudFlowRunner):
def run_task(self, *args, **kwargs):
with prefect.utilities.configuration.set_temporary_config({
"logging.log_to_cloud":
False,
"cloud.auth_token":
""
}):
return super().run_task(*args, **kwargs)
@prefect.task(result_handler=JSONResultHandler())
def log_stuff():
logger = prefect.context.get("logger")
logger.critical("important log right here")
return (
prefect.context.config.special_key,
prefect.context.config.cloud.auth_token,
)
calls = []
class Client(MagicMock):
def write_run_logs(self, *args, **kwargs):
calls.append(args)
def set_task_run_state(self, *args, **kwargs):
return kwargs.get("state")
def get_flow_run_info(self, *args, **kwargs):
return MagicMock(
id="flow_run_id",
task_runs=[MagicMock(task_slug=log_stuff.slug, id="TESTME")],
)
monkeypatch.setattr("prefect.client.Client", Client)
monkeypatch.setattr("prefect.engine.cloud.task_runner.Client", Client)
monkeypatch.setattr("prefect.engine.cloud.flow_runner.Client", Client)
prefect.utilities.logging.prefect_logger.handlers[-1].client = Client()
with prefect.utilities.configuration.set_temporary_config({
"logging.log_to_cloud":
True,
"special_key":
42,
"cloud.auth_token":
"original",
}):
# captures config at init
runner = CustomFlowRunner(flow=prefect.Flow("test", tasks=[log_stuff]))
flow_state = runner.run(
return_tasks=[log_stuff],
task_contexts={log_stuff: dict(special_key=99)})
assert flow_state.is_successful()
assert flow_state.result[log_stuff].result == (42, "original")
time.sleep(0.75)
logs = [log for call in calls for log in call[0]]
assert len(logs) >= 6 # actual number of logs
loggers = [c["name"] for c in logs]
assert set(loggers) == {
"prefect.CloudTaskRunner",
"prefect.CustomFlowRunner",
"prefect.log_stuff",
}
task_run_ids = [c["task_run_id"] for c in logs if c["task_run_id"]]
assert set(task_run_ids) == {"TESTME"}
def test_results_are_different_if_handled(self):
r = Result("3", result_handler=JSONResultHandler())
s = Result("3", result_handler=JSONResultHandler())
s.store_safe_value()
assert s != r
