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"""
cost: str
endpoint_id: int
endpoint: str
reqnum: int
request_idx: int
signature: str
+3 -7
View File
@@ -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,
+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
lib~=4.0
nltk~=3.9
+11 -14
View File
@@ -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:
endpoint_id = endpoint_info["id"]
endpoint_api_key = endpoint_info["api_key"]
if endpoint_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} ")
+1 -1
View File
@@ -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)