fix cases of synchronous code blocking event loop (#755)

Author: olliestanley

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

@@ -17,7 +17,7 @@ class AsyncModelEndpointInferenceGateway(ABC):
     """
 
     @abstractmethod
-    def create_task(
+    async def create_task(
         self,
         topic: str,
         predict_request: EndpointPredictV1Request,

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

@@ -1,4 +1,6 @@
 # This is the abstract class defining putting and retrieving tasks into a queue.
+import asyncio
+import functools
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
 
@@ -32,6 +34,36 @@ def send_task(
         Returns: The unique identifier for the task.
         """
 
+    async def send_task_async(
+        self,
+        task_name: str,
+        queue_name: str,
+        args: Optional[List[Any]] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+        expires: Optional[int] = None,
+    ) -> CreateAsyncTaskV1Response:
+        """
+        Non-blocking version of send_task that runs in a thread executor
+        to avoid blocking the event loop.
+
+        Note: This is a workaround for Celery's synchronous API. Ideally the
+        gateway interface (including get_task) would be natively async, but
+        Celery lacks first-class asyncio support, so we use run_in_executor
+        as a pragmatic bridge.
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            functools.partial(
+                self.send_task,
+                task_name=task_name,
+                queue_name=queue_name,
+                args=args,
+                kwargs=kwargs,
+                expires=expires,
+            ),
+        )
+
     @abstractmethod
     def get_task(self, task_id: str) -> GetAsyncTaskV1Response:
         """

model-engine/model_engine_server/domain/use_cases/async_inference_use_cases.py

@@ -66,7 +66,7 @@ async def execute(
         task_name = model_endpoint.record.current_model_bundle.celery_task_name()
 
         inference_gateway = self.model_endpoint_service.get_async_model_endpoint_inference_gateway()
-        return inference_gateway.create_task(
+        return await inference_gateway.create_task(
             topic=model_endpoint.record.destination,
             predict_request=request,
             task_timeout_seconds=DEFAULT_TASK_TIMEOUT_SECONDS,

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

@@ -23,7 +23,7 @@ class LiveAsyncModelEndpointInferenceGateway(AsyncModelEndpointInferenceGateway)
     def __init__(self, task_queue_gateway: TaskQueueGateway):
         self.task_queue_gateway = task_queue_gateway
 
-    def create_task(
+    async def create_task(
         self,
         topic: str,
         predict_request: EndpointPredictV1Request,
@@ -35,7 +35,7 @@ def create_task(
         # key in some fields, and root overriding only reflects in the json() output.
         predict_args = json.loads(predict_request.json())
 
-        send_task_response = self.task_queue_gateway.send_task(
+        send_task_response = await self.task_queue_gateway.send_task_async(
             task_name=task_name,
             queue_name=topic,
             args=[predict_args, datetime.now(), predict_request.return_pickled],

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

@@ -51,7 +51,7 @@ def __init__(
         self.resource_gateway = resource_gateway
         self.task_queue_gateway = task_queue_gateway
 
-    def create_model_endpoint_infra(
+    async def create_model_endpoint_infra(
         self,
         *,
         model_endpoint_record: ModelEndpointRecord,
@@ -105,7 +105,7 @@ def create_model_endpoint_infra(
             default_callback_url=default_callback_url,
             default_callback_auth=default_callback_auth,
         )
-        response = self.task_queue_gateway.send_task(
+        response = await self.task_queue_gateway.send_task_async(
             task_name=BUILD_TASK_NAME,
             queue_name=get_service_builder_queue(SERVICE_IDENTIFIER, SERVICE_BUILDER_QUEUE),
             # celery request is required to be JSON serializables
@@ -196,7 +196,6 @@ async def update_model_endpoint_infra(
             default_callback_url = endpoint_config.default_callback_url
         if default_callback_auth is None and endpoint_config is not None:
             default_callback_auth = endpoint_config.default_callback_auth
-
         aws_role = infra_state.aws_role
         results_s3_bucket = infra_state.results_s3_bucket
 
@@ -225,7 +224,7 @@ async def update_model_endpoint_infra(
             default_callback_url=default_callback_url,
             default_callback_auth=default_callback_auth,
         )
-        response = self.task_queue_gateway.send_task(
+        response = await self.task_queue_gateway.send_task_async(
             task_name=BUILD_TASK_NAME,
             queue_name=get_service_builder_queue(SERVICE_IDENTIFIER, SERVICE_BUILDER_QUEUE),
             kwargs=dict(build_endpoint_request_json=build_endpoint_request.dict()),

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

@@ -18,7 +18,7 @@ class ModelEndpointInfraGateway(ABC):
     """
 
     @abstractmethod
-    def create_model_endpoint_infra(
+    async def create_model_endpoint_infra(
         self,
         *,
         model_endpoint_record: ModelEndpointRecord,

model-engine/model_engine_server/infra/services/live_batch_job_orchestration_service.py

@@ -267,10 +267,10 @@ async def _read_or_submit_tasks(
     async def _submit_tasks(
         self, queue_name: str, input_path: str, task_name: str
     ) -> List[BatchEndpointInProgressTask]:
-        def _create_task(
+        async def _create_task(
             predict_request: BatchEndpointInferencePrediction,
         ) -> BatchEndpointInProgressTask:
-            response = self.async_model_endpoint_inference_gateway.create_task(
+            response = await self.async_model_endpoint_inference_gateway.create_task(
                 topic=queue_name,
                 predict_request=predict_request.request,
                 task_timeout_seconds=DEFAULT_TASK_TIMEOUT_SECONDS,
@@ -301,9 +301,8 @@ def _create_task(
                     BatchEndpointInferencePrediction(request=request, reference_id=reference_id)
                 )
 
-        executor = ThreadPoolExecutor()
-        task_ids = list(executor.map(_create_task, inputs))
-        return task_ids
+        task_ids = await asyncio.gather(*[_create_task(inp) for inp in inputs])
+        return list(task_ids)
 
     def _poll_tasks(
         self,

model-engine/model_engine_server/infra/services/live_model_endpoint_service.py

@@ -187,7 +187,7 @@ async def create_model_endpoint(
                 public_inference=public_inference,
             )
         )
-        creation_task_id = self.model_endpoint_infra_gateway.create_model_endpoint_infra(
+        creation_task_id = await self.model_endpoint_infra_gateway.create_model_endpoint_infra(
             model_endpoint_record=model_endpoint_record,
             min_workers=min_workers,
             max_workers=max_workers,

model-engine/tests/unit/conftest.py

@@ -1071,7 +1071,7 @@ def __init__(
     def _get_deployment_name(user_id: str, model_endpoint_name: str) -> str:
         return f"{user_id}-{model_endpoint_name}"
 
-    def create_model_endpoint_infra(
+    async def create_model_endpoint_infra(
         self,
         *,
         model_endpoint_record: ModelEndpointRecord,
@@ -1638,7 +1638,7 @@ class FakeAsyncModelEndpointInferenceGateway(AsyncModelEndpointInferenceGateway)
     def __init__(self):
         self.tasks = []
 
-    def create_task(
+    async def create_task(
         self,
         topic: str,
         predict_request: EndpointPredictV1Request,

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

@@ -13,11 +13,11 @@ def fake_live_async_model_inference_gateway(fake_task_queue_gateway):
 
 
 @pytest.mark.asyncio
-def test_task_create_get_url(
+async def test_task_create_get_url(
     fake_live_async_model_inference_gateway: LiveAsyncModelEndpointInferenceGateway,
     endpoint_predict_request_1,
 ):
-    create_response = fake_live_async_model_inference_gateway.create_task(
+    create_response = await fake_live_async_model_inference_gateway.create_task(
         "test_topic", endpoint_predict_request_1[0], 60
     )
     task_id = create_response.task_id
@@ -41,11 +41,11 @@ def test_task_create_get_url(
 
 
 @pytest.mark.asyncio
-def test_task_create_get_args_callback(
+async def test_task_create_get_args_callback(
     fake_live_async_model_inference_gateway: LiveAsyncModelEndpointInferenceGateway,
     endpoint_predict_request_2,
 ):
-    create_response = fake_live_async_model_inference_gateway.create_task(
+    create_response = await fake_live_async_model_inference_gateway.create_task(
         "test_topic", endpoint_predict_request_2[0], 60
     )
     task_id = create_response.task_id