def test_get_flow(coefficients, unused_tcp_port):
with tmp(Path(SOURCE_DIR) / "test-data/local/prefect_test_case"):
config = parse_config("config.yml")
config.update(
{
"config_path": os.getcwd(),
ids.REALIZATIONS: 2,
ids.EXECUTOR: "local",
}
)
inputs = {}
coeffs_trans = coefficient_transmitters(
coefficients, config.get(ids.STORAGE)["storage_path"]
)
script_trans = script_transmitters(config)
for iens in range(2):
inputs[iens] = {**coeffs_trans[iens], **script_trans[iens]}
config.update(
{
"inputs": inputs,
"outputs": output_transmitters(config),
}
)
server_config = EvaluatorServerConfig(unused_tcp_port)
for permuted_steps in permutations(config["steps"]):
permuted_config = copy.deepcopy(config)
permuted_config["steps"] = permuted_steps
permuted_config["dispatch_uri"] = server_config.dispatch_uri
ensemble = PrefectEnsemble(permuted_config)
for iens in range(2):
with prefect.context(
url=server_config.url,
token=server_config.token,
cert=server_config.cert,
):
flow = ensemble.get_flow(ensemble._ee_id, [iens])
# Get the ordered tasks and retrieve their step ids.
flow_steps = [
task.get_step()
for task in flow.sorted_tasks()
if isinstance(task, UnixTask)
]
assert len(flow_steps) == 4
realization_steps = list(
ensemble.get_reals()[iens].get_steps_sorted_topologically()
)
# Testing realization steps
for step_ordering in [realization_steps, flow_steps]:
mapping = {
step._name: idx for idx, step in enumerate(step_ordering)
}
assert mapping["second_degree"] < mapping["zero_degree"]
assert mapping["zero_degree"] < mapping["add_coeffs"]
assert mapping["first_degree"] < mapping["add_coeffs"]
assert mapping["second_degree"] < mapping["add_coeffs"]
def test_run_prefect_ensemble_exception(unused_tcp_port, coefficients):
with tmp(os.path.join(SOURCE_DIR, "test-data/local/prefect_test_case")):
config = parse_config("config.yml")
config.update(
{
"config_path": os.getcwd(),
"realizations": 2,
"executor": "local",
}
)
inputs = {}
coeffs_trans = coefficient_transmitters(
coefficients, config.get(ids.STORAGE)["storage_path"]
)
script_trans = script_transmitters(config)
for iens in range(2):
inputs[iens] = {**coeffs_trans[iens], **script_trans[iens]}
config.update(
{
"inputs": inputs,
"outputs": output_transmitters(config),
}
)
service_config = EvaluatorServerConfig(unused_tcp_port)
config["config_path"] = Path(config["config_path"])
config["run_path"] = Path(config["run_path"])
config["storage"]["storage_path"] = Path(config["storage"]["storage_path"])
config["dispatch_uri"] = service_config.dispatch_uri
ensemble = PrefectEnsemble(config)
ensemble.get_flow = dummy_get_flow
evaluator = EnsembleEvaluator(ensemble, service_config, 0, ee_id="1")
with evaluator.run() as mon:
for event in mon.track():
if event.data is not None and event.data.get("status") in [
"Failed",
"Stopped",
]:
mon.signal_done()
assert evaluator._ensemble.get_status() == "Failed"