[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

@@ -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

.gitignore

@@ -21,7 +21,7 @@ todesktop.json
 khoj_assistant.egg-info
 /config/khoj*.yml
 .pytest_cache
-khoj.log
+*.log
 static
 
 # Obsidian plugin artifacts

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:

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"

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" ]

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"]
 

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

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)

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:

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:

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"))

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