feat: support Celery auto-scale (langgenius#6249)

Co-authored-by: takatost <[email protected]>

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

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.

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}