From 2741835605275d11a58edfa5725cc379ca2a4676 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 5 Jan 2024 19:03:11 +0530 Subject: [PATCH] build(Dockerfile): move prisma build to dockerfile Seems to solve - https://github.com/BerriAI/litellm/issues/1321 --- .gitignore | 1 + Dockerfile | 17 +++++++++-- docker/.env.example | 4 +-- litellm/proxy/proxy_server.py | 53 ++++++++++++++++++++++++++++------- litellm/proxy/utils.py | 50 ++++++++++++++++++--------------- retry_push.sh | 28 ++++++++++++++++++ schema.prisma | 33 ++++++++++++++++++++++ 7 files changed, 149 insertions(+), 37 deletions(-) create mode 100644 retry_push.sh create mode 100644 schema.prisma diff --git a/.gitignore b/.gitignore index 29c296915..618e3d874 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,4 @@ proxy_server_config_@.yaml .gitignore proxy_server_config_2.yaml litellm/proxy/secret_managers/credentials.json +hosted_config.yaml diff --git a/Dockerfile b/Dockerfile index b76aaf1d1..180bde57a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,7 +3,6 @@ ARG LITELLM_BUILD_IMAGE=python:3.9 # Runtime image ARG LITELLM_RUNTIME_IMAGE=python:3.9-slim - # Builder stage FROM $LITELLM_BUILD_IMAGE as builder @@ -35,8 +34,12 @@ RUN pip wheel --no-cache-dir --wheel-dir=/wheels/ -r requirements.txt # Runtime stage FROM $LITELLM_RUNTIME_IMAGE as runtime +ARG with_database WORKDIR /app +# Copy the current directory contents into the container at /app +COPY . . +RUN ls -la /app # Copy the built wheel from the builder stage to the runtime stage; assumes only one wheel file is present COPY --from=builder /app/dist/*.whl . @@ -45,9 +48,17 @@ COPY --from=builder /wheels/ /wheels/ # Install the built wheel using pip; again using a wildcard if it's the only file RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ && rm -f *.whl && rm -rf /wheels +# Check if the with_database argument is set to 'true' +RUN echo "Value of with_database is: ${with_database}" +# If true, execute the following instructions +RUN if [ "$with_database" = "true" ]; then \ + prisma generate; \ + chmod +x /app/retry_push.sh; \ + /app/retry_push.sh; \ + fi -EXPOSE 4000/tcp +EXPOSE 8000/tcp # Set your entrypoint and command ENTRYPOINT ["litellm"] -CMD ["--port", "4000"] \ No newline at end of file +CMD ["--config", "./hosted_config.yaml", "--port", "8000", "--num_workers", "8"] \ No newline at end of file diff --git a/docker/.env.example b/docker/.env.example index 91934506a..613f99706 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -6,10 +6,10 @@ LITELLM_MASTER_KEY="sk-1234" ############ -# Database - You can change these to any PostgreSQL database that has logical replication enabled. +# Database - You can change these to any PostgreSQL database. ############ -# LITELLM_DATABASE_URL="your-postgres-db-url" +LITELLM_DATABASE_URL="your-postgres-db-url" ############ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2a97917b1..c34a701c6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -519,16 +519,12 @@ class ProxyConfig: user_config_file_path = config_file_path # Load existing config ## Yaml - if os.path.exists(f"{file_path}"): - with open(f"{file_path}", "r") as config_file: - config = yaml.safe_load(config_file) - else: - config = { - "model_list": [], - "general_settings": {}, - "router_settings": {}, - "litellm_settings": {}, - } + if file_path is not None: + if os.path.exists(f"{file_path}"): + with open(f"{file_path}", "r") as config_file: + config = yaml.safe_load(config_file) + else: + raise Exception(f"File not found! - {file_path}") ## DB if ( @@ -2328,6 +2324,21 @@ async def update_config(config_info: ConfigYAML): raise HTTPException(status_code=500, detail=f"An error occurred - {str(e)}") +@router.get( + "/config/get", + tags=["config.yaml"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_config(): + """ + Master key only. + + Returns the config. Mainly used for testing. + """ + global proxy_config + return await proxy_config.get_config() + + @router.get("/config/yaml", tags=["config.yaml"]) async def config_yaml_endpoint(config_info: ConfigYAML): """ @@ -2416,6 +2427,28 @@ async def health_endpoint( } +@router.get("/health/readiness", tags=["health"]) +async def health_readiness(): + """ + Unprotected endpoint for checking if worker can receive requests + """ + global prisma_client + if prisma_client is not None: # if db passed in, check if it's connected + if prisma_client.db.is_connected() == True: + return {"status": "healthy"} + else: + return {"status": "healthy"} + raise HTTPException(status_code=503, detail="Service Unhealthy") + + +@router.get("/health/liveliness", tags=["health"]) +async def health_liveliness(): + """ + Unprotected endpoint for checking if worker is alive + """ + return "I'm alive!" + + @router.get("/") async def home(request: Request): return "LiteLLM: RUNNING" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5670d90da..514653295 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -250,30 +250,36 @@ def on_backoff(details): class PrismaClient: def __init__(self, database_url: str, proxy_logging_obj: ProxyLogging): - print_verbose( - "LiteLLM: DATABASE_URL Set in config, trying to 'pip install prisma'" - ) - ## init logging object - self.proxy_logging_obj = proxy_logging_obj - os.environ["DATABASE_URL"] = database_url - # Save the current working directory - original_dir = os.getcwd() - # set the working directory to where this script is - abspath = os.path.abspath(__file__) - dname = os.path.dirname(abspath) - os.chdir(dname) - + ### Check if prisma client can be imported (setup done in Docker build) try: - subprocess.run(["prisma", "generate"]) - subprocess.run( - ["prisma", "db", "push", "--accept-data-loss"] - ) # this looks like a weird edge case when prisma just wont start on render. we need to have the --accept-data-loss - finally: - os.chdir(original_dir) - # Now you can import the Prisma Client - from prisma import Client # type: ignore + from prisma import Client # type: ignore - self.db = Client() # Client to connect to Prisma db + self.db = Client() # Client to connect to Prisma db + except: # if not - go through normal setup process + print_verbose( + "LiteLLM: DATABASE_URL Set in config, trying to 'pip install prisma'" + ) + ## init logging object + self.proxy_logging_obj = proxy_logging_obj + os.environ["DATABASE_URL"] = database_url + # Save the current working directory + original_dir = os.getcwd() + # set the working directory to where this script is + abspath = os.path.abspath(__file__) + dname = os.path.dirname(abspath) + os.chdir(dname) + + try: + subprocess.run(["prisma", "generate"]) + subprocess.run( + ["prisma", "db", "push", "--accept-data-loss"] + ) # this looks like a weird edge case when prisma just wont start on render. we need to have the --accept-data-loss + finally: + os.chdir(original_dir) + # Now you can import the Prisma Client + from prisma import Client # type: ignore + + self.db = Client() # Client to connect to Prisma db def hash_token(self, token: str): # Hash the string using SHA-256 diff --git a/retry_push.sh b/retry_push.sh new file mode 100644 index 000000000..5c41d72a0 --- /dev/null +++ b/retry_push.sh @@ -0,0 +1,28 @@ +#!/bin/bash + +retry_count=0 +max_retries=3 +exit_code=1 + +until [ $retry_count -ge $max_retries ] || [ $exit_code -eq 0 ] +do + retry_count=$((retry_count+1)) + echo "Attempt $retry_count..." + + # Run the Prisma db push command + prisma db push --accept-data-loss + + exit_code=$? + + if [ $exit_code -ne 0 ] && [ $retry_count -lt $max_retries ]; then + echo "Retrying in 10 seconds..." + sleep 10 + fi +done + +if [ $exit_code -ne 0 ]; then + echo "Unable to push database changes after $max_retries retries." + exit 1 +fi + +echo "Database push successful!" \ No newline at end of file diff --git a/schema.prisma b/schema.prisma new file mode 100644 index 000000000..d12cac8f2 --- /dev/null +++ b/schema.prisma @@ -0,0 +1,33 @@ +datasource client { + provider = "postgresql" + url = env("DATABASE_URL") +} + +generator client { + provider = "prisma-client-py" +} + +model LiteLLM_UserTable { + user_id String @unique + max_budget Float? + spend Float @default(0.0) + user_email String? +} + +// required for token gen +model LiteLLM_VerificationToken { + token String @unique + spend Float @default(0.0) + expires DateTime? + models String[] + aliases Json @default("{}") + config Json @default("{}") + user_id String? + max_parallel_requests Int? + metadata Json @default("{}") +} + +model LiteLLM_Config { + param_name String @id + param_value Json? +} \ No newline at end of file