Logging post inference hook implementation (#428)

* add logging hook

* plumb other columns through endpoint config

* add args to unit tests

* change to labels

* add config and assume role

* new storage gateway

* add stream name config

* fix test

* undo conftest

* handle error

* fix test

* fake streaming storage gateway

* move client to fn

* change to fake gateway

* PR comments

* catch err

* update test

* remove error in response

* try small test

* add more tests

* fix test

---------

Co-authored-by: Sai Atmakuri <[email protected]>

Author: tiffzhao5

model-engine/model_engine_server/api/dependencies.py

@@ -43,6 +43,12 @@
     LLMModelEndpointService,
     ModelEndpointService,
 )
+from model_engine_server.inference.domain.gateways.streaming_storage_gateway import (
+    StreamingStorageGateway,
+)
+from model_engine_server.inference.infra.gateways.firehose_streaming_storage_gateway import (
+    FirehoseStreamingStorageGateway,
+)
 from model_engine_server.infra.gateways import (
     CeleryTaskQueueGateway,
     FakeMonitoringMetricsGateway,
@@ -137,6 +143,7 @@ class ExternalInterfaces:
     cron_job_gateway: CronJobGateway
     monitoring_metrics_gateway: MonitoringMetricsGateway
     tokenizer_repository: TokenizerRepository
+    streaming_storage_gateway: StreamingStorageGateway
 
 
 def get_default_monitoring_metrics_gateway() -> MonitoringMetricsGateway:
@@ -265,6 +272,8 @@ def _get_external_interfaces(
 
     tokenizer_repository = LiveTokenizerRepository(llm_artifact_gateway=llm_artifact_gateway)
 
+    streaming_storage_gateway = FirehoseStreamingStorageGateway()
+
     external_interfaces = ExternalInterfaces(
         docker_repository=docker_repository,
         model_bundle_repository=model_bundle_repository,
@@ -287,6 +296,7 @@ def _get_external_interfaces(
         cron_job_gateway=cron_job_gateway,
         monitoring_metrics_gateway=monitoring_metrics_gateway,
         tokenizer_repository=tokenizer_repository,
+        streaming_storage_gateway=streaming_storage_gateway,
     )
     return external_interfaces
 

model-engine/model_engine_server/common/constants.py

@@ -2,6 +2,7 @@
 
 BILLING_POST_INFERENCE_HOOK: str = "billing"
 CALLBACK_POST_INFERENCE_HOOK: str = "callback"
+LOGGING_POST_INFERENCE_HOOK: str = "logging"
 READYZ_FPATH: str = "/tmp/readyz"
 DEFAULT_CELERY_TASK_NAME: str = "hosted_model_inference.inference.async_inference.tasks.predict"
 LIRA_CELERY_TASK_NAME: str = "ml_serve.celery_service.exec_func"

model-engine/model_engine_server/core/config.py

@@ -42,6 +42,8 @@ class InfraConfig:
     profile_ml_worker: str = "default"
     profile_ml_inference_worker: str = "default"
     identity_service_url: Optional[str] = None
+    firehose_role_arn: Optional[str] = None
+    firehose_stream_name: Optional[str] = None
 
     @classmethod
     def from_yaml(cls, yaml_path) -> "InfraConfig":

model-engine/model_engine_server/domain/entities/model_endpoint_entity.py

@@ -88,6 +88,10 @@ class ModelEndpointConfig(BaseModel):
     billing_tags: Optional[Dict[str, Any]] = None
     default_callback_url: Optional[str] = None
     default_callback_auth: Optional[CallbackAuth]
+    endpoint_id: Optional[str] = None
+    endpoint_type: Optional[ModelEndpointType]
+    bundle_id: Optional[str] = None
+    labels: Optional[Dict[str, str]] = None
 
     def serialize(self) -> str:
         return python_json_to_b64(dict_not_none(**self.dict()))

model-engine/model_engine_server/domain/exceptions.py

@@ -164,3 +164,9 @@ class TriggerNameAlreadyExistsException(DomainException):
     """
     Thrown if the requested name already exists in the trigger repository
     """
+
+
+class StreamPutException(DomainException):
+    """
+    Thrown if the streaming storage gateway fails to put a record.
+    """

model-engine/model_engine_server/inference/async_inference/tasks.py

@@ -18,6 +18,9 @@
 from model_engine_server.inference.infra.gateways.datadog_inference_monitoring_metrics_gateway import (
     DatadogInferenceMonitoringMetricsGateway,
 )
+from model_engine_server.inference.infra.gateways.firehose_streaming_storage_gateway import (
+    FirehoseStreamingStorageGateway,
+)
 from model_engine_server.inference.post_inference_hooks import PostInferenceHooksHandler
 
 logger = make_logger(logger_name())
@@ -46,6 +49,11 @@ def init_worker_global():
         default_callback_url=endpoint_config.default_callback_url,
         default_callback_auth=endpoint_config.default_callback_auth,
         monitoring_metrics_gateway=DatadogInferenceMonitoringMetricsGateway(),
+        endpoint_id=endpoint_config.endpoint_id,
+        endpoint_type=endpoint_config.endpoint_type,
+        bundle_id=endpoint_config.bundle_id,
+        labels=endpoint_config.labels,
+        streaming_storage_gateway=FirehoseStreamingStorageGateway(),
     )
     # k8s health check
     with open(READYZ_FPATH, "w") as f:

model-engine/model_engine_server/inference/domain/gateways/streaming_storage_gateway.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class StreamingStorageGateway(ABC):
+    """
+    Base class for a gateway that stores data through a streaming mechanism.
+    """
+
+    @abstractmethod
+    def put_record(self, stream_name: str, record: Dict[str, Any]) -> None:
+        """
+        Put a record into a streaming storage mechanism.
+
+        Args:
+            stream_name: The name of the stream.
+            record: The record to put into the stream.
+        """
+        pass

model-engine/model_engine_server/inference/forwarding/forwarding.py

@@ -14,6 +14,9 @@
 from model_engine_server.inference.infra.gateways.datadog_inference_monitoring_metrics_gateway import (
     DatadogInferenceMonitoringMetricsGateway,
 )
+from model_engine_server.inference.infra.gateways.firehose_streaming_storage_gateway import (
+    FirehoseStreamingStorageGateway,
+)
 from model_engine_server.inference.post_inference_hooks import PostInferenceHooksHandler
 
 __all__: Sequence[str] = (
@@ -279,6 +282,11 @@ def endpoint(route: str) -> str:
             default_callback_url=endpoint_config.default_callback_url,
             default_callback_auth=endpoint_config.default_callback_auth,
             monitoring_metrics_gateway=DatadogInferenceMonitoringMetricsGateway(),
+            endpoint_id=endpoint_config.endpoint_id,
+            endpoint_type=endpoint_config.endpoint_type,
+            bundle_id=endpoint_config.bundle_id,
+            labels=endpoint_config.labels,
+            streaming_storage_gateway=FirehoseStreamingStorageGateway(),
         )
 
         return Forwarder(
@@ -451,6 +459,11 @@ def endpoint(route: str) -> str:
             default_callback_url=endpoint_config.default_callback_url,
             default_callback_auth=endpoint_config.default_callback_auth,
             monitoring_metrics_gateway=DatadogInferenceMonitoringMetricsGateway(),
+            endpoint_id=endpoint_config.endpoint_id,
+            endpoint_type=endpoint_config.endpoint_type,
+            bundle_id=endpoint_config.bundle_id,
+            labels=endpoint_config.labels,
+            streaming_storage_gateway=FirehoseStreamingStorageGateway(),
         )
 
         return StreamingForwarder(

model-engine/model_engine_server/inference/infra/gateways/firehose_streaming_storage_gateway.py

@@ -0,0 +1,63 @@
+import json
+from typing import Any, Dict
+
+import boto3
+from model_engine_server.core.config import infra_config
+from model_engine_server.core.loggers import logger_name, make_logger
+from model_engine_server.domain.exceptions import StreamPutException
+from model_engine_server.inference.domain.gateways.streaming_storage_gateway import (
+    StreamingStorageGateway,
+)
+
+logger = make_logger(logger_name())
+
+
+class FirehoseStreamingStorageGateway(StreamingStorageGateway):
+    """
+    A gateway that stores data through the AWS Kinesis Firehose streaming mechanism.
+    """
+
+    def __init__(self):
+        pass
+
+    """
+    Creates a new firehose client.
+
+    Streams with Snowflake as a destination and the AWS profile live in different 
+    accounts. Firehose doesn't support resource-based policies, so we need to assume 
+    a new role to write to the stream.
+    """
+
+    def _get_firehose_client(self):
+        sts_client = boto3.client("sts", region_name=infra_config().default_region)
+        assumed_role_object = sts_client.assume_role(
+            RoleArn=infra_config().firehose_role_arn,
+            RoleSessionName="AssumeMlLoggingRoleSession",
+        )
+        credentials = assumed_role_object["Credentials"]
+        session = boto3.Session(
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+        )
+        firehose_client = session.client("firehose", region_name=infra_config().default_region)
+        return firehose_client
+
+    def put_record(self, stream_name: str, record: Dict[str, Any]) -> None:
+        """
+        Put a record into a Firehose stream.
+
+        Args:
+            stream_name: The name of the stream.
+            record: The record to put into the stream.
+        """
+        firehose_response = self._get_firehose_client().put_record(
+            DeliveryStreamName=stream_name, Record={"Data": json.dumps(record).encode("utf-8")}
+        )
+        if firehose_response["ResponseMetadata"]["HTTPStatusCode"] != 200:
+            raise StreamPutException(
+                f"Failed to put record into firehose stream {stream_name}. Record content: {record}"
+            )
+        logger.info(
+            f"Logged to firehose stream {stream_name}. Record content: {record}, Record ID: {firehose_response['RecordId']}"
+        )

model-engine/model_engine_server/inference/post_inference_hooks.py

@@ -4,13 +4,22 @@
 
 import requests
 from fastapi.responses import JSONResponse
-from model_engine_server.common.constants import CALLBACK_POST_INFERENCE_HOOK
+from model_engine_server.common.constants import (
+    CALLBACK_POST_INFERENCE_HOOK,
+    LOGGING_POST_INFERENCE_HOOK,
+)
 from model_engine_server.common.dtos.tasks import EndpointPredictV1Request
+from model_engine_server.core.config import infra_config
 from model_engine_server.core.loggers import logger_name, make_logger
 from model_engine_server.domain.entities import CallbackAuth, CallbackBasicAuth
+from model_engine_server.domain.entities.model_endpoint_entity import ModelEndpointType
+from model_engine_server.domain.exceptions import StreamPutException
 from model_engine_server.inference.domain.gateways.inference_monitoring_metrics_gateway import (
     InferenceMonitoringMetricsGateway,
 )
+from model_engine_server.inference.domain.gateways.streaming_storage_gateway import (
+    StreamingStorageGateway,
+)
 from tenacity import Retrying, stop_after_attempt, wait_exponential
 
 logger = make_logger(logger_name())
@@ -76,6 +85,61 @@ def handle(
                 assert 200 <= res.status_code < 300
 
 
+class LoggingHook(PostInferenceHook):
+    def __init__(
+        self,
+        endpoint_name: str,
+        bundle_name: str,
+        user_id: str,
+        endpoint_id: Optional[str],
+        endpoint_type: Optional[ModelEndpointType],
+        bundle_id: Optional[str],
+        labels: Optional[Dict[str, str]],
+        streaming_storage_gateway: StreamingStorageGateway,
+    ):
+        super().__init__(endpoint_name, bundle_name, user_id)
+        self._endpoint_id = endpoint_id
+        self._endpoint_type = endpoint_type
+        self._bundle_id = bundle_id
+        self._labels = labels
+        self._streaming_storage_gateway = streaming_storage_gateway
+
+    def handle(
+        self,
+        request_payload: EndpointPredictV1Request,
+        response: Dict[str, Any],
+        task_id: Optional[str],
+    ):
+        if (
+            not self._endpoint_id
+            or not self._endpoint_type
+            or not self._bundle_id
+            or not self._labels
+        ):
+            logger.warning(
+                "No endpoint_id, endpoint_type, bundle_id, or labels specified for request."
+            )
+            return
+        response["task_id"] = task_id
+        data_record = {
+            "REQUEST_BODY": request_payload.json(),
+            "RESPONSE_BODY": response,
+            "ENDPOINT_ID": self._endpoint_id,
+            "ENDPOINT_NAME": self._endpoint_name,
+            "ENDPOINT_TYPE": self._endpoint_type.value,
+            "BUNDLE_ID": self._bundle_id,
+            "LABELS": self._labels,
+        }
+        stream_name = infra_config().firehose_stream_name
+        if stream_name is None:
+            logger.warning("No firehose stream name specified. Logging hook will not be executed.")
+            return
+        try:
+            self._streaming_storage_gateway.put_record(stream_name=stream_name, record=data_record)
+        except StreamPutException as e:
+            logger.error(f"Error in logging hook {e}")
+
+
 class PostInferenceHooksHandler:
     def __init__(
         self,
@@ -88,6 +152,11 @@ def __init__(
         default_callback_auth: Optional[CallbackAuth],
         post_inference_hooks: Optional[List[str]],
         monitoring_metrics_gateway: InferenceMonitoringMetricsGateway,
+        endpoint_id: Optional[str],
+        endpoint_type: Optional[ModelEndpointType],
+        bundle_id: Optional[str],
+        labels: Optional[Dict[str, str]],
+        streaming_storage_gateway: StreamingStorageGateway,
     ):
         self._monitoring_metrics_gateway = monitoring_metrics_gateway
         self._hooks: Dict[str, PostInferenceHook] = {}
@@ -104,6 +173,17 @@ def __init__(
                         default_callback_url,
                         default_callback_auth,
                     )
+                elif hook_lower == LOGGING_POST_INFERENCE_HOOK:
+                    self._hooks[hook_lower] = LoggingHook(
+                        endpoint_name,
+                        bundle_name,
+                        user_id,
+                        endpoint_id,
+                        endpoint_type,
+                        bundle_id,
+                        labels,
+                        streaming_storage_gateway,
+                    )
                 else:
                     raise ValueError(f"Hook {hook_lower} is currently not supported.")