Merge pull request #237 from lsst-dm/tickets/DM-43082-3

DM-43082: Knative Keda co-existence

Author: dspeck1

python/activator/activator.py

@@ -27,16 +27,20 @@
 import json
 import logging
 import os
+import signal
+import socket
 import sys
 import time
-import signal
-import yaml
 import uuid
+import yaml
 
 import boto3
 from botocore.handlers import validate_bucket_name
 import cloudevents.http
 import confluent_kafka as kafka
+from confluent_kafka.serialization import SerializationContext, MessageField
+from confluent_kafka.schema_registry import SchemaRegistryClient
+from confluent_kafka.schema_registry.avro import AvroDeserializer
 import flask
 
 from .config import PipelinesConfig
@@ -53,7 +57,8 @@
 from .repo_tracker import LocalRepoTracker
 from .visit import FannedOutVisit
 
-
+# Platform that prompt processing will run on
+platform = os.environ["PLATFORM"].lower()
 # The short name for the instrument.
 instrument_name = os.environ["RUBIN_INSTRUMENT"]
 # The skymap to use in the central repo
@@ -76,6 +81,31 @@
 kafka_group_id = str(uuid.uuid4())
 # The topic on which to listen to updates to image_bucket
 bucket_topic = os.environ.get("BUCKET_TOPIC", "rubin-prompt-processing")
+# Offset for Kafka bucket notification.
+bucket_notification_kafka_offset_reset = os.environ.get("BUCKET_NOTIFICATION_KAFKA_OFFSET_RESET", "latest")
+
+# Conditionally load keda environment variables
+if platform == "keda":
+    # Kafka Schema Registry URL for next visit fan out messages
+    fan_out_schema_registry_url = os.environ["FAN_OUT_SCHEMA_REGISTRY_URL"]
+    # Kafka cluster with next visit fanned out messages.
+    fan_out_kafka_cluster = os.environ["FAN_OUT_KAFKA_CLUSTER"]
+    # Kafka group for next visit fan out messages.
+    fan_out_kafka_group_id = os.environ["FAN_OUT_KAFKA_GROUP_ID"]
+    # Kafka topic for next visit fan out messages.
+    fan_out_kafka_topic = os.environ["FAN_OUT_KAFKA_TOPIC"]
+    # Kafka topic offset for next visit fan out messages.
+    fan_out_kafka_topic_offset = os.environ["FAN_OUT_KAFKA_TOPIC_OFFSET"]
+    # Kafka Fan Out SASL Mechansim.
+    fan_out_kafka_sasl_mechanism = os.environ["FAN_OUT_KAFKA_SASL_MECHANISM"]
+    # Kafka Fan Out Security Protocol.
+    fan_out_kafka_security_protocol = os.environ["FAN_OUT_KAFKA_SECURITY_PROTOCOL"]
+    # Kafka Fan Out Consumer Username.
+    fan_out_kafka_sasl_username = os.environ["FAN_OUT_KAFKA_SASL_USERNAME"]
+    # Kafka Fan Out Consumer Password.
+    fan_out_kafka_sasl_password = os.environ["FAN_OUT_KAFKA_SASL_PASSWORD"]
+    # Time to wait for fanned out messages before spawning new pod.
+    fanned_out_msg_listen_timeout = int(os.environ.get("FANNED_OUT_MSG_LISTEN_TIMEOUT", 300))
 
 _log = logging.getLogger("lsst." + __name__)
 _log.setLevel(logging.DEBUG)
@@ -127,7 +157,7 @@ def _get_consumer():
     return kafka.Consumer({
         "bootstrap.servers": kafka_cluster,
         "group.id": kafka_group_id,
-        "auto.offset.reset": "latest",  # default, but make explicit
+        "auto.offset.reset": bucket_notification_kafka_offset_reset,
     })
 
 
@@ -195,6 +225,107 @@ def create_app():
         sys.exit(3)
 
 
+def dict_to_fanned_out_visit(obj, ctx):
+    """
+    Converts object literal(dict) to a Fanned Out instance.
+    Args:
+        ctx (SerializationContext): Metadata pertaining to the serialization
+            operation.
+        obj (dict): Object literal(dict)
+    """
+
+    if obj is None:
+        return None
+
+    return FannedOutVisit(**obj)
+
+
+def keda_start():
+
+    try:
+        setup_usdf_logger(
+            labels={"instrument": instrument_name},
+        )
+
+        # Initialize local registry
+        registry = LocalRepoTracker.get()
+        registry.init_tracker()
+
+        # Check initialization and abort early
+        _get_consumer()
+        _get_storage_client()
+        _get_central_butler()
+        _get_local_repo()
+
+        _log.info("Worker ready to handle requests.")
+
+    except Exception as e:
+        _log.critical("Failed to start worker; aborting.")
+        _log.exception(e)
+        sys.exit(1)
+
+    # Initialize schema registry for fan out
+    fan_out_schema_registry_conf = {'url': fan_out_schema_registry_url}
+    fan_out_schema_registry_client = SchemaRegistryClient(fan_out_schema_registry_conf)
+
+    fan_out_avro_deserializer = AvroDeserializer(schema_registry_client=fan_out_schema_registry_client,
+                                                 from_dict=dict_to_fanned_out_visit)
+    fan_out_consumer_conf = {
+        "bootstrap.servers": fan_out_kafka_cluster,
+        "group.id": fan_out_kafka_group_id,
+        "auto.offset.reset": fan_out_kafka_topic_offset,
+        "sasl.mechanism": fan_out_kafka_sasl_mechanism,
+        "security.protocol": fan_out_kafka_security_protocol,
+        "sasl.username": fan_out_kafka_sasl_username,
+        "sasl.password": fan_out_kafka_sasl_password,
+        'enable.auto.commit': False
+    }
+
+    _log.info("starting fan out consumer")
+    fan_out_consumer = kafka.Consumer(fan_out_consumer_conf, logger=_log)
+    fan_out_consumer.subscribe([fan_out_kafka_topic])
+    fan_out_listen_start_time = time.time()
+
+    try:
+        while time.time() - fan_out_listen_start_time < fanned_out_msg_listen_timeout:
+
+            fan_out_message = fan_out_consumer.poll(timeout=5)
+            if fan_out_message is None:
+                continue
+            if fan_out_message.error():
+                _log.warning("Fanned out consumer error: %s", fan_out_message.error())
+            else:
+                deserialized_fan_out_visit = fan_out_avro_deserializer(fan_out_message.value(),
+                                                                       SerializationContext(
+                                                                       fan_out_message.topic(),
+                                                                       MessageField.VALUE))
+                _log.info("Unpacked message as %r.", deserialized_fan_out_visit)
+
+                # Calculate time to load knative and receive message based on timestamp in Kafka message
+                _log.debug("Message timestamp %r", fan_out_message.timestamp())
+                fan_out_kafka_msg_timestamp = fan_out_message.timestamp()
+                fan_out_to_prompt_time = int(time.time() * 1000) - fan_out_kafka_msg_timestamp[1]
+                _log.debug("Seconds since fan out message delivered %r", fan_out_to_prompt_time/1000)
+
+                # Commit message and close client
+                fan_out_consumer.commit(message=fan_out_message, asynchronous=False)
+                fan_out_consumer.close()
+
+                try:
+                    # Process fan out visit
+                    process_visit(deserialized_fan_out_visit)
+                except Exception as e:
+                    _log.critical("Process visit failed; aborting.")
+                    _log.exception(e)
+                finally:
+                    _log.info("Processing completed for %s", socket.gethostname())
+                    break
+
+    finally:
+        # TODO Handle local registry unregistration on DM-47975
+        _log.info("Finished listening for fanned out messages")
+
+
 def _graceful_shutdown(signum: int, stack_frame):
     """Signal handler for cases where the service should gracefully shut down.
 
@@ -586,10 +717,18 @@ def server_error(e: Exception) -> tuple[str, int]:
 
 
 def main():
-    # This function is only called in test environments. Container
-    # deployments call `create_app()()` through Gunicorn.
-    app = create_app()
-    app.run(host="127.0.0.1", port=8080, debug=True)
+    # Knative deployments call `create_app()()` through Gunicorn.
+    # Keda deployments invoke main.
+    if platform == "knative":
+        _log.info("starting standalone Flask app")
+        app = create_app()
+        app.run(host="127.0.0.1", port=8080, debug=True)
+    # starts keda instance of the application
+    elif platform == "keda":
+        _log.info("starting keda instance")
+        keda_start()
+    else:
+        _log.info("no platform defined")
 
 
 if __name__ == "__main__":