diff --git a/builder/requirements.txt b/builder/requirements.txt index a93caea..7e8cf28 100644 --- a/builder/requirements.txt +++ b/builder/requirements.txt @@ -2,7 +2,7 @@ hf_transfer ray pandas pyarrow -runpod==1.5.3 +runpod==1.6.2 huggingface-hub packaging typing-extensions==4.7.1 diff --git a/openai/server.py b/openai/server.py deleted file mode 100644 index 3c7986e..0000000 --- a/openai/server.py +++ /dev/null @@ -1,70 +0,0 @@ -from fastapi import FastAPI, Request, HTTPException, BackgroundTasks -from fastapi.responses import StreamingResponse, JSONResponse -import requests -import json -import asyncio -import os - -app = FastAPI() - -ENDPOINT_ID = "pst10x7hvwoz2k" -RUNPOD_API_KEY = os.environ.get('RUNPOD_API_KEY') -BASE_URL = f'https://api.runpod.ai/v2/{ENDPOINT_ID}' -MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1" - -def transform_request(openai_request): - # minimal - transformed_request = { - "messages": openai_request.get("messages", []), - "batch_size": 30, - "apply_chat_template": True, - "sampling_params": { - "max_tokens": openai_request.get("max_tokens", 100), - "temperature": 0, - "ignore_eos": True, - "n": 1 - }, - "stream": True, - "use_openai_format": True - } - return transformed_request - - -async def stream_data(run_url, headers, test_payload): - response = requests.post(run_url, headers=headers, data=json.dumps({"input": test_payload})) - if response.status_code != 200: - raise HTTPException(status_code=500, detail="Error initializing the stream") - - job_id = response.json().get("id") - stream_url = BASE_URL + f'/stream/{job_id}' - response = {} - - while response.get('status') != 'COMPLETED': - response = requests.get(stream_url, headers=headers).json() - stream = response.get('stream', []) - if stream: - for chunk in stream: - yield chunk["output"] - await asyncio.sleep(0.1) - -@app.post("/v1/chat/completions") -async def create_chat_completion(request: Request, background_tasks: BackgroundTasks): - openai_request = await request.json() - transformed_request = transform_request(openai_request) - - headers = {'Authorization': f'Bearer {RUNPOD_API_KEY}', 'Content-Type': 'application/json'} - run_url = BASE_URL + "/run" - - return StreamingResponse(stream_data(run_url, headers, transformed_request)) - -@app.get("/v1/models") -async def get_model(): - """ - Endpoint to return the model name. - """ - return JSONResponse(content={"model": MODEL_NAME}) - - -if __name__ == "__main__": - import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8000) diff --git a/openai/test_openai_stream.py b/openai/test_openai_stream.py deleted file mode 100644 index 23761d1..0000000 --- a/openai/test_openai_stream.py +++ /dev/null @@ -1,45 +0,0 @@ -""" -Move to SRC directory -""" - -import os -from utils import JobInput -from engine import vLLMEngine - -vllm_engine = vLLMEngine() - -async def handler(job): - job_input = JobInput(job["input"]) - results_generator = vllm_engine.generate(job_input) - async for batch in results_generator: - yield batch - -test_payload = { - "input": { - "messages": [ - {"role": "user", "content": "Write me a 3000 word long and detailed essay about how the french revolution impacted the rest of europe over the 18th century."}, - ], - "batch_size": 2, - "apply_chat_template": True, - "sampling_params": { - "max_tokens": 10, - "temperature": 0, - "ignore_eos": True, - "n":1 - }, - "stream": True, - "use_openai_format": True - } -} - -async def test_handler(): - print("Start of output") - print("=" *50) - async for batch in handler(test_payload): - print(batch, end="") - print("=" *50) - print("End of output") - -import asyncio - -asyncio.run(test_handler()) \ No newline at end of file diff --git a/openai/test_openai_worker.py b/openai/test_openai_worker.py deleted file mode 100644 index 169c8f6..0000000 --- a/openai/test_openai_worker.py +++ /dev/null @@ -1,45 +0,0 @@ -import requests -import json -import time -import os - -ENDPOINT_ID = "pst10x7hvwoz2k" -RUNPOD_API_KEY = os.environ.get('RUNPOD_API_KEY') - -test_payload = { - "messages": [ - {"role": "user", "content": "Write me a 3000 word long and detailed essay about how the french revolution impacted the rest of europe over the 18th century."}, - ], - "batch_size": 2, # How many tokens to yield per batch - "apply_chat_template": True, - "sampling_params": { - "max_tokens": 4, - "temperature": 0, - "ignore_eos": True, - "n":1 - }, - "stream": True, - "use_openai_format": True -} - -base_url = f'https://api.runpod.ai/v2/{ENDPOINT_ID}' -run_url = base_url + "/run" -headers = { - 'Authorization': f'Bearer {RUNPOD_API_KEY}', - 'Content-Type': 'application/json', -} - -job_id = requests.post(run_url, headers=headers, data=json.dumps({"input":test_payload})).json()["id"] - -stream_url = base_url + f'/stream/{job_id}' -response = {} - -while response.get('status') != 'COMPLETED': - response = requests.get(stream_url, headers=headers).json() - stream = response.get('stream', []) - if stream: - batch_data = "" - for chunk in stream: - batch_data += chunk["output"] - print(batch_data, end="") - time.sleep(0.1) \ No newline at end of file diff --git a/openai/vllm_openai.py b/openai/vllm_openai.py deleted file mode 100644 index e3d8141..0000000 --- a/openai/vllm_openai.py +++ /dev/null @@ -1,25 +0,0 @@ -import os -from openai import OpenAI -client = OpenAI( - api_key="XYZ", - base_url="http://0.0.0.0:8000/v1" -) - - -stream = client.chat.completions.create( - model="mistralai/Mistral-7B-Instruct-v0.1", - messages=[ - {"role": "user", "content": "Who won the world series in 2020?"}, - {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}, - {"role": "user", "content": "Where was it played? Write a 1000 word essay about this"} - ], - max_tokens=300, - stream=True, -) - - -for chunk in stream: - print(chunk.choices[0].delta.content, end="") - -# When you send a request to /openai/v1/chat/completions or to /openai/v1/completions -# and you set the stream parameter to true, we need SSE streaming of the content \ No newline at end of file diff --git a/src/config.py b/src/config.py index 6027c5b..3176248 100644 --- a/src/config.py +++ b/src/config.py @@ -1,6 +1,5 @@ import os from dotenv import load_dotenv -from typing import Tuple, Optional from utils import count_physical_cores from torch.cuda import device_count diff --git a/src/constants.py b/src/constants.py index 55608d1..ce056b5 100644 --- a/src/constants.py +++ b/src/constants.py @@ -1,7 +1,9 @@ from typing import Union -DEFAULT_BATCH_SIZE = 30 +DEFAULT_BATCH_SIZE = 50 DEFAULT_MAX_CONCURRENCY = 300 +DEFAULT_BATCH_SIZE_GROWTH_FACTOR = 3 +DEFAULT_MIN_BATCH_SIZE = 1 SAMPLING_PARAM_TYPES = { "n": int, @@ -25,4 +27,4 @@ SAMPLING_PARAM_TYPES = { "skip_special_tokens": bool, "spaces_between_special_tokens": bool, "include_stop_str_in_output": bool -} +} \ No newline at end of file diff --git a/src/engine.py b/src/engine.py index c0d42dc..22483e0 100644 --- a/src/engine.py +++ b/src/engine.py @@ -4,16 +4,15 @@ import json from dotenv import load_dotenv from torch.cuda import device_count -from transformers import AutoTokenizer -from typing import Union, AsyncGenerator +from typing import AsyncGenerator from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams from vllm.entrypoints.openai.serving_chat import OpenAIServingChat from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion -from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse, CompletionRequest +from vllm.entrypoints.openai.protocol import ChatCompletionRequest, CompletionRequest, ErrorResponse -from utils import DummyRequest, OpenAIRequest, JobInput, BatchSize -from constants import DEFAULT_MAX_CONCURRENCY, DEFAULT_BATCH_SIZE +from utils import DummyRequest, JobInput, BatchSize, create_error_response +from constants import DEFAULT_MAX_CONCURRENCY, DEFAULT_BATCH_SIZE, DEFAULT_BATCH_SIZE_GROWTH_FACTOR, DEFAULT_MIN_BATCH_SIZE from tokenizer import TokenizerWrapper from config import EngineConfig @@ -26,14 +25,13 @@ class vLLMEngine: self.llm = self._initialize_llm() if engine is None else engine self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY)) self.default_batch_size = int(os.getenv("DEFAULT_BATCH_SIZE", DEFAULT_BATCH_SIZE)) - self.batch_size_growth_factor = int(os.getenv("BATCH_SIZE_GROWTH_FACTOR", 1)) - self.min_batch_size = int(os.getenv("MIN_BATCH_SIZE", 1)) + self.batch_size_growth_factor = int(os.getenv("BATCH_SIZE_GROWTH_FACTOR", DEFAULT_BATCH_SIZE_GROWTH_FACTOR)) + self.min_batch_size = int(os.getenv("MIN_BATCH_SIZE", DEFAULT_MIN_BATCH_SIZE)) def dynamic_batch_size(self, current_batch_size, growth_factor): return min(current_batch_size*growth_factor, self.default_batch_size) async def generate(self, job_input: JobInput): - # Adjust to use attributes from JobInput directly async for batch in self._generate_vllm( llm_input=job_input.llm_input, validated_sampling_params=job_input.validated_sampling_params, @@ -115,6 +113,7 @@ class OpenAIvLLMEngine: self.default_batch_size = vllm_engine.default_batch_size self.batch_size_growth_factor, self.min_batch_size = vllm_engine.batch_size_growth_factor, vllm_engine.min_batch_size self._initialize_engines() + self.raw_openai_output = bool(int(os.getenv("RAW_OPENAI_OUTPUT", 0))) def _initialize_engines(self): self.chat_engine = OpenAIServingChat( @@ -122,36 +121,38 @@ class OpenAIvLLMEngine: ) self.completion_engine = OpenAIServingCompletion(self.llm, self.config["model"]) - async def generate(self, openai_request: OpenAIRequest): - if openai_request.route == "/models": - yield self._handle_model_request() - elif openai_request.route in ["/chat/completions", "/completions"]: + async def generate(self, openai_request: JobInput): + if openai_request.openai_route == "/v1/models": + yield await self._handle_model_request() + elif openai_request.openai_route in ["/v1/chat/completions", "/v1/completions"]: async for response in self._handle_chat_or_completion_request(openai_request): yield response else: - raise ValueError("Invalid route") + yield create_error_response("Invalid route").model_dump() - def _handle_model_request(self): - return self.config["model"] + async def _handle_model_request(self): + models = await self.chat_engine.show_available_models() + return models.model_dump() - async def _handle_chat_or_completion_request(self, openai_request: OpenAIRequest): - if openai_request.route == "/chat/completions": + async def _handle_chat_or_completion_request(self, openai_request: JobInput): + if openai_request.openai_route == "/v1/chat/completions": request_class = ChatCompletionRequest generator_function = self.chat_engine.create_chat_completion - else: # "/completions" + elif openai_request.openai_route == "/v1/completions": request_class = CompletionRequest generator_function = self.completion_engine.create_completion - request = request_class( - model=self.config["model"], - **openai_request.inputs - ) + try: + request = request_class( + **openai_request.openai_input + ) + except Exception as e: + yield create_error_response(str(e)).model_dump() + return response_generator = await generator_function(request, DummyRequest()) - if isinstance(response_generator, ErrorResponse): - raise ValueError(response_generator.model_dump()) - elif not openai_request.inputs.get("stream"): + if not openai_request.openai_input.get("stream") or isinstance(response_generator, ErrorResponse): yield response_generator.model_dump() else: batch = [] @@ -159,13 +160,24 @@ class OpenAIvLLMEngine: batch_size = BatchSize(self.default_batch_size, self.min_batch_size, self.batch_size_growth_factor) async for chunk_str in response_generator: - if "data" in chunk_str and not "[DONE]" in chunk_str: - batch.append(json.loads(chunk_str.removeprefix("data: ").rstrip("\n\n"))) + if "data" in chunk_str: + if self.raw_openai_output: + data = chunk_str + elif "[DONE]" in chunk_str: + continue + else: + data = json.loads(chunk_str.removeprefix("data: ").rstrip("\n\n")) if not self.raw_openai_output else chunk_str + batch.append(data) batch_token_counter += 1 if batch_token_counter >= batch_size.current_batch_size: + if self.raw_openai_output: + batch = "".join(batch) yield batch batch = [] batch_token_counter = 0 batch_size.update() if batch: - yield batch \ No newline at end of file + if self.raw_openai_output: + batch = "".join(batch) + yield batch + \ No newline at end of file diff --git a/src/handler.py b/src/handler.py index 3646e5f..176ec7e 100644 --- a/src/handler.py +++ b/src/handler.py @@ -1,26 +1,18 @@ import os import runpod -from utils import JobInput, OpenAIRequest +from utils import JobInput from engine import vLLMEngine, OpenAIvLLMEngine -if not os.getenv("MODEL_NAME"): - os.environ["MODEL_NAME"] = "facebook/opt-125m" - os.environ["CUSTOM_CHAT_TEMPLATE"] = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" - vllm_engine = vLLMEngine() OpenAIvLLMEngine = OpenAIvLLMEngine(vllm_engine) async def handler(job): - if "openai" in job: - openai_request = OpenAIRequest(job) - results_generator = OpenAIvLLMEngine.generate(openai_request) - else: - job_input = JobInput(job["input"]) - results_generator = vllm_engine.generate(job_input) - + job_input = JobInput(job["input"]) + engine = OpenAIvLLMEngine if job_input.openai_route else vllm_engine + results_generator = engine.generate(job_input) async for batch in results_generator: yield batch - + runpod.serverless.start( { "handler": handler, diff --git a/src/utils.py b/src/utils.py index 10096a1..5663000 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,7 +1,10 @@ import logging +from http import HTTPStatus from typing import Any, Dict +from constants import SAMPLING_PARAM_TYPES from vllm.utils import random_uuid -from constants import SAMPLING_PARAM_TYPES, DEFAULT_BATCH_SIZE +from vllm.entrypoints.openai.protocol import ErrorResponse + logging.basicConfig(level=logging.INFO) @@ -50,11 +53,8 @@ class JobInput: self.growth_factor = float(growth_factor) if growth_factor else None min_batch_size = job.get("min_batch_size") self.min_batch_size = int(min_batch_size) if min_batch_size else None - -class OpenAIRequest: - def __init__(self, request): - self.route = request["openai"]["route"] - self.inputs = request["input"] + self.openai_route = job.get("openai_route") + self.openai_input = job.get("openai_input") class DummyRequest: async def is_disconnected(self): @@ -74,4 +74,8 @@ class BatchSize: def update(self): if self.is_dynamic: self.current_batch_size = min(self.current_batch_size*self.growth_factor, self.max_batch_size) - \ No newline at end of file + +def create_error_response(message: str, err_type: str = "BadRequestError", status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> ErrorResponse: + return ErrorResponse(message=message, + type=err_type, + code=status_code.value) \ No newline at end of file