updating the llm worker

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 

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)

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(