def _evaluate(self, dispatch_url, ee_id):
asyncio.set_event_loop(asyncio.get_event_loop())
try:
with Client(dispatch_url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_STARTED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
self.run_flow(ee_id)
with Client(dispatch_url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_STOPPED,
"source": f"/ert/ee/{self._ee_id}",
"datacontenttype": "application/octet-stream",
},
cloudpickle.dumps(self.config["outputs"]),
)
c.send(to_json(event).decode())
except Exception as e:
logger.exception(
"An exception occurred while starting the ensemble evaluation",
exc_info=True,
)
with Client(dispatch_url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_FAILED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
def _evaluate(self, dispatch_url, ee_id):
try:
input_files = self._fetch_input_files()
with Client(dispatch_url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_STARTED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
self.run_flow(ee_id, dispatch_url, input_files)
with Client(dispatch_url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_STOPPED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
except Exception:
logger.exception(
"An exception occurred while starting the ensemble evaluation",
exc_info=True,
)
with Client(dispatch_url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_FAILED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
def run(self, expected_res=None):
run_path = self._run_path / str(self._iens)
storage = self.storage_driver(run_path)
self.retrieve_resources(expected_res, storage)
with Client(self._url) as ee_client:
event = CloudEvent(
{
"type": ids.EVTYPE_FM_STEP_START,
"source":
f"/ert/ee/{self._ee_id}/real/{self._iens}/stage/{self._stage_id}/step/{self._step_id}",
"datacontenttype": "application/json",
}, )
ee_client.send(to_json(event).decode())
outputs = []
self.run_jobs(ee_client, run_path)
for output in self._outputs:
if not (run_path / output).exists():
raise FileNotFoundError(
f"Output file {output} was not generated!")
outputs.append(storage.store(output, self._iens))
event = CloudEvent(
{
"type": ids.EVTYPE_FM_STEP_SUCCESS,
"source":
f"/ert/ee/{self._ee_id}/real/{self._iens}/stage/{self._stage_id}/step/{self._step_id}",
"datacontenttype": "application/json",
}, )
ee_client.send(to_json(event).decode())
return {"iens": self._iens, "outputs": outputs}
def run_jobs(self, client, run_path):
for index, job in enumerate(self._job_list):
self.logger.info(f"Running command {self._cmd} {job['name']}")
event = CloudEvent(
{
"type": ids.EVTYPE_FM_JOB_START,
"source":
f"/ert/ee/{self._ee_id}/real/{self._iens}/stage/{self._stage_id}/step/{self._step_id}/job/{job['id']}",
"datacontenttype": "application/json",
}, )
client.send(to_json(event).decode())
shell_cmd = [self._cmd, job["executable"], *job["args"]]
cmd_exec = subprocess.run(
shell_cmd,
universal_newlines=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
cwd=run_path,
)
self.logger.info(cmd_exec.stdout)
if cmd_exec.returncode != 0:
self.logger.error(cmd_exec.stderr)
event = CloudEvent(
{
"type": ids.EVTYPE_FM_JOB_FAILURE,
"source":
f"/ert/ee/{self._ee_id}/real/{self._iens}/stage/{self._stage_id}/step/{self._step_id}/job/{job['id']}",
"datacontenttype": "application/json",
},
{
"stderr": cmd_exec.stderr,
"stdout": cmd_exec.stdout
},
)
client.send(to_json(event).decode())
raise RuntimeError(
f"Script {job['name']} failed with exception {cmd_exec.stderr}"
)
event = CloudEvent(
{
"type": ids.EVTYPE_FM_JOB_SUCCESS,
"source":
f"/ert/ee/{self._ee_id}/real/{self._iens}/stage/{self._stage_id}/step/{self._step_id}/job/{job['id']}",
"datacontenttype": "application/json",
},
{"stdout": cmd_exec.stdout},
)
client.send(to_json(event).decode())
def _evaluate(self):
get_event_loop()
try:
with Client(
self._ee_config.dispatch_uri,
self._ee_config.token,
self._ee_config.cert,
) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_STARTED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
with prefect.context(
url=self._ee_config.dispatch_uri,
token=self._ee_config.token,
cert=self._ee_config.cert,
):
self.run_flow(self._ee_id)
with Client(
self._ee_config.dispatch_uri,
self._ee_config.token,
self._ee_config.cert,
) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_STOPPED,
"source": f"/ert/ee/{self._ee_id}",
"datacontenttype": "application/octet-stream",
},
cloudpickle.dumps(self._outputs),
)
c.send(to_json(event).decode())
except Exception as e:
logger.exception(
"An exception occurred while starting the ensemble evaluation",
exc_info=True,
)
with Client(
self._ee_config.dispatch_uri,
self._ee_config.token,
self._ee_config.cert,
) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_ENSEMBLE_FAILED,
"source": f"/ert/ee/{self._ee_id}",
}, )
c.send(to_json(event).decode())
async def _publish_changes(ee_id, changes, websocket):
events = [
JobQueue._translate_change_to_cloudevent(ee_id, real_id, status)
for real_id, status in changes.items()
]
for event in events:
await websocket.send(to_json(event))
def _publish_event(self):
with Client(self._evaluator_url, self._token, self._cert) as client:
while True:
event = self._event_queue.get()
if event is None:
return
client.send(to_json(event).decode())
async def test_run_and_cancel_legacy_ensemble(tmpdir, unused_tcp_port,
make_ensemble_builder):
num_reals = 10
conf_file = Path(tmpdir / CONFIG_FILE)
with tmpdir.as_cwd():
with open(conf_file, "w") as f:
f.write(f'port: "{unused_tcp_port}"\n')
ensemble = make_ensemble_builder(tmpdir, num_reals, 2).build()
config = load_config(conf_file)
evaluator = EnsembleEvaluator(ensemble, config, ee_id="1")
thread = threading.Thread(
name="test_eval",
target=evaluator.run_and_get_successful_realizations,
args=(),
)
thread.start()
# Wait for evaluator to start
await wait_for_ws(config["url"], 10)
# Send termination request to the evaluator
async with websockets.connect(config["client_url"]) as websocket:
out_cloudevent = CloudEvent({
"type": identifiers.EVTYPE_EE_USER_CANCEL,
"source": "/ert/test/0",
"id": "ID",
})
await websocket.send(to_json(out_cloudevent))
thread.join()
assert evaluator._snapshot.get_status() == "Cancelled"
def send_dispatch_event(client, event_type, source, event_id, data):
event1 = CloudEvent({
"type": event_type,
"source": source,
"id": event_id
}, data)
client.send(to_json(event1))
def terminate_message(self):
out_cloudevent = CloudEvent({
"type": identifiers.EVTYPE_EE_TERMINATED,
"source": f"/ert/ee/{self._ee_id}",
"id": self.event_index(),
})
message = to_json(out_cloudevent).decode()
return message
def test_json_can_talk_to_itself(specversion):
event = CloudEvent(test_attributes, test_data)
event_json = to_json(event)
event = from_json(event_json)
for key, val in test_attributes.items():
assert event[key] == val
assert event.data == test_data
async def send_cloudevent(self, url, event, retries=10):
for retry in range(retries):
try:
async with websockets.connect(url) as websocket:
await websocket.send(to_json(event))
return
except ConnectionRefusedError:
await asyncio.sleep(1)
raise IOError(f"Could not send event {event} to url {url}")
def test_to_json(specversion):
event = CloudEvent(test_attributes, test_data)
event_json = to_json(event)
event_dict = json.loads(event_json)
for key, val in test_attributes.items():
assert event_dict[key] == val
assert event_dict["data"] == test_data
def to_json(
self,
event: CloudEvent,
data_marshaller: types.MarshallerType = None) -> Union[str, bytes]:
if not data_marshaller:
datacontenttype = (event[_DATACONTENTTYPE] if _DATACONTENTTYPE
in event else _DEFAULT_DATACONTETYPE)
data_marshaller = self._marshallers.get(datacontenttype)
return to_json(event, data_marshaller=data_marshaller)
async def send_cloudevent(self,
url,
event,
token=None,
cert=None,
retries=1):
client = Client(url, token, cert)
await client._send(
to_json(event, data_marshaller=serialization.evaluator_marshaller))
await client.websocket.close()
async def _send():
async with websockets.connect(
self._client_uri,
ssl=self._ssl_context,
extra_headers=self._extra_headers,
) as websocket:
message = to_json(
cloud_event,
data_marshaller=serialization.evaluator_marshaller)
await websocket.send(message)
def _evaluate(self, client_url, dispatch_url, ee_id):
super()._evaluate(dispatch_url, ee_id)
with Client(client_url) as client:
client.send(
to_json(
CloudEvent({
"type": identifiers.EVTYPE_EE_USER_DONE,
"source": f"/ert/ee/{ee_id}",
"id": f"event-user-done",
})))
def test_json_can_talk_to_itself_base64(specversion):
data = b"test123"
event = CloudEvent(test_attributes, data)
event_json = to_json(event)
event = from_json(event_json)
for key, val in test_attributes.items():
assert event[key] == val
assert event.data == data
def create_snapshot_msg(ee_id, snapshot, event_index):
data = snapshot.to_dict()
out_cloudevent = CloudEvent(
{
"type": identifiers.EVTYPE_EE_SNAPSHOT,
"source": f"/ert/ee/{ee_id}",
"id": event_index,
},
data,
)
return to_json(out_cloudevent).decode()
def _send_event(self, cloud_event: CloudEvent) -> None:
with ExitStack() as stack:
duplexer = self._ws_duplexer
if not duplexer:
duplexer = SyncWebsocketDuplexer(self._client_uri,
self._base_uri, self._cert,
self._token)
stack.callback(duplexer.stop)
duplexer.send(
to_json(cloud_event,
data_marshaller=serialization.evaluator_marshaller))
def terminate_message(self):
out_cloudevent = CloudEvent(
{
"type": identifiers.EVTYPE_EE_TERMINATED,
"source": f"/ert/ee/{self._ee_id}",
"id": str(self.event_index()),
}
)
message = to_json(
out_cloudevent, data_marshaller=serialization.evaluator_marshaller
).decode()
return message
def _on_task_failure(task, state, url):
if prefect_context.task_run_count > task.max_retries:
with Client(url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_FM_STEP_FAILURE,
"source": task.get_step()._source,
"datacontenttype": "application/json",
},
{"error_msg": state.message},
)
c.send(to_json(event).decode())
def _on_task_failure(task, state, url):
if prefect_context.task_run_count > task.max_retries:
with Client(url) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_FM_STEP_FAILURE,
"source":
f"/ert/ee/{task.get_ee_id()}/real/{task.get_iens()}/stage/{task.get_stage_id()}/step/{task.get_step_id()}",
"datacontenttype": "application/json",
},
{"stderr": state.message},
)
c.send(to_json(event).decode())
async def send_cloudevent( # pylint: disable=too-many-arguments
self,
url: str,
event: CloudEvent,
token: Optional[str] = None,
cert: Optional[Union[str, bytes]] = None,
retries: int = 1,
) -> None:
client = Client(url, token, cert)
await client._send(to_json(event,
data_marshaller=evaluator_marshaller))
assert client.websocket # mypy
await client.websocket.close()
def create_snapshot_msg(ee_id, iter_, snapshot, event_index):
data = snapshot.to_dict()
data["iter"] = iter_
out_cloudevent = CloudEvent(
{
"type": identifiers.EVTYPE_EE_SNAPSHOT,
"source": f"/ert/ee/{ee_id}",
"id": event_index,
},
data,
)
return to_json(
out_cloudevent, data_marshaller=serialization.evaluator_marshaller
).decode()
def _from_cncf_events(event):
"""This takes in a CNCF cloudevent and returns a dictionary.
If cloud events library is not installed, the event is returned back.
"""
try:
from cloudevents.http import to_json
return json.loads(to_json(event))
except (AttributeError, ImportError):
# means this is not a CNCF event
return event
except Exception as err: # pylint: disable=broad-except
msg = """Failed to serialize the event. Please ensure your
CloudEvents is correctly formatted (https://pypi.org/project/cloudevents/)"""
raise_with_traceback(ValueError, msg, err)
async def _send_snapshot_update(self, snapshot_mutate_event):
self._snapshot.merge_event(snapshot_mutate_event)
out_cloudevent = CloudEvent(
{
"type": identifiers.EVTYPE_EE_SNAPSHOT_UPDATE,
"source": f"/ert/ee/{self._ee_id}",
"id": self.event_index(),
},
snapshot_mutate_event.to_dict(),
)
out_msg = to_json(out_cloudevent).decode()
if out_msg and self._clients:
await asyncio.wait(
[client.send(out_msg) for client in self._clients])
def test_to_json_base64(specversion):
data = b"test123"
event = CloudEvent(test_attributes, data)
event_json = to_json(event)
event_dict = json.loads(event_json)
for key, val in test_attributes.items():
assert event_dict[key] == val
# test data was properly marshalled into data_base64
data_base64 = event_dict["data_base64"].encode()
test_data_base64 = base64.b64encode(data)
assert data_base64 == test_data_base64
def _on_task_failure(self, task, state):
if prefect_context.task_run_count > task.max_retries:
url = prefect_context.url
token = prefect_context.token
cert = prefect_context.cert
with Client(url, token, cert) as c:
event = CloudEvent(
{
"type": ids.EVTYPE_FM_STEP_FAILURE,
"source": task.get_step().get_source(self._ee_id),
"datacontenttype": "application/json",
},
{"error_msg": state.message},
)
c.send(to_json(event).decode())
def _on_task_failure(
task: Union[UnixTask, FunctionTask], state: State, ee_id: str
) -> None:
if prefect_context.task_run_count > task.max_retries:
url = prefect_context.url
token = prefect_context.token
cert = prefect_context.cert
with Client(url, token, cert) as c:
event = CloudEvent(
{
"type": EVTYPE_FM_STEP_FAILURE,
"source": task.step.source(ee_id),
"datacontenttype": "application/json",
},
{"error_msg": state.message},
)
c.send(to_json(event).decode())
