Update celery forwarder to use greenlets instead of processes

Author: dmchoiboi

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

@@ -3,7 +3,9 @@
 from datetime import datetime, timedelta
 from typing import Any, Dict, Optional, TypedDict, Union
 
+from aiohttp import ClientConnectionError
 from celery import Celery, Task, states
+from gevent import monkey
 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
@@ -23,7 +25,8 @@
 from model_engine_server.inference.infra.gateways.datadog_inference_monitoring_metrics_gateway import (
     DatadogInferenceMonitoringMetricsGateway,
 )
-from requests import ConnectionError
+
+monkey.patch_all()
 
 logger = make_logger(logger_name())
 
@@ -132,9 +135,9 @@ def after_return(
         base=ErrorHandlingTask,
         name=LIRA_CELERY_TASK_NAME,
         track_started=True,
-        autoretry_for=(ConnectionError,),
+        autoretry_for=(ClientConnectionError,),
     )
-    def exec_func(payload, arrival_timestamp, *ignored_args, **ignored_kwargs):
+    async 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:
@@ -144,7 +147,7 @@ def exec_func(payload, arrival_timestamp, *ignored_args, **ignored_kwargs):
             # Don't fail the celery task even if there's a status code
             # (otherwise we can't really control what gets put in the result attribute)
             # in the task (https://docs.celeryq.dev/en/stable/reference/celery.result.html#celery.result.AsyncResult.status)
-            result = forwarder(payload)
+            result = await forwarder.forward(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)
@@ -177,12 +180,7 @@ def start_celery_service(
         concurrency=concurrency,
         loglevel="INFO",
         optimization="fair",
-        # Don't use pool="solo" so we can send multiple concurrent requests over
-        # Historically, pool="solo" argument fixes the known issues of celery and some of the libraries.
-        # Particularly asyncio and torchvision transformers. This isn't relevant since celery-forwarder
-        # is quite lightweight
-        # TODO: we should probably use eventlet or gevent for the pool, since
-        # the forwarder is nearly the most extreme example of IO bound.
+        pool="gevent",
     )
     worker.start()