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())