testing vllm

Dockerfile

@@ -23,4 +23,5 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 # Add src files (Worker Template)
 ADD src .
 
+CMD python -m vllm.entrypoints.openai.api_server --model facebook/opt-125m --port 443 --host 127.0.0.1
 CMD python -u /handler.py

builder/requirements.txt

@@ -1,4 +1,5 @@
 # Required Python packages get listed here, one per line.
 # Reccomended to lock the version number to avoid unexpected changes.
 
-runpod==0.9.12
+runpod==git+https://github.com/runpod/runpod-python@vllm#egg=runpod
+vllm==0.1.1

src/handler.py

@@ -1,21 +1,37 @@
 #!/usr/bin/env python
 ''' Contains the handler function that will be called by the serverless. '''
-
 import runpod
-
-# Load models into VRAM here so they can be warm between requests
-
+import requests
+import json
+from typing import Dict
 
 def handler(event):
     '''
     This is the handler function that will be called by the serverless.
     '''
-    print(event)
+    # Prepare the request
+    llm_endpoint_data: Dict = event['llm_endpoint']
+    llm_body_data: Dict = event['llm_body']
 
-    # do the things
+    llm_endpoint_url: str = llm_endpoint_data['url']
+    llm_endpoint_request_type: str = llm_endpoint_data['request_type']
 
-    # return the output that you want to be returned like pre-signed URLs to output artifacts
-    return "Hello World"
+    url = "http://127.0.0.1:443/" + llm_endpoint_url
+    headers = {
+        "Content-Type": "application/json"
+    }
+
+    # Process the request
+    if llm_endpoint_request_type.lower() == 'post':
+        response = requests.post(url, headers=headers,
+                                 data=json.dumps(llm_body_data))
+
+    elif llm_endpoint_request_type.lower() == 'get':
+        response = requests.get(url, headers=headers)
+
+    response_data = response.json()
+
+    return response_data
 
 
-runpod.serverless.start({"handler": handler})
+runpod.serverless.start({"handler": handler}, serverless_llm=True)