MRG: Merge pull request #669 from octue/speed-up-event-replaying

Speed up event replaying

Author: cortadocodes

docs/source/inter_service_compatibility.rst

@@ -63,6 +63,10 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._subscriber = MockSubscriber()
 
+    @property
+    def subscriber(self):
+        return self._subscriber
+
     def create(self, allow_existing=False):
         """Register the subscription in the global subscriptions dictionary.
 
@@ -78,11 +82,14 @@ def create(self, allow_existing=False):
             self._created = True
 
     def delete(self):
-        """Do nothing.
+        """Delete the subscription from the global subscriptions dictionary.
 
         :return None:
         """
-        pass
+        try:
+            SUBSCRIPTIONS.remove(self.name)
+        except KeyError:
+            pass
 
     def exists(self, timeout=5):
         """Check if the subscription exists in the global subscriptions dictionary.
@@ -310,7 +317,6 @@ def __init__(self, backend, service_id=None, run_function=None, children=None, *
         super().__init__(backend, service_id, run_function, *args, **kwargs)
         self.children = children or {}
         self._publisher = MockPublisher()
-        self.subscriber = MockSubscriber()
 
     @property
     def publisher(self):

octue/cloud/emulators/_pub_sub.py

@@ -39,7 +39,8 @@ class AbstractEventHandler:
     :param dict|None event_handlers: a mapping of event type names to callables that handle each type of event. The handlers must not mutate the events.
     :param dict schema: the JSON schema to validate events against
     :param bool include_service_metadata_in_logs: if `True`, include the SRUIDs and question UUIDs of the service revisions involved in the question to the start of the log message
-    :param bool only_handle_result: if `True`, skip handling non-result events and only handle the "result" event when received
+    :param bool only_handle_result: if `True`, skip handling non-result events and only handle the "result" event when received (turning this on speeds up event handling)
+    :param bool validate_events: if `True`, validate events before attempting to handle them (turning this off speeds up event handling)
     :return None:
     """
 
@@ -51,12 +52,14 @@ def __init__(
         schema=SERVICE_COMMUNICATION_SCHEMA,
         include_service_metadata_in_logs=True,
         only_handle_result=False,
+        validate_events=True,
     ):
         self.handle_monitor_message = handle_monitor_message
         self.record_events = record_events
         self.schema = schema
         self.include_service_metadata_in_logs = include_service_metadata_in_logs
         self.only_handle_result = only_handle_result
+        self.validate_events = validate_events
 
         self.handled_events = []
         self._start_time = None
@@ -113,7 +116,7 @@ def _extract_and_validate_event(self, container):
         recipient = attributes.get("recipient")
         child_sdk_version = attributes.get("sender_sdk_version")
 
-        if not is_event_valid(
+        if self.validate_events and not is_event_valid(
             event=event,
             attributes=attributes,
             recipient=recipient,
@@ -125,8 +128,8 @@ def _extract_and_validate_event(self, container):
 
         logger.debug(
             "%r: Received an event related to question %r.",
-            attributes["recipient"],
-            attributes["question_uuid"],
+            attributes.get("recipient"),
+            attributes.get("question_uuid"),
         )
 
         return (event, attributes)

octue/cloud/events/handler.py

@@ -15,7 +15,8 @@ class EventReplayer(AbstractEventHandler):
     :param dict|None event_handlers: a mapping of event type names to callables that handle each type of event. The handlers must not mutate the events.
     :param dict|str schema: the JSON schema to validate events against
     :param bool include_service_metadata_in_logs: if `True`, include the SRUIDs and question UUIDs of the service revisions involved in the question to the start of the log message
-    :param bool only_handle_result: if `True`, skip non-result events and only handle the "result" event if present
+    :param bool only_handle_result: if `True`, skip non-result events and only handle the "result" event if present (turning this on speeds up event handling)
+    :param bool validate_events: if `True`, validate events before attempting to handle them (turning this off speeds up event handling)
     :return None:
     """
 
@@ -27,6 +28,7 @@ def __init__(
         schema=SERVICE_COMMUNICATION_SCHEMA,
         include_service_metadata_in_logs=True,
         only_handle_result=False,
+        validate_events=True,
     ):
         event_handlers = event_handlers or {
             "question": self._handle_question,
@@ -45,6 +47,7 @@ def __init__(
             schema=schema,
             include_service_metadata_in_logs=include_service_metadata_in_logs,
             only_handle_result=only_handle_result,
+            validate_events=validate_events,
         )
 
     def handle_events(self, events):
@@ -56,6 +59,10 @@ def handle_events(self, events):
         super().handle_events()
 
         for event in events:
+            # Skip validation and handling of other event kinds if only the result event is wanted.
+            if self.only_handle_result and event.get("event", {}).get("kind") != "result":
+                continue
+
             event, attributes = self._extract_and_validate_event(event)
 
             # Skip the event if it fails validation.

octue/cloud/events/replayer.py

@@ -3,6 +3,7 @@
 import logging
 import time
 from datetime import datetime, timedelta
+from functools import cached_property
 
 from google.api_core import retry
 from google.cloud.pubsub_v1 import SubscriberClient
@@ -91,11 +92,20 @@ def __init__(
             include_service_metadata_in_logs=include_service_metadata_in_logs,
         )
 
-        self._subscriber = SubscriberClient()
         self._heartbeat_checker = None
         self._last_heartbeat = None
         self._alive = True
 
+    @cached_property
+    def subscriber(self):
+        """Get or instantiate the subscriber client. The client isn't instantiated until this property is called for the
+        first time. This allows checking for the `GOOGLE_APPLICATION_CREDENTIALS` environment variable to be put off
+        until it's needed.
+
+        :return google.cloud.pubsub_v1.SubscriberClient:
+        """
+        return SubscriberClient()
+
     @property
     def total_run_time(self):
         """The amount of time elapsed since `self.handle_events` was called. If it hasn't been called yet, this is
@@ -156,7 +166,7 @@ def handle_events(self, timeout=60, maximum_heartbeat_interval=300):
 
         finally:
             self._heartbeat_checker.cancel()
-            self._subscriber.close()
+            self.subscriber.close()
 
         if self.handled_events:
             last_event = self.handled_events[-1]
@@ -221,7 +231,7 @@ def _pull_available_events(self, timeout):
         while self._alive:
             logger.debug("Pulling events from Google Pub/Sub: attempt %d.", attempt)
 
-            pull_response = self._subscriber.pull(
+            pull_response = self.subscriber.pull(
                 request={"subscription": self.subscription.path, "max_messages": MAX_SIMULTANEOUS_MESSAGES_PULL},
                 retry=retry.Retry(),
             )
@@ -240,7 +250,7 @@ def _pull_available_events(self, timeout):
         if not pull_response.received_messages:
             return []
 
-        self._subscriber.acknowledge(
+        self.subscriber.acknowledge(
             request={
                 "subscription": self.subscription.path,
                 "ack_ids": [message.ack_id for message in pull_response.received_messages],

octue/cloud/pub_sub/events.py

@@ -84,8 +84,6 @@ def __init__(self, backend, service_id=None, run_function=None, service_registri
 
         self._pub_sub_id = convert_service_id_to_pub_sub_form(self.id)
         self._local_sdk_version = importlib.metadata.version("octue")
-        self._publisher = None
-        self._services_topic = None
         self._event_handler = None
 
     def __repr__(self):
@@ -95,23 +93,20 @@ def __repr__(self):
         """
         return f"<{type(self).__name__}({self.id!r})>"
 
-    @property
+    @functools.cached_property
     def publisher(self):
         """Get or instantiate the publisher client for the service. No publisher is instantiated until this property is
         called for the first time. This allows checking for the `GOOGLE_APPLICATION_CREDENTIALS` environment variable to
         be put off until it's needed.
 
         :return google.cloud.pubsub_v1.PublisherClient:
         """
-        if not self._publisher:
-            self._publisher = pubsub_v1.PublisherClient(
-                batch_settings=BATCH_SETTINGS,
-                publisher_options=pubsub_v1.types.PublisherOptions(enable_message_ordering=True),
-            )
-
-        return self._publisher
+        return pubsub_v1.PublisherClient(
+            batch_settings=BATCH_SETTINGS,
+            publisher_options=pubsub_v1.types.PublisherOptions(enable_message_ordering=True),
+        )
 
-    @property
+    @functools.cached_property
     def services_topic(self):
         """Get the Octue services topic that all events in the project are published to. No topic is instantiated until
         this property is called for the first time. This allows checking for the `GOOGLE_APPLICATION_CREDENTIALS`
@@ -120,15 +115,12 @@ def services_topic(self):
         :raise octue.exceptions.ServiceNotFound: if the topic doesn't exist in the project
         :return octue.cloud.pub_sub.topic.Topic: the Octue services topic for the project
         """
-        if not self._services_topic:
-            topic = Topic(name=OCTUE_SERVICES_PREFIX, project_name=self.backend.project_name)
-
-            if not topic.exists():
-                raise octue.exceptions.ServiceNotFound(f"{topic!r} cannot be found.")
+        topic = Topic(name=OCTUE_SERVICES_PREFIX, project_name=self.backend.project_name)
 
-            self._services_topic = topic
+        if not topic.exists():
+            raise octue.exceptions.ServiceNotFound(f"{topic!r} cannot be found.")
 
-        return self._services_topic
+        return topic
 
     @property
     def received_events(self):
@@ -533,7 +525,7 @@ def _emit_event(
         attributes.update(
             {
                 "uuid": str(uuid.uuid4()),
-                "datetime": datetime.datetime.utcnow().isoformat(),
+                "datetime": datetime.datetime.now(tz=datetime.timezone.utc).isoformat(),
                 "question_uuid": question_uuid,
                 "parent_question_uuid": parent_question_uuid,
                 "originator_question_uuid": originator_question_uuid,

octue/cloud/pub_sub/service.py

@@ -1,4 +1,5 @@
 import logging
+from functools import cached_property
 
 import google.api_core.exceptions
 from google.cloud.pubsub_v1 import SubscriberClient
@@ -72,9 +73,18 @@ def __init__(
 
         self.push_endpoint = push_endpoint
         self.enable_message_ordering = enable_message_ordering
-        self._subscriber = SubscriberClient()
         self._created = False
 
+    @cached_property
+    def subscriber(self):
+        """Get or instantiate the subscriber client. The client isn't instantiated until this property is called for the
+        first time. This allows checking for the `GOOGLE_APPLICATION_CREDENTIALS` environment variable to be put off
+        until it's needed.
+
+        :return google.cloud.pubsub_v1.SubscriberClient:
+        """
+        return SubscriberClient()
+
     @property
     def creation_triggered_locally(self):
         """Was the subscription successfully created by calling `self.create` locally? This is `False` if its creation
@@ -116,13 +126,13 @@ def create(self, allow_existing=False):
         subscription = self._create_proto_message_subscription()
 
         if not allow_existing:
-            subscription = self._subscriber.create_subscription(request=subscription)
+            subscription = self.subscriber.create_subscription(request=subscription)
             self._created = True
             self._log_creation()
             return subscription
 
         try:
-            subscription = self._subscriber.create_subscription(request=subscription)
+            subscription = self.subscriber.create_subscription(request=subscription)
             self._created = True
         except google.api_core.exceptions.AlreadyExists:
             pass
@@ -135,7 +145,7 @@ def update(self):
 
         :return None:
         """
-        self._subscriber.update_subscription(
+        self.subscriber.update_subscription(
             request=UpdateSubscriptionRequest(
                 mapping=None,
                 subscription=self._create_proto_message_subscription(),  # noqa
@@ -156,7 +166,7 @@ def delete(self):
 
         :return None:
         """
-        self._subscriber.delete_subscription(subscription=self.path)
+        self.subscriber.delete_subscription(subscription=self.path)
         logger.info("Subscription %r deleted.", self.path)
 
     def exists(self, timeout=5):
@@ -166,7 +176,7 @@ def exists(self, timeout=5):
         :return bool:
         """
         try:
-            self._subscriber.get_subscription(subscription=self.path, timeout=timeout)
+            self.subscriber.get_subscription(subscription=self.path, timeout=timeout)
             return True
         except google.api_core.exceptions.NotFound:
             return False

octue/cloud/pub_sub/subscription.py

@@ -1,6 +1,7 @@
 import logging
 import time
 from datetime import datetime
+from functools import cached_property
 
 import google.api_core.exceptions
 from google.cloud.pubsub_v1 import PublisherClient
@@ -24,9 +25,18 @@ def __init__(self, name, project_name):
         self.name = name
         self.project_name = project_name
         self.path = self.generate_topic_path(self.project_name, self.name)
-        self._publisher = PublisherClient()
         self._created = False
 
+    @cached_property
+    def publisher(self):
+        """Get or instantiate the publisher client. The client isn't instantiated until this property is called for the
+        first time. This allows checking for the `GOOGLE_APPLICATION_CREDENTIALS` environment variable to be put off
+        until it's needed.
+
+        :return google.cloud.pubsub_v1.PublisherClient:
+        """
+        return PublisherClient()
+
     @property
     def creation_triggered_locally(self):
         """Was the topic successfully created by calling `self.create` locally? This is `False` if its creation was
@@ -52,15 +62,15 @@ def create(self, allow_existing=False):
         posix_timestamp_with_no_decimals = str(datetime.now().timestamp()).split(".")[0]
 
         if not allow_existing:
-            self._publisher.create_topic(
+            self.publisher.create_topic(
                 request=Topic_(name=self.path, labels={"created": posix_timestamp_with_no_decimals})
             )
             self._created = True
             self._log_creation()
             return
 
         try:
-            self._publisher.create_topic(
+            self.publisher.create_topic(
                 request=Topic_(name=self.path, labels={"created": posix_timestamp_with_no_decimals})
             )
             self._created = True
@@ -74,14 +84,14 @@ def get_subscriptions(self):
 
         :return list(str):
         """
-        return list(self._publisher.list_topic_subscriptions(topic=self.path))
+        return list(self.publisher.list_topic_subscriptions(topic=self.path))
 
     def delete(self):
         """Delete the topic from Google Pub/Sub.
 
         :return None:
         """
-        self._publisher.delete_topic(topic=self.path)
+        self.publisher.delete_topic(topic=self.path)
         logger.info("Topic %r deleted.", self.path)
 
     def exists(self, timeout=10):
@@ -94,7 +104,7 @@ def exists(self, timeout=10):
 
         while time.time() - start_time <= timeout:
             try:
-                self._publisher.get_topic(topic=self.path)
+                self.publisher.get_topic(topic=self.path)
                 return True
             except google.api_core.exceptions.NotFound:
                 time.sleep(1)