Metrics for stuck async requests (#471)

Author: squeakymouse

model-engine/model_engine_server/core/celery/__init__.py

@@ -1,11 +1,19 @@
 from typing import Sequence
 
-from .app import TaskVisibility, celery_app, get_all_db_indexes, get_redis_host_port, inspect_app
+from .app import (
+    DEFAULT_TASK_VISIBILITY_SECONDS,
+    TaskVisibility,
+    celery_app,
+    get_all_db_indexes,
+    get_redis_host_port,
+    inspect_app,
+)
 
 __all__: Sequence[str] = (
     "celery_app",
     "get_all_db_indexes",
     "get_redis_host_port",
     "inspect_app",
     "TaskVisibility",
+    "DEFAULT_TASK_VISIBILITY_SECONDS",
 )

model-engine/model_engine_server/core/celery/app.py

@@ -31,6 +31,9 @@
 ] = "model_engine_server.core.celery.abs:AzureBlockBlobBackend"
 
 
+DEFAULT_TASK_VISIBILITY_SECONDS = 86400
+
+
 @unique
 class TaskVisibility(IntEnum):
     """

model-engine/model_engine_server/domain/gateways/monitoring_metrics_gateway.py

@@ -24,14 +24,12 @@ class MonitoringMetricsGateway(ABC):
     def emit_attempted_build_metric(self):
         """
         Service builder attempted metric
-
         """
 
     @abstractmethod
     def emit_successful_build_metric(self):
         """
         Service builder succeeded metric
-
         """
 
     @abstractmethod
@@ -44,49 +42,42 @@ def emit_build_time_metric(self, duration_seconds: float):
     def emit_image_build_cache_hit_metric(self, image_type: str):
         """
         Service builder image build cache hit metric
-
         """
 
     @abstractmethod
     def emit_image_build_cache_miss_metric(self, image_type: str):
         """
         Service builder image build cache miss metric
-
         """
 
     @abstractmethod
     def emit_docker_failed_build_metric(self):
         """
         Service builder docker build failed metric
-
         """
 
     @abstractmethod
     def emit_database_cache_hit_metric(self):
         """
         Successful database cache metric
-
         """
 
     @abstractmethod
     def emit_database_cache_miss_metric(self):
         """
         Missed database cache metric
-
         """
 
     @abstractmethod
     def emit_route_call_metric(self, route: str, metadata: MetricMetadata):
         """
         Route call metric
-
         """
         pass
 
     @abstractmethod
     def emit_token_count_metrics(self, token_usage: TokenUsage, metadata: MetricMetadata):
         """
         Token count metrics
-
         """
         pass

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

@@ -30,3 +30,23 @@ def emit_successful_post_inference_hook(self, hook: str):
         Args:
             hook: The name of the hook
         """
+
+    @abstractmethod
+    def emit_async_task_received_metric(self, queue_name: str):
+        """
+        Async task received metric
+
+        Args:
+            queue_name: The name of the Celery queue
+        """
+        pass
+
+    @abstractmethod
+    def emit_async_task_stuck_metric(self, queue_name: str):
+        """
+        Async task stuck metric
+
+        Args:
+            queue_name: The name of the Celery queue
+        """
+        pass

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

@@ -1,12 +1,17 @@
 import argparse
 import json
+from datetime import datetime, timedelta
 from typing import Any, Dict, Optional, TypedDict, Union
 
 from celery import Celery, Task, states
 from model_engine_server.common.constants import DEFAULT_CELERY_TASK_NAME, LIRA_CELERY_TASK_NAME
 from model_engine_server.common.dtos.model_endpoints import BrokerType
 from model_engine_server.common.dtos.tasks import EndpointPredictV1Request
-from model_engine_server.core.celery import TaskVisibility, celery_app
+from model_engine_server.core.celery import (
+    DEFAULT_TASK_VISIBILITY_SECONDS,
+    TaskVisibility,
+    celery_app,
+)
 from model_engine_server.core.config import infra_config
 from model_engine_server.core.loggers import logger_name, make_logger
 from model_engine_server.core.utils.format import format_stacktrace
@@ -15,6 +20,9 @@
     LoadForwarder,
     load_named_config,
 )
+from model_engine_server.inference.infra.gateways.datadog_inference_monitoring_metrics_gateway import (
+    DatadogInferenceMonitoringMetricsGateway,
+)
 
 logger = make_logger(logger_name())
 
@@ -68,6 +76,8 @@ def create_celery_service(
         backend_protocol=backend_protocol,
     )
 
+    monitoring_metrics_gateway = DatadogInferenceMonitoringMetricsGateway()
+
     class ErrorHandlingTask(Task):
         """Sets a 'custom' field with error in the Task response for FAILURE.
 
@@ -112,13 +122,18 @@ def after_return(
     # See documentation for options:
     # https://docs.celeryproject.org/en/stable/userguide/tasks.html#list-of-options
     @app.task(base=ErrorHandlingTask, name=LIRA_CELERY_TASK_NAME, track_started=True)
-    def exec_func(payload, *ignored_args, **ignored_kwargs):
+    def exec_func(payload, arrival_timestamp, *ignored_args, **ignored_kwargs):
         if len(ignored_args) > 0:
             logger.warning(f"Ignoring {len(ignored_args)} positional arguments: {ignored_args=}")
         if len(ignored_kwargs) > 0:
             logger.warning(f"Ignoring {len(ignored_kwargs)} keyword arguments: {ignored_kwargs=}")
         try:
-            return forwarder(payload)
+            monitoring_metrics_gateway.emit_async_task_received_metric(queue_name)
+            result = forwarder(payload)
+            request_duration = datetime.now() - arrival_timestamp
+            if request_duration > timedelta(seconds=DEFAULT_TASK_VISIBILITY_SECONDS):
+                monitoring_metrics_gateway.emit_async_task_stuck_metric(queue_name)
+            return result
         except Exception:
             logger.exception("Celery service failed to respond to request.")
             raise
@@ -131,8 +146,8 @@ def exec_func(payload, *ignored_args, **ignored_kwargs):
         name=DEFAULT_CELERY_TASK_NAME,
         track_started=True,
     )
-    def exec_func_pre_lira(payload, *ignored_args, **ignored_kwargs):
-        return exec_func(payload, *ignored_args, **ignored_kwargs)
+    def exec_func_pre_lira(payload, arrival_timestamp, *ignored_args, **ignored_kwargs):
+        return exec_func(payload, arrival_timestamp, *ignored_args, **ignored_kwargs)
 
     return app
 

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

@@ -10,3 +10,11 @@ def emit_attempted_post_inference_hook(self, hook: str):
 
     def emit_successful_post_inference_hook(self, hook: str):
         statsd.increment(f"scale_launch.post_inference_hook.{hook}.success")
+
+    def emit_async_task_received_metric(self, queue_name: str):
+        statsd.increment(
+            "scale_launch.async_task.received.count", tags=[f"queue_name:{queue_name}"]
+        )  # pragma: no cover
+
+    def emit_async_task_stuck_metric(self, queue_name: str):
+        statsd.increment("scale_launch.async_task.stuck.count", tags=[f"queue_name:{queue_name}"])

model-engine/model_engine_server/infra/gateways/live_async_model_endpoint_inference_gateway.py

@@ -1,4 +1,5 @@
 import json
+from datetime import datetime
 
 from model_engine_server.common.constants import DEFAULT_CELERY_TASK_NAME
 from model_engine_server.common.dtos.tasks import (
@@ -37,7 +38,7 @@ def create_task(
         send_task_response = self.task_queue_gateway.send_task(
             task_name=task_name,
             queue_name=topic,
-            args=[predict_args, predict_request.return_pickled],
+            args=[predict_args, datetime.now(), predict_request.return_pickled],
             expires=task_timeout_seconds,
         )
         return CreateAsyncTaskV1Response(task_id=send_task_response.task_id)

model-engine/tests/unit/infra/gateways/test_live_async_model_inference_gateway.py

@@ -1,4 +1,5 @@
 import json
+from datetime import datetime, timedelta
 from typing import Any
 
 import pytest
@@ -22,10 +23,11 @@ def test_task_create_get_url(
     task_id = create_response.task_id
     task_queue_gateway: Any = fake_live_async_model_inference_gateway.task_queue_gateway
     assert len(task_queue_gateway.queue) == 1
-    assert task_queue_gateway.queue[task_id]["args"] == [
-        endpoint_predict_request_1[0].dict(),
-        endpoint_predict_request_1[0].return_pickled,
-    ]
+    assert task_queue_gateway.queue[task_id]["args"][0] == endpoint_predict_request_1[0].dict()
+    assert (datetime.now() - task_queue_gateway.queue[task_id]["args"][1]) < timedelta(seconds=1)
+    assert (
+        task_queue_gateway.queue[task_id]["args"][2] == endpoint_predict_request_1[0].return_pickled
+    )
 
     get_response_1 = fake_live_async_model_inference_gateway.get_task(task_id)
     assert get_response_1 == GetAsyncTaskV1Response(task_id=task_id, status=TaskStatus.PENDING)
@@ -49,17 +51,18 @@ def test_task_create_get_args_callback(
     task_id = create_response.task_id
     task_queue_gateway: Any = fake_live_async_model_inference_gateway.task_queue_gateway
     assert len(task_queue_gateway.queue) == 1
-    assert task_queue_gateway.queue[task_id]["args"] == [
-        {
-            "args": endpoint_predict_request_2[0].args.__root__,
-            "url": None,
-            "cloudpickle": None,
-            "callback_auth": json.loads(endpoint_predict_request_2[0].callback_auth.json()),
-            "callback_url": endpoint_predict_request_2[0].callback_url,
-            "return_pickled": endpoint_predict_request_2[0].return_pickled,
-        },
-        endpoint_predict_request_2[0].return_pickled,
-    ]
+    assert task_queue_gateway.queue[task_id]["args"][0] == {
+        "args": endpoint_predict_request_2[0].args.__root__,
+        "url": None,
+        "cloudpickle": None,
+        "callback_auth": json.loads(endpoint_predict_request_2[0].callback_auth.json()),
+        "callback_url": endpoint_predict_request_2[0].callback_url,
+        "return_pickled": endpoint_predict_request_2[0].return_pickled,
+    }
+    assert (datetime.now() - task_queue_gateway.queue[task_id]["args"][1]) < timedelta(seconds=1)
+    assert (
+        task_queue_gateway.queue[task_id]["args"][2] == endpoint_predict_request_2[0].return_pickled
+    )
 
     get_response_1 = fake_live_async_model_inference_gateway.get_task(task_id)
     assert get_response_1 == GetAsyncTaskV1Response(task_id=task_id, status=TaskStatus.PENDING)