diff --git a/Dockerfile b/Dockerfile index 8b21393..0286ee2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -28,10 +28,10 @@ RUN --mount=type=cache,target=/root/.cache/pip \ rm /requirements.txt # Add src files (Worker Template) -ADD src . +ADD src . # Quick temporary updates -RUN pip install git+https://github.com/runpod/runpod-python@main#egg=runpod --compile +RUN pip install git+https://github.com/runpod/runpod-python@experimental#egg=runpod --compile # Prepare the models inside the docker image ARG HUGGING_FACE_HUB_TOKEN= @@ -58,4 +58,4 @@ RUN mkdir -p /model RUN MODEL_NAME=$MODEL_NAME MODEL_REVISION=$MODEL_REVISION MODEL_BASE_PATH=$MODEL_BASE_PATH HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN python -u /download_model.py # Start the handler -CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER python -u /handler.py +CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER python -u /handler.py diff --git a/src/benchmark.py b/src/benchmark.py new file mode 100644 index 0000000..9b966cc --- /dev/null +++ b/src/benchmark.py @@ -0,0 +1,54 @@ +import concurrent.futures +import requests +import time +import os + +RUNPOD_ENDPOINT = os.environ('RUNPOD_ENDPOINT') +RUNPOD_API_KEY = os.environ('RUNPOD_API_KEY') + +url = "https://api.runpod.ai/v2/{RUNPOD_ENDPOINT}/run" + +headers = { + "Authorization": RUNPOD_API_KEY, + "Content-Type": "application/json" +} + +prompt = """ +Write me an essay about how the french revolution impacted the rest of europe over the 18th century. +""" + +payload = { + "input": { + "stream": True, + "prompt": prompt, + "sampling_params": { + "max_tokens": 1000, + "n": "1", + "presence_penalty": "0.2", + "frequency_penalty": "0.7", + "temperature": "0.3", + } + } +} + + +def make_request(url, headers, payload): + response = requests.post(url, headers=headers, json=payload) + return response + +while True: + # Number of concurrent requests to make per second. + num_requests = 100 + + with concurrent.futures.ThreadPoolExecutor(max_workers=num_requests) as executor: + futures = [executor.submit(make_request, url, headers, payload) + for _ in range(num_requests)] + + # Wait for all requests to complete + for future in concurrent.futures.as_completed(futures): + response = future.result() + # Handle response as needed + print(response.status_code) + + # Sleep for 1 second before starting the next iteration + time.sleep(1) diff --git a/src/handler.py b/src/handler.py index 3f90f38..64962e8 100644 --- a/src/handler.py +++ b/src/handler.py @@ -14,6 +14,7 @@ MODEL_BASE_PATH = os.environ.get('MODEL_BASE_PATH', '/runpod-volume/') STREAMING = os.environ.get('STREAMING', False) == 'True' TOKENIZER = os.environ.get('TOKENIZER', None) + if not MODEL_NAME: print("Error: The model has not been provided.") @@ -38,13 +39,16 @@ engine_args = AsyncEngineArgs( # Create the vLLM asynchronous engine llm = AsyncLLMEngine.from_engine_args(engine_args) - def concurrency_controller() -> bool: # Compute pending sequences total_pending_sequences = len( llm.engine.scheduler.waiting) + len(llm.engine.scheduler.swapped) - return total_pending_sequences > 0 + return total_pending_sequences > 30 +# Execute engine step indefinitely +# async def engine_step(): +# while True: +# await llm.engine_step() # Validation def validate_sampling_params(sampling_params): @@ -200,7 +204,7 @@ async def handler(job): if STREAMING: print("Starting the vLLM serverless worker with streaming enabled.") runpod.serverless.start( - {"handler": handler_streaming, "concurrency_controller": concurrency_controller}) + {"handler": handler_streaming, "concurrency_controller": concurrency_controller, "return_aggregate_stream": True}) else: print("Starting the vLLM serverless worker with streaming disabled.") runpod.serverless.start(