Compare commits

..

3 Commits

Author SHA1 Message Date
Edgar Lin fdd50a2aaa correct version pin for aiohttp 2025-12-10 19:14:34 -08:00
LucasArmandVast 70f8a8f534 Merge pull request #72 from vast-ai/hotfix-pin-pycares
Hotfix: pin pycares
2025-12-10 20:41:44 -05:00
Lucas Armand 7be8aa6397 pin pycares 2025-12-10 17:38:03 -08:00
5 changed files with 20 additions and 25 deletions
+1 -1
View File
@@ -66,7 +66,7 @@ class AuthData:
"""data used to authenticate requester""" """data used to authenticate requester"""
cost: str cost: str
endpoint_id: int endpoint: str
reqnum: int reqnum: int
request_idx: int request_idx: int
signature: str signature: str
+3 -7
View File
@@ -75,7 +75,6 @@ def print_truncate_res(res: str):
@dataclass @dataclass
class ClientState: class ClientState:
endpoint_group_name: str endpoint_group_name: str
endpoint_id: int
api_key: str api_key: str
server_url: str server_url: str
worker_endpoint: str worker_endpoint: str
@@ -96,7 +95,7 @@ class ClientState:
self.status = ClientStatus.Error self.status = ClientStatus.Error
return return
route_payload = { route_payload = {
"endpoint_id": self.endpoint_id, "endpoint": self.endpoint_group_name,
"api_key": self.api_key, "api_key": self.api_key,
"cost": self.payload.count_workload(), "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 = threading.Thread(target=print_state, args=(clients, num_requests))
print_thread.daemon = True # makes threads get killed on program exit print_thread.daemon = True # makes threads get killed on program exit
print_thread.start() 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 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") log.debug(f"Endpoint {endpoint_group_name} not found for API key")
return return
endpoint_id = endpoint_info["id"]
endpoint_api_key = endpoint_info["api_key"]
try: try:
for _ in range(num_requests): for _ in range(num_requests):
client = ClientState( client = ClientState(
endpoint_group_name=endpoint_group_name, endpoint_group_name=endpoint_group_name,
endpoint_id=endpoint_id,
api_key=endpoint_api_key, api_key=endpoint_api_key,
server_url=server_url, server_url=server_url,
worker_endpoint=worker_endpoint, worker_endpoint=worker_endpoint,
+3 -1
View File
@@ -1,4 +1,6 @@
aiohttp[speedups]==3.10.1 aiohttp==3.10.1
aiodns~=3.6.0
pycares~=4.11.0
anyio~=4.4 anyio~=4.4
lib~=4.0 lib~=4.0
nltk~=3.9 nltk~=3.9
+11 -14
View File
@@ -13,11 +13,11 @@ from vastai import Serverless
ENDPOINT_NAME = "my-comfyui-endpoint" ENDPOINT_NAME = "my-comfyui-endpoint"
COST = 100 # Use a constant cost for image generation 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" WORKER_ENDPOINT = "/prompt"
COST = 100 COST = 100
route_payload = { route_payload = {
"endpoint_id": endpoint_id, "endpoint": endpoint_group_name,
"api_key": api_key, "api_key": api_key,
"cost": COST, "cost": COST,
} }
@@ -32,7 +32,7 @@ def call_default_workflow(endpoint_id: int, api_key: str, server_url: str) -> No
auth_data = dict( auth_data = dict(
signature=message["signature"], signature=message["signature"],
cost=message["cost"], cost=message["cost"],
endpoint_id=message["endpoint_id"], endpoint=message["endpoint"],
reqnum=message["reqnum"], reqnum=message["reqnum"],
url=message["url"], 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( 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: ) -> None:
WORKER_ENDPOINT = "/custom-workflow" WORKER_ENDPOINT = "/custom-workflow"
COST = 100 COST = 100
route_payload = { route_payload = {
"endpoint_id": endpoint_id, "endpoint": endpoint_group_name,
"api_key": api_key, "api_key": api_key,
"cost": COST, "cost": COST,
} }
@@ -72,7 +72,7 @@ def call_custom_workflow_for_sd3(
auth_data = dict( auth_data = dict(
signature=message["signature"], signature=message["signature"],
cost=message["cost"], cost=message["cost"],
endpoint_id=message["endpoint_id"], endpoint=message["endpoint"],
reqnum=message["reqnum"], reqnum=message["reqnum"],
url=message["url"], url=message["url"],
request_idx=message["request_idx"], request_idx=message["request_idx"],
@@ -146,28 +146,25 @@ def call_custom_workflow_for_sd3(
if __name__ == "__main__": if __name__ == "__main__":
from lib.test_utils import test_args from lib.test_utils import test_args
log = logging.getLogger(__name__)
args = test_args.parse_args() 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, endpoint_name=args.endpoint_group_name,
account_api_key=args.api_key, account_api_key=args.api_key,
instance=args.instance, instance=args.instance,
) )
if endpoint_info: if endpoint_api_key:
endpoint_id = endpoint_info["id"]
endpoint_api_key = endpoint_info["api_key"]
try: try:
call_default_workflow( call_default_workflow(
endpoint_id=endpoint_id,
api_key=endpoint_api_key, api_key=endpoint_api_key,
endpoint_group_name=args.endpoint_group_name,
server_url=args.server_url, server_url=args.server_url,
) )
call_custom_workflow_for_sd3( call_custom_workflow_for_sd3(
endpoint_id=endpoint_id,
api_key=endpoint_api_key, api_key=endpoint_api_key,
endpoint_group_name=args.endpoint_group_name,
server_url=args.server_url, server_url=args.server_url,
) )
except Exception as e: except Exception as e:
log.error(f"Error during API call: {e}") log.error(f"Error during API call: {e}")
else: 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} ")
+1 -1
View File
@@ -60,7 +60,7 @@ def do_one(endpoint_name: str,
worker_session): worker_session):
try: try:
workload = payload.count_workload() 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}"} headers = {"Authorization": f"Bearer {endpoint_api_key}"}
start = time.time() start = time.time()
r0 = route_session.post(urljoin(server_url, "/route/"), json=route_payload, headers=headers, timeout=4) r0 = route_session.post(urljoin(server_url, "/route/"), json=route_payload, headers=headers, timeout=4)