From 9ed51e791b2c82ffda88767bc64a0f856f37c572 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 8 Mar 2024 15:24:08 -0800
Subject: [PATCH] (fix) default num workers

---
 litellm/proxy/proxy_cli.py                    |  9 +++-
 .../proxy_load_test/litellm_proxy_config.yaml |  6 +++
 litellm/proxy/proxy_load_test/locustfile.py   | 27 ++++++++++
 .../proxy/proxy_load_test/openai_endpoint.py  | 50 +++++++++++++++++++
 4 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 litellm/proxy/proxy_load_test/litellm_proxy_config.yaml
 create mode 100644 litellm/proxy/proxy_load_test/locustfile.py
 create mode 100644 litellm/proxy/proxy_load_test/openai_endpoint.py

diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index f7eba02ecb8..367bbbb7008 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -16,6 +16,13 @@ from importlib import resources
 import shutil
 
 telemetry = None
+default_num_workers = 1
+try:
+    default_num_workers = os.cpu_count() or 1
+    if default_num_workers is not None and default_num_workers > 0:
+        default_num_workers -= 1
+except:
+    pass
 
 
 def append_query_params(url, params):
@@ -57,7 +64,7 @@ def is_port_in_use(port):
 @click.option("--port", default=8000, help="Port to bind the server to.", envvar="PORT")
 @click.option(
     "--num_workers",
-    default=1,
+    default=default_num_workers,
     help="Number of gunicorn workers to spin up",
     envvar="NUM_WORKERS",
 )
diff --git a/litellm/proxy/proxy_load_test/litellm_proxy_config.yaml b/litellm/proxy/proxy_load_test/litellm_proxy_config.yaml
new file mode 100644
index 00000000000..2e107d36689
--- /dev/null
+++ b/litellm/proxy/proxy_load_test/litellm_proxy_config.yaml
@@ -0,0 +1,6 @@
+model_list:
+  - model_name: gpt-3.5-turbo
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: http://0.0.0.0:8090
\ No newline at end of file
diff --git a/litellm/proxy/proxy_load_test/locustfile.py b/litellm/proxy/proxy_load_test/locustfile.py
new file mode 100644
index 00000000000..2cd2e2fccee
--- /dev/null
+++ b/litellm/proxy/proxy_load_test/locustfile.py
@@ -0,0 +1,27 @@
+from locust import HttpUser, task, between
+
+
+class MyUser(HttpUser):
+    wait_time = between(1, 5)
+
+    @task
+    def chat_completion(self):
+        headers = {
+            "Content-Type": "application/json",
+            # Include any additional headers you may need for authentication, etc.
+        }
+
+        # Customize the payload with "model" and "messages" keys
+        payload = {
+            "model": "gpt-3.5-turbo",
+            "messages": [
+                {"role": "system", "content": "You are a chat bot."},
+                {"role": "user", "content": "Hello, how are you?"},
+            ],
+            # Add more data as necessary
+        }
+
+        # Make a POST request to the "chat/completions" endpoint
+        response = self.client.post("chat/completions", json=payload, headers=headers)
+
+        # Print or log the response if needed
diff --git a/litellm/proxy/proxy_load_test/openai_endpoint.py b/litellm/proxy/proxy_load_test/openai_endpoint.py
new file mode 100644
index 00000000000..b3291ce7098
--- /dev/null
+++ b/litellm/proxy/proxy_load_test/openai_endpoint.py
@@ -0,0 +1,50 @@
+# import sys, os
+# sys.path.insert(
+#     0, os.path.abspath("../")
+# )  # Adds the parent directory to the system path
+from fastapi import FastAPI, Request, status, HTTPException, Depends
+from fastapi.responses import StreamingResponse
+from fastapi.security import OAuth2PasswordBearer
+from fastapi.middleware.cors import CORSMiddleware
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# for completion
+@app.post("/chat/completions")
+@app.post("/v1/chat/completions")
+async def completion(request: Request):
+    return {
+        "id": "chatcmpl-123",
+        "object": "chat.completion",
+        "created": 1677652288,
+        "model": "gpt-3.5-turbo-0125",
+        "system_fingerprint": "fp_44709d6fcb",
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": "\n\nHello there, how may I assist you today?",
+                },
+                "logprobs": None,
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21},
+    }
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    # run this on 8090, 8091, 8092 and 8093
+    uvicorn.run(app, host="0.0.0.0", port=8090)