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_prefect_no_retries(unused_tcp_port, coefficients, tmpdir, function_config):
def function_that_fails_once(coeffs):
run_path = Path("ran_once")
if not run_path.exists():
run_path.touch()
raise RuntimeError("This is an expected ERROR")
run_path.unlink()
return []
with tmpdir.as_cwd():
pickle_func = cloudpickle.dumps(function_that_fails_once)
config = function_config
coeffs_trans = coefficient_transmitters(
coefficients, config.get(ids.STORAGE)["storage_path"]
)
service_config = EvaluatorServerConfig(unused_tcp_port)
config["realizations"] = len(coefficients)
config["executor"] = "local"
config["max_retries"] = 0
config["retry_delay"] = 1
config["steps"][0]["jobs"][0]["executable"] = pickle_func
config["inputs"] = {
iens: coeffs_trans[iens] for iens in range(len(coefficients))
}
config["outputs"] = output_transmitters(config)
config["dispatch_uri"] = service_config.dispatch_uri
ensemble = PrefectEnsemble(config)
evaluator = EnsembleEvaluator(ensemble, service_config, 0, ee_id="1")
event_list = []
with evaluator.run() as mon:
for event in mon.track():
event_list.append(event)
if event.data is not None and event.data.get("status") in [
state.ENSEMBLE_STATE_FAILED,
state.ENSEMBLE_STATE_STOPPED,
]:
mon.signal_done()
step_failed = False
job_failed = False
for real in ensemble.snapshot.get_reals().values():
for step in real.steps.values():
for job in step.jobs.values():
if job.status == state.JOB_STATE_FAILURE:
job_failed = True
assert job.error == "This is an expected ERROR"
if step.status == state.STEP_STATE_FAILURE:
step_failed = True
assert ensemble.get_status() == state.ENSEMBLE_STATE_FAILED
assert job_failed, f"Events: {event_list}"
assert step_failed, f"Events: {event_list}"
def evaluate(
workspace_root: Path,
evaluation_name: str,
input_records: MultiEnsembleRecord,
ensemble_config: EnsembleConfig,
stages_config: StagesConfig,
) -> MultiEnsembleRecord:
evaluation_tmp_dir = _create_evaluator_tmp_dir(workspace_root,
evaluation_name)
config = EvaluatorServerConfig()
ee_config = _build_ee_config(
evaluation_tmp_dir,
ensemble_config,
stages_config,
input_records,
config.dispatch_uri,
)
ensemble = PrefectEnsemble(ee_config) # type: ignore
ee = EnsembleEvaluator(ensemble=ensemble, config=config, iter_=0)
result = _run(ee)
output_records = _prepare_output_records(result)
return output_records
def test_prefect_no_retries(unused_tcp_port, coefficients, tmpdir,
function_config):
def function_that_fails_once(coeffs):
run_path = Path("ran_once")
if not run_path.exists():
run_path.touch()
raise RuntimeError("This is an expected ERROR")
run_path.unlink()
return []
with tmpdir.as_cwd():
pickle_func = cloudpickle.dumps(function_that_fails_once)
config = function_config
coeffs_trans = coefficient_transmitters(
coefficients,
config.get(ids.STORAGE)["storage_path"])
service_config = EvaluatorServerConfig(unused_tcp_port)
config["realizations"] = len(coefficients)
config["executor"] = "local"
config["max_retries"] = 0
config["retry_delay"] = 1
config["steps"][0]["jobs"][0]["executable"] = pickle_func
config["inputs"] = {
iens: coeffs_trans[iens]
for iens in range(len(coefficients))
}
config["outputs"] = output_transmitters(config)
config["dispatch_uri"] = service_config.dispatch_uri
ensemble = PrefectEnsemble(config)
evaluator = EnsembleEvaluator(ensemble, service_config, 0, ee_id="1")
step_failed = False
job_failed = False
with evaluator.run() as mon:
for event in mon.track():
# Capture the job error messages
if event.data is not None and "This is an expected ERROR" in str(
event.data):
for real in event.data["reals"].values():
for step in real["steps"].values():
for job in step["jobs"].values():
if job["status"] == "Failed":
job_failed = True
if step["status"] == "Failed":
step_failed = True
if event.data is not None and event.data.get("status") in [
"Failed",
"Stopped",
]:
mon.signal_done()
assert evaluator._ensemble.get_status() == "Failed"
assert job_failed
assert step_failed
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"
def test_prefect_reties(unused_tcp_port, coefficients, tmpdir,
function_config):
def function_that_fails_once(coeffs):
run_path = Path("ran_once")
if not run_path.exists():
run_path.touch()
raise RuntimeError("This is an expected ERROR")
run_path.unlink()
return []
with tmpdir.as_cwd():
pickle_func = cloudpickle.dumps(function_that_fails_once)
config = function_config
coeffs_trans = coefficient_transmitters(
coefficients,
config.get(ids.STORAGE)["storage_path"])
service_config = EvaluatorServerConfig(unused_tcp_port)
config["realizations"] = len(coefficients)
config["executor"] = "local"
config["max_retries"] = 2
config["retry_delay"] = 1
config["steps"][0]["jobs"][0]["executable"] = pickle_func
config["inputs"] = {
iens: coeffs_trans[iens]
for iens in range(len(coefficients))
}
config["outputs"] = output_transmitters(config)
config["dispatch_uri"] = service_config.dispatch_uri
ensemble = PrefectEnsemble(config)
evaluator = EnsembleEvaluator(ensemble, service_config, 0, ee_id="1")
error_event_reals = []
with evaluator.run() as mon:
for event in mon.track():
# Caputure the job error messages
if event.data is not None and "This is an expected ERROR" in str(
event.data):
error_event_reals.append(event.data["reals"])
if event.data is not None and event.data.get("status") in [
"Failed",
"Stopped",
]:
mon.signal_done()
assert evaluator._snapshot.get_status() == "Stopped"
successful_realizations = evaluator._snapshot.get_successful_realizations(
)
assert successful_realizations == config["realizations"]
# Check we get only one job error message per realization
assert len(error_event_reals) == config["realizations"]
for idx, reals in enumerate(error_event_reals):
assert len(reals) == 1
assert str(idx) in reals
def test_cancel_run_prefect_ensemble(unused_tcp_port, coefficients):
with tmp(Path(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)
evaluator = EnsembleEvaluator(ensemble, service_config, 0, ee_id="2")
with evaluator.run() as mon:
cancel = True
for _ in mon.track():
if cancel:
mon.signal_cancel()
cancel = False
assert evaluator._ensemble.get_status() == "Cancelled"
def evaluate(workspace_root, evaluation_name, input_records, ensemble_config,
stages_config):
evaluation_tmp_dir = _create_evaluator_tmp_dir(workspace_root,
evaluation_name)
config = EvaluatorServerConfig()
ee_config = _build_ee_config(
evaluation_tmp_dir,
ensemble_config,
stages_config,
input_records,
config.dispatch_uri,
)
ensemble = PrefectEnsemble(ee_config)
ee = EnsembleEvaluator(ensemble=ensemble, config=config, iter_=0)
result = _run(ee)
responses = _prepare_responses(result)
return responses
def test_run_prefect_ensemble(unused_tcp_port, coefficients):
test_path = Path(SOURCE_DIR) / "test-data/local/prefect_test_case"
with tmp(test_path):
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["dispatch_uri"] = service_config.dispatch_uri
ensemble = PrefectEnsemble(config)
evaluator = EnsembleEvaluator(ensemble, service_config, 0, ee_id="1")
with evaluator.run() as mon:
for event in mon.track():
if isinstance(event.data, dict) and event.data.get("status") in [
"Failed",
"Stopped",
]:
mon.signal_done()
assert evaluator._ensemble.get_status() == "Stopped"
successful_realizations = evaluator._ensemble.get_successful_realizations()
assert successful_realizations == config["realizations"]
def test_run_prefect_for_function_defined_outside_py_environment(
unused_tcp_port, coefficients, tmpdir, function_config):
with tmpdir.as_cwd():
# Create temporary module that defines a function `bar`
# 'bar' returns a call to different function 'internal_call' defined in the same python file
module_path = Path(tmpdir) / "foo"
module_path.mkdir()
init_file = module_path / "__init__.py"
init_file.touch()
file_path = module_path / "bar.py"
file_path.write_text(
"def bar(coeffs):\n return internal_call(coeffs)\n"
"def internal_call(coeffs):\n return [sum(coeffs.values())]\n")
spec = importlib.util.spec_from_file_location("foo", str(file_path))
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
func = getattr(module, "bar")
pickle_func = cloudpickle.dumps(func)
init_file.unlink()
file_path.unlink()
# Check module is not in the python environment
with pytest.raises(ModuleNotFoundError):
import foo.bar
config = function_config
coeffs_trans = coefficient_transmitters(
coefficients,
config.get(ids.STORAGE)["storage_path"])
config["realizations"] = len(coefficients)
config["executor"] = "local"
config["steps"][0]["jobs"][0]["executable"] = pickle_func
config["inputs"] = {iens: coeffs_trans[iens] for iens in range(2)}
config["outputs"] = output_transmitters(config)
service_config = EvaluatorServerConfig(unused_tcp_port)
config["dispatch_uri"] = service_config.dispatch_uri
ensemble = PrefectEnsemble(config)
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",
]:
results = mon.get_result()
mon.signal_done()
assert evaluator._snapshot.get_status() == "Stopped"
successful_realizations = evaluator._snapshot.get_successful_realizations(
)
assert successful_realizations == config["realizations"]
expected_results = [
pickle.loads(pickle_func)(coeffs) for coeffs in coefficients
]
transmitter_futures = [
res["function_output"].load() for res in results.values()
]
results = asyncio.get_event_loop().run_until_complete(
asyncio.gather(*transmitter_futures))
assert expected_results == [res.data for res in results]