From 5f3f6b7c61532bbc7456022fa2442ae16d8998e3 Mon Sep 17 00:00:00 2001 From: sabaimran <65192171+sabaimran@users.noreply.github.com> Date: Thu, 26 Oct 2023 13:15:31 -0700 Subject: [PATCH] [Multi-User Part 5]: Add a production Docker file and use a gunicorn configuration with it (#514) - Add a productionized setup for the Khoj server using `gunicorn` with multiple workers for handling requests - Add a new Dockerfile meant for production config at `ghcr.io/khoj-ai/khoj:prod`; the existing Docker config should remain the same --- .github/workflows/dockerize_production.yml | 48 ++++++++++++++++++++++ .gitignore | 2 +- docker-compose.yml | 4 ++ gunicorn-config.py | 10 +++++ prod.Dockerfile | 30 ++++++++++++++ pyproject.toml | 3 +- src/khoj/configure.py | 3 ++ src/khoj/main.py | 9 +++- src/khoj/routers/auth.py | 5 ++- src/khoj/utils/cli.py | 9 +++- tests/test_cli.py | 4 +- tests/test_text_search.py | 1 - 12 files changed, 117 insertions(+), 11 deletions(-) create mode 100644 .github/workflows/dockerize_production.yml create mode 100644 gunicorn-config.py create mode 100644 prod.Dockerfile diff --git a/.github/workflows/dockerize_production.yml b/.github/workflows/dockerize_production.yml new file mode 100644 index 00000000..97fc876d --- /dev/null +++ b/.github/workflows/dockerize_production.yml @@ -0,0 +1,48 @@ +name: dockerize-prod + +on: + pull_request: + push: + tags: + - "*" + branches: + - master + paths: + - src/khoj/** + - config/** + - pyproject.toml + - prod.Dockerfile + - .github/workflows/dockerize_production.yml + workflow_dispatch: + +env: + DOCKER_IMAGE_TAG: 'prod' + +jobs: + build: + name: Build Production Docker Image, Push to Container Registry + runs-on: ubuntu-latest + steps: + - name: Checkout Code + uses: actions/checkout@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v2 + + - name: Login to GitHub Container Registry + uses: docker/login-action@v2 + with: + registry: ghcr.io + username: ${{ github.repository_owner }} + password: ${{ secrets.PAT }} + + - name: 📦 Build and Push Docker Image + uses: docker/build-push-action@v2 + with: + context: . + file: prod.Dockerfile + platforms: linux/amd64 + push: true + tags: ghcr.io/${{ github.repository }}:${{ env.DOCKER_IMAGE_TAG }} + build-args: | + PORT=42110 diff --git a/.gitignore b/.gitignore index e3e93428..35315263 100644 --- a/.gitignore +++ b/.gitignore @@ -21,7 +21,7 @@ todesktop.json khoj_assistant.egg-info /config/khoj*.yml .pytest_cache -khoj.log +*.log static # Obsidian plugin artifacts diff --git a/docker-compose.yml b/docker-compose.yml index d6048916..c75aa4fc 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -39,6 +39,7 @@ services: - ./tests/data/embeddings/:/root/.khoj/content/ - ./tests/data/models/:/root/.khoj/search/ - khoj_config:/root/.khoj/ + - khoj_models:/root/.cache/torch/sentence_transformers # Use 0.0.0.0 to explicitly set the host ip for the service on the container. https://pythonspeed.com/articles/docker-connection-refused/ environment: - POSTGRES_DB=postgres @@ -46,9 +47,12 @@ services: - POSTGRES_PASSWORD=postgres - POSTGRES_HOST=database - POSTGRES_PORT=5432 + - GOOGLE_CLIENT_SECRET=bar + - GOOGLE_CLIENT_ID=foo command: --host="0.0.0.0" --port=42110 -vv volumes: khoj_config: khoj_db: + khoj_models: diff --git a/gunicorn-config.py b/gunicorn-config.py new file mode 100644 index 00000000..1760ae38 --- /dev/null +++ b/gunicorn-config.py @@ -0,0 +1,10 @@ +import multiprocessing + +bind = "0.0.0.0:42110" +workers = 4 +worker_class = "uvicorn.workers.UvicornWorker" +timeout = 120 +keep_alive = 60 +accesslog = "access.log" +errorlog = "error.log" +loglevel = "debug" diff --git a/prod.Dockerfile b/prod.Dockerfile new file mode 100644 index 00000000..3cf6a600 --- /dev/null +++ b/prod.Dockerfile @@ -0,0 +1,30 @@ +# Use Nvidia's latest Ubuntu 22.04 image as the base image +FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 + +LABEL org.opencontainers.image.source https://github.com/khoj-ai/khoj + +# Install System Dependencies +RUN apt update -y && apt -y install python3-pip git + +WORKDIR /app + +# Install Application +COPY pyproject.toml . +COPY README.md . +RUN sed -i 's/dynamic = \["version"\]/version = "0.0.0"/' pyproject.toml && \ + TMPDIR=/home/cache/ pip install --cache-dir=/home/cache/ -e . + +# Copy Source Code +COPY . . + +RUN apt install vim -y + +# Set the PYTHONPATH environment variable in order for it to find the Django app. +ENV PYTHONPATH=/app/src:$PYTHONPATH + +# Run the Application +# There are more arguments required for the application to run, +# but these should be passed in through the docker-compose.yml file. +ARG PORT +EXPOSE ${PORT} +ENTRYPOINT [ "gunicorn", "-c", "gunicorn-config.py", "src.khoj.main:app" ] diff --git a/pyproject.toml b/pyproject.toml index f9ef020c..d5b7f0ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ dependencies = [ "schedule == 1.1.0", "sentence-transformers == 2.2.2", "transformers >= 4.28.0", - "torch >= 2.0.1", + "torch == 2.0.1", "uvicorn == 0.17.6", "aiohttp == 3.8.5", "langchain >= 0.0.187", @@ -70,6 +70,7 @@ dependencies = [ "psycopg2-binary == 2.9.9", "google-auth == 2.23.3", "python-multipart == 0.0.6", + "gunicorn == 21.2.0", ] dynamic = ["version"] diff --git a/src/khoj/configure.py b/src/khoj/configure.py index 56585328..67ca3543 100644 --- a/src/khoj/configure.py +++ b/src/khoj/configure.py @@ -103,6 +103,9 @@ def configure_server( user: KhojUser = None, ): # Update Config + if config == None: + logger.info(f"🚨 Khoj is not configured.\nInitializing it with a default config.") + config = FullConfig() state.config = config # Initialize Search Models from Config and initialize content diff --git a/src/khoj/main.py b/src/khoj/main.py index 8fe40e76..d434c461 100644 --- a/src/khoj/main.py +++ b/src/khoj/main.py @@ -65,7 +65,7 @@ logging.basicConfig(handlers=[rich_handler]) logger = logging.getLogger("khoj") -def run(): +def run(should_start_server=True): # Turn Tokenizers Parallelism Off. App does not support it. os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -107,7 +107,10 @@ def run(): configure_middleware(app) initialize_server(args.config) - start_server(app, host=args.host, port=args.port, socket=args.socket) + + # If the server is started through gunicorn (external to the script), don't start the server + if should_start_server: + start_server(app, host=args.host, port=args.port, socket=args.socket) def set_state(args): @@ -139,3 +142,5 @@ def poll_task_scheduler(): if __name__ == "__main__": run() +else: + run(should_start_server=False) diff --git a/src/khoj/routers/auth.py b/src/khoj/routers/auth.py index ab1964b8..5c375bd0 100644 --- a/src/khoj/routers/auth.py +++ b/src/khoj/routers/auth.py @@ -36,16 +36,17 @@ else: @auth_router.get("/login") async def login_get(request: Request): - redirect_uri = request.url_for("auth") + redirect_uri = str(request.app.url_path_for("auth")) return await oauth.google.authorize_redirect(request, redirect_uri) @auth_router.post("/login") async def login(request: Request): - redirect_uri = request.url_for("auth") + redirect_uri = str(request.app.url_path_for("auth")) return await oauth.google.authorize_redirect(request, redirect_uri) +@auth_router.post("/redirect") @auth_router.post("/token") @requires(["authenticated"], redirect="login_page") async def generate_token(request: Request, token_name: Optional[str] = None) -> str: diff --git a/src/khoj/utils/cli.py b/src/khoj/utils/cli.py index c72320a1..5090e399 100644 --- a/src/khoj/utils/cli.py +++ b/src/khoj/utils/cli.py @@ -3,6 +3,9 @@ import argparse import pathlib from importlib.metadata import version import os +import logging + +logger = logging.getLogger(__name__) # Internal Packages from khoj.utils.helpers import resolve_absolute_path @@ -17,7 +20,7 @@ def cli(args=None): # Setup Argument Parser for the Commandline Interface parser = argparse.ArgumentParser(description="Start Khoj; An AI personal assistant for your Digital Brain") parser.add_argument( - "--config-file", "-c", default="~/.khoj/khoj.yml", type=pathlib.Path, help="YAML file to configure Khoj" + "--config-file", default="~/.khoj/khoj.yml", type=pathlib.Path, help="YAML file to configure Khoj" ) parser.add_argument( "--regenerate", @@ -42,7 +45,9 @@ def cli(args=None): help="Run Khoj in anonymous mode. This does not require any login for connecting users.", ) - args = parser.parse_args(args) + args, remaining_args = parser.parse_known_args(args) + + logger.debug(f"Ignoring unknown commandline args: {remaining_args}") args.version_no = version("khoj-assistant") if args.version: diff --git a/tests/test_cli.py b/tests/test_cli.py index cff2a7f3..e3daa2c6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -25,7 +25,7 @@ def test_cli_invalid_config_file_path(): non_existent_config_file = f"non-existent-khoj-{random()}.yml" # Act - actual_args = cli([f"-c={non_existent_config_file}"]) + actual_args = cli([f"--config-file={non_existent_config_file}"]) # Assert assert actual_args.config_file == resolve_absolute_path(non_existent_config_file) @@ -35,7 +35,7 @@ def test_cli_invalid_config_file_path(): # ---------------------------------------------------------------------------------------------------- def test_cli_config_from_file(): # Act - actual_args = cli(["-c=tests/data/config.yml", "--regenerate", "-vvv"]) + actual_args = cli(["--config-file=tests/data/config.yml", "--regenerate", "-vvv"]) # Assert assert actual_args.config_file == resolve_absolute_path(Path("tests/data/config.yml")) diff --git a/tests/test_text_search.py b/tests/test_text_search.py index ec8034ef..aeeaa85f 100644 --- a/tests/test_text_search.py +++ b/tests/test_text_search.py @@ -17,7 +17,6 @@ from khoj.utils.fs_syncer import collect_files, get_org_files from database.models import LocalOrgConfig, KhojUser, Embeddings, GithubConfig logger = logging.getLogger(__name__) -from khoj.utils.rawconfig import ContentConfig, SearchConfig # Test