jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: support Celery auto-scale (#6249)

Browse source 
Co-authored-by: takatost <takatost@gmail.com>
This commit is contained in:
Nam Vu 2024-07-31 13:34:44 +07:00 • committed by GitHub
commit 8eb0d0fddd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 43 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

15
api/docker/entrypoint.sh
View file

@ -8,8 +8,21 @@ if [[ "${MIGRATION_ENABLED}" == "true" ]]; then
fi
if [[ "${MODE}" == "worker" ]]; then
exec celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} -c ${CELERY_WORKER_AMOUNT:-1} --loglevel INFO \
# Get the number of available CPU cores
if [ "${CELERY_AUTO_SCALE,,}" = "true" ]; then
# Set MAX_WORKERS to the number of available cores if not specified
AVAILABLE_CORES=$(nproc)
MAX_WORKERS=${CELERY_MAX_WORKERS:-$AVAILABLE_CORES}
MIN_WORKERS=${CELERY_MIN_WORKERS:-1}
CONCURRENCY_OPTION="--autoscale=${MAX_WORKERS},${MIN_WORKERS}"
else
CONCURRENCY_OPTION="-c ${CELERY_WORKER_AMOUNT:-1}"
fi
exec celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} $CONCURRENCY_OPTION --loglevel INFO \
-Q ${CELERY_QUEUES:-dataset,generation,mail,ops_trace,app_deletion}
elif [[ "${MODE}" == "beat" ]]; then
exec celery -A app.celery beat --loglevel INFO
else