Skip to content

Commit

Permalink
add agent resource sync API endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
fregataa committed May 28, 2024
1 parent afdcb6f commit fdd53c9
Showing 1 changed file with 36 additions and 1 deletion.
37 changes: 36 additions & 1 deletion src/ai/backend/manager/api/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,11 @@
import trafaret as t
from aiohttp import hdrs, web
from dateutil.tz import tzutc
from pydantic import BaseModel, Field
from pydantic import (
AliasChoices,
BaseModel,
Field,
)
from redis.asyncio import Redis
from sqlalchemy.orm import noload, selectinload
from sqlalchemy.sql.expression import null, true
Expand Down Expand Up @@ -967,6 +971,36 @@ async def sync_agent_registry(request: web.Request, params: Any) -> web.StreamRe
return web.json_response({}, status=200)


class SyncAgentResourceRequestModel(BaseModel):
agent_id: AgentId = Field(
validation_alias=AliasChoices("agent_id", "agent"),
description="Target agent id to sync resource.",
)


@server_status_required(ALL_ALLOWED)
@auth_required
@pydantic_params_api_handler(SyncAgentResourceRequestModel)
async def sync_agent_resource(
request: web.Request, params: SyncAgentResourceRequestModel
) -> web.Response:
root_ctx: RootContext = request.app["_root.context"]
requester_access_key, owner_access_key = await get_access_key_scopes(request)

agent_id = params.agent_id
log.info(
"SYNC_AGENT_RESOURCE (ak:{}/{}, a:{})", requester_access_key, owner_access_key, agent_id
)

async with root_ctx.db.begin() as db_conn:
try:
await root_ctx.registry.sync_agent_resource(db_conn, [agent_id])
except BackendError:
log.exception("SYNC_AGENT_RESOURCE: exception")
raise
return web.Response(status=204)


@server_status_required(ALL_ALLOWED)
@auth_required
@check_api_params(
Expand Down Expand Up @@ -2274,6 +2308,7 @@ def create_app(
cors.add(app.router.add_route("POST", "/_/create-cluster", create_cluster))
cors.add(app.router.add_route("GET", "/_/match", match_sessions))
cors.add(app.router.add_route("POST", "/_/sync-agent-registry", sync_agent_registry))
cors.add(app.router.add_route("POST", "/_/sync-agent-resource", sync_agent_resource))
session_resource = cors.add(app.router.add_resource(r"/{session_name}"))
cors.add(session_resource.add_route("GET", get_info))
cors.add(session_resource.add_route("PATCH", restart))
Expand Down

0 comments on commit fdd53c9

Please sign in to comment.