def get_step(step_name, inputs, outputs, jobs, type_="unix"):
input_map: Dict[str, ert.data.RecordTransmitter] = {}
output_map: Dict[str, ert.data.RecordTransmitter] = {}
real_source = "/real/0"
step_source = "/real/0/step/0"
step_builder = (ee.create_step_builder().set_parent_source(
source=real_source).set_id("0").set_name(step_name).set_type(type_))
for idx, (name, executable, args) in enumerate(jobs):
step_builder.add_job(ee.create_job_builder().set_id(
str(idx)).set_name(name).set_executable(executable).set_args(
args).set_parent_source(step_source))
for name, path, mime, factory in inputs:
step_builder.add_input(
ee.create_file_io_builder().set_name(name).set_path(
Path(path)).set_mime(mime).set_transformation(
ert.data.ExecutableRecordTransformation(
)).set_transmitter_factory(factory))
for name, path, mime, factory in outputs:
step_builder.add_output(
ee.create_file_io_builder().set_name(name).set_path(path).set_mime(
mime).set_transmitter_factory(factory))
for input_ in step_builder._inputs:
input_map[input_._name] = input_.transmitter_factory()()
for output in step_builder._outputs:
output_map[output._name] = output.transmitter_factory()()
return step_builder.build(), input_map, output_map
def sum_coeffs_step(test_data_path, transmitter_factory):
step_builder = ee.create_step_builder().set_name("add_coeffs").set_type(
"unix")
step_builder.add_input(ee.create_file_io_builder(
).set_name("sum_up").set_path(
Path("sum_coeffs.py")).set_mime("text/x-python").set_transformation(
ert.data.ExecutableRecordTransformation()).set_transmitter_factory(
partial(
create_script_transmitter,
"sum_up",
test_data_path / "sum_coeffs.py",
transmitter_factory=transmitter_factory,
)))
step_builder.add_input(
ee.create_file_io_builder().set_name("input0").set_path(
Path("poly_0.out")).set_mime(
"application/json").set_transmitter_factory(
partial(transmitter_factory, "input0")))
step_builder.add_input(
ee.create_file_io_builder().set_name("input1").set_path(
Path("poly_1.out")).set_mime(
"application/json").set_transmitter_factory(
partial(transmitter_factory, "input1")))
step_builder.add_input(
ee.create_file_io_builder().set_name("input2").set_path(
Path("poly_2.out")).set_mime(
"application/json").set_transmitter_factory(
partial(transmitter_factory, "input2")))
step_builder.add_output(
ee.create_file_io_builder().set_name("sum_output").set_path(
Path("poly_sum.out")).set_mime(
"application/json").set_transmitter_factory(
partial(transmitter_factory, "sum_output")))
step_builder.add_job(
ee.create_job_builder().set_name("sum_up").set_executable(
Path("sum_coeffs.py")).set_args([]))
return step_builder
def add_commands(
transportable_commands: Tuple[ert3.config.TransportableCommand, ...],
storage_type: str,
storage_path: str,
step: StepBuilder,
) -> None:
def command_location(name: str) -> FilePath:
return next(
(cmd.location for cmd in transportable_commands if cmd.name == name),
pathlib.Path(name),
)
async def transform_output(
transmitter: ert.data.RecordTransmitter, mime: str, location: pathlib.Path
) -> None:
transformation = ert.data.ExecutableRecordTransformation()
record = await transformation.transform_output(mime, location)
await transmitter.transmit_record(record)
for command in transportable_commands:
transmitter: ert.data.RecordTransmitter
if storage_type == "shared_disk":
transmitter = ert.data.SharedDiskRecordTransmitter(
name=command.name,
storage_path=pathlib.Path(storage_path),
)
elif storage_type == "ert_storage":
transmitter = ert.storage.StorageRecordTransmitter(
name=command.name, storage_url=storage_path
)
else:
raise ValueError(f"Unsupported transmitter type: {storage_type}")
get_event_loop().run_until_complete(
transform_output(
transmitter=transmitter,
mime="application/octet-stream",
location=command.location,
)
)
step.add_input(
create_file_io_builder()
.set_name(command.name)
.set_path(command_location(command.name))
.set_mime("application/octet-stream")
.set_transformation(ert.data.ExecutableRecordTransformation())
# cast necessary due to https://github.com/python/mypy/issues/9656
.set_transmitter_factory(
lambda _t=transmitter: cast(ert.data.RecordTransmitter, _t)
)
)
def function_ensemble_builder_factory(
ensemble_size,
transmitter_factory,
coefficients,
):
job_builder = ee.create_job_builder().set_name("user_defined_function")
step_builder = (ee.create_step_builder().set_name(
"function_evaluation").set_type("function"))
coeffs_input = (ee.create_file_io_builder().set_name("coeffs").set_path(
"coeffs").set_mime("application/json"))
for iens, values in enumerate(coefficients):
transmitter = create_input_transmitter(values,
transmitter_factory("coeffs"))
coeffs_input.set_transmitter_factory(lambda _t=transmitter: _t, iens)
step_builder.add_input(coeffs_input)
step_builder.add_output(
ee.create_file_io_builder().set_name("function_output").set_path(
"output").set_mime("application/json").set_transmitter_factory(
partial(transmitter_factory, "function_output")))
step_builder.add_job(job_builder)
real_builder = ee.create_realization_builder().active(True).add_step(
step_builder)
builder = (ee.create_ensemble_builder().set_custom_port_range(
custom_port_range=range(1024, 65535)).set_ensemble_size(
ensemble_size).set_max_running(6).set_max_retries(2).set_executor(
"local").set_forward_model(real_builder))
def build(pickled_function):
job_builder.set_executable(pickled_function)
return builder
return build
def get_degree_step(degree, degree_spelled, transmitter_factory,
test_data_path, coefficients):
step_builder = (ee.create_step_builder().set_name(
f"{degree_spelled}_degree").set_type("unix"))
input_name = f"generate_{degree_spelled}_degree"
step_builder.add_input(
ee.create_file_io_builder().set_name(input_name).set_path(
Path("evaluate_coeffs.py")
).set_mime("text/x-python").set_transformation(
ert.data.ExecutableRecordTransformation()).set_transmitter_factory(
partial(
create_script_transmitter,
input_name,
test_data_path / "evaluate_coeffs.py",
transmitter_factory=transmitter_factory,
)))
coeffs_input = (ee.create_file_io_builder().set_name("coeffs").set_path(
Path("coeffs.json")).set_mime("application/json"))
for iens, values in enumerate(coefficients):
transmitter = create_input_transmitter(values,
transmitter_factory("coeffs"))
coeffs_input.set_transmitter_factory(lambda _t=transmitter: _t, iens)
step_builder.add_input(coeffs_input)
output_name = f"input{degree}"
step_builder.add_output(
ee.create_file_io_builder().set_name(output_name).set_path(
Path(f"poly_{degree}.out")).set_mime(
"application/json").set_transmitter_factory(
partial(transmitter_factory, output_name)))
step_builder.add_job(ee.create_job_builder().set_name(
f"generate_{degree_spelled}_degree").set_executable(
Path("evaluate_coeffs.py")).set_args([f"{degree}"]))
return step_builder
def zero_degree_step(transmitter_factory, test_data_path, coefficients):
step_builder = get_degree_step(
degree=0,
degree_spelled="zero",
transmitter_factory=transmitter_factory,
test_data_path=test_data_path,
coefficients=coefficients,
)
step_builder.add_input(
ee.create_file_io_builder().set_name("input2").set_path(
Path("poly_2.out")).set_mime(
"application/json").set_transmitter_factory(
partial(transmitter_factory, "input2")))
return step_builder
def add_step_outputs(
storage_type: str,
step_config: ert3.config.Step,
storage_path: str,
ensemble_size: int,
step: StepBuilder,
) -> None:
for record_name, output in step_config.output.items():
transformation = _get_output_recordtransformation(output)
output = (
create_file_io_builder()
.set_name(record_name)
.set_path(pathlib.Path(output.location))
.set_mime(output.mime)
.set_transformation(transformation)
)
for iens in range(0, ensemble_size):
factory: Callable[
[Type[ert.data.RecordTransmitter]], ert.data.RecordTransmitter
]
if storage_type == "shared_disk":
factory = partial(
ert.data.SharedDiskRecordTransmitter,
name=record_name,
storage_path=pathlib.Path(storage_path),
)
elif storage_type == "ert_storage":
factory = partial(
ert.storage.StorageRecordTransmitter,
name=record_name,
storage_url=storage_path,
iens=iens,
)
else:
raise ValueError(
f"unexpected storage type{storage_type} for {record_name} record"
)
output.set_transmitter_factory(factory, iens)
step.add_output(output)
def add_step_inputs(
inputs: Tuple[ert3.config.LinkedInput, ...],
transmitters: Dict[int, Dict[str, ert.data.RecordTransmitter]],
step: StepBuilder,
) -> None:
for input_ in inputs:
transformation = _get_input_recordtransformation(input_)
step_input = (
create_file_io_builder()
.set_name(input_.name)
.set_mime(input_.dest_mime)
.set_path(pathlib.Path(input_.dest_location))
.set_transformation(transformation)
)
for iens, io_to_transmitter in transmitters.items():
trans = io_to_transmitter[input_.name]
# cast necessary due to https://github.com/python/mypy/issues/9656
step_input.set_transmitter_factory(
lambda _t=trans: cast(ert.data.RecordTransmitter, _t), iens
)
step.add_input(step_input)