cleanup

.gitignore

@@ -158,3 +158,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+runpod.toml

Dockerfile

@@ -1,7 +1,9 @@
 # Base image
 FROM runpod/base:0.4.2-cuda12.1.0
 
-ARG HUGGING_FACE_HUB_TOKEN
+ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets"
+ENV HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub"
+ENV TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub"
 
 # Install Python dependencies (Worker Template)
 COPY builder/requirements.txt /requirements.txt
@@ -10,40 +12,20 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     python3.11 -m pip install --upgrade -r /requirements.txt --no-cache-dir && \
     rm /requirements.txt
 
-# Add src files (Worker Template)
-ADD src .
-
 # Prepare argument for the model and tokenizer
 ARG MODEL_NAME=""
 ENV MODEL_NAME=$MODEL_NAME
-ARG MODEL_REVISION="main"
-ENV MODEL_REVISION=$MODEL_REVISION
-ARG MODEL_BASE_PATH="/runpod-volume/"
-ENV MODEL_BASE_PATH=$MODEL_BASE_PATH
+
+# Tokenizer is a path ot a huggingface repo
 ARG TOKENIZER=
 ENV TOKENIZER=$TOKENIZER
-ARG STREAMING=
-ENV STREAMING=$STREAMING
+
+# A string awq or squeezellm
 ARG QUANTIZATION=
 ENV QUANTIZATION=$QUANTIZATION
-ARG MAX_CONCURRENCY=
-ENV MAX_CONCURRENCY=$MAX_CONCURRENCY
 
-ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets"
-ENV HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub"
-ENV TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub"
-
-# Download the models
-RUN mkdir -p /model
-
-# Set environment variables
-ENV MODEL_NAME=$MODEL_NAME \
-    MODEL_REVISION=$MODEL_REVISION \
-    MODEL_BASE_PATH=$MODEL_BASE_PATH \
-    HUGGING_FACE_HUB_TOKEN=$HUGGING_FACE_HUB_TOKEN
-
-# Set the entrypoint
-ENTRYPOINT ["/entrypoint.sh"]
+# Add source files
+ADD src .
 
 # Start the handler
-CMD STREAMING=$STREAMING MODEL_NAME=$MODEL_NAME MODEL_BASE_PATH=$MODEL_BASE_PATH TOKENIZER=$TOKENIZER QUANTIZATION=$QUANTIZATION  python3.11 /handler.py
+CMD python3.11 /handler.py

entrypoint.sh

@@ -1,26 +0,0 @@
-#!/bin/bash
-
-# Export specific ENV variables to /etc/rp_environment
-echo "Exporting environment variables..."
-printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
-echo 'source /etc/rp_environment' >> ~/.bashrc
-
-if [[ $PUBLIC_KEY ]]
-then
-    mkdir -p ~/.ssh
-    chmod 700 ~/.ssh
-    echo $PUBLIC_KEY >> ~/.ssh/authorized_keys
-    chmod 700 -R ~/.ssh
-    # Start the SSH service in the background
-    service ssh start
-else
-    echo "No PUBLIC_KEY ENV variable provided, not starting openSSH daemon"
-fi
-
-# Conditional execution based on TESTING
-if [ "$TESTING" = "True" ]; then
-    sleep infinity
-else
-    # Execute the passed arguments (CMD)
-    exec "$@"
-fi

run_multiple.sh

@@ -1,43 +0,0 @@
-#!/bin/bash
-
-# Number of concurrent requests
-num_instances=40
-
-# Start time
-start_time=$(date +%s)
-
-# Function to run the command
-run_command() {
-    curl -v --request POST \
-         --url https://api.runpod.ai/v2/u4f0txz9dssz6q/runsync \
-         --header 'accept: application/json' \
-         --header 'authorization: UGXLPIBFYJXUCKOBNR8CCJ3KP8GWJNI97F91QT6Y' \
-         --header 'content-type: application/json' \
-         --data '
-    {
-      "input": {
-        "prompt": "Who is the president of the United States?",
-        "sampling_params": {
-          "max_tokens": "300",
-          "ignore_eos": true
-        }
-      }
-    }
-    '
-}
-
-# Run instances in the background
-for ((i = 1; i <= num_instances; i++)); do
-    run_command &
-done
-
-# Wait for all instances to finish
-wait
-
-# End time
-end_time=$(date +%s)
-
-# Calculate total time
-total_time=$((end_time - start_time))
-
-echo "Total time taken for $num_instances instances (1k tokens each): $total_time seconds"

src/benchmark.py

@@ -1,61 +0,0 @@
-import requests
-import json
-import time
-import os
-
-url = "https://api.runpod.ai/v2/4hlrhh430u5tz7/runsync"
-
-headers = {
-    "Authorization":"UGXLPIBFYJXUCKOBNR8CCJ3KP8GWJNI97F91QT6Y",
-    "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",
-        }
-    }
-}
-
-import concurrent.futures
-import time
-
-def make_request(url, headers, payload):
-    response = requests.post(url, headers=headers, json=payload)
-    return response
-
-url = "https://api.runpod.ai/v2/4hlrhh430u5tz7/status/1a73bdc5-aede-4c0a-8fa2-2ac2fb3010dc"
-
-while True:
-    # Number of concurrent requests to make
-    num_requests = 1
-
-    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.json())
-        print(response.status_code)
-
-    # Sleep for 1 second before starting the next iteration
-    time.sleep(1)
-
-
-# response_json = json.loads(response.text)
-
-# get_status = requests.get(status_url, headers=headers)
-# print(get_status.text)

src/entrypoint.sh

@@ -1,26 +0,0 @@
-#!/bin/bash
-
-# Export specific ENV variables to /etc/rp_environment
-echo "Exporting environment variables..."
-printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
-echo 'source /etc/rp_environment' >> ~/.bashrc
-
-if [[ $PUBLIC_KEY ]]
-then
-    mkdir -p ~/.ssh
-    chmod 700 ~/.ssh
-    echo $PUBLIC_KEY >> ~/.ssh/authorized_keys
-    chmod 700 -R ~/.ssh
-    # Start the SSH service in the background
-    service ssh start
-else
-    echo "No PUBLIC_KEY ENV variable provided, not starting openSSH daemon"
-fi
-
-# Conditional execution based on TESTING
-if [ "$TESTING" = "True" ]; then
-    sleep infinity
-else
-    # Execute the passed arguments (CMD)
-    exec "$@"
-fi

src/handler.py

@@ -1,50 +1,31 @@
 #!/usr/bin/env python
 ''' Contains the handler function that will be called by the serverless worker. '''
 
-# Start the vLLM serving layer on our RunPod worker.
-from typing import Generator
-from metrics import vllm_log_system_stats
-from templates import DEFAULT_TEMPLATE, LLAMA2_TEMPLATE
-from vllm import AsyncLLMEngine, SamplingParams, AsyncEngineArgs
-from vllm.utils import random_uuid
-import runpod
 import os
+from typing import Generator
+
+import runpod
+
+from metrics import vllm_log_system_stats
+from vllm import AsyncLLMEngine, SamplingParams, AsyncEngineArgs, utils
+
+
+NUM_GPU_SHARD = int(os.environ.get('NUM_GPU_SHARD', 1))  # Number of GPUs to shard the model across
 
 # Prepare the model and tokenizer
-MODEL_NAME = os.environ.get('MODEL_NAME')
-STREAMING = os.environ.get('STREAMING', False) == 'True'
-TOKENIZER = os.environ.get('TOKENIZER', None)
-DTYPE = "auto"
+MODEL_NAME = os.environ["MODEL_NAME"]
+TOKENIZER = os.environ.get('TOKENIZER', MODEL_NAME)
 
-MODEL_BASE_PATH = os.environ.get('MODEL_BASE_PATH')
-if not os.path.exists(BASE_VOLUME):
-    os.makedirs(BASE_VOLUME)
 
-USE_FULL_METRICS = os.environ.get('USE_FULL_METRICS', True)
-MAX_CONCURRENCY = os.environ.get('MAX_CONCURRENCY', 200)
-TOTAL_RUNNING_JOBS = 0
+MODEL_BASE_PATH = os.environ.get('MODEL_BASE_PATH', "/runpod-volume/")
+if not os.path.exists(MODEL_BASE_PATH):
+    os.makedirs(MODEL_BASE_PATH)
+
+USE_FULL_METRICS = os.environ.get('USE_FULL_METRICS', True)  # From the SDK, need to review later.
 
 # Set up quantization-related parameters
 QUANTIZATION = os.environ.get('QUANTIZATION', None)
-
-if type(QUANTIZATION) is str and QUANTIZATION.lower() in ["awq", "squeezellm"]:
-    QUANTIZATION = None
-    print("Invalid quantization parameter. Using default value of None.")
-else:
-    DTYPE = "half"
-
-if not MODEL_NAME:
-    print("Error: The model has not been provided.")
-
-if not TOKENIZER or len(TOKENIZER) == 0:
-    print("Error: The tokenizer has not been provided. Defaulting to MODEL_NAME.")
-
-# Tensor parallelism
-try:
-    NUM_GPU_SHARD = int(os.environ.get('NUM_GPU_SHARD', 1))
-except ValueError:
-    print("Error: NUM_GPU_SHARD should be an integer. Using default value of 1.")
-    NUM_GPU_SHARD = 1
+DTYPE = "auto" if str(QUANTIZATION).lower() not in ['squeezellm', 'awq'] else "half"
 
 
 # Prepare the engine's arguments
@@ -67,11 +48,6 @@ llm.engine._log_system_stats = lambda x, y: vllm_log_system_stats(
     llm.engine, x, y)
 
 
-def concurrency_controller() -> int:
-    global TOTAL_RUNNING_JOBS
-    return max(0, MAX_CONCURRENCY - TOTAL_RUNNING_JOBS)
-
-
 def prepare_metrics() -> dict:
     # The vLLM metrics are updated every 5 seconds, see metrics.py for the _LOGGING_INTERVAL_SEC field.
     if hasattr(llm.engine, 'metrics'):
@@ -147,232 +123,78 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
     '''
     This is the handler function that will be called by the serverless worker.
     '''
-    print("Job received by handler: {}".format(job))
-    global TOTAL_RUNNING_JOBS
-    # Retrieve the job input.
+    print(f"Job received by handler: {job}")
+
     job_input = job['input']
-
-    # Utilize the built-in llama2 template if a llama2 base model is being employed.
-    llama_models = ["llama-2-7b-chat-hf", "llama-2-13b-chat-hf", "llama-2-70b-chat-hf", "elinas/chronos-13b-v2"]
-
-    # Create the prompt using the template.
     prompt = job_input['prompt']
+    sampling_params = job_input.get('sampling_params', None)
 
-    # Validate and set sampling parameters
-    sampling_params = validate_and_set_sampling_params(job_input.get('sampling_params', None))
-
-    # Print job input and sampling parameters
-    print("Job Input:", job_input)
-    print("Sampling Parameters:", sampling_params)
+    # Might be able to remove this later
+    sampling_params = validate_and_set_sampling_params(sampling_params)
 
     # Send request to VLLM
-    request_id = random_uuid()
-    TOTAL_RUNNING_JOBS += 1
+    request_id = utils.random_uuid()
     results_generator = llm.generate(prompt, sampling_params, request_id)
 
-    # Keep track of the stream's information to perform the appropriate chunking.
-    class Tracker():
-        def __init__(self):
-            self.positions = None
-            self.stream_index = 0
-
-        def inc_stream_idx(self):
-            self.stream_index += 1
-
-    tracker = Tracker()
-
-    def extract_next_chunk(request_output):
-        """
-        Extracts and processes generated chunks and token counts from the request output.
-
-        Args:
-            request_output (CompletionOutput): The output of a language model request.
-
-        Returns:
-            tuple: A tuple containing two lists - chunk_outputs (extracted chunks) and num_output_tokens (generated token counts).
-        """
-        chunk_outputs = []  # List to store extracted chunks
-        num_output_tokens = []  # List to store generated token counts
-
-        # Iterate over each completion in the request output
-        for idx, completion in enumerate(request_output.outputs):
-            # Extract the current chunk position from the tracker
-            chunk_pos = tracker.positions[idx]['chunk_pos']
-
-            # Append the chunk to the output
-            chunk_outputs.append(completion.text[chunk_pos:])
-
-            # Update the chunk position in the tracker
-            tracker.positions[idx]['chunk_pos'] = len(completion.text)
-
-            # Calculate the number of generated tokens in the current completion
-            num_generated_tokens = len(completion.token_ids) - tracker.positions[idx]['token_pos']
-
-            # Append the token count to the output
-            num_output_tokens.append(num_generated_tokens)
-
-            # Update the token position in the tracker
-            tracker.positions[idx]['token_pos'] = len(completion.token_ids)
-
-        return chunk_outputs, num_output_tokens
+    stream_index = 0
+    chunk_positions = None
+    aggregate_text = []
+    aggregate_metrics = {'input_tokens': 0, 'output_tokens': 0}
 
     async for request_output in results_generator:
-        # Initialize chunk positions if not already done
-        if tracker.positions is None:
-            tracker.positions = [{'chunk_pos': 0, 'token_pos': 0}] * len(request_output.outputs)
+        if chunk_positions is None:
+            chunk_positions = [0] * len(request_output.outputs)
+
+        text_outputs, output_tokens = [], []
+        for idx, output in enumerate(request_output.outputs):
+            chunk_pos = chunk_positions[idx]
+            text_outputs.append(output.text[chunk_pos:])
+            chunk_positions[idx] = len(output.text)
+            output_tokens.append(len(output.token_ids) - chunk_pos)
 
-        # Metrics for the vLLM serverless worker
         runpod_metrics = prepare_metrics() if USE_FULL_METRICS else {}
-
-        # Number of generated sequences
-        num_seqs = sampling_params.n
-
-        # Extract the next chunk from the output
-        text_outputs, output_tokens = extract_next_chunk(request_output)
-
-        # Record job input and token counts
-        # runpod_metrics['job_input'] = job_input
-
-        # Only include the input_tokens count for the very first stream response. This is to avoid duplicate counting.
-        if tracker.stream_index == 0:
+        if stream_index == 0:
             input_tokens_count = len(request_output.prompt_token_ids)
-            runpod_metrics['input_tokens'] = sum([input_tokens_count] * num_seqs)
         else:
-            runpod_metrics['input_tokens'] = sum([0] * num_seqs)
+            input_tokens_count = 0
 
-        # Include the output tokens count [#, #, #, ...]
-        runpod_metrics['output_tokens'] = sum(output_tokens)
+        runpod_metrics.update({
+            "input_tokens": input_tokens_count,
+            "output_tokens": sum(output_tokens),
+            "scenario": 'stream',
+            "stream_index": stream_index
+        })
 
-        # Store the scenario type and stream index
-        runpod_metrics['scenario'] = 'stream'
-        runpod_metrics['stream_index'] = tracker.stream_index
+        stream_index += 1
 
-        # Increment the index within the stream
-        tracker.inc_stream_idx()
-
-        ret = {
+        yield {
             "text": text_outputs,
             "input_tokens": runpod_metrics['input_tokens'],
             "output_tokens": runpod_metrics['output_tokens']
         }
 
-        # Include metrics for the job.
-        runpod.serverless.modules.rp_metrics.metrics_collector.push_metrics_internal(
-            job_id=job['id'],
-            metrics=runpod_metrics
-        )
+        # Aggregate text and metrics
+        if not aggregate_text:
+            aggregate_text = [""] * len(text_outputs)
+        for idx, text in enumerate(text_outputs):
+            aggregate_text[idx] += text
+        aggregate_metrics['input_tokens'] += runpod_metrics['input_tokens']
+        aggregate_metrics['output_tokens'] += runpod_metrics['output_tokens']
 
-        # Keep track of the final output
-        final_output = request_output
-
-        # Include metrics in the highest level for the job output for aggregrate.
-        def aggregate_function(streamed_outputs):
-            aggregate_output = [""] * len(streamed_outputs[0]['text'])
-            for stream in streamed_outputs:
-                for id, seq in enumerate(stream['text']):
-                    aggregate_output[id] += seq
-
-            # Number of generated sequences
-            num_seqs = sampling_params.n
-
-            # Aggregate metrics to expose to the user
-            input_tokens = len(final_output.prompt_token_ids) * num_seqs
-            output_tokens = sum([len(output.token_ids) for output in final_output.outputs])
-
-            return {
-                "text": aggregate_output,
-                "input_tokens": input_tokens,
-                "output_tokens": output_tokens,
+    yield {
+        "text": aggregate_text,
+        "input_tokens": aggregate_metrics['input_tokens'],
+        "output_tokens": aggregate_metrics['output_tokens']
     }
 
-        # Update the aggregate transformation function
-        runpod.serverless.modules.rp_metrics.metrics_collector.update_stream_aggregate(
-            job_id=job['id'],
-            aggregate_function=aggregate_function
-        )
 
-        # Yield the output
-        yield ret
-
-        TOTAL_RUNNING_JOBS -= 1
-
-
-async def handler(job: dict) -> dict[str, list]:
-    '''
-    This is the handler function that will be called by the serverless worker.
-    '''
-    print("Job received by handler: {}".format(job))
-    global TOTAL_RUNNING_JOBS
-    # Retrieve the job input.
-    job_input = job['input']
-    # Create the prompt using the template.
-    prompt = job_input['prompt']
-
-    # Validate and set sampling parameters
-    sampling_params = validate_and_set_sampling_params(job_input.get('sampling_params', None))
-
-    # Print job input and sampling parameters
-    print("Job Input:", job_input)
-    print("Sampling Parameters:", sampling_params)
-
-    # Send request to VLLM
-    request_id = random_uuid()
-    TOTAL_RUNNING_JOBS += 1
-    results_generator = llm.generate(prompt, sampling_params, request_id)
-
-    # Get the final generated output
-    final_output = None
-    async for request_output in results_generator:
-        final_output = request_output
-
-    # Extract prompt and text outputs
-    prompt = final_output.prompt
-    text_outputs = [output.text for output in final_output.outputs]
-
-    # Number of generated sequences
-    num_seqs = sampling_params.n
-
-    # Prepare metrics if full metrics are enabled
-    runpod_metrics = prepare_metrics() if USE_FULL_METRICS else {}
-
-    # Record job input and token counts
-    # runpod_metrics['job_input'] = job_input
-
-    runpod_metrics['input_tokens'] = len(final_output.prompt_token_ids) * num_seqs
-    runpod_metrics['output_tokens'] = sum([len(output.token_ids) for output in final_output.outputs])
-
-    # Store the scenario type
-    runpod_metrics['scenario'] = 'batch'
-
-    # Include metrics for the job.
-    runpod.serverless.modules.rp_metrics.metrics_collector.push_metrics_internal(
-        job_id=job['id'],
-        metrics=runpod_metrics
-    )
-
-    ret = {
-        "text": text_outputs,
-        "input_tokens": runpod_metrics['input_tokens'],
-        "output_tokens": runpod_metrics['output_tokens']
-    }
-
-    TOTAL_RUNNING_JOBS -= 1
-
-    return ret
+def concurrency_modifier() -> int:
+    return 100
 
 
 # Start the serverless worker with appropriate settings
-if STREAMING:
-    print("Starting the vLLM serverless worker with streaming enabled.")
-    runpod.serverless.start({
+runpod.serverless.start({
     "handler": handler_streaming,
-        "concurrency_controller": concurrency_controller,
+    "concurrency_modifier": concurrency_modifier,
     "return_aggregate_stream": True
-    })
-else:
-    print("Starting the vLLM serverless worker with streaming disabled.")
-    runpod.serverless.start({
-        "handler": handler,
-        "concurrency_controller":
-            concurrency_controller
-    })
+})