feat: support Celery auto-scale

api/docker/entrypoint.sh

@@ -8,8 +8,16 @@ if [[ "${MIGRATION_ENABLED}" == "true" ]]; then
 fi
 
 if [[ "${MODE}" == "worker" ]]; then
-  celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} -c ${CELERY_WORKER_AMOUNT:-1} --loglevel INFO \
-    -Q ${CELERY_QUEUES:-dataset,generation,mail,ops_trace,app_deletion}
+  if [ "${CELERY_AUTO_SCALE,,}" = "true" ]; then
+    MAX_WORKERS=${CELERY_MAX_WORKERS:-4}
+    MIN_WORKERS=${CELERY_MIN_WORKERS:-1}
+    CONCURRENCY_OPTION="--autoscale=${MAX_WORKERS},${MIN_WORKERS}"
+  else
+    CONCURRENCY_OPTION="-c ${CELERY_WORKER_AMOUNT:-1}"
+  fi
+
+  celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} $CONCURRENCY_OPTION --loglevel INFO \
+    -Q ${CELERY_QUEUES:-dataset,generation,mail}
 elif [[ "${MODE}" == "beat" ]]; then
   celery -A app.celery beat --loglevel INFO
 else

docker/.env.example

@@ -124,6 +124,31 @@ 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=
+
 # ------------------------------
 # Database Configuration
 # The database uses PostgreSQL. Please use the public schema.

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:-}
   DB_USERNAME: ${DB_USERNAME:-postgres}
   DB_PASSWORD: ${DB_PASSWORD:-difyai123456}
   DB_HOST: ${DB_HOST:-db}