diff --git a/changes/9640.feature.md b/changes/9640.feature.md new file mode 100644 index 00000000000..825cd5bc75b --- /dev/null +++ b/changes/9640.feature.md @@ -0,0 +1 @@ +Add REST API endpoints for prometheus query preset CRUD and execution (BEP-1050). diff --git a/src/ai/backend/common/dto/clients/prometheus/__init__.py b/src/ai/backend/common/dto/clients/prometheus/__init__.py index 290010f811c..abe10e9c246 100644 --- a/src/ai/backend/common/dto/clients/prometheus/__init__.py +++ b/src/ai/backend/common/dto/clients/prometheus/__init__.py @@ -1,3 +1,4 @@ +from .defs import PROMETHEUS_DURATION_PATTERN from .request import QueryTimeRange from .response import ( LabelValueResponse, @@ -9,6 +10,7 @@ ) __all__ = [ + "PROMETHEUS_DURATION_PATTERN", "QueryTimeRange", "LabelValueResponse", "MetricResponse", diff --git a/src/ai/backend/common/dto/clients/prometheus/defs.py b/src/ai/backend/common/dto/clients/prometheus/defs.py new file mode 100644 index 00000000000..53e2a1e45c4 --- /dev/null +++ b/src/ai/backend/common/dto/clients/prometheus/defs.py @@ -0,0 +1,4 @@ +"""Constants for Prometheus DTOs.""" + +# Prometheus duration format: e.g. "5m", "1h30m", "10s", "1d" +PROMETHEUS_DURATION_PATTERN = r"^(?:\d+(?:ms|s|m|h|d|w|y))+$" diff --git a/src/ai/backend/common/dto/manager/prometheus_query_preset/__init__.py b/src/ai/backend/common/dto/manager/prometheus_query_preset/__init__.py new file mode 100644 index 00000000000..bc78d602bdc --- /dev/null +++ b/src/ai/backend/common/dto/manager/prometheus_query_preset/__init__.py @@ -0,0 +1,71 @@ +""" +Prometheus Query Definition DTOs for Manager API. +""" + +from .path import ( + QueryDefinitionIdPathParam, +) +from .request import ( + CreateQueryDefinitionOptionsRequest, + CreateQueryDefinitionRequest, + ExecuteQueryDefinitionOptionsRequest, + ExecuteQueryDefinitionRequest, + MetricLabelEntry, + ModifyQueryDefinitionOptionsRequest, + ModifyQueryDefinitionRequest, + QueryDefinitionFilter, + SearchQueryDefinitionsRequest, +) +from .response import ( + CreateQueryDefinitionResponse, + DeleteQueryDefinitionResponse, + ExecuteQueryDefinitionResponse, + GetQueryDefinitionResponse, + MetricLabelEntryDTO, + MetricValueDTO, + ModifyQueryDefinitionResponse, + PaginationInfo, + QueryDefinitionDTO, + QueryDefinitionExecuteData, + QueryDefinitionMetricResult, + QueryDefinitionOptionsDTO, + SearchQueryDefinitionsResponse, +) +from .types import ( + OrderDirection, + QueryDefinitionOrder, + QueryDefinitionOrderField, +) + +__all__ = ( + # Path DTOs + "QueryDefinitionIdPathParam", + # Request DTOs + "CreateQueryDefinitionOptionsRequest", + "CreateQueryDefinitionRequest", + "ExecuteQueryDefinitionOptionsRequest", + "ExecuteQueryDefinitionRequest", + "MetricLabelEntry", + "ModifyQueryDefinitionOptionsRequest", + "ModifyQueryDefinitionRequest", + "QueryDefinitionFilter", + "SearchQueryDefinitionsRequest", + # Response DTOs + "CreateQueryDefinitionResponse", + "DeleteQueryDefinitionResponse", + "ExecuteQueryDefinitionResponse", + "GetQueryDefinitionResponse", + "MetricLabelEntryDTO", + "MetricValueDTO", + "ModifyQueryDefinitionResponse", + "PaginationInfo", + "QueryDefinitionDTO", + "QueryDefinitionExecuteData", + "QueryDefinitionMetricResult", + "QueryDefinitionOptionsDTO", + "SearchQueryDefinitionsResponse", + # Types + "OrderDirection", + "QueryDefinitionOrder", + "QueryDefinitionOrderField", +) diff --git a/src/ai/backend/common/dto/manager/prometheus_query_preset/path.py b/src/ai/backend/common/dto/manager/prometheus_query_preset/path.py new file mode 100644 index 00000000000..0d657d91c8e --- /dev/null +++ b/src/ai/backend/common/dto/manager/prometheus_query_preset/path.py @@ -0,0 +1,20 @@ +""" +Path parameter DTOs for Prometheus Query Definition API endpoints. +Shared between Client SDK and Manager API. +""" + +from __future__ import annotations + +from uuid import UUID + +from pydantic import Field + +from ai.backend.common.api_handlers import BaseRequestModel + +__all__ = ("QueryDefinitionIdPathParam",) + + +class QueryDefinitionIdPathParam(BaseRequestModel): + """Path parameter for query definition ID.""" + + id: UUID = Field(description="The query definition ID") diff --git a/src/ai/backend/common/dto/manager/prometheus_query_preset/request.py b/src/ai/backend/common/dto/manager/prometheus_query_preset/request.py new file mode 100644 index 00000000000..f40f384f04a --- /dev/null +++ b/src/ai/backend/common/dto/manager/prometheus_query_preset/request.py @@ -0,0 +1,142 @@ +""" +Request DTOs for Prometheus Query Definition API endpoints. +Shared between Client SDK and Manager API. +""" + +from __future__ import annotations + +import re + +from pydantic import Field, field_validator + +from ai.backend.common.api_handlers import SENTINEL, BaseRequestModel, Sentinel +from ai.backend.common.dto.clients.prometheus.defs import PROMETHEUS_DURATION_PATTERN +from ai.backend.common.dto.clients.prometheus.request import QueryTimeRange +from ai.backend.common.dto.manager.defs import DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT +from ai.backend.common.dto.manager.query import StringFilter + +from .types import QueryDefinitionOrder + +__all__ = ( + "CreateQueryDefinitionOptionsRequest", + "CreateQueryDefinitionRequest", + "ExecuteQueryDefinitionOptionsRequest", + "ExecuteQueryDefinitionRequest", + "MetricLabelEntry", + "ModifyQueryDefinitionOptionsRequest", + "ModifyQueryDefinitionRequest", + "QueryDefinitionFilter", + "SearchQueryDefinitionsRequest", +) + + +class CreateQueryDefinitionOptionsRequest(BaseRequestModel): + """Options for a prometheus query definition.""" + + filter_labels: list[str] = Field(description="Allowed filter label keys") + group_labels: list[str] = Field(description="Allowed group-by label keys") + + +class CreateQueryDefinitionRequest(BaseRequestModel): + """Request to create a prometheus query definition.""" + + name: str = Field(description="Human-readable name") + metric_name: str = Field(description="Prometheus metric name") + query_template: str = Field(description="PromQL template with placeholders") + time_window: str | None = Field( + default=None, pattern=PROMETHEUS_DURATION_PATTERN, description="Default time window" + ) + options: CreateQueryDefinitionOptionsRequest = Field(description="Query definition options") + + +class ModifyQueryDefinitionOptionsRequest(BaseRequestModel): + """Options for modifying a prometheus query definition. + + Each field is optional — only provided fields are updated. + """ + + filter_labels: list[str] | None = Field(default=None, description="Allowed filter label keys") + group_labels: list[str] | None = Field(default=None, description="Allowed group-by label keys") + + +class ModifyQueryDefinitionRequest(BaseRequestModel): + """Request to modify a prometheus query definition. + + Only ``time_window`` uses ``Sentinel`` because it is the only nullable DB column; + all other fields are non-nullable, so ``None`` simply means "do not update". + """ + + name: str | None = Field(default=None, description="Human-readable name") + metric_name: str | None = Field(default=None, description="Prometheus metric name") + query_template: str | None = Field( + default=None, description="PromQL template with placeholders" + ) + time_window: str | Sentinel | None = Field(default=SENTINEL, description="Default time window") + options: ModifyQueryDefinitionOptionsRequest | None = Field( + default=None, description="Query definition options" + ) + + @field_validator("time_window", mode="after") + @classmethod + def _validate_time_window(cls, v: str | Sentinel | None) -> str | Sentinel | None: + if isinstance(v, str) and not re.match(PROMETHEUS_DURATION_PATTERN, v): + raise ValueError(f"Invalid Prometheus duration format: {v!r}") + return v + + +class QueryDefinitionFilter(BaseRequestModel): + """Filter for prometheus query definition search.""" + + name: StringFilter | None = Field(default=None, description="Filter by name") + + +class SearchQueryDefinitionsRequest(BaseRequestModel): + """Request body for searching prometheus query definitions with filters, orders, and pagination.""" + + filter: QueryDefinitionFilter | None = Field(default=None, description="Filter conditions") + order: list[QueryDefinitionOrder] | None = Field( + default=None, description="Order specifications" + ) + offset: int = Field(default=0, ge=0, description="Number of items to skip") + limit: int = Field( + default=DEFAULT_PAGE_LIMIT, + ge=1, + le=MAX_PAGE_LIMIT, + description="Maximum items to return", + ) + + +class MetricLabelEntry(BaseRequestModel): + """A key-value label entry for executing a query definition.""" + + key: str = Field(description="Label key") + value: str = Field(description="Label value") + + +class ExecuteQueryDefinitionOptionsRequest(BaseRequestModel): + """Execution options for a prometheus query definition.""" + + filter_labels: list[MetricLabelEntry] = Field( + default_factory=list, + description="Filter labels as key-value pairs", + ) + group_labels: list[str] = Field( + default_factory=list, + description="Group-by labels", + ) + + +class ExecuteQueryDefinitionRequest(BaseRequestModel): + """Request to execute a prometheus query definition.""" + + options: ExecuteQueryDefinitionOptionsRequest = Field( + default_factory=ExecuteQueryDefinitionOptionsRequest, + description="Execution options (filter and group labels)", + ) + window: str | None = Field( + default=None, pattern=PROMETHEUS_DURATION_PATTERN, description="Time window override" + ) + time_range: QueryTimeRange | None = Field( + default=None, + description="Time range for the query; if omitted, executes an instant query at the current time", + ) diff --git a/src/ai/backend/common/dto/manager/prometheus_query_preset/response.py b/src/ai/backend/common/dto/manager/prometheus_query_preset/response.py new file mode 100644 index 00000000000..a6a59241104 --- /dev/null +++ b/src/ai/backend/common/dto/manager/prometheus_query_preset/response.py @@ -0,0 +1,123 @@ +""" +Response DTOs for Prometheus Query Definition API endpoints. +Shared between Client SDK and Manager API. +""" + +from __future__ import annotations + +from datetime import datetime +from uuid import UUID + +from pydantic import BaseModel, Field + +from ai.backend.common.api_handlers import BaseResponseModel + +__all__ = ( + "CreateQueryDefinitionResponse", + "DeleteQueryDefinitionResponse", + "ExecuteQueryDefinitionResponse", + "GetQueryDefinitionResponse", + "MetricLabelEntryDTO", + "MetricValueDTO", + "ModifyQueryDefinitionResponse", + "PaginationInfo", + "QueryDefinitionDTO", + "QueryDefinitionExecuteData", + "QueryDefinitionMetricResult", + "QueryDefinitionOptionsDTO", + "SearchQueryDefinitionsResponse", +) + + +class QueryDefinitionOptionsDTO(BaseModel): + """Options DTO for a prometheus query definition.""" + + filter_labels: list[str] = Field(description="Allowed filter label keys") + group_labels: list[str] = Field(description="Allowed group-by label keys") + + +class QueryDefinitionDTO(BaseModel): + """DTO for prometheus query definition data.""" + + id: UUID = Field(description="Query definition ID") + name: str = Field(description="Human-readable name") + metric_name: str = Field(description="Prometheus metric name") + query_template: str = Field(description="PromQL template") + time_window: str | None = Field(default=None, description="Default time window") + options: QueryDefinitionOptionsDTO = Field(description="Query definition options") + created_at: datetime = Field(description="Creation timestamp") + updated_at: datetime = Field(description="Last update timestamp") + + +class PaginationInfo(BaseModel): + """Pagination information.""" + + total: int = Field(description="Total count of items") + offset: int = Field(description="Current offset") + limit: int = Field(description="Current limit") + + +class CreateQueryDefinitionResponse(BaseResponseModel): + """Response for creating a query definition.""" + + item: QueryDefinitionDTO + + +class GetQueryDefinitionResponse(BaseResponseModel): + """Response for getting a query definition by ID.""" + + item: QueryDefinitionDTO + + +class SearchQueryDefinitionsResponse(BaseResponseModel): + """Response for searching query definitions.""" + + items: list[QueryDefinitionDTO] + pagination: PaginationInfo + + +class ModifyQueryDefinitionResponse(BaseResponseModel): + """Response for modifying a query definition.""" + + item: QueryDefinitionDTO + + +class DeleteQueryDefinitionResponse(BaseResponseModel): + """Response for deleting a query definition.""" + + id: UUID = Field(description="Deleted query definition ID") + + +class MetricLabelEntryDTO(BaseModel): + """A key-value label entry in the execute response.""" + + key: str + value: str + + +class MetricValueDTO(BaseModel): + """A single (timestamp, value) data point from Prometheus.""" + + timestamp: float + value: str + + +class QueryDefinitionMetricResult(BaseModel): + """A single metric result from query definition execution.""" + + metric: list[MetricLabelEntryDTO] + values: list[MetricValueDTO] + + +class QueryDefinitionExecuteData(BaseModel): + """Data field of the execute response.""" + + result_type: str + result: list[QueryDefinitionMetricResult] + + +class ExecuteQueryDefinitionResponse(BaseResponseModel): + """Response for executing a query definition.""" + + status: str + data: QueryDefinitionExecuteData diff --git a/src/ai/backend/common/dto/manager/prometheus_query_preset/types.py b/src/ai/backend/common/dto/manager/prometheus_query_preset/types.py new file mode 100644 index 00000000000..b264a4c5aac --- /dev/null +++ b/src/ai/backend/common/dto/manager/prometheus_query_preset/types.py @@ -0,0 +1,39 @@ +""" +Shared types for Prometheus Query Definition DTOs. +""" + +from __future__ import annotations + +from enum import StrEnum + +from pydantic import Field + +from ai.backend.common.api_handlers import BaseRequestModel + +__all__ = ( + "OrderDirection", + "QueryDefinitionOrder", + "QueryDefinitionOrderField", +) + + +class OrderDirection(StrEnum): + """Order direction for sorting.""" + + ASC = "asc" + DESC = "desc" + + +class QueryDefinitionOrderField(StrEnum): + """Fields available for ordering prometheus query definitions.""" + + NAME = "name" + CREATED_AT = "created_at" + UPDATED_AT = "updated_at" + + +class QueryDefinitionOrder(BaseRequestModel): + """Order specification for prometheus query definitions.""" + + field: QueryDefinitionOrderField = Field(description="Field to order by") + direction: OrderDirection = Field(default=OrderDirection.ASC, description="Order direction") diff --git a/src/ai/backend/manager/api/rest/prometheus_query_preset/__init__.py b/src/ai/backend/manager/api/rest/prometheus_query_preset/__init__.py new file mode 100644 index 00000000000..e1950a5b006 --- /dev/null +++ b/src/ai/backend/manager/api/rest/prometheus_query_preset/__init__.py @@ -0,0 +1,7 @@ +from .handler import PrometheusQueryPresetHandler +from .registry import register_prometheus_query_preset_routes + +__all__ = [ + "PrometheusQueryPresetHandler", + "register_prometheus_query_preset_routes", +] diff --git a/src/ai/backend/manager/api/rest/prometheus_query_preset/adapter.py b/src/ai/backend/manager/api/rest/prometheus_query_preset/adapter.py new file mode 100644 index 00000000000..ff7c550901b --- /dev/null +++ b/src/ai/backend/manager/api/rest/prometheus_query_preset/adapter.py @@ -0,0 +1,142 @@ +"""Adapter for converting Prometheus Query Definition domain data to DTOs.""" + +from __future__ import annotations + +from uuid import UUID + +from ai.backend.common.api_handlers import Sentinel +from ai.backend.common.dto.clients.prometheus.response import MetricResponse +from ai.backend.common.dto.manager.prometheus_query_preset import ( + MetricLabelEntryDTO, + MetricValueDTO, + ModifyQueryDefinitionRequest, + OrderDirection, + QueryDefinitionDTO, + QueryDefinitionFilter, + QueryDefinitionMetricResult, + QueryDefinitionOptionsDTO, + QueryDefinitionOrder, + QueryDefinitionOrderField, + SearchQueryDefinitionsRequest, +) +from ai.backend.manager.api.rest.adapter import BaseFilterAdapter +from ai.backend.manager.data.prometheus_query_preset import PrometheusQueryPresetData +from ai.backend.manager.models.prometheus_query_preset import PrometheusQueryPresetRow +from ai.backend.manager.repositories.base import ( + BatchQuerier, + OffsetPagination, + QueryCondition, + QueryOrder, + Updater, +) +from ai.backend.manager.repositories.prometheus_query_preset.options import ( + PrometheusQueryPresetConditions, + PrometheusQueryPresetOrders, +) +from ai.backend.manager.repositories.prometheus_query_preset.updaters import ( + PrometheusQueryPresetUpdaterSpec, +) +from ai.backend.manager.types import OptionalState, TriState + + +class PrometheusQueryPresetAdapter(BaseFilterAdapter): + """Adapter for converting between query definition domain data and DTOs.""" + + def convert_to_dto(self, data: PrometheusQueryPresetData) -> QueryDefinitionDTO: + """Convert domain data to query definition DTO.""" + return QueryDefinitionDTO( + id=data.id, + name=data.name, + metric_name=data.metric_name, + query_template=data.query_template, + time_window=data.time_window, + options=QueryDefinitionOptionsDTO( + filter_labels=data.filter_labels, + group_labels=data.group_labels, + ), + created_at=data.created_at, + updated_at=data.updated_at, + ) + + def convert_metric_response(self, response: MetricResponse) -> QueryDefinitionMetricResult: + """Convert a Prometheus MetricResponse to a QueryDefinitionMetricResult DTO.""" + metric_labels = [ + MetricLabelEntryDTO(key=key, value=str(val)) + for key, val in response.metric.model_dump(exclude_none=True).items() + ] + values = [MetricValueDTO(timestamp=ts, value=v) for ts, v in response.values] + return QueryDefinitionMetricResult(metric=metric_labels, values=values) + + def build_updater( + self, request: ModifyQueryDefinitionRequest, preset_id: UUID + ) -> Updater[PrometheusQueryPresetRow]: + """Build an Updater from a modify request.""" + spec = PrometheusQueryPresetUpdaterSpec( + name=( + OptionalState.update(request.name) + if request.name is not None + else OptionalState.nop() + ), + metric_name=( + OptionalState.update(request.metric_name) + if request.metric_name is not None + else OptionalState.nop() + ), + query_template=( + OptionalState.update(request.query_template) + if request.query_template is not None + else OptionalState.nop() + ), + time_window=TriState.nop() + if isinstance(request.time_window, Sentinel) + else TriState.nullify() + if request.time_window is None + else TriState.update(request.time_window), + filter_labels=( + OptionalState.update(request.options.filter_labels) + if request.options is not None and request.options.filter_labels is not None + else OptionalState.nop() + ), + group_labels=( + OptionalState.update(request.options.group_labels) + if request.options is not None and request.options.group_labels is not None + else OptionalState.nop() + ), + ) + return Updater(spec=spec, pk_value=preset_id) + + def build_querier(self, request: SearchQueryDefinitionsRequest) -> BatchQuerier: + """Build a BatchQuerier from search request.""" + conditions = self._convert_filter(request.filter) if request.filter else [] + orders = [self._convert_order(o) for o in request.order] if request.order else [] + return BatchQuerier( + conditions=conditions, + orders=orders, + pagination=OffsetPagination(limit=request.limit, offset=request.offset), + ) + + def _convert_filter(self, filter_req: QueryDefinitionFilter) -> list[QueryCondition]: + """Convert query definition filter to list of query conditions.""" + conditions: list[QueryCondition] = [] + if filter_req.name is not None: + condition = self.convert_string_filter( + filter_req.name, + contains_factory=PrometheusQueryPresetConditions.by_name_contains, + equals_factory=PrometheusQueryPresetConditions.by_name_equals, + starts_with_factory=PrometheusQueryPresetConditions.by_name_starts_with, + ends_with_factory=PrometheusQueryPresetConditions.by_name_ends_with, + ) + if condition is not None: + conditions.append(condition) + return conditions + + def _convert_order(self, order: QueryDefinitionOrder) -> QueryOrder: + """Convert query definition order specification to query order.""" + ascending = order.direction == OrderDirection.ASC + if order.field == QueryDefinitionOrderField.NAME: + return PrometheusQueryPresetOrders.name(ascending=ascending) + if order.field == QueryDefinitionOrderField.CREATED_AT: + return PrometheusQueryPresetOrders.created_at(ascending=ascending) + if order.field == QueryDefinitionOrderField.UPDATED_AT: + return PrometheusQueryPresetOrders.updated_at(ascending=ascending) + raise ValueError(f"Unknown order field: {order.field}") diff --git a/src/ai/backend/manager/api/rest/prometheus_query_preset/handler.py b/src/ai/backend/manager/api/rest/prometheus_query_preset/handler.py new file mode 100644 index 00000000000..53f8fa389ca --- /dev/null +++ b/src/ai/backend/manager/api/rest/prometheus_query_preset/handler.py @@ -0,0 +1,183 @@ +"""REST API handler for Prometheus Query Preset operations.""" + +from __future__ import annotations + +from http import HTTPStatus + +from ai.backend.common.api_handlers import APIResponse, BodyParam, PathParam +from ai.backend.common.dto.manager.prometheus_query_preset import ( + CreateQueryDefinitionRequest, + CreateQueryDefinitionResponse, + DeleteQueryDefinitionResponse, + ExecuteQueryDefinitionRequest, + ExecuteQueryDefinitionResponse, + GetQueryDefinitionResponse, + ModifyQueryDefinitionRequest, + ModifyQueryDefinitionResponse, + PaginationInfo, + QueryDefinitionExecuteData, + QueryDefinitionIdPathParam, + SearchQueryDefinitionsRequest, + SearchQueryDefinitionsResponse, +) +from ai.backend.manager.data.prometheus_query_preset import ( + ExecutePresetOptions, +) +from ai.backend.manager.models.prometheus_query_preset import PrometheusQueryPresetRow +from ai.backend.manager.repositories.base import ( + Creator, +) +from ai.backend.manager.repositories.prometheus_query_preset.creators import ( + PrometheusQueryPresetCreatorSpec, +) +from ai.backend.manager.services.prometheus_query_preset.actions.create import ( + CreatePresetAction, +) +from ai.backend.manager.services.prometheus_query_preset.actions.delete import ( + DeletePresetAction, +) +from ai.backend.manager.services.prometheus_query_preset.actions.execute_preset import ( + ExecutePresetAction, +) +from ai.backend.manager.services.prometheus_query_preset.actions.get import ( + GetPresetAction, +) +from ai.backend.manager.services.prometheus_query_preset.actions.modify import ( + ModifyPresetAction, +) +from ai.backend.manager.services.prometheus_query_preset.actions.search import ( + SearchPresetsAction, +) +from ai.backend.manager.services.prometheus_query_preset.processors import ( + PrometheusQueryPresetProcessors, +) + +from .adapter import PrometheusQueryPresetAdapter + + +class PrometheusQueryPresetHandler: + """REST API handler for Prometheus Query Preset CRUD and execution.""" + + def __init__( + self, + *, + processor: PrometheusQueryPresetProcessors, + ) -> None: + self._processor = processor + self._adapter = PrometheusQueryPresetAdapter() + + async def create_preset( + self, + body: BodyParam[CreateQueryDefinitionRequest], + ) -> APIResponse: + """Create a new prometheus query preset.""" + creator: Creator[PrometheusQueryPresetRow] = Creator( + spec=PrometheusQueryPresetCreatorSpec( + name=body.parsed.name, + metric_name=body.parsed.metric_name, + query_template=body.parsed.query_template, + time_window=body.parsed.time_window, + filter_labels=body.parsed.options.filter_labels, + group_labels=body.parsed.options.group_labels, + ) + ) + action_result = await self._processor.create_preset.wait_for_complete( + CreatePresetAction(creator=creator) + ) + resp = CreateQueryDefinitionResponse( + item=self._adapter.convert_to_dto(action_result.preset) + ) + return APIResponse.build(status_code=HTTPStatus.CREATED, response_model=resp) + + async def search_presets( + self, + body: BodyParam[SearchQueryDefinitionsRequest], + ) -> APIResponse: + """Search presets with filters, orders, and pagination.""" + querier = self._adapter.build_querier(body.parsed) + action_result = await self._processor.search_presets.wait_for_complete( + SearchPresetsAction(querier=querier) + ) + resp = SearchQueryDefinitionsResponse( + items=[ + self._adapter.convert_to_dto(preset_data) for preset_data in action_result.items + ], + pagination=PaginationInfo( + total=action_result.total_count, + offset=body.parsed.offset, + limit=body.parsed.limit, + ), + ) + return APIResponse.build(status_code=HTTPStatus.OK, response_model=resp) + + async def get_preset( + self, + path: PathParam[QueryDefinitionIdPathParam], + ) -> APIResponse: + """Get a preset by ID.""" + action_result = await self._processor.get_preset.wait_for_complete( + GetPresetAction(preset_id=path.parsed.id) + ) + resp = GetQueryDefinitionResponse(item=self._adapter.convert_to_dto(action_result.preset)) + return APIResponse.build(status_code=HTTPStatus.OK, response_model=resp) + + async def modify_preset( + self, + path: PathParam[QueryDefinitionIdPathParam], + body: BodyParam[ModifyQueryDefinitionRequest], + ) -> APIResponse: + """Modify a preset.""" + updater = self._adapter.build_updater(body.parsed, path.parsed.id) + action_result = await self._processor.modify_preset.wait_for_complete( + ModifyPresetAction(preset_id=path.parsed.id, updater=updater) + ) + resp = ModifyQueryDefinitionResponse( + item=self._adapter.convert_to_dto(action_result.preset) + ) + return APIResponse.build(status_code=HTTPStatus.OK, response_model=resp) + + async def delete_preset( + self, + path: PathParam[QueryDefinitionIdPathParam], + ) -> APIResponse: + """Delete a preset.""" + action_result = await self._processor.delete_preset.wait_for_complete( + DeletePresetAction(preset_id=path.parsed.id) + ) + resp = DeleteQueryDefinitionResponse(id=action_result.preset_id) + return APIResponse.build(status_code=HTTPStatus.OK, response_model=resp) + + async def execute_preset( + self, + path: PathParam[QueryDefinitionIdPathParam], + body: BodyParam[ExecuteQueryDefinitionRequest], + ) -> APIResponse: + """Execute a preset with given parameters.""" + filter_labels = {entry.key: entry.value for entry in body.parsed.options.filter_labels} + options = ExecutePresetOptions( + filter_labels=filter_labels, + group_labels=body.parsed.options.group_labels, + ) + + action_result = await self._processor.execute_preset.wait_for_complete( + ExecutePresetAction( + preset_id=path.parsed.id, + options=options, + window=body.parsed.window, + time_range=body.parsed.time_range, + ) + ) + + prom_response = action_result.response + result_items = [ + self._adapter.convert_metric_response(metric_response) + for metric_response in prom_response.data.result + ] + resp = ExecuteQueryDefinitionResponse( + status=prom_response.status, + data=QueryDefinitionExecuteData( + result_type=prom_response.data.result_type, + result=result_items, + ), + ) + return APIResponse.build(status_code=HTTPStatus.OK, response_model=resp) diff --git a/src/ai/backend/manager/api/rest/prometheus_query_preset/registry.py b/src/ai/backend/manager/api/rest/prometheus_query_preset/registry.py new file mode 100644 index 00000000000..3343f8ad5a9 --- /dev/null +++ b/src/ai/backend/manager/api/rest/prometheus_query_preset/registry.py @@ -0,0 +1,37 @@ +"""Prometheus Query Definition module registrar.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ai.backend.manager.api.rest.middleware.auth import superadmin_required +from ai.backend.manager.api.rest.routing import RouteRegistry + +from .handler import PrometheusQueryPresetHandler + +if TYPE_CHECKING: + from ai.backend.manager.api.rest.types import RouteDeps + + +def register_prometheus_query_preset_routes( + handler: PrometheusQueryPresetHandler, route_deps: RouteDeps +) -> RouteRegistry: + """Build the prometheus query definition sub-application.""" + reg = RouteRegistry.create("resource/prometheus-query-definitions", route_deps.cors_options) + + # CRUD endpoints (superadmin only) + reg.add("POST", "", handler.create_preset, middlewares=[superadmin_required]) + reg.add("POST", "/search", handler.search_presets, middlewares=[superadmin_required]) + reg.add("GET", "/{id}", handler.get_preset, middlewares=[superadmin_required]) + reg.add("PATCH", "/{id}", handler.modify_preset, middlewares=[superadmin_required]) + reg.add("DELETE", "/{id}", handler.delete_preset, middlewares=[superadmin_required]) + + # Execute endpoint (superadmin only) + reg.add( + "POST", + "/{id}/execute", + handler.execute_preset, + middlewares=[superadmin_required], + ) + + return reg diff --git a/src/ai/backend/manager/api/rest/tree.py b/src/ai/backend/manager/api/rest/tree.py index 52c3db3cdeb..790a886f9a5 100644 --- a/src/ai/backend/manager/api/rest/tree.py +++ b/src/ai/backend/manager/api/rest/tree.py @@ -102,6 +102,8 @@ def build_api_routes( from .notification.registry import register_notification_routes from .object_storage.handler import ObjectStorageHandler from .object_storage.registry import register_object_storage_routes + from .prometheus_query_preset import PrometheusQueryPresetHandler + from .prometheus_query_preset.registry import register_prometheus_query_preset_routes from .quota_scope.handler import QuotaScopeHandler from .quota_scope.registry import register_quota_scope_routes from .ratelimit.registry import register_ratelimit_routes @@ -276,6 +278,10 @@ def build_api_routes( error_monitor=error_monitor, ) + # Prometheus query preset handler + prometheus_processor = processors.prometheus_query_preset + prometheus_query_preset_handler = PrometheusQueryPresetHandler(processor=prometheus_processor) + # 3. Build all registries return [ register_auth_routes(auth_handler, route_deps), @@ -348,4 +354,5 @@ def build_api_routes( register_export_routes(export_handler, route_deps), register_agent_routes(agent_handler, route_deps), register_resource_slot_routes(resource_slot_handler, route_deps), + register_prometheus_query_preset_routes(prometheus_query_preset_handler, route_deps), ] diff --git a/src/ai/backend/manager/repositories/prometheus_query_preset/options.py b/src/ai/backend/manager/repositories/prometheus_query_preset/options.py index c3e91e309bc..7650c3f29bb 100644 --- a/src/ai/backend/manager/repositories/prometheus_query_preset/options.py +++ b/src/ai/backend/manager/repositories/prometheus_query_preset/options.py @@ -72,13 +72,6 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner - @staticmethod - def by_metric_name_equals(metric_name: str) -> QueryCondition: - def inner() -> sa.sql.expression.ColumnElement[bool]: - return PrometheusQueryPresetRow.metric_name == metric_name - - return inner - @staticmethod def by_cursor_forward(cursor_id: str) -> QueryCondition: cursor_uuid = uuid.UUID(cursor_id) diff --git a/src/ai/backend/manager/services/prometheus_query_preset/actions/execute_preset.py b/src/ai/backend/manager/services/prometheus_query_preset/actions/execute_preset.py index 75861cd249b..1d48e5a6f45 100644 --- a/src/ai/backend/manager/services/prometheus_query_preset/actions/execute_preset.py +++ b/src/ai/backend/manager/services/prometheus_query_preset/actions/execute_preset.py @@ -17,7 +17,7 @@ class ExecutePresetAction(PrometheusQueryPresetAction): preset_id: UUID options: ExecutePresetOptions window: str | None - time_range: QueryTimeRange + time_range: QueryTimeRange | None @override def entity_id(self) -> str | None: diff --git a/src/ai/backend/manager/services/prometheus_query_preset/service.py b/src/ai/backend/manager/services/prometheus_query_preset/service.py index 390ad3b6eba..12858f998e0 100644 --- a/src/ai/backend/manager/services/prometheus_query_preset/service.py +++ b/src/ai/backend/manager/services/prometheus_query_preset/service.py @@ -2,7 +2,7 @@ from ai.backend.common.clients.prometheus.client import PrometheusClient from ai.backend.common.clients.prometheus.preset import MetricPreset -from ai.backend.common.exception import PrometheusQueryPresetInvalidLabel +from ai.backend.common.exception import InvalidAPIParameters, PrometheusQueryPresetInvalidLabel from ai.backend.logging.utils import BraceStyleAdapter from ai.backend.manager.data.prometheus_query_preset import ( ExecutePresetOptions, @@ -101,6 +101,9 @@ async def execute_preset(self, action: ExecutePresetAction) -> ExecutePresetActi group_by=set(action.options.group_labels), window=window, ) + if action.time_range is None: + # TODO: Implement instant query execution (query_instant) in the Prometheus client and use it here. + raise InvalidAPIParameters("time_range is required for preset execution") response = await self._prometheus_client.query_range( preset=metric_preset, time_range=action.time_range,