Files
privateGPT/scripts/worker_entrypoint
Javier Martinez 183cd03857 feat!: PrivateGPT revamp v1 (#2230)
* feat!: PrivateGPT revamp v1

* chore(docs): update nodejs
2026-06-02 16:55:46 +02:00

105 lines
3.4 KiB
Bash
Executable File

#!/bin/bash
: "${PGPT_WORKER_MODE:=mixed}" # mixed, worker, flower
: "${PGPT_FLOWER_PORT:=5555}"
: "${PGPT_FLOWER_PASSWORD:=flower}"
: "${PGPT_FLOWER_URL_PREFIX:=}"
: "${PGPT_FLOWER_PERSISTENT:=true}"
: "${PGPT_FLOWER_PERSISTENT_DB:=celery_flower.db}"
: "${PGPT_FLOWER_MAXIMUM_TASKS:=1000}"
: "${PGPT_CELERY_LOG_LEVEL:=info}"
: "${PGPT_CELERY_QUEUES:=}"
: "${PGPT_CELERY_HOSTNAME:=default}"
: "${PGPT_CELERY_POOL:=prefork}"
: "${PGPT_CELERY_CONCURRENCY:=}"
: "${PGPT_CELERY_TIME_LIMIT:=}"
: "${PGPT_WORKER_APP_MODULE:=private_gpt}"
: "${API_ENABLED:=true}"
: "${API_PORT:=8090}"
CELERY_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery"
HEALTHCHECK_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery.healthcheck:app"
# Create directory for PID files
mkdir -p /tmp/private_gpt_pids
# Add signal handling
cleanup() {
echo "Shutting down services..."
jobs -p | xargs -r kill
exit 0
}
trap cleanup EXIT INT TERM
# Run flower on the specified port
if [[ "$PGPT_WORKER_MODE" == "flower" || "$PGPT_WORKER_MODE" == "mixed" ]]; then
echo "Starting flower on port $PGPT_FLOWER_PORT"
FLOWER_ARGS=()
if [[ -n "$PGPT_FLOWER_URL_PREFIX" ]]; then
FLOWER_ARGS+=("--url_prefix=$PGPT_FLOWER_URL_PREFIX")
fi
if [[ -n "$PGPT_FLOWER_PORT" ]]; then
FLOWER_ARGS+=("--port=$PGPT_FLOWER_PORT")
fi
if [[ -n "$PGPT_FLOWER_PASSWORD" ]]; then
FLOWER_ARGS+=("--basic_auth=flower:$PGPT_FLOWER_PASSWORD")
fi
if [[ "$PGPT_WORKER_MODE" == "mixed" ]]; then
FLOWER_ARGS+=("--logging=none")
fi
if [[ "$PGPT_FLOWER_PERSISTENT" == "true" ]]; then
FLOWER_ARGS+=("--persistent=True")
if [[ -n "$PGPT_FLOWER_PERSISTENT_DB" ]]; then
FLOWER_ARGS+=("--db=$PGPT_FLOWER_PERSISTENT_DB")
fi
if [[ -n "$PGPT_FLOWER_MAXIMUM_TASKS" ]]; then
FLOWER_ARGS+=("--max-tasks=$PGPT_FLOWER_MAXIMUM_TASKS")
fi
fi
# print all args except password
args=("${FLOWER_ARGS[@]}")
args=("${args[@]//:$PGPT_FLOWER_PASSWORD/}")
echo "Starting flower with args: ${args[*]}"
python -m celery --app "$CELERY_APP_MODULE" flower "${FLOWER_ARGS[@]}" &
fi
# Run celery worker using the config
if [[ "$PGPT_WORKER_MODE" == "worker" || "$PGPT_WORKER_MODE" == "mixed" ]]; then
WORKER_ARGS=()
if [[ -n "$PGPT_CELERY_LOG_LEVEL" ]]; then
WORKER_ARGS+=("--loglevel=$PGPT_CELERY_LOG_LEVEL")
fi
if [[ -n "$PGPT_CELERY_QUEUES" ]]; then
WORKER_ARGS+=("--queues=$PGPT_CELERY_QUEUES")
if [[ -n "$PGPT_CELERY_HOSTNAME" ]]; then
PGPT_CELERY_HOSTNAME="${PGPT_CELERY_QUEUES}"
fi
fi
if [[ -n "$PGPT_CELERY_HOSTNAME" ]]; then
WORKER_ARGS+=("--hostname=${PGPT_CELERY_HOSTNAME}%h")
fi
if [[ -n "$PGPT_CELERY_POOL" ]]; then
WORKER_ARGS+=("--pool=$PGPT_CELERY_POOL")
fi
if [[ -n "$PGPT_CELERY_CONCURRENCY" ]]; then
WORKER_ARGS+=("--concurrency=$PGPT_CELERY_CONCURRENCY")
fi
if [[ -n "$PGPT_CELERY_TIME_LIMIT" ]]; then
WORKER_ARGS+=("--time-limit=$PGPT_CELERY_TIME_LIMIT")
fi
echo "Starting celery worker with args: ${WORKER_ARGS[*]}"
python -m celery --app "$CELERY_APP_MODULE" worker "${WORKER_ARGS[@]}" &
fi
# Start the Healthcheck server
if [[ "$API_ENABLED" == "true" ]]; then
echo "Starting healthcheck server on port $API_PORT"
python -m uvicorn "$HEALTHCHECK_APP_MODULE" --host 0.0.0.0 --port $API_PORT --no-access-log --log-level critical &
fi
# Wait for all background processes
wait