def test_unix_step(unused_tcp_port):
host = "localhost"
url = f"ws://{host}:{unused_tcp_port}"
messages = []
mock_ws_thread = threading.Thread(target=partial(_mock_ws,
messages=messages),
args=(host, unused_tcp_port))
mock_ws_thread.start()
def _on_task_failure(task, state):
raise Exception(state.message)
with tmp(Path(SOURCE_DIR) / "test-data/local/prefect_test_case"):
config = parse_config("config.yml")
storage = storage_driver_factory(config=config.get("storage"),
run_path=".")
resource = storage.store("unix_test_script.py")
jobs = [{
"id": "0",
"name": "test_script",
"executable": "unix_test_script.py",
"args": ["vas"],
}]
stage_task = UnixStep(
resources=[resource],
outputs=["output.out"],
job_list=jobs,
iens=1,
cmd="python3",
url=url,
step_id="step_id_0",
stage_id="stage_id_0",
ee_id="ee_id_0",
on_failure=_on_task_failure,
run_path=config.get("run_path"),
storage_config=config.get("storage"),
max_retries=1,
retry_delay=timedelta(seconds=2),
)
flow = Flow("testing")
flow.add_task(stage_task)
flow_run = flow.run()
# Stop the mock evaluator WS server
with Client(url) as c:
c.send("stop")
mock_ws_thread.join()
task_result = flow_run.result[stage_task]
assert task_result.is_successful()
assert flow_run.is_successful()
assert len(task_result.result["outputs"]) == 1
expected_path = storage.get_storage_path(1) / "output.out"
output_path = flow_run.result[stage_task].result["outputs"][0]
assert expected_path == output_path
assert output_path.exists()
def gen_coef(parameters, real, config):
data = {}
paths = {}
storage = storage_driver_factory(config.get("storage"), "coeff")
for iens in range(real):
for name, elements in parameters.items():
for element in elements:
start, end = element["args"]
data[element["name"]] = uniform(start, end)
Path("coeff").mkdir(exist_ok=True)
file_name = "coeffs.json"
file_path = Path("coeff") / file_name
with open(file_path, "w") as f:
json.dump(data, f)
paths[iens] = (storage.store(file_name, iens), )
return paths
def store_resources(self, resources):
storage = storage_driver_factory(self.config.get("storage"),
self.config.get("config_path", "."))
stored_resources = [storage.store(res) for res in resources]
return stored_resources
def storage_driver(self, run_path):
Path(run_path).mkdir(parents=True, exist_ok=True)
return storage_driver_factory(self._storage_config, run_path)
def test_storage_driver(tmpdir):
storage_path = Path(".my_storage_test_path")
test_config = {"type": "shared_disk", "storage_path": storage_path}
with tmpdir.as_cwd():
file_name = "test_file.txt"
with open(file_name, "w") as f:
f.write("...")
storage = storage_driver_factory(config=test_config,
run_path=tmpdir.strpath)
assert storage is not None
# Check storage path
global_storage_path = storage.get_storage_path(None)
expected_global_storage_path = storage_path / "global"
assert expected_global_storage_path == global_storage_path
storage_path_iens_1 = storage.get_storage_path(1)
expected_storage_path = storage_path / "1"
assert expected_storage_path == storage_path_iens_1
# Store global file
global_resource_uri = storage.store(file_name)
expected_uri = storage_path / "global" / file_name
assert expected_uri == global_resource_uri
assert global_resource_uri.is_file()
# Store file for realization 1
real_resource = storage.store(file_name, 1)
expected_uri = storage_path / "1" / file_name
assert expected_uri == real_resource
assert real_resource.is_file()
# Retrieve resource from global storage
storage.retrieve(global_resource_uri, target="global_resource")
assert Path("global_resource").is_file()
# Retrieve resource from realization storage
storage.retrieve(real_resource, target="realization_resource")
assert Path("realization_resource").is_file()
# Retrieve non-existing file
file_uri = global_storage_path / "foo"
assert not file_uri.is_file()
with pytest.raises(ValueError,
match=f"Storage does not contain file: {file_uri}"):
storage.retrieve(file_uri)
# Retrieve a file that is not in storage, but exists outside of it.
file_uri = Path("foo.txt")
with open(file_uri, "w") as f:
f.write("bar")
assert file_uri.is_file()
with pytest.raises(ValueError,
match=f"Storage does not contain file: {file_uri}"):
storage.retrieve(file_uri)
# Retrieve an existing file that is not in storage using a relative uri.
file_uri = storage_path / ".." / "foo.txt"
with open(file_uri, "w") as f:
f.write("bar")
assert file_uri.is_file()
with pytest.raises(ValueError,
match=f"Storage does not contain file: {file_uri}"):
storage.retrieve(file_uri)
def test_on_task_failure(unused_tcp_port):
host = "localhost"
url = f"ws://{host}:{unused_tcp_port}"
messages = []
mock_ws_thread = threading.Thread(target=partial(_mock_ws,
messages=messages),
args=(host, unused_tcp_port))
mock_ws_thread.start()
with tmp(Path(SOURCE_DIR) / "test-data/local/prefect_test_case", False):
config = parse_config("config.yml")
storage = storage_driver_factory(config=config.get("storage"),
run_path=".")
resource = storage.store("unix_test_retry_script.py")
jobs = [{
"id": "0",
"name": "test_script",
"executable": "unix_test_retry_script.py",
"args": [],
}]
stage_task = UnixStep(
resources=[resource],
outputs=[],
job_list=jobs,
iens=1,
cmd="python3",
url=url,
step_id="step_id_0",
stage_id="stage_id_0",
ee_id="ee_id_0",
on_failure=partial(PrefectEnsemble._on_task_failure, url=url),
run_path=config.get("run_path"),
storage_config=config.get("storage"),
max_retries=3,
retry_delay=timedelta(seconds=1),
)
flow = Flow("testing")
flow.add_task(stage_task)
flow_run = flow.run()
# Stop the mock evaluator WS server
with Client(url) as c:
c.send("stop")
mock_ws_thread.join()
task_result = flow_run.result[stage_task]
assert task_result.is_successful()
assert flow_run.is_successful()
fail_job_messages = [
msg for msg in messages if ids.EVTYPE_FM_JOB_FAILURE in msg
]
fail_step_messages = [
msg for msg in messages if ids.EVTYPE_FM_STEP_FAILURE in msg
]
expected_job_failed_messages = 2
expected_step_failed_messages = 0
assert expected_job_failed_messages == len(fail_job_messages)
assert expected_step_failed_messages == len(fail_step_messages)