def valid_fetcher_event(benchmark_event, benchmark_doc: BenchmarkDoc) -> FetcherBenchmarkEvent:
payload = FetcherPayload(
toml=benchmark_doc,
# We don't care about datasets here
datasets=[],
)
return create_from_object(FetcherBenchmarkEvent, benchmark_event, payload=payload)
def test_create_from_object(base_event_as_dict):
event = BenchmarkEvent.from_json(json.dumps(base_event_as_dict))
obj = create_from_object(FetcherBenchmarkEvent,
event,
payload=FetcherPayload(toml="", datasets=[]))
assert type(obj) == FetcherBenchmarkEvent
assert obj.action_id == event.action_id
def expected_executor_event(
fetcher_event: FetcherBenchmarkEvent) -> ExecutorBenchmarkEvent:
job_payload = ExecutorPayload.create_from_fetcher_payload(
fetcher_event.payload, BENCHMARK_JOB)
return create_from_object(ExecutorBenchmarkEvent,
fetcher_event,
payload=job_payload)
def mock_scheduled_run_fetcher_event():
scheduled_run_benchmark_event = create_benchmark_event(
scheduled_run_benchmark())
doc = scheduled_run_benchmark_event.payload.toml
return create_from_object(FetcherBenchmarkEvent,
scheduled_run_benchmark_event,
payload=FetcherPayload(datasets=[], toml=doc))
def fetcher_event(benchmark_event: BenchmarkEvent) -> FetcherBenchmarkEvent:
toml = BenchmarkDoc(contents={"var": "val"},
doc="DONTCARE",
sha1="DONTCARE")
return create_from_object(FetcherBenchmarkEvent,
benchmark_event,
payload=FetcherPayload(datasets=[], toml=toml))
def create_fetcher_benchmark_event_from_dict(
benchmark_event_dummy_payload: BenchmarkEvent,
toml_dict: Dict[str, Any]) -> FetcherBenchmarkEvent:
doc = BenchmarkDoc(toml_dict, "var = val",
"") # We don't care about the initial TOML
fetch_payload = FetcherPayload(toml=doc, datasets=[])
return create_from_object(
FetcherBenchmarkEvent,
benchmark_event_dummy_payload,
payload=fetch_payload,
action_id=benchmark_event_dummy_payload.action_id.replace("_", "-"),
)
def handle_event(self, event: FetcherBenchmarkEvent,
kafka_service: KafkaService):
# Only handle single run benchmarks
if not ExecutorEventHandler.is_single_run(event):
logging.debug(f"Ignoring non single-run event: {event}")
return
engine_id = ExecutorEventHandler.get_engine_id(event)
engine = self.execution_engines.get(engine_id)
if not engine:
# Ok. I've failed, but may be another service can have this engine
if engine_id in self.valid_execution_engines:
logging.info(
f"{engine_id} is whitelisted, but not present here. Nothing to do"
)
else:
# It's really something weird
logging.warning(f"Unknown engine {engine_id}")
return
try:
kafka_service.send_status_message_event(
event, Status.PENDING,
"Processing benchmark submission request...")
job = engine.run(event)
except ExecutionEngineException as e:
logger.exception("Engine throws exception")
kafka_service.send_status_message_event(event, Status.ERROR,
str(e))
raise KafkaServiceCallbackException from e
payload = ExecutorPayload.create_from_fetcher_payload(
event.payload, job)
response_event = create_from_object(ExecutorBenchmarkEvent,
event,
payload=payload)
kafka_service.send_status_message_event(
response_event, Status.SUCCEEDED,
"Benchmark successfully created...")
kafka_service.send_event(response_event, topic=self.producer_topic)
def test_producer(benchmark_event: BenchmarkEvent, kafka_service_config: KafkaServiceConfig):
consumer, producer = create_kafka_consumer_producer(kafka_service_config, SERVICE_NAME)
expected_job = BenchmarkJob(id=JOB_ID, k8s_yaml="")
expected_payload = ExecutorPayload.create_from_fetcher_payload(benchmark_event.payload, job=expected_job)
expected_event = create_from_object(ExecutorBenchmarkEvent, benchmark_event, payload=expected_payload)
producer.send(kafka_service_config.producer_topic, benchmark_event, key=benchmark_event.client_id)
filter_event = get_message_is_the_response(benchmark_event)
is_expected_event = get_event_equals(expected_event)
while True:
records = consumer.poll(POLL_TIMEOUT_MS)
for topic, recs in records.items():
for msg in recs:
if filter_event(msg.value):
assert is_expected_event(msg.value)
consumer.close()
return
def test_create_from_object_when_specifying_field_with_wrong_type(base_event):
with pytest.raises(ValueError):
create_from_object(FetcherBenchmarkEvent,
base_event,
payload="payload")
def test_create_from_object_when_input_object_is_not_a_dataclass():
with pytest.raises(ValueError):
create_from_object(FetcherBenchmarkEvent, object())
def test_create_from_object_when_desired_class_is_not_a_dataclass(base_event):
with pytest.raises(ValueError):
create_from_object(object, base_event)
def handle_event(self, event: CommandRequestEvent,
kafka_service: KafkaService):
# Safety check
payload = event.payload
if not payload or not isinstance(payload, CommandRequestPayload):
logger.info("Wrong payload passed")
return
command = event.payload.command
args = event.payload.args
logger.info(f"Command {command} called with {args}")
if not hasattr(self.cmd_object, command):
# May be another callback in the chain will be able to handle it
logger.info("No command method found")
return
method = getattr(self.cmd_object, command)
if callable(method):
sig = signature(method)
result = None
code = KafkaCommandCallback.CODE_SUCCESS
msg = self._get_message(KafkaCommandCallback.OK, command)
pos_args = []
kw_args = {}
if isinstance(args, list):
pos_args = tuple(args)
elif isinstance(args, dict):
kw_args = dict(args)
if "event" in sig.parameters:
kw_args["event"] = event
if "kafka_service" in sig.parameters:
kw_args["kafka_service"] = kafka_service
params = OrderedDict(sig.parameters)
# Remove positionals
orderparams = list(params)
for i in range(0, min(len(pos_args), len(orderparams))):
del params[orderparams[i]]
# Remove named
for k in kw_args:
if k in params:
del params[k]
from_event = vars(event)
for key in from_event.keys() & params.keys():
kw_args[key] = from_event[key]
try:
logger.info(f"executing {command}({pos_args}, {kw_args})")
result = method(*pos_args, **kw_args)
except TypeError as e:
logger.exception("Method invocation failed")
code = KafkaCommandCallback.CODE_CLIENT_ERROR
msg = self._get_message(KafkaCommandCallback.INVALID_ARGS,
command, str(e))
except Exception as e:
logger.exception("Command failed")
code = KafkaCommandCallback.CODE_COMMAND_ERROR
msg = self._get_message(KafkaCommandCallback.UNKNOWN_ERROR,
command, str(e))
response_payload = CommandResponsePayload(code, msg, event, result)
response_event = create_from_object(CommandResponseEvent,
event,
payload=response_payload)
kafka_service.send_event(response_event, self.cmd_return_topic)
else:
logger.info(f"Uncallable {command} member requested")
