build(Dockerfile): move prisma build to dockerfile

Seems to solve - BerriAI#1321

Author: krrishdholakia

.gitignore

@@ -31,3 +31,4 @@ proxy_server_config_@.yaml
 .gitignore
 proxy_server_config_2.yaml
 litellm/proxy/secret_managers/credentials.json
+hosted_config.yaml

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"]
+CMD ["--config", "./hosted_config.yaml", "--port", "8000", "--num_workers", "8"]

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

litellm/proxy/proxy_server.py

@@ -519,16 +519,12 @@ async def get_config(self, config_file_path: Optional[str] = None) -> dict:
             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"

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
-
-        self.db = Client()  # Client to connect to Prisma db
+            from prisma import Client  # type: ignore
+
+            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

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

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?
+}