Merge pull request #82 from hyscale-lab/feat/tbt-calculation

joeltancej · web-flow · commit 83b74efc5d15 · 2025-09-17T13:40:25.000+08:00
TBT calculation for non-streaming inferences
diff --git a/README.md b/README.md
@@ -46,13 +46,13 @@ PERPLEXITY_AI_API="your-perplexity-ai-api-key"
 HYPERBOLIC_API="your-hyperbolic-api-key"
 GROQ_API_KEY="your-groq-api-key"
 GEMINI_API_KEY="your-gemini-api-key"
-AZURE_LLAMA_8B_API="your-azure-llama-8b-api-key"
-AZURE_LLAMA_3.1_70B_API="your-azure-llama-70b-api-key"
 MISTRAL_LARGE_API="your-mistral-large-api-key"
 AWS_BEDROCK_ACCESS_KEY_ID="your-aws-bedrock-access-key-id"
 AWS_BEDROCK_SECRET_ACCESS_KEY="your-aws-bedrock-secret-key"
 AWS_BEDROCK_REGION="your-aws-bedrock-region"
 DYNAMODB_ENDPOINT_URL="your-dynamodb-endpoint-url"
+AZURE_AI_ENDPOINT="your-azure-ai-endpoint"
+AZURE_AI_API_KEY="your-azure-ai-api-key"
 ```
 
 ## **Usage**
diff --git a/providers/anthropic_provider.py b/providers/anthropic_provider.py
@@ -68,9 +68,20 @@ def perform_inference(self, model, prompt, max_output=100, verbosity=True):
                 timeout=500,
             )
             elapsed = timer() - start
+
+            usage = getattr(response, "usage", None)
+            total_tokens = (getattr(usage, "output_tokens", 0) or 0) if usage else 0
+
+            tbt = elapsed / max(total_tokens, 1)
+            tps = (total_tokens / elapsed)
+
             self.log_metrics(model, "response_times", elapsed)
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
             # Process and display the response
             if verbosity:
+                print(f"Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
                 self.display_response(response, elapsed)
             return elapsed
 
diff --git a/providers/aws_provider.py b/providers/aws_provider.py
@@ -72,8 +72,18 @@ def perform_inference(self, model, prompt, max_output=100, verbosity=True):
             model_response = json.loads(response["body"].read())
             generated_text = model_response.get("generation", "")
 
+            total_tokens = model_response.get("generation_token_count") or 0
+
+            tbt = total_time / max(total_tokens, 1)
+            tps = (total_tokens / total_time)
+
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
+
             if verbosity:
                 print(f"[INFO] Total response time: {total_time:.4f} seconds")
+                print(f"[INFO] Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
                 print("[INFO] Generated response:")
                 print(generated_text)
 
diff --git a/providers/azure_provider.py b/providers/azure_provider.py
@@ -1,83 +1,93 @@
 import os
-import requests
 import numpy as np
 from providers.base_provider import ProviderInterface
 from time import perf_counter as timer
-import re
+from azure.ai.inference import ChatCompletionsClient
+from azure.ai.inference.models import SystemMessage, UserMessage
+from azure.core.credentials import AzureKeyCredential
 
 
 class Azure(ProviderInterface):
     def __init__(self):
         """Initialize AzureProvider with required API information."""
         super().__init__()
 
+        self.endpoint = os.getenv("AZURE_AI_ENDPOINT")
+        self.api_key = os.getenv("AZURE_AI_API_KEY")
+
         # Map model names to Azure model IDs
         self.model_map = {
             # "mistral-7b-instruct-v0.1": "mistral-7b-instruct-v0.1",
-            "meta-llama-3.1-8b-instruct": "Meta-Llama-3-1-8B-Instruct-fyp",
+            "meta-llama-3.1-8b-instruct": "Meta-Llama-3.1-8B-Instruct-fyp",
             "meta-llama-3.1-70b-instruct": "Meta-Llama-3-1-70B-Instruct-fyp",
             "mistral-large": "Mistral-Large-2411-yatcd",
             "common-model": "Mistral-Large-2411-yatcd",
         }
 
-        # Define API keys for each model
-        self.model_api_keys = {
-            # "mistral-7b-instruct-v0.1": os.environ.get("MISTRAL_API_KEY"),
-            "meta-llama-3.1-8b-instruct": os.environ.get("AZURE_LLAMA_8B_API"),
-            "meta-llama-3.1-70b-instruct": os.environ.get("AZURE_LLAMA_3.1_70B_API"),
-            "mistral-large": os.environ.get("MISTRAL_LARGE_API"),
-            "common-model": os.environ.get("MISTRAL_LARGE_API")
-        }
+        self._client = None
+
+    def _ensure_client(self):
+        """
+        Create the Azure client only when first used.
+        Raise a clear error if env vars are missing.
+        """
+        if self._client is not None:
+            return
+
+        if not self.api_key or not isinstance(self.api_key, str):
+            raise RuntimeError(
+                "Azure provider misconfigured: AZURE_AI_API_KEY is missing or not a string."
+            )
+        if not self.endpoint:
+            raise RuntimeError(
+                "Azure provider misconfigured: AZURE_AI_ENDPOINT is missing."
+            )
+
+        credential = AzureKeyCredential(self.api_key)
+        self._client = ChatCompletionsClient(
+            endpoint=self.endpoint,
+            credential=credential,
+            api_version="2024-05-01-preview",
+        )
 
     def get_model_name(self, model):
         """Retrieve the model name based on the input key."""
         return self.model_map.get(model, None)
 
-    def get_model_api_key(self, model):
-        """Retrieve the API key for a specific model."""
-        api_key = self.model_api_keys.get(model)
-        if not api_key:
-            raise ValueError(
-                f"No API key found for model '{model}'. Ensure it is set in environment variables."
-            )
-        return api_key
-
     def perform_inference(self, model, prompt, max_output=100, verbosity=True):
         """Performs non-streaming inference request to Azure."""
         try:
+            self._ensure_client()
+            client = self._client
             model_id = self.get_model_name(model)
-            api_key = self.get_model_api_key(model)
             if model_id is None:
                 print(f"Model {model} not available.")
                 return None
             start_time = timer()
-            endpoint = f"https://{model_id}.eastus.models.ai.azure.com/chat/completions"
-            response = requests.post(
-                f"{endpoint}",
-                headers={
-                    "Authorization": f"Bearer {api_key}",
-                    "Content-Type": "application/json",
-                },
-                json={
-                    "messages": [
-                        {"role": "system", "content": self.system_prompt},
-                        {"role": "user", "content": prompt},
-                    ],
-                    "max_tokens": max_output,
-                },
-                timeout=500,
+            response = client.complete(
+                messages=[
+                    SystemMessage(content=self.system_prompt),
+                    UserMessage(content=prompt),
+                ],
+                max_tokens=max_output,
+                model=model_id
             )
             elapsed = timer() - start_time
-            if response.status_code != 200:
-                print(f"Error: {response.status_code} - {response.text}")
-                return None
 
-            # Parse and display response
-            inference = response.json()
+            usage = response.get("usage")
+            total_tokens = usage.get("completion_tokens") or 0
+            tbt = elapsed / max(total_tokens, 1)
+            tps = (total_tokens / elapsed)
+
             self.log_metrics(model, "response_times", elapsed)
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
+
             if verbosity:
-                print(f"Response: {inference['choices'][0]['message']['content']}")
-            return inference
+                print(f"Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
+                print(f"Response: {response['choices'][0]['message']['content']}")
+            return response
         
         except Exception as e:
             print(f"[ERROR] Inference failed for model '{model}': {e}")
@@ -87,72 +97,46 @@ def perform_inference_streaming(
         self, model, prompt, max_output=100, verbosity=True
     ):
         """Performs streaming inference request to Azure."""
+        self._ensure_client()
+        client = self._client
         model_id = self.get_model_name(model)
-        api_key = self.get_model_api_key(model)
         if model_id is None:
             print(f"Model {model} not available.")
             return None
 
         inter_token_latencies = []
-        endpoint = f"https://{model_id}.eastus.models.ai.azure.com/chat/completions"
         start_time = timer()
         try:
-            response = requests.post(
-                f"{endpoint}",
-                headers={
-                    "Authorization": f"Bearer {api_key}",
-                    "Content-Type": "application/json",
-                },
-                json={
-                    "messages": [
-                        # {"role": "system", "content": self.system_prompt + "\nThe number appended at the end is not important."},
-                        # {"role": "user", "content": prompt + " " + str(timer())},
-                        {"role": "system", "content": self.system_prompt},
-                        {"role": "user", "content": prompt},
-                    ],
-                    "max_tokens": max_output,
-                    "stream": True,
-                },
-                stream=True,
-                timeout=500,
-            )
-
             first_token_time = None
-            for line in response.iter_lines():
-                if line:
-                    # print(line)
-                    if first_token_time is None:
-                        # print(line)
-                        first_token_time = timer()
-                        ttft = first_token_time - start_time
-                        prev_token_time = first_token_time
-                        if verbosity:
+            with client.complete(
+                stream=True,
+                messages=[
+                    SystemMessage(content=self.system_prompt),
+                    UserMessage(content=prompt),
+                ],
+                max_tokens=max_output,
+                model=model_id
+            ) as response:
+                for event in response:
+                    if not event.choices or not event.choices[0].delta:
+                        continue
+
+                    delta = event.choices[0].delta
+                    if delta.content:
+                        if first_token_time is None:
+                            first_token_time = timer()
+                            ttft = first_token_time - start_time
+                            prev_token_time = first_token_time
                             print(f"##### Time to First Token (TTFT): {ttft:.4f} seconds\n")
 
-                    line_str = line.decode("utf-8").strip()
-                    
-                    if line_str == "data: [DONE]":
-                        # print(line_str)
-                        # print("here")
-                        total_time = timer() - start_time
-                        break
-
-                    # Capture token timing
-                    time_to_next_token = timer()
-                    inter_token_latency = time_to_next_token - prev_token_time
-                    prev_token_time = time_to_next_token
-                    inter_token_latencies.append(inter_token_latency)
-
-                    # Display token if verbosity is enabled
-                    match = re.search(r'"content"\s*:\s*"(.*?)"', line_str)
-                    if match:
-                        print(match.group(1), end="")
-                    # if verbosity:
-                    #     if len(inter_token_latencies) < 20:
-                    #         print(line_str[19:].split('"')[5], end="")
-                    #     elif len(inter_token_latencies) == 20:
-                    #         print("...")
+                        time_to_next_token = timer()
+                        inter_token_latency = time_to_next_token - prev_token_time
+                        prev_token_time = time_to_next_token
+                        inter_token_latencies.append(inter_token_latency)
+
+                        print(delta.content, end="", flush=True)
 
+            total_time = timer() - start_time
             # Calculate total metrics
 
             if verbosity:
diff --git a/providers/base_provider.py b/providers/base_provider.py
@@ -37,8 +37,22 @@ def perform_inference(self, model, prompt, max_output=100, verbosity=True):
                 timeout=(1, 2)
             )
             elapsed = timer() - start
+
+            usage = getattr(response, "usage", None)
+            total_tokens = 0
+            if usage:
+                total_tokens = getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", None) or 0
+
+            tbt = elapsed / max(total_tokens, 1)
+            tps = (total_tokens / elapsed)
+
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
             self.log_metrics(model, "response_times", elapsed)
+
             if verbosity:
+                print(f"Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
                 self.display_response(response, elapsed)
             return elapsed
             
diff --git a/providers/cloudflare_provider.py b/providers/cloudflare_provider.py
@@ -61,14 +61,25 @@ def perform_inference(self, model, prompt, max_output=100, verbosity=True):
             )
 
             elapsed = timer() - start_time
-            # print("request sucess")
-            # log response times metric
-            self.log_metrics(model, "response_times", elapsed)
 
             inference = response.json()
+
+            meta = inference.get("result", {})
+            usage = meta.get("usage", {})
+            total_tokens = usage.get("completion_tokens") or 0
+
+            tbt = elapsed / max(total_tokens, 1)
+            tps = (total_tokens / elapsed)
+
+            self.log_metrics(model, "response_times", elapsed)
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
+            
             print(inference)
             # logging.debug(inference["result"]["response"])
             if verbosity:
+                print(f"Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
                 print(inference["result"]["response"][:50])
 
                 print(f"#### _Generated in *{elapsed:.2f}* seconds_")
diff --git a/providers/google_provider.py b/providers/google_provider.py
@@ -60,8 +60,19 @@ def perform_inference(self, model, prompt, max_output=100, verbosity=True):
             )
             elapsed = timer() - start_time
 
+            usage = getattr(response, "usage_metadata", None)
+            total_tokens = (getattr(usage, "candidates_token_count", 0) or 0) if usage else 0
+
+            tbt = elapsed / max(total_tokens, 1)
+            tps = (total_tokens / elapsed)
+
             self.log_metrics(model, "response_times", elapsed)
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
+
             if verbosity:
+                print(f"Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
                 print(response.text)
                 print(f"\nGenerated in {elapsed:.2f} seconds")
             return elapsed
diff --git a/providers/vllm_provider.py b/providers/vllm_provider.py
@@ -49,10 +49,20 @@ def perform_inference(self, model, prompt, vllm_ip, max_output=100, verbosity=Tr
             )
             elapsed = timer() - start_time
 
-            # Log response times metric
+            data = response.json()
+            usage = data.get("usage") or {}
+            total_tokens = usage.get("completion_tokens")
+
+            tbt = elapsed / max(total_tokens, 1)
+            tps = (total_tokens / elapsed)
+
+            self.log_metrics(model, "totaltokens", total_tokens)
+            self.log_metrics(model, "timebetweentokens", tbt)
+            self.log_metrics(model, "tps", tps)
             self.log_metrics(model, "response_times", elapsed)
 
             if verbosity:
+                print(f"Tokens: {total_tokens}, Avg TBT: {tbt:.4f}s, TPS: {tps:.2f}")
                 print(f"#### _Generated in *{elapsed:.2f}* seconds_")
             
             print(response)
diff --git a/requirements.txt b/requirements.txt
@@ -18,4 +18,5 @@ groq==0.13.0
 google-generativeai==0.8.3
 fastapi==0.115.6
 uvicorn==0.32.1
-pytest-asyncio==0.25.0
+pytest-asyncio==0.25.0
+azure-ai-inference==1.0.0b9
diff --git a/test_files/providers/base_provider_test.py b/test_files/providers/base_provider_test.py
diff --git a/test_files/providers/google_provider_test.py b/test_files/providers/google_provider_test.py
