131 lines
5.3 KiB
Python
131 lines
5.3 KiB
Python
"""Vast.ai lease inventory and administration route registration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from aiohttp import web
|
|
|
|
if __package__:
|
|
from .route_context import RouteContext
|
|
from .vast_api import VastApiClient
|
|
from .vast_leases import VastLeaseManager
|
|
from .vast_service import VastService
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from route_context import RouteContext
|
|
from vast_api import VastApiClient
|
|
from vast_leases import VastLeaseManager
|
|
from vast_service import VastService
|
|
|
|
|
|
def _vast_api_client() -> VastApiClient:
|
|
"""Build a Vast API client from the current environment."""
|
|
api_key = str(os.getenv("VAST_API_KEY") or "").strip()
|
|
if not api_key:
|
|
raise RuntimeError("Set VAST_API_KEY first.")
|
|
base_url = str(os.getenv("COMFY_MODAL_VAST_API_BASE_URL") or "").strip()
|
|
return VastApiClient(api_key, **({"base_url": base_url} if base_url else {}))
|
|
|
|
|
|
def register_vast_routes(prompt_server: Any, ctx: RouteContext) -> None:
|
|
"""Register Vast status, credential verification, reap, and destroy routes."""
|
|
registry = ctx.vast_registry
|
|
|
|
if hasattr(prompt_server.routes, "get"):
|
|
|
|
@prompt_server.routes.get("/remote/vast/status")
|
|
async def vast_status(request: web.Request) -> web.Response:
|
|
"""Return refreshed credential-free managed Vast lease inventory."""
|
|
del request
|
|
if registry is None:
|
|
return web.json_response(
|
|
{
|
|
"configured": False,
|
|
"leases": [],
|
|
"error": "ComfyUI user directory unavailable.",
|
|
},
|
|
status=503,
|
|
)
|
|
try:
|
|
if not os.getenv("VAST_API_KEY"):
|
|
state = await asyncio.to_thread(registry.load)
|
|
leases = state.leases
|
|
else:
|
|
manager = VastLeaseManager.for_inventory(
|
|
api_client=_vast_api_client(),
|
|
registry=registry,
|
|
owner_id=ctx.settings.app_name,
|
|
)
|
|
leases = await manager.refresh_owned_leases()
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
return web.json_response(
|
|
{
|
|
"configured": bool(os.getenv("VAST_API_KEY")),
|
|
"leases": [],
|
|
"error": str(exc),
|
|
},
|
|
status=502,
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"configured": bool(os.getenv("VAST_API_KEY")),
|
|
"image_configured": bool(os.getenv("COMFY_MODAL_VAST_IMAGE")),
|
|
"leases": [lease.to_dict() for lease in leases],
|
|
}
|
|
)
|
|
|
|
@prompt_server.routes.post("/remote/vast/verify")
|
|
async def vast_verify(request: web.Request) -> web.Response:
|
|
"""Verify the configured Vast credential without returning it."""
|
|
del request
|
|
try:
|
|
account = await _vast_api_client().verify_credentials()
|
|
except RuntimeError as exc:
|
|
status = 400 if not os.getenv("VAST_API_KEY") else 502
|
|
return web.json_response(
|
|
{"verified": False, "error": str(exc)}, status=status
|
|
)
|
|
except (OSError, ValueError) as exc:
|
|
return web.json_response(
|
|
{"verified": False, "error": str(exc)}, status=502
|
|
)
|
|
return web.json_response({"verified": True, "account": account})
|
|
|
|
@prompt_server.routes.post("/remote/vast/reap")
|
|
async def vast_reap(request: web.Request) -> web.Response:
|
|
"""Destroy only owned idle leases whose configured deadline has expired."""
|
|
del request
|
|
try:
|
|
service = VastService.from_environment(
|
|
ctx.settings,
|
|
repo_root=Path(__file__).resolve().parent,
|
|
)
|
|
destroyed = await service.lease_manager.destroy_expired()
|
|
except (OSError, RuntimeError, ValueError) as exc:
|
|
return web.json_response({"error": str(exc)}, status=502)
|
|
return web.json_response({"destroyed_instance_ids": list(destroyed)})
|
|
|
|
@prompt_server.routes.post("/remote/vast/destroy")
|
|
async def vast_destroy(request: web.Request) -> web.Response:
|
|
"""Destroy one exact idle registry-owned lease after ownership checks."""
|
|
try:
|
|
payload = await request.json()
|
|
instance_id = int(payload.get("instance_id"))
|
|
if registry is None:
|
|
raise RuntimeError("ComfyUI user directory is unavailable.")
|
|
manager = VastLeaseManager.for_inventory(
|
|
api_client=_vast_api_client(),
|
|
registry=registry,
|
|
owner_id=ctx.settings.app_name,
|
|
)
|
|
destroyed = await manager.destroy_owned_lease(
|
|
instance_id,
|
|
allow_active_work=payload.get("force") is True,
|
|
)
|
|
except (OSError, RuntimeError, TypeError, ValueError) as exc:
|
|
return web.json_response({"error": str(exc)}, status=502)
|
|
return web.json_response({"instance_id": instance_id, "destroyed": destroyed})
|