add chatgpt-api-response-timeout-secs flag, set this to 20 mins in test

e49924e1b9ad4fba710e7d2a87f43c8c98e54b23 · 2024-07-20 00:27:14 -0700 · Alex Cheema
Files touched

M .github/workflows/test.yml
M exo/api/chatgpt_api.py
M main.py
Diff

commit e49924e1b9ad4fba710e7d2a87f43c8c98e54b23
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sat Jul 20 00:27:14 2024 -0700

    add chatgpt-api-response-timeout-secs flag, set this to 20 mins in test
---
 .github/workflows/test.yml | 31 ++++++++++++++++++++++++-------
 exo/api/chatgpt_api.py     |  4 ++--
 main.py                    |  3 ++-
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 9d76f005..9a3f311b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -78,11 +78,11 @@ jobs:
     - name: Run chatgpt api integration test
       run: |
         # Start first instance
-        DEBUG_DISCOVERY=9 DEBUG=9 python3 main.py --listen-port 5678 --broadcast-port 5679 --chatgpt-api-port 8000 > output1.log 2>&1 &
+        DEBUG_DISCOVERY=9 DEBUG=9 python3 main.py --listen-port 5678 --broadcast-port 5679 --chatgpt-api-port 8000 --chatgpt-api-response-timeout-secs 1200 > output1.log 2>&1 &
         PID1=$!
 
         # Start second instance
-        DEBUG_DISCOVERY=9 DEBUG=9 python3 main.py --listen-port 5679 --broadcast-port 5678 --chatgpt-api-port 8001 > output2.log 2>&1 &
+        DEBUG_DISCOVERY=9 DEBUG=9 python3 main.py --listen-port 5679 --broadcast-port 5678 --chatgpt-api-port 8001 --chatgpt-api-response-timeout-secs 1200 > output2.log 2>&1 &
         PID2=$!
 
         # Wait for discovery
@@ -96,22 +96,39 @@ jobs:
               "messages": [{"role": "user", "content": "Placeholder to load model..."}],
               "temperature": 0.7
             }'
+        curl -s http://localhost:8001/v1/chat/completions \
+            -H "Content-Type: application/json" \
+            -d '{
+              "model": "llama-3-8b",
+              "messages": [{"role": "user", "content": "Placeholder to load model..."}],
+              "temperature": 0.7
+            }'
+
+        response_1=$(curl -s http://localhost:8000/v1/chat/completions \
+          -H "Content-Type: application/json" \
+          -d '{
+            "model": "llama-3-8b",
+            "messages": [{"role": "user", "content": "Who was the king of pop?"}],
+            "temperature": 0.7
+          }')
+        echo "Response 1: $response_1"
 
-        response=$(curl -s http://localhost:8000/v1/chat/completions \
+        response_2=$(curl -s http://localhost:8000/v1/chat/completions \
           -H "Content-Type: application/json" \
           -d '{
             "model": "llama-3-8b",
             "messages": [{"role": "user", "content": "Who was the king of pop?"}],
             "temperature": 0.7
           }')
-        echo "Response: $response"
+        echo "Response 2: $response_2"
 
-        if ! echo "$response" | grep -q "Michael Jackson"; then
+        if ! echo "$response_1" | grep -q "Michael Jackson" || ! echo "$response_2" | grep -q "Michael Jackson"; then
           echo "Test failed: Response does not contain 'Michael Jackson'"
-          echo "Response: $response"
+          echo "Response 1: $response_1"
+          echo "Response 2: $response_2"
           exit 1
         else
-          echo "Test passed: Response contains 'Michael Jackson'"
+          echo "Test passed: Response from both nodes contains 'Michael Jackson'"
         fi
 
         # Stop both instances
diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index 65b57b01..f6675ee3 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -117,10 +117,10 @@ def build_prompt(tokenizer, messages: List[Message]):
 
 
 class ChatGPTAPI:
-    def __init__(self, node: Node, inference_engine_classname: str):
+    def __init__(self, node: Node, inference_engine_classname: str, response_timeout_secs: int = 90):
         self.node = node
         self.inference_engine_classname = inference_engine_classname
-        self.response_timeout_secs = 90
+        self.response_timeout_secs = response_timeout_secs
         self.app = web.Application()
         self.prev_token_lens: Dict[str, int] = {}
         self.stream_tasks: Dict[str, asyncio.Task] = {}
diff --git a/main.py b/main.py
index 89f3f2fd..2758473a 100644
--- a/main.py
+++ b/main.py
@@ -22,6 +22,7 @@ parser.add_argument("--listen-port", type=int, default=5678, help="Listening por
 parser.add_argument("--broadcast-port", type=int, default=5678, help="Broadcast port for discovery")
 parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of peers to wait to connect to before starting")
 parser.add_argument("--chatgpt-api-port", type=int, default=8000, help="ChatGPT API port")
+parser.add_argument("--chatgpt-api-response-timeout-secs", type=int, default=90, help="ChatGPT API response timeout in seconds")
 parser.add_argument("--inference-engine", type=str, default=None, help="Inference engine to use")
 args = parser.parse_args()
 
@@ -57,7 +58,7 @@ discovery = GRPCDiscovery(args.node_id, args.node_port, args.listen_port, args.b
 node = StandardNode(args.node_id, None, inference_engine, discovery, partitioning_strategy=RingMemoryWeightedPartitioningStrategy(), chatgpt_api_endpoint=f"http://localhost:{args.chatgpt_api_port}/v1/chat/completions", web_chat_url=f"http://localhost:{args.chatgpt_api_port}")
 server = GRPCServer(node, args.node_host, args.node_port)
 node.server = server
-api = ChatGPTAPI(node, inference_engine.__class__.__name__)
+api = ChatGPTAPI(node, inference_engine.__class__.__name__, response_timeout_secs=args.chatgpt_api_response_timeout_secs)
 
 node.on_token.register("main_log").on_next(lambda _, tokens , __: print(inference_engine.tokenizer.decode(tokens) if hasattr(inference_engine, "tokenizer") else tokens))
← 7dd7ccab do one request to load the model then another to check the r · back to Exo · async model downloading fixes #30 a4cc6677 →