mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-16 17:00:12 +00:00
feat: resumable chat worker + tool worker + async tokenizer (#2298)
* fix: avoid to block the loop * fix: blocks in expansion * fix: remove maximum concurrent users ... * fix: multiplexer * fix: readers * fix: more fixes ... * fix: impl * feat: tool scheduler * feat: add adaptative * feat: add chat worker * fix: max * feat: add chat/tools workers * fix: mypy * feat: add generic scheduler * fix: get result * feat: do serializable the tool executor * fix: tools * fix: config * fix: config * fix: args * fix: config * fix: serializer * Revert "fix: blocks in expansion" This reverts commita2110f94a8. * fix: unify all logic * feat: add ingestion scheduler * fix: settings * fix: config * feat: add arq worker to chat * fix: arq worker * fix: add nest * fix: mypy * fix: await * fix: script stress * fix: tokenizer * fix: chat scheduler * fix: mypy * fix: add async tokenizer * fix: improve condense * fix: tool scheduler * feat: add initial real async chat worker * fix: mypy * fix: do resumable local executor ... ... ... fix: revert usleess changes fix: remove parent chat job fix: refactor fix: loop ref: rename models fix: chat engine fix: mypy ... ... ... fix: fix deps * fix: tests * fix: tests * ... * fix: stream * fix: config * fix: scheduler * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Handle PGPT_WORKER_MODE=celery in health check worker status * fix: cancel * fix: arch * fix: test ingestion * fix: deserialization of chat messages * fix: broken results * fix: mypy * fix: test * fix: config * fix: remove arq tool worker * fix: output cls * fix: preserve early resumable tool callbacks * fix: preserve async tool result order * refactor: address worker PR review comments * fix: mypy * test: colocate ARQ chat enqueue coverage * fix: remove redis from tests * test: isolate chat mocks and cancellation timing * fix: tests (cherry picked from commit218b599c66) # Conflicts: # tests/server/chat/anthropic/test_anthropic_client.py # tests/server/chat/anthropic/test_langchain_anthropic.py # tests/server/chat/test_chat_knowledge_revamp.py # tests/server/chat/test_chat_routes.py # tests/server/chat/test_chat_routes_skills_integration.py * fix: tests (cherry picked from commitfc5ec0f72a) * fix: ruff * fix: test * fix: worker config (cherry picked from commit1371c275a1) * fix: principal * test: remove flaky chat cancellation assertion --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
This commit is contained in:
Javier Martinez
GitHub
@@ -1,104 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
: "${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_MODE:?PGPT_WORKER_MODE is required}"
|
||||
: "${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
|
||||
exec python -m "${PGPT_WORKER_APP_MODULE}.worker" "$@"
|
||||
|
||||
Reference in New Issue
Block a user