diff --git a/Dockerfile b/Dockerfile index 7d03455..578ef87 100644 --- a/Dockerfile +++ b/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 diff --git a/builder/requirements.txt b/builder/requirements.txt index 70e05ec..6552836 100644 --- a/builder/requirements.txt +++ b/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 diff --git a/src/handler.py b/src/handler.py index dc0c957..70dc81a 100644 --- a/src/handler.py +++ b/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)