download vast.ai's root certificate in order to make pyworker requests

* Addresses breaking change in test_utils.py

Endpoint.get_endpoint_api_key() now requires instance

Moves the call to this function out of the APIClient and into main

* Ensure make_benchmark_payload has a value to calculate the workload

---------

Co-authored-by: Nader Arbabian <nader@vast.ai>

lib/backend.py

@@ -5,7 +5,7 @@ import base64
 import subprocess
 import dataclasses
 import logging
-from asyncio import sleep, gather, Semaphore
+from asyncio import wait, sleep, gather, Semaphore, FIRST_COMPLETED, create_task
 from typing import Tuple, Awaitable, NoReturn, List, Union, Callable, Optional
 from functools import cached_property
 from distutils.util import strtobool
@@ -123,6 +123,12 @@ class Backend:
             return web.json_response(dict(error="invalid JSON"), status=422)
         workload = payload.count_workload()
 
+        async def cancel_api_call_if_disconnected() -> web.Response:
+            await request.wait_for_disconnection()
+            log.debug(f"request with reqnum: {auth_data.reqnum} was canceled")
+            self.metrics._request_canceled(workload=workload, reqnum=auth_data.reqnum)
+            return web.Response(status=500)
+
         async def make_request() -> Union[web.Response, web.StreamResponse]:
             log.debug(f"got request, {auth_data.reqnum}")
             self.metrics._request_start(workload=workload, reqnum=auth_data.reqnum)
@@ -168,7 +174,15 @@ class Backend:
             return web.Response(status=401)
 
         try:
-            return await make_request()
+            done, pending = await wait(
+                [
+                    create_task(make_request()),
+                    create_task(cancel_api_call_if_disconnected()),
+                ],
+                return_when=FIRST_COMPLETED,
+            )
+            [task.cancel() for task in pending]
+            return done.pop().result()
         except Exception as e:
             log.debug(f"Exception in main handler loop {e}")
             return web.Response(status=500)

lib/server.py

@@ -27,7 +27,7 @@ def start_server(backend: Backend, routes: List[web.RouteDef], **kwargs):
         log.debug("starting server...")
         app = web.Application()
         app.add_routes(routes)
-        runner = web.AppRunner(app, handler_cancellation=True)
+        runner = web.AppRunner(app)
         await runner.setup()
         site = web.TCPSite(
             runner,

lib/test_utils.py

@@ -10,6 +10,7 @@ from collections import Counter
 from dataclasses import dataclass, field, asdict
 from urllib.parse import urljoin
 from utils.endpoint_util import Endpoint
+from utils.ssl import get_cert_file_path
 import requests
 
 from lib.data_types import AuthData, ApiPayload
@@ -120,9 +121,11 @@ class ClientState:
         self.url = worker_address
         url = urljoin(worker_address, self.worker_endpoint)
         self.status = ClientStatus.Generating
+
         response = requests.post(
             url,
             json=req_data,
+            verify=get_cert_file_path(),
         )
         if response.status_code != 200:
             self.infer_error.append(

requirements.txt

@@ -1,4 +1,4 @@
-aiohttp~=3.11
+aiohttp==3.10.1
 anyio~=4.4
 lib~=4.0
 nltk~=3.9

utils/ssl.py

@@ -0,0 +1,15 @@
+import tempfile
+from functools import cache
+
+import requests
+
+
+@cache
+def get_cert_file_path():
+    cert_url = "https://console.vast.ai/static/jvastai_root.cer"
+    response = requests.get(cert_url)
+    response.raise_for_status()
+    # Use a temporary file that is not deleted on close
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".cer", mode="wb") as f:
+        f.write(response.content)
+        return f.name

workers/comfyui/client.py

@@ -5,6 +5,7 @@ import requests
 
 from lib.test_utils import print_truncate_res
 from utils.endpoint_util import Endpoint
+from utils.ssl import get_cert_file_path
 
 """
 NOTE: this client example uses a custom comfy workflow compatible with SD3 only
@@ -51,6 +52,7 @@ def call_default_workflow(
     response = requests.post(
         url,
         json=req_data,
+        verify=get_cert_file_path(),
     )
     response.raise_for_status()
     print_truncate_res(str(response.json()))
@@ -141,6 +143,7 @@ def call_custom_workflow_for_sd3(
     response = requests.post(
         url,
         json=req_data,
+        verify=get_cert_file_path(),
     )
     response.raise_for_status()
     print_truncate_res(str(response.json()))

workers/openai/client.py

@@ -6,6 +6,7 @@ from urllib.parse import urljoin
 from typing import Dict, Any, Optional, Iterator, Union, List
 import requests
 from utils.endpoint_util import Endpoint
+from utils.ssl import get_cert_file_path
 from .data_types.client import CompletionConfig, ChatCompletionConfig
 
 logging.basicConfig(
@@ -19,40 +20,37 @@ COMPLETIONS_PROMPT = "the capital of USA is"
 CHAT_PROMPT = "Think step by step: Tell me about the Python programming language."
 TOOLS_PROMPT = "Can you list the files in the current working directory and tell me what you see? What do you think this directory might be for?"
 
+
 class APIClient:
     """Lightweight client focused solely on API communication"""
-    
+
     # Remove the generic WORKER_ENDPOINT since we're now going direct
     DEFAULT_COST = 100
     DEFAULT_TIMEOUT = 4
-    
-    def __init__(self, endpoint_group_name: str, api_key: str, server_url: str):
+
+    def __init__(
+        self,
+        endpoint_group_name: str,
+        api_key: str,
+        server_url: str,
+        endpoint_api_key: str,
+    ):
         self.endpoint_group_name = endpoint_group_name
         self.api_key = api_key
         self.server_url = server_url
-        self.endpoint_api_key = self._get_endpoint_api_key()
-    
-    def _get_endpoint_api_key(self) -> Optional[str]:
-        """Get the endpoint API key"""
-        endpoint_api_key = Endpoint.get_endpoint_api_key(
-            endpoint_name=self.endpoint_group_name,
-            account_api_key=self.api_key,
-        )
-        if not endpoint_api_key:
-            log.error(f"Failed to get API key for endpoint {self.endpoint_group_name}")
-        return endpoint_api_key
-    
+        self.endpoint_api_key = endpoint_api_key
+
     def _get_worker_url(self, cost: int = DEFAULT_COST) -> Dict[str, Any]:
         """Get worker URL and auth data from routing service"""
         if not self.endpoint_api_key:
             raise ValueError("No valid endpoint API key available")
-            
+
         route_payload = {
             "endpoint": self.endpoint_group_name,
             "api_key": self.endpoint_api_key,
             "cost": cost,
         }
-        
+
         response = requests.post(
             urljoin(self.server_url, "/route/"),
             json=route_payload,
@@ -60,7 +58,7 @@ class APIClient:
         )
         response.raise_for_status()
         return response.json()
-    
+
     def _create_auth_data(self, message: Dict[str, Any]) -> Dict[str, Any]:
         """Create auth data from routing response"""
         return {
@@ -70,42 +68,46 @@ class APIClient:
             "reqnum": message["reqnum"],
             "url": message["url"],
         }
-    
-    def _make_request(self, payload: Dict[str, Any], endpoint: str, method: str = "POST", 
-                     stream: bool = False) -> Union[Dict[str, Any], Iterator[str]]:
+
+    def _make_request(
+        self,
+        payload: Dict[str, Any],
+        endpoint: str,
+        method: str = "POST",
+        stream: bool = False,
+    ) -> Union[Dict[str, Any], Iterator[str]]:
         """Make request directly to the specific worker endpoint"""
         # Get worker URL and auth data
-        cost = payload.get('max_tokens')
+        cost = payload.get("max_tokens", self.DEFAULT_COST)
         message = self._get_worker_url(cost=cost)
         worker_url = message["url"]
         auth_data = self._create_auth_data(message)
-        
-        req_data = {
-            "payload": {
-                "input": payload
-            },
-            "auth_data": auth_data
-        }
+
+        req_data = {"payload": {"input": payload}, "auth_data": auth_data}
 
         url = urljoin(worker_url, endpoint)
         log.debug(f"Making direct request to: {url}")
         log.debug(f"Payload: {req_data}")
-        
+
         # Make the request using the specified method
         if method.upper() == "POST":
-            response = requests.post(url, json=req_data, stream=stream)
+            response = requests.post(
+                url, json=req_data, stream=stream, verify=get_cert_file_path()
+            )
         elif method.upper() == "GET":
-            response = requests.get(url, params=req_data, stream=stream)
+            response = requests.get(
+                url, params=req_data, stream=stream, verify=get_cert_file_path()
+            )
         else:
             raise ValueError(f"Unsupported HTTP method: {method}")
-            
+
         response.raise_for_status()
-        
+
         if stream:
             return self._handle_streaming_response(response)
         else:
             return response.json()
-    
+
     def _handle_streaming_response(self, response: requests.Response) -> Iterator[str]:
         """Handle streaming response and yield tokens"""
         try:
@@ -124,61 +126,60 @@ class APIClient:
             log.error(f"Error handling streaming response: {e}")
             raise
 
-    
-    def call_completions(self, config: CompletionConfig) -> Union[Dict[str, Any], Iterator[str]]:
-        payload = config.to_dict()
-        
-        return self._make_request(
-            payload=payload,
-            endpoint="/v1/completions",
-            stream=config.stream
-        )
-    
-    def call_chat_completions(self, config: ChatCompletionConfig) -> Union[Dict[str, Any], Iterator[str]]:
+    def call_completions(
+        self, config: CompletionConfig
+    ) -> Union[Dict[str, Any], Iterator[str]]:
         payload = config.to_dict()
 
         return self._make_request(
-            payload=payload,
-            endpoint="/v1/chat/completions",
-            stream=config.stream
+            payload=payload, endpoint="/v1/completions", stream=config.stream
+        )
+
+    def call_chat_completions(
+        self, config: ChatCompletionConfig
+    ) -> Union[Dict[str, Any], Iterator[str]]:
+        payload = config.to_dict()
+
+        return self._make_request(
+            payload=payload, endpoint="/v1/chat/completions", stream=config.stream
         )
 
 
 class ToolManager:
     """Handles tool definitions and execution"""
-    
+
     @staticmethod
     def list_files() -> str:
         """Execute ls on current directory"""
         try:
-            result = subprocess.run(['ls', '-la', '.'], capture_output=True, text=True, timeout=10)
+            result = subprocess.run(
+                ["ls", "-la", "."], capture_output=True, text=True, timeout=10
+            )
             if result.returncode == 0:
                 return result.stdout
             else:
                 return f"Error: {result.stderr}"
         except Exception as e:
             return f"Error running ls: {e}"
-    
+
     @staticmethod
     def get_ls_tool_definition() -> List[Dict[str, Any]]:
         """Get the ls tool definition"""
-        return [{
-            "type": "function",
-            "function": {
-                "name": "list_files",
-                "description": "List files and directories in the cwd",
-                "parameters": {
-                    "type": "object",
-                    "properties": {},
-                    "required": []
-                }
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "list_files",
+                    "description": "List files and directories in the cwd",
+                    "parameters": {"type": "object", "properties": {}, "required": []},
+                },
             }
-        }]
-    
+        ]
+
     def execute_tool_call(self, tool_call: Dict[str, Any]) -> str:
         """Execute a tool call and return the result"""
         function_name = tool_call["function"]["name"]
-        
+
         if function_name == "list_files":
             return self.list_files()
         else:
@@ -187,13 +188,17 @@ class ToolManager:
 
 class APIDemo:
     """Demo and testing functionality for the API client"""
-    
-    def __init__(self, client: APIClient, model: str, tool_manager: ToolManager = None):
+
+    def __init__(
+        self, client: APIClient, model: str, tool_manager: Optional[ToolManager] = None
+    ):
         self.client = client
         self.model = model
         self.tool_manager = tool_manager or ToolManager()
-    
-    def handle_streaming_response(self, response_stream, show_reasoning: bool = True) -> str:
+
+    def handle_streaming_response(
+        self, response_stream, show_reasoning: bool = True
+    ) -> str:
         """
         Handle streaming chat response and display all output.
         """
@@ -260,178 +265,181 @@ class APIDemo:
 
         return full_response
 
-
     def test_tool_support(self) -> bool:
         """Test if the endpoint supports function calling"""
         log.debug("Testing endpoint tool calling support...")
-        
+
         # Try a simple request with minimal tools to test support
         messages = [{"role": "user", "content": "Hello"}]
-        minimal_tool = [{
-            "type": "function",
-            "function": {
-                "name": "test_function",
-                "description": "Test function"
+        minimal_tool = [
+            {
+                "type": "function",
+                "function": {"name": "test_function", "description": "Test function"},
             }
-        }]
-        
+        ]
+
         config = ChatCompletionConfig(
             model=self.model,
             messages=messages,
             max_tokens=10,
             tools=minimal_tool,
-            tool_choice="none"  # Don't actually call the tool
+            tool_choice="none",  # Don't actually call the tool
         )
-        
+
         try:
             response = self.client.call_chat_completions(config)
             return True
         except Exception as e:
             log.error(f"Error: Endpoint does not support tool calling: {e}")
             return False
-    
+
     def demo_completions(self) -> None:
         """Demo: test basic completions endpoint"""
         print("=" * 60)
         print("COMPLETIONS DEMO")
         print("=" * 60)
-        
+
         config = CompletionConfig(
-            model=self.model,
-            prompt=COMPLETIONS_PROMPT,
-            stream=False
+            model=self.model, prompt=COMPLETIONS_PROMPT, stream=False
+        )
+
+        log.info(
+            f"Testing completions with model '{self.model}' and prompt: '{config.prompt}'"
         )
-        
-        log.info(f"Testing completions with model '{self.model}' and prompt: '{config.prompt}'")
         response = self.client.call_completions(config)
-        
+
         if isinstance(response, dict):
             print("\nResponse:")
             print(json.dumps(response, indent=2))
         else:
             log.error("Unexpected response format")
-    
+
     def demo_chat(self, use_streaming: bool = True) -> None:
         """
         Demo: test chat completions endpoint with optional streaming
         """
         print("=" * 60)
-        print(f"CHAT COMPLETIONS DEMO {'(STREAMING)' if use_streaming else '(NON-STREAMING)'}")
+        print(
+            f"CHAT COMPLETIONS DEMO {'(STREAMING)' if use_streaming else '(NON-STREAMING)'}"
+        )
         print("=" * 60)
-        
+
         config = ChatCompletionConfig(
             model=self.model,
             messages=[{"role": "user", "content": CHAT_PROMPT}],
             stream=use_streaming,
         )
-        
+
         log.info(f"Testing chat completions with model '{self.model}'...")
         response = self.client.call_chat_completions(config)
-        
+
         if use_streaming:
             try:
                 self.handle_streaming_response(response, show_reasoning=True)
             except Exception as e:
                 log.error(f"\nError during streaming: {e}")
                 import traceback
+
                 traceback.print_exc()
                 return
-            
+
         else:
             if isinstance(response, dict):
                 choice = response.get("choices", [{}])[0]
                 message = choice.get("message", {})
                 content = message.get("content", "")
-                reasoning = message.get("reasoning_content", "") or message.get("reasoning", "")
-                
+                reasoning = message.get("reasoning_content", "") or message.get(
+                    "reasoning", ""
+                )
+
                 if reasoning:
                     print(f"\n🧠 Reasoning: \033[90m{reasoning}\033[0m")
-                
+
                 print(f"\n💬 Assistant: {content}")
                 print(f"\nFull Response:")
                 print(json.dumps(response, indent=2))
             else:
                 log.error("Unexpected response format")
 
-
-    
     def demo_ls_tool(self) -> None:
         """Demo: ask LLM to list files in the current directory and describe what it sees"""
         print("=" * 60)
         print("TOOL USE DEMO: List Directory Contents")
         print("=" * 60)
-        
+
         # Test if tools are supported first
         if not self.test_tool_support():
             return
-        
+
         # Request with tool available
-        messages = [
-            {"role": "user", "content": TOOLS_PROMPT}
-        ]
-        
+        messages = [{"role": "user", "content": TOOLS_PROMPT}]
+
         config = ChatCompletionConfig(
             model=self.model,
             messages=messages,
             tools=self.tool_manager.get_ls_tool_definition(),
-            tool_choice="auto"
+            tool_choice="auto",
         )
-        
+
         log.info(f"Making initial request with tool using model '{self.model}'...")
         response = self.client.call_chat_completions(config)
-        
+
         if not isinstance(response, dict):
             raise ValueError("Expected dict response for tool use")
-        
+
         choice = response.get("choices", [{}])[0]
         message = choice.get("message", {})
-        
+
         print(f"Assistant response: {message.get('content', 'No content')}")
-        
+
         # Check for tool calls
         tool_calls = message.get("tool_calls")
         if not tool_calls:
-            raise ValueError("No tool calls made - model may not support function calling")
-        
+            raise ValueError(
+                "No tool calls made - model may not support function calling"
+            )
+
         print(f"Tool calls detected: {len(tool_calls)}")
-        
+
         # Execute the tool call
         for tool_call in tool_calls:
             function_name = tool_call["function"]["name"]
             print(f"Executing tool: {function_name}")
-            
+
             tool_result = self.tool_manager.execute_tool_call(tool_call)
             print(f"Tool result:\n{tool_result}")
-            
+
             # Add tool result and continue conversation
             messages.append(message)  # Add assistant's message with tool call
-            messages.append({
-                "role": "tool",
-                "tool_call_id": tool_call["id"],
-                "content": tool_result
-            })
-            
+            messages.append(
+                {
+                    "role": "tool",
+                    "tool_call_id": tool_call["id"],
+                    "content": tool_result,
+                }
+            )
+
             # Get final response
             final_config = ChatCompletionConfig(
                 model=self.model,
                 messages=messages,
-                tools=self.tool_manager.get_ls_tool_definition()
+                tools=self.tool_manager.get_ls_tool_definition(),
             )
-            
+
             print("Getting final response...")
             final_response = self.client.call_chat_completions(final_config)
-            
+
             if isinstance(final_response, dict):
                 final_choice = final_response.get("choices", [{}])[0]
                 final_message = final_choice.get("message", {})
                 final_content = final_message.get("content", "")
-                
+
                 print("\n" + "=" * 60)
                 print("FINAL LLM ANALYSIS:")
                 print("=" * 60)
                 print(final_content)
                 print("=" * 60)
-    
+
     def interactive_chat(self) -> None:
         """Interactive chat session with streaming"""
         print("=" * 60)
@@ -440,40 +448,39 @@ class APIDemo:
         print(f"Using model: {self.model}")
         print("Type 'quit' to exit, 'clear' to clear history")
         print()
-        
+
         messages = []
-        
+
         while True:
             try:
                 user_input = input("You: ").strip()
-                
-                if user_input.lower() == 'quit':
+
+                if user_input.lower() == "quit":
                     print("👋 Goodbye!")
                     break
-                elif user_input.lower() == 'clear':
+                elif user_input.lower() == "clear":
                     messages = []
                     print("Chat history cleared")
                     continue
                 elif not user_input:
                     continue
-                
+
                 messages.append({"role": "user", "content": user_input})
-                
+
                 config = ChatCompletionConfig(
-                    model=self.model,
-                    messages=messages,
-                    stream=True,
-                    temperature=0.7
+                    model=self.model, messages=messages, stream=True, temperature=0.7
                 )
-                
+
                 print("Assistant: ", end="", flush=True)
-                
+
                 response = self.client.call_chat_completions(config)
-                assistant_content = self.handle_streaming_response(response, show_reasoning=True)
-                
+                assistant_content = self.handle_streaming_response(
+                    response, show_reasoning=True
+                )
+
                 # Add assistant response to conversation history
                 messages.append({"role": "assistant", "content": assistant_content})
-                
+
             except KeyboardInterrupt:
                 print("\n👋 Chat interrupted. Goodbye!")
                 break
@@ -485,50 +492,49 @@ class APIDemo:
 def main():
     """Main function with CLI switches for different tests"""
     from lib.test_utils import test_args
-    
+
     # Add mandatory model argument
     test_args.add_argument(
-        "--model", 
-        required=True,
-        help="Model to use for requests (required)"
+        "--model", required=True, help="Model to use for requests (required)"
     )
-    
+
     # Add test mode arguments
     test_args.add_argument(
-        "--completion", 
-        action="store_true",
-        help="Test completions endpoint"
+        "--completion", action="store_true", help="Test completions endpoint"
     )
     test_args.add_argument(
-        "--chat", 
+        "--chat",
         action="store_true",
-        help="Test chat completions endpoint (non-streaming)"
+        help="Test chat completions endpoint (non-streaming)",
     )
     test_args.add_argument(
-        "--chat-stream", 
+        "--chat-stream",
         action="store_true",
-        help="Test chat completions endpoint with streaming"
+        help="Test chat completions endpoint with streaming",
     )
     test_args.add_argument(
-        "--tools", 
+        "--tools",
         action="store_true",
-        help="Test function calling with ls tool (non-streaming)"
+        help="Test function calling with ls tool (non-streaming)",
     )
     test_args.add_argument(
-        "--interactive", 
+        "--interactive",
         action="store_true",
-        help="Start interactive streaming chat session"
+        help="Start interactive streaming chat session",
     )
-    
+
     args = test_args.parse_args()
-    
+
     # Check that only one test mode is selected
     test_modes = [
-        args.completion, args.chat, args.chat_stream, 
-        args.tools, args.interactive
+        args.completion,
+        args.chat,
+        args.chat_stream,
+        args.tools,
+        args.interactive,
     ]
     selected_count = sum(test_modes)
-    
+
     if selected_count == 0:
         print("Please specify exactly one test mode:")
         print("  --completion    : Test completions endpoint")
@@ -536,27 +542,42 @@ def main():
         print("  --chat-stream   : Test chat completions endpoint with streaming")
         print("  --tools         : Test function calling with ls tool (non-streaming)")
         print("  --interactive   : Start interactive streaming chat session")
-        print(f"\nExample: python {sys.argv[0]} --model Qwen/Qwen3-8B --chat-stream -k YOUR_KEY -e YOUR_ENDPOINT")
+        print(
+            f"\nExample: python {sys.argv[0]} --model Qwen/Qwen3-8B --chat-stream -k YOUR_KEY -e YOUR_ENDPOINT"
+        )
         sys.exit(1)
     elif selected_count > 1:
         print("Please specify exactly one test mode")
         sys.exit(1)
-    
+
     try:
+        endpoint_api_key = Endpoint.get_endpoint_api_key(
+            endpoint_name=args.endpoint_group_name,
+            account_api_key=args.api_key,
+            instance=args.instance,
+        )
+
+        if not endpoint_api_key:
+            log.error(
+                f"Could not retrieve API key for endpoint '{args.endpoint_group_name}'. Exiting."
+            )
+            sys.exit(1)
+
         # Create the core API client
         client = APIClient(
             endpoint_group_name=args.endpoint_group_name,
             api_key=args.api_key,
-            server_url=args.server_url
+            server_url=args.server_url,
+            endpoint_api_key=endpoint_api_key,
         )
-        
+
         # Create tool manager and demo (passing the model parameter)
         tool_manager = ToolManager()
         demo = APIDemo(client, args.model, tool_manager)
-        
+
         print(f"Using model: {args.model}")
         print("=" * 60)
-        
+
         # Run the selected test
         if args.completion:
             demo.demo_completions()
@@ -568,11 +589,11 @@ def main():
             demo.demo_ls_tool()
         elif args.interactive:
             demo.interactive_chat()
-        
+
     except Exception as e:
         log.error(f"Error during test: {e}", exc_info=True)
         sys.exit(1)
 
 
 if __name__ == "__main__":
-    main()
+    main()

workers/openai/data_types/client.py

@@ -3,11 +3,13 @@ from dataclasses import dataclass, field, fields, is_dataclass
 from typing import Optional, List, Dict, Any
 
 
-class SerializableDataclass:    
+class SerializableDataclass:
     def _serialize_recursive(self, obj: Any) -> Any:
         if is_dataclass(obj):
-            return {field.name: self._serialize_recursive(getattr(obj, field.name)) 
-                    for field in fields(obj)}
+            return {
+                field.name: self._serialize_recursive(getattr(obj, field.name))
+                for field in fields(obj)
+            }
         elif isinstance(obj, dict):
             return {key: self._serialize_recursive(value) for key, value in obj.items()}
         elif isinstance(obj, (list, tuple)):
@@ -16,10 +18,10 @@ class SerializableDataclass:
             return [self._serialize_recursive(item) for item in obj]
         else:
             return obj
-    
+
     def to_dict(self) -> Dict[str, Any]:
         return self._serialize_recursive(self)
-    
+
     def to_json(self, indent: int = 2) -> str:
         return json.dumps(self.to_dict(), indent=indent)
 
@@ -27,6 +29,7 @@ class SerializableDataclass:
 @dataclass
 class CompletionConfig(SerializableDataclass):
     """Configuration for completion requests"""
+
     model: str
     prompt: str = "Hello"
     max_tokens: int = 256
@@ -39,8 +42,9 @@ class CompletionConfig(SerializableDataclass):
 @dataclass
 class ChatCompletionConfig(SerializableDataclass):
     """Configuration for chat completion requests"""
+
     model: str
-    messages: list = None
+    messages: list = field(default_factory=list)
     max_tokens: int = 2096
     temperature: float = 0.7
     top_k: int = 20
@@ -48,7 +52,7 @@ class ChatCompletionConfig(SerializableDataclass):
     stream: bool = False
     tools: Optional[List[Dict[str, Any]]] = field(default_factory=list)
     tool_choice: str = "auto"
-    
+
     def __post_init__(self):
         if self.messages is None:
-            self.messages = [{"role": "user", "content": "Hello"}]
+            self.messages = [{"role": "user", "content": "Hello"}]

workers/openai/data_types/server.py

@@ -2,7 +2,7 @@ import os, json, random
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from lib.data_types import EndpointHandler, ApiPayload, JsonDataException
-from typing import Union, Type, Dict, Any
+from typing import Union, Type, Dict, Any, Optional
 from aiohttp import web, ClientResponse
 import nltk
 import logging
@@ -10,41 +10,39 @@ import logging
 nltk.download("words")
 WORD_LIST = nltk.corpus.words.words()
 log = logging.getLogger(__name__)
-            
+
 """
 Generic dataclass accepts any dictionary in input.
 """
+
+
 @dataclass
 class GenericData(ApiPayload, ABC):
     input: Dict[str, Any]
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "GenericData":
-        return cls(
-            input=data["input"]
-        )
-    
+        return cls(input=data["input"])
+
     @classmethod
     def from_json_msg(cls, json_msg: Dict[str, Any]) -> "GenericData":
         errors = {}
-        
+
         # Validate required parameters
         required_params = ["input"]
         for param in required_params:
             if param not in json_msg:
                 errors[param] = "missing parameter"
-        
+
         if errors:
             raise JsonDataException(errors)
-        
+
         try:
             # Create clean data dict and delegate to from_dict
-            clean_data = {
-                "input": json_msg["input"]
-            }
-            
+            clean_data = {"input": json_msg["input"]}
+
             return cls.from_dict(clean_data)
-            
+
         except (json.JSONDecodeError, JsonDataException) as e:
             errors["parameters"] = str(e)
             raise JsonDataException(errors)
@@ -59,7 +57,8 @@ class GenericData(ApiPayload, ABC):
 
     def count_workload(self) -> int:
         return self.input.get("max_tokens", 0)
-    
+
+
 @dataclass
 class GenericHandler(EndpointHandler[GenericData], ABC):
 
@@ -67,10 +66,10 @@ class GenericHandler(EndpointHandler[GenericData], ABC):
     @abstractmethod
     def endpoint(self) -> str:
         pass
-    
+
     @property
-    def healthcheck_endpoint(self) -> str:
-        return os.environ.get('MODEL_HEALTH_ENDPOINT')
+    def healthcheck_endpoint(self) -> Optional[str]:
+        return os.environ.get("MODEL_HEALTH_ENDPOINT")
 
     @classmethod
     def payload_cls(cls) -> Type[GenericData]:
@@ -82,17 +81,17 @@ class GenericHandler(EndpointHandler[GenericData], ABC):
 
     async def generate_client_response(
         self, client_request: web.Request, model_response: ClientResponse
-            ) -> Union[web.Response, web.StreamResponse]:
+    ) -> Union[web.Response, web.StreamResponse]:
         match model_response.status:
             case 200:
                 # Check if the response is actually streaming based on response headers/content-type
                 is_streaming_response = (
-                    model_response.content_type == "text/event-stream" or
-                    model_response.content_type == "application/x-ndjson" or
-                    model_response.headers.get("Transfer-Encoding") == "chunked" or
-                    "stream" in model_response.content_type.lower()
+                    model_response.content_type == "text/event-stream"
+                    or model_response.content_type == "application/x-ndjson"
+                    or model_response.headers.get("Transfer-Encoding") == "chunked"
+                    or "stream" in model_response.content_type.lower()
                 )
-                
+
                 if is_streaming_response:
                     log.debug("Detected streaming response...")
                     res = web.StreamResponse()
@@ -109,12 +108,13 @@ class GenericHandler(EndpointHandler[GenericData], ABC):
                     return web.Response(
                         body=content,
                         status=200,
-                        content_type=model_response.content_type
+                        content_type=model_response.content_type,
                     )
             case code:
                 log.debug("SENDING RESPONSE: ERROR: unknown code")
                 return web.Response(status=code)
 
+
 @dataclass
 class CompletionsData(GenericData):
     @classmethod
@@ -123,55 +123,60 @@ class CompletionsData(GenericData):
         model = os.environ.get("MODEL_NAME")
         if not model:
             raise ValueError("MODEL_NAME environment variable not set")
-        
+
         test_input = {
             "model": model,
             "prompt": prompt,
-            "temperature": 0.7
+            "temperature": 0.7,
+            "max_tokens": 500,
         }
         return cls(input=test_input)
-    
+
+
 @dataclass
 class CompletionsHandler(GenericHandler):
     @property
     def endpoint(self) -> str:
         return "/v1/completions"
-    
+
     @classmethod
     def payload_cls(cls) -> Type[CompletionsData]:
         return CompletionsData
-    
+
     def make_benchmark_payload(self) -> CompletionsData:
         return CompletionsData.for_test()
 
+
 @dataclass
 class ChatCompletionsData(GenericData):
     """Chat completions-specific data implementation"""
-    
+
     @classmethod
     def for_test(cls) -> "ChatCompletionsData":
         prompt = " ".join(random.choices(WORD_LIST, k=int(250)))
         model = os.environ.get("MODEL_NAME")
         if not model:
             raise ValueError("MODEL_NAME environment variable not set")
-        
+
         # Chat completions use messages format instead of prompt
         test_input = {
             "model": model,
             "messages": [{"role": "user", "content": prompt}],
-            "temperature": 0.7
+            "temperature": 0.7,
+            "max_tokens": 500,
         }
         return cls(input=test_input)
-    
-@dataclass  
-class ChatCompletionsHandler(GenericHandler): 
+
+
+@dataclass
+class ChatCompletionsHandler(GenericHandler):
     @property
     def endpoint(self) -> str:
         return "/v1/chat/completions"
-    
+
     @classmethod
     def payload_cls(cls) -> Type[ChatCompletionsData]:
         return ChatCompletionsData
-        
+
     def make_benchmark_payload(self) -> ChatCompletionsData:
         return ChatCompletionsData.for_test()

workers/openai/server.py

@@ -7,20 +7,20 @@ from lib.server import start_server
 
 # This line indicates that the inference server is listening
 MODEL_SERVER_START_LOG_MSG = [
-    "Application startup complete.", # vLLM
-    "llama runner started", # Ollama
-    '"message":"Connected","target":"text_generation_router"', # TGI
-    '"message":"Connected","target":"text_generation_router::server"', # TGI
+    "Application startup complete.",  # vLLM
+    "llama runner started",  # Ollama
+    '"message":"Connected","target":"text_generation_router"',  # TGI
+    '"message":"Connected","target":"text_generation_router::server"',  # TGI
 ]
 
 MODEL_SERVER_ERROR_LOG_MSGS = [
-    "INFO exited: vllm", # vLLM
-    "RuntimeError: Engine", # vLLM
-    "Error: pull model manifest:", # Ollama
-    "stalled; retrying", # Ollama
-    "Error: WebserverFailed", # TGI
-    "Error: DownloadError", # TGI
-    "Error: ShardCannotStart", #TGI
+    "INFO exited: vllm",  # vLLM
+    "RuntimeError: Engine",  # vLLM
+    "Error: pull model manifest:",  # Ollama
+    "stalled; retrying",  # Ollama
+    "Error: WebserverFailed",  # TGI
+    "Error: DownloadError",  # TGI
+    "Error: ShardCannotStart",  # TGI
 ]
 
 logging.basicConfig(
@@ -31,8 +31,8 @@ logging.basicConfig(
 log = logging.getLogger(__file__)
 
 backend = Backend(
-    model_server_url=os.environ.get("MODEL_SERVER_URL"),
-    model_log_file=os.environ.get("MODEL_LOG"),
+    model_server_url=os.environ["MODEL_SERVER_URL"],
+    model_log_file=os.environ["MODEL_LOG"],
     allow_parallel_requests=True,
     benchmark_handler=CompletionsHandler(benchmark_runs=3, benchmark_words=256),
     log_actions=[
@@ -45,9 +45,11 @@ backend = Backend(
     ],
 )
 
+
 async def handle_ping(_):
     return web.Response(body="pong")
 
+
 routes = [
     web.post("/v1/completions", backend.create_handler(CompletionsHandler())),
     web.post("/v1/chat/completions", backend.create_handler(ChatCompletionsHandler())),

workers/openai/test_load.py

@@ -7,22 +7,22 @@ WORKER_ENDPOINT = "/v1/completions"
 if __name__ == "__main__":
     # Check if MODEL_NAME environment variable is set
     model_name_set = os.environ.get("MODEL_NAME") is not None
-    
+
     # Add model argument - required only if MODEL_NAME is not set
     test_args.add_argument(
-        "--model", 
+        "--model",
         dest="model",
         required=not model_name_set,
-        help="Model to use for completions request (required if MODEL_NAME env var not set)"
+        help="Model to use for completions request (required if MODEL_NAME env var not set)",
     )
-    
+
     # Parse known args to get model early, before test_load_cmd adds its args
     known_args, _ = test_args.parse_known_args()
-    
+
     # Set environment variable if model was provided
-    if hasattr(known_args, 'model') and known_args.model:
+    if hasattr(known_args, "model") and known_args.model:
         os.environ["MODEL_NAME"] = known_args.model
         print(f"Set MODEL_NAME environment variable to: {known_args.model}")
-    
+
     # Now call test_load_cmd normally - it will add its own args and re-parse
-    test_load_cmd(CompletionsData, WORKER_ENDPOINT, arg_parser=test_args)
+    test_load_cmd(CompletionsData, WORKER_ENDPOINT, arg_parser=test_args)

workers/tgi/client.py

@@ -4,6 +4,7 @@ import json
 from urllib.parse import urljoin
 import requests
 from utils.endpoint_util import Endpoint
+from utils.ssl import get_cert_file_path
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -42,7 +43,11 @@ def call_generate(endpoint_group_name: str, api_key: str, server_url: str) -> No
     req_data = dict(payload=payload, auth_data=auth_data)
     url = urljoin(url, WORKER_ENDPOINT)
     print(f"url: {url}")
-    response = requests.post(url, json=req_data)
+    response = requests.post(
+        url,
+        json=req_data,
+        verify=get_cert_file_path(),
+    )
     response.raise_for_status()
     res = response.json()
     print(res)