Skip to content

Update celery forwarder to use greenlets instead of processes #689

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Update celery forwarder to use greenlets instead of processes #689

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1f06017
Update celery forwarder to use greenlets instead of processes
dmchoiboi Mar 1, 2025
f491aae
Add imports
dmchoiboi Mar 4, 2025
2fb2264
fix gevent import order
dmchoiboi Mar 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 6 additions & 8 deletions model-engine/model_engine_server/inference/forwarding/celery_forwarder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from typing import Any, Dict, Optional, TypedDict, Union

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
Expand All @@ -23,7 +24,9 @@
from model_engine_server.inference.infra.gateways.datadog_inference_monitoring_metrics_gateway import (
DatadogInferenceMonitoringMetricsGateway,
)
from requests import ConnectionError
from request import ConnectionError

monkey.patch_all()

logger = make_logger(logger_name())

Expand Down Expand Up @@ -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 = 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)
Expand Down Expand Up @@ -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()

Expand Down