diff --git a/api/docker/entrypoint.sh b/api/docker/entrypoint.sh index a53d84c6e9..9cf5c505d1 100755 --- a/api/docker/entrypoint.sh +++ b/api/docker/entrypoint.sh @@ -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 diff --git a/docker/.env.example b/docker/.env.example index fa0b0d6950..b4a55bbba5 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -124,10 +124,36 @@ GUNICORN_TIMEOUT=360 # The number of Celery workers. The default is 1, and can be set as needed. CELERY_WORKER_AMOUNT= +# Flag indicating whether to enable autoscaling of Celery workers. +# +# Autoscaling is useful when tasks are CPU intensive and can be dynamically +# allocated and deallocated based on the workload. +# +# When autoscaling is enabled, the maximum and minimum number of workers can +# be specified. The autoscaling algorithm will dynamically adjust the number +# of workers within the specified range. +# +# Default is false (i.e., autoscaling is disabled). +# +# Example: +# CELERY_AUTO_SCALE=true +CELERY_AUTO_SCALE=false + +# The maximum number of Celery workers that can be autoscaled. +# This is optional and only used when autoscaling is enabled. +# Default is not set. +CELERY_MAX_WORKERS= + +# The minimum number of Celery workers that can be autoscaled. +# This is optional and only used when autoscaling is enabled. +# Default is not set. +CELERY_MIN_WORKERS= + # API Tool configuration API_TOOL_DEFAULT_CONNECT_TIMEOUT=10 API_TOOL_DEFAULT_READ_TIMEOUT=60 + # ------------------------------ # Database Configuration # The database uses PostgreSQL. Please use the public schema. diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 6a3d61b309..de7afa8efb 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -22,6 +22,9 @@ x-shared-env: &shared-api-worker-env CELERY_WORKER_CLASS: ${CELERY_WORKER_CLASS:-} GUNICORN_TIMEOUT: ${GUNICORN_TIMEOUT:-360} CELERY_WORKER_AMOUNT: ${CELERY_WORKER_AMOUNT:-} + CELERY_AUTO_SCALE: ${CELERY_AUTO_SCALE:-false} + CELERY_MAX_WORKERS: ${CELERY_MAX_WORKERS:-} + CELERY_MIN_WORKERS: ${CELERY_MIN_WORKERS:-} API_TOOL_DEFAULT_CONNECT_TIMEOUT: ${API_TOOL_DEFAULT_CONNECT_TIMEOUT:-10} API_TOOL_DEFAULT_READ_TIMEOUT: ${API_TOOL_DEFAULT_READ_TIMEOUT:-60} DB_USERNAME: ${DB_USERNAME:-postgres}