correct version pin for aiohttp

Hotfix: pin pycares

lib/data_types.py

@@ -66,7 +66,7 @@ class AuthData:
     """data used to authenticate requester"""
 
     cost: str
-    endpoint_id: int
+    endpoint: str
     reqnum: int
     request_idx: int
     signature: str

lib/test_utils.py

@@ -75,7 +75,6 @@ def print_truncate_res(res: str):
 @dataclass
 class ClientState:
     endpoint_group_name: str
-    endpoint_id: int
     api_key: str
     server_url: str
     worker_endpoint: str
@@ -96,7 +95,7 @@ class ClientState:
             self.status = ClientStatus.Error
             return
         route_payload = {
-            "endpoint_id": self.endpoint_id,
+            "endpoint": self.endpoint_group_name,
             "api_key": self.api_key,
             "cost": self.payload.count_workload(),
         }
@@ -245,19 +244,16 @@ def run_test(
     print_thread = threading.Thread(target=print_state, args=(clients, num_requests))
     print_thread.daemon = True  # makes threads get killed on program exit
     print_thread.start()
-    endpoint_info = Endpoint.get_endpoint_info(
+    endpoint_api_key = Endpoint.get_endpoint_api_key(
         endpoint_name=endpoint_group_name, account_api_key=api_key, instance=instance
     )
-    if not endpoint_info:
+    if not endpoint_api_key:
         log.debug(f"Endpoint {endpoint_group_name} not found for API key")
         return
-    endpoint_id = endpoint_info["id"]
-    endpoint_api_key = endpoint_info["api_key"]
     try:
         for _ in range(num_requests):
             client = ClientState(
                 endpoint_group_name=endpoint_group_name,
-                endpoint_id=endpoint_id,
                 api_key=endpoint_api_key,
                 server_url=server_url,
                 worker_endpoint=worker_endpoint,

requirements.txt

@@ -1,4 +1,6 @@
-aiohttp[speedups]==3.10.1
+aiohttp==3.10.1
+aiodns~=3.6.0
+pycares~=4.11.0
 anyio~=4.4
 lib~=4.0
 nltk~=3.9
@@ -8,4 +10,4 @@ Requests~=2.32
 transformers~=4.52
 utils==1.0.*
 hf_transfer>=0.1.9
-vastai-sdk>=0.2.0
+vastai-sdk>=0.2.0

workers/comfyui/client.py

@@ -13,11 +13,11 @@ from vastai import Serverless
 ENDPOINT_NAME = "my-comfyui-endpoint"
 COST = 100 # Use a constant cost for image generation
 
-def call_default_workflow(endpoint_id: int, api_key: str, server_url: str) -> None:
+def call_default_workflow(client: Serverless) -> None:
     WORKER_ENDPOINT = "/prompt"
     COST = 100
     route_payload = {
-        "endpoint_id": endpoint_id,
+        "endpoint": endpoint_group_name,
         "api_key": api_key,
         "cost": COST,
     }
@@ -32,7 +32,7 @@ def call_default_workflow(endpoint_id: int, api_key: str, server_url: str) -> No
     auth_data = dict(
         signature=message["signature"],
         cost=message["cost"],
-        endpoint_id=message["endpoint_id"],
+        endpoint=message["endpoint"],
         reqnum=message["reqnum"],
         url=message["url"],
     )
@@ -52,12 +52,12 @@ def call_default_workflow(endpoint_id: int, api_key: str, server_url: str) -> No
 
 
 def call_custom_workflow_for_sd3(
-    endpoint_id: int, api_key: str, server_url: str
+    endpoint_group_name: str, api_key: str, server_url: str
 ) -> None:
     WORKER_ENDPOINT = "/custom-workflow"
     COST = 100
     route_payload = {
-        "endpoint_id": endpoint_id,
+        "endpoint": endpoint_group_name,
         "api_key": api_key,
         "cost": COST,
     }
@@ -72,7 +72,7 @@ def call_custom_workflow_for_sd3(
     auth_data = dict(
         signature=message["signature"],
         cost=message["cost"],
-        endpoint_id=message["endpoint_id"],
+        endpoint=message["endpoint"],
         reqnum=message["reqnum"],
         url=message["url"],
         request_idx=message["request_idx"],
@@ -146,28 +146,25 @@ def call_custom_workflow_for_sd3(
 if __name__ == "__main__":
     from lib.test_utils import test_args
 
-    log = logging.getLogger(__name__)
     args = test_args.parse_args()
-    endpoint_info = Endpoint.get_endpoint_info(
+    endpoint_api_key = Endpoint.get_endpoint_api_key(
         endpoint_name=args.endpoint_group_name,
         account_api_key=args.api_key,
         instance=args.instance,
     )
-    if endpoint_info:
+    if endpoint_api_key:
-        endpoint_id = endpoint_info["id"]
-        endpoint_api_key = endpoint_info["api_key"]
         try:
             call_default_workflow(
-                endpoint_id=endpoint_id,
                 api_key=endpoint_api_key,
+                endpoint_group_name=args.endpoint_group_name,
                 server_url=args.server_url,
             )
             call_custom_workflow_for_sd3(
-                endpoint_id=endpoint_id,
                 api_key=endpoint_api_key,
+                endpoint_group_name=args.endpoint_group_name,
                 server_url=args.server_url,
             )
         except Exception as e:
             log.error(f"Error during API call: {e}")
     else:
-        log.error(f"Failed to get endpoint info for {args.endpoint_group_name}")
+        log.error(f"Failed to get API key for endpoint {args.endpoint_group_name} ")

workers/openai/test_load.py

@@ -60,7 +60,7 @@ def do_one(endpoint_name: str,
            worker_session):
     try:
         workload = payload.count_workload()
-        route_payload = {"endpoint_id": endpoint_id, "api_key": endpoint_api_key, "cost": workload}
+        route_payload = {"endpoint": endpoint_name, "api_key": endpoint_api_key, "cost": workload}
         headers = {"Authorization": f"Bearer {endpoint_api_key}"}
         start = time.time()
         r0 = route_session.post(urljoin(server_url, "/route/"), json=route_payload, headers=headers, timeout=4)