
import os
import subprocess
import json
import time
from termcolor import colored
from tenacity import retry, wait_random_exponential, stop_after_attempt
from transformers import AutoTokenizer

# Model should be available on Hugging Face or stored locally
# If private, be sure to add HUGGING_FACE_ACCESS_TOKEN to environment variables
model = 'casperhansen/mixtral-instruct-awq'

# For Runpod with TGI. Replace <POD_ID> with your Runpod Pod ID
api_endpoint = "https://<POD-ID>-8080.proxy.runpod.net"

tgi_api_base = api_endpoint + '/generate'

tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)

# # ## Manually adjust the prompt. Not Recommended. Here is Vicuna 1.1 prompt
## format. System messages not supported.
# tokenizer.chat_template = "{% set sep = ' ' %}{% set sep2 = '</s>' %}{{ 'A chat between a curious user and an artificial intelligence assistant.\n\nThe assistant gives helpful, detailed, and polite answers to user questions.\n\n' }}{% if messages[0]['role'] == 'system' %}{{ '' }}{% set start_index = 1 %}{% else %}{% set start_index = 0 %}{% endif %}{% for i in range(start_index, messages|length) %}{% if messages[i]['role'] == 'user' %}{{ 'USER:\n' + messages[i]['content'].strip() + (sep if i % 2 == start_index else sep2) }}{% elif messages[i]['role'] == 'assistant' %}{{ 'ASSISTANT:\n' + messages[i]['content'].strip() + (sep if i % 2 == start_index else sep2) }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT:\n' }}{% endif %}"

# # OPTION TO MANUALLY FORMAT MESSAGES (INSTEAD OF USING tokenizer.apply_chat_template)
# B_SYS = "<<SYS>>\n"
# E_SYS = "\n<</SYS>>\n\n"
# B_INST = "[INST] "
# E_INST = " [/INST]\n\n"
# BOS_token = "<s>"
# EOS_token = "</s>"

# def format_messages(messages):
    # formatted_string = ''
    # formatted_string += BOS_token
    # formatted_string += B_INST

    # for message in messages:
    #     if message['role'] == 'system':
    #         formatted_string += B_SYS
    #         formatted_string += message['content']
    #         formatted_string += E_SYS
    #     elif message['role'] in ['user']:
    #         formatted_string += message['content']
    #         formatted_string += E_INST
    #     elif message['role'] in ['assistant']:
    #         formatted_string += message['content']
    #         formatted_string += EOS_token
    #         formatted_string += BOS_token
    #         formatted_string += B_INST

    # return formatted_string

@retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(3))
def chat_completion_request_runpod(messages):
    # formatted_messages = format_messages(messages)

    formatted_messages = tokenizer.apply_chat_template(messages, tokenize=False, \
                                                       add_generation_prompt=True)

    # print(formatted_messages)

    # Properly escape the string for JSON
    json_payload = json.dumps({
        "inputs": formatted_messages,
        "parameters": {
            "max_new_tokens": 500,
            "do_sample": False,
            # "stop": ["<step>"] #required for codellama 70b
            }})

    start_time = time.time()  # Start timing

    try:
        # Execute the curl command
        curl_command = f"""
        curl -s {tgi_api_base} \
            -X POST \
            -d '{json_payload}' \
            -H 'Content-Type: application/json'
        """

        response = subprocess.run(curl_command, shell=True, check=True, \
                                  stdout=subprocess.PIPE)
        response_time = time.time() - start_time  # Calculate response time

        response = response.stdout.decode()
        response = json.loads(response).get("generated_text", "No generated text found")

        # # Log the first and last 25 characters and the response time
        # print(f"Response Time: {response_time} seconds")
        # print(f"Start of Response: {response[:25]}")
        # print(f"End of Response: {response[-25:]}")

        # Calculate tokens per second
        tokens_generated = len(response)/4  # Assuming each word is a token
        tokens_per_second = tokens_generated / response_time if response_time > 0 else 0
        prompt_tokens = chat_response.usage.prompt_tokens if completion_text else 0

        # Print promt and generated tokens, time taken and tokens per second
        print(f"Total Time: {response_time:.2f} seconds")
        print(f"Prompt Tokens: {prompt_tokens:.2f}")
        print(f"Tokens Generated: {tokens_generated:.2f}")
        print(f"Tokens per Second: {tokens_per_second:.2f}")

        return response
    except subprocess.CalledProcessError as e:
        print("Unable to generate ChatCompletion response")
        print(f"Exception: {e}")
        return str(e)

def pretty_print_conversation(messages):
    role_to_color = {
        "system": "red",
        "user": "green",
        "assistant": "blue",
        "tool": "magenta",
    }

    for message in messages:
        if message["role"] == "system":
            print(colored(f"system: {message['content']}\n", role_to_color[message["role"]]))
        elif message["role"] == "user":
            print(colored(f"user: {message['content']}\n", role_to_color[message["role"]]))
        elif message["role"] == "assistant" and message.get("function_call"):
            print(colored(f"assistant: {message['function_call']}\n", \
                          role_to_color[message["role"]]))
        elif message["role"] == "assistant" and not message.get("function_call"):
            print(colored(f"assistant: {message['content']}\n", role_to_color[message["role"]]))
        elif message["role"] == "tool":
            print(colored(f"function ({message['name']}): {message['content']}\n", \
                          role_to_color[message["role"]]))

# Chat
messages = []
# messages.append({"role": "system", "content": "You are a helpful assistant."})
messages.append({"role": "user", "content": "Write a long essay on the topic of spring."})
# messages.append({"role": "user", "content": "Write a short piece of python code to add up the first 10 prime fibonacci numbers."})

chat_response = chat_completion_request_runpod(messages)
messages.append({"role": "assistant", "content": chat_response})

pretty_print_conversation(messages)


import threading
import os
import json
import time
import requests
from tenacity import retry, wait_random_exponential, stop_after_attempt
from transformers import AutoTokenizer

# Model should be available on Hugging Face
# If private, be sure to add HUGGING_FACE_ACCESS_TOKEN to environment variables
model = 'casperhansen/mixtral-instruct-awq'

# For Runpod with TGI. Replace <POD_ID> with your Runpod Pod ID
api_endpoint = "https://<POD-ID>-8080.proxy.runpod.net"

tgi_api_base = api_endpoint + '/generate'

tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)

# # Manually adjust the prompt. Not Recommended. Here is Vicuna 1.1 prompt format. System messages not supported.
# tokenizer.chat_template = "{% set sep = ' ' %}{% set sep2 = '</s>' %}{{ 'A chat between a curious user and an artificial intelligence assistant.\n\nThe assistant gives helpful, detailed, and polite answers to user questions.\n\n' }}{% if messages[0]['role'] == 'system' %}{{ '' }}{% set start_index = 1 %}{% else %}{% set start_index = 0 %}{% endif %}{% for i in range(start_index, messages|length) %}{% if messages[i]['role'] == 'user' %}{{ 'USER:\n' + messages[i]['content'].strip() + (sep if i % 2 == start_index else sep2) }}{% elif messages[i]['role'] == 'assistant' %}{{ 'ASSISTANT:\n' + messages[i]['content'].strip() + (sep if i % 2 == start_index else sep2) }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT:\n' }}{% endif %}"

# # OPTION TO MANUALLY FORMAT MESSAGES (INSTEAD OF USING tokenizer.apply_chat_template)
# B_SYS = "<<SYS>>\n"
# E_SYS = "\n<</SYS>>\n\n"
# B_INST = "[INST] "
# E_INST = " [/INST]\n\n"
# BOS_token = "<s>"
# EOS_token = "</s>"

# def format_messages(messages):
    # formatted_string = ''
    # formatted_string += BOS_token
    # formatted_string += B_INST

    # for message in messages:
    #     if message['role'] == 'system':
    #         formatted_string += B_SYS
    #         formatted_string += message['content']
    #         formatted_string += E_SYS
    #     elif message['role'] in ['user']:
    #         formatted_string += message['content']
    #         formatted_string += E_INST
    #     elif message['role'] in ['assistant']:
    #         formatted_string += message['content']
    #         formatted_string += EOS_token
    #         formatted_string += BOS_token
    #         formatted_string += B_INST

    # return formatted_string

# @retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(3))
def chat_completion_request_threaded(messages, request_number):
    # formatted_messages = format_messages(messages)

    formatted_messages = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    json_payload = {"inputs": formatted_messages, "parameters": {"max_new_tokens": 500, "do_sample": False}}

    start_time = time.time()  # Start timing

    try:
        response = requests.post(tgi_api_base, json=json_payload)
        response_time = time.time() - start_time  # Calculate response time

        if response.status_code == 200:
            response_content = response.json().get("generated_text", "No generated text found")
        else:
            raise Exception(f"Request failed with status code {response.status_code}")

        # print(response_content)

        # Calculate tokens per second
        tokens_generated = len(response_content) / 4
        tokens_per_second = tokens_generated / response_time if response_time > 0 else 0

        # Print time taken and tokens per second for each request
        print(f"Request #{request_number}: Total Time: {response_time:.2f} seconds, Tokens per Second: {tokens_per_second:.2f}")

        return response_content
    except Exception as e:
        print(f"Unable to generate ChatCompletion response for Request #{request_number}")
        print(f"Exception: {e}")
        return str(e)

def send_request_every_x_seconds(interval, total_requests):
    for i in range(total_requests):
        threading.Timer(interval * i, send_request, args=(i+1,)).start()

def send_request(request_number):
    messages = [
        {"role": "user", "content": "Write a long essay on the topic of spring."}
    ]
    chat_completion_request_threaded(messages, request_number)

# Start sending requests every x seconds
send_request_every_x_seconds(0.125, 12)  # Modify as needed for your use case


from openai import OpenAI
import os
import time
from dotenv import load_dotenv
from termcolor import colored

model = 'casperhansen/mixtral-instruct-awq'

# For Runpod with vLLM. Replace <POD_ID> with your Runpod Pod ID
api_endpoint = "https://<POD-ID>-8000.proxy.runpod.net"

openai_api_base = api_endpoint + '/v1'

# Initialize the OpenAI client
client = OpenAI(
    api_key="EMPTY",  # Replace with your actual API key if required
    base_url=openai_api_base,
)

def chat_completion_request_openai(messages, client):
    start_time = time.time()  # Start timing

    # Create chat completions using the OpenAI client
    chat_response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        max_tokens=500
    )

    response_time = time.time() - start_time  # Calculate response time

    # Extract the completion text from the response
    if chat_response.choices:
        completion_text = chat_response.choices[0].message.content
    else:
        completion_text = None

    # Calculate tokens per second
    prompt_tokens = chat_response.usage.prompt_tokens if completion_text else 0
    tokens_generated = chat_response.usage.completion_tokens if completion_text else 0
    tokens_per_second = tokens_generated / response_time if response_time > 0 else 0

    # print(chat_response)

    # Print time taken and tokens per second
    print(f"Total Time: {response_time:.2f} seconds")
    print(f"Prompt Tokens: {prompt_tokens:.2f}")
    print(f"Tokens Generated: {tokens_generated:.2f}")
    print(f"Tokens per Second: {tokens_per_second:.2f}")

    return completion_text

def pretty_print_conversation(messages):
    role_to_color = {
        "system": "red",
        "user": "green",
        "assistant": "blue",
        "tool": "magenta",
    }

    for message in messages:
        color = role_to_color.get(message["role"], "grey")
        print(colored(f"{message['role']}: {message['content']}\n", color))

# Test the function
messages = [
    {"role": "user", "content": "Write a long essay on the topic of spring."}
]

chat_response = chat_completion_request_openai(messages, client)
messages.append({"role": "assistant", "content": chat_response})

pretty_print_conversation(messages)


from openai import OpenAI
import os
import time
import threading
# from termcolor import colored  # Uncomment if you wish to use colored output

# Model should be available on Hugging Face or stored locally
# If private, be sure to add HUGGING_FACE_ACCESS_TOKEN to environment variables
model = 'casperhansen/mixtral-instruct-awq'

# For Runpod with TGI. Replace <POD_ID> with your Runpod Pod ID
api_endpoint = "https://<POD-ID>-8080.proxy.runpod.net"

openai_api_base = api_endpoint + '/v1'

# Initialize the OpenAI client
client = OpenAI(
    api_key="EMPTY",  # Replace with your actual API key if required
    base_url=openai_api_base,
)

def chat_completion_request_openai(messages, client, request_number):
    start_time = time.time()  # Start timing

    # Create chat completions using the OpenAI client
    chat_response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        max_tokens=500
    )

    response_time = time.time() - start_time  # Calculate response time

    # Extract the completion text from the response
    if chat_response.choices:
        completion_text = chat_response.choices[0].message.content
    else:
        completion_text = None

    # Calculate tokens per second
    prompt_tokens = chat_response.usage.prompt_tokens if completion_text else 0
    tokens_generated = chat_response.usage.completion_tokens if completion_text else 0
    tokens_per_second = tokens_generated / response_time if response_time > 0 else 0

    # Print header and response details
    print(f"\n---------- Request #{request_number} ----------")
    print(f"Total Time Taken: {response_time:.2f} seconds")
    print(f"Prompt tokens: {prompt_tokens:.2f}")
    print(f"Tokens generated: {tokens_generated:.2f}")
    print(f"Tokens per Second: {tokens_per_second:.2f}\n")

    return completion_text

def send_request_every_x_seconds():
    for i in range(12):
        threading.Timer(0.125 * i, send_request, args=(i+1,)).start()

def send_request(request_number):
    messages = [
        {"role": "user", "content": "Write a long essay on the topic of spring."}
    ]

    chat_completion_request_openai(messages, client, request_number)

# Start sending requests every x seconds
send_request_every_x_seconds()

GPU	Latency (s)	Tokens Per Second
1x RTX A6000 (48GB VRAM)	29.05	17.21
1x A100 (80GB VRAM)	25.20	19.84
1x H100 PCIe (80GB VRAM)	25.76	19.41
1x H100 SXM5 (80GB VRAM)	24.64	20.29

GPU	Concurrent Requests	Average Latency (s)	Average Tokens Per Second	Cost ($/hr)
1x RTX A6000 (48GB)	25	46.51	11.87	0.79
1x RTX A6000 (48GB)	50	82.87	6.71	0.79
1x RTX A6000 (48GB)	100	Timeout Error	Timeout Error	0.79
1x A100 (80GB)	25	43.67	12.84	1.89
1x A100 (80GB)	50	57.15	9.66	1.89
1x A100 (80GB)	100	85.80	6.51	1.89
1x H100 PCIe (80GB)	25	45.69	12.04	3.89
1x H100 PCIe (80GB)	50	62.77	8.76	3.89
1x H100 PCIe (80GB)	100	99.53	5.56	3.89
1x H100 SXM5 (80GB)	25	35.43	15.71	4.69
1x H100 SXM5 (80GB)	50	48.66	11.38	4.69
1x H100 SXM5 (80GB)	100	72.40	7.59	4.69

GPU	Latency (s)	Tokens Per Second
1x RTX A6000 (48GB VRAM)	30.09	16.62
1x A100 (80GB VRAM)	29.12	17.17
1x H100 SXM5 (80GB VRAM)	27.96	17.89

GPU	Concurrent Requests	Average Latency (s)	Average Tokens Per Second	Cost ($/hr)
1x RTX A6000 (48GB)	25	48.63	10.32	0.79
1x RTX A6000 (48GB)	50	82.87	6.71	0.79
1x RTX A6000 (48GB)	100	502 Error	502 Error	0.79
1x A100 (80GB)	25	45.44	10.86	1.89
1x A100 (80GB)	50	72.12	6.89	1.89
1x A100 (80GB)	100	502 Error	502 Error	1.89
1x H100 SXM5 (80GB)	25	39.25	12.57	4.69
1x H100 SXM5 (80GB)	50	56.42	8.83	4.69
1x H100 SXM5 (80GB)	100	91.05	5.55	4.69

Model Serving: How to Set Up an Inference API and Measure Performance¶

Why should you read this notebook?¶

Source Code¶

Pre-requisites¶

Table of Contents¶

Considerations for Model Serving¶

1. Privacy / Customization¶

2. Model size¶

3. Inference toolkits¶

4. GPU selection¶

5. Cloud providers¶

Text Generation Inference (TGI)¶

About¶

Setup¶

Selecting a GPU¶

Using the Docker image¶

Querying the model¶

TGI Speed Test Script (Python)¶

TGI Speed Test Results¶

Test Setup¶

Test Results¶

TGI Concurrency Test Script¶

Test Setup¶

TGI Concurrency Test Results¶

vLLM¶

About¶

Setup¶

Selecting a GPU¶

Using the Docker image¶

Querying the model¶

vLLM Speed Test Script (Python)¶

vLLM Speed Test Results¶

Test Setup¶

Test Results¶

vLLM Concurrency Test Script¶

vLLM Concurrency Test Results¶