refactor: use endpoint_id instead of endpoint name for routing

- Update route_payload to use endpoint_id instead of endpoint name
- Update AuthData to expect endpoint_id (int) instead of endpoint (str)
- Update ClientState to track endpoint_id
- Update comfyui client functions to use endpoint_id
- Fetch endpoint info (id + api_key) instead of just api_key

This aligns with the autoscaler changes in AUTO-848 that switched
to ID-based endpoint lookups for improved security and consistency.
This commit is contained in:
Colter Downing
2025-12-06 14:46:41 -08:00
parent 138fc3ac47
commit 68d8ce4bfd
4 changed files with 23 additions and 16 deletions
+14 -11
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(client: Serverless) -> None:
def call_default_workflow(endpoint_id: int, api_key: str, server_url: str) -> None:
WORKER_ENDPOINT = "/prompt"
COST = 100
route_payload = {
"endpoint": endpoint_group_name,
"endpoint_id": endpoint_id,
"api_key": api_key,
"cost": COST,
}
@@ -32,7 +32,7 @@ def call_default_workflow(client: Serverless) -> None:
auth_data = dict(
signature=message["signature"],
cost=message["cost"],
endpoint=message["endpoint"],
endpoint_id=message["endpoint_id"],
reqnum=message["reqnum"],
url=message["url"],
)
@@ -52,12 +52,12 @@ def call_default_workflow(client: Serverless) -> None:
def call_custom_workflow_for_sd3(
endpoint_group_name: str, api_key: str, server_url: str
endpoint_id: int, api_key: str, server_url: str
) -> None:
WORKER_ENDPOINT = "/custom-workflow"
COST = 100
route_payload = {
"endpoint": endpoint_group_name,
"endpoint_id": endpoint_id,
"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=message["endpoint"],
endpoint_id=message["endpoint_id"],
reqnum=message["reqnum"],
url=message["url"],
request_idx=message["request_idx"],
@@ -146,25 +146,28 @@ 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_api_key = Endpoint.get_endpoint_api_key(
endpoint_info = Endpoint.get_endpoint_info(
endpoint_name=args.endpoint_group_name,
account_api_key=args.api_key,
instance=args.instance,
)
if endpoint_api_key:
if endpoint_info:
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 API key for endpoint {args.endpoint_group_name} ")
log.error(f"Failed to get endpoint info for {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": endpoint_name, "api_key": endpoint_api_key, "cost": workload}
route_payload = {"endpoint_id": endpoint_id, "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)