diff --git a/CHANGELOG.md b/CHANGELOG.md index ead66098..a3f8233b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,8 @@ These are the section headers that we use: - Added `GET /api/v1/workspaces/:workspace_id/users` endpoint to get the users of a workspace. ([#153](https://github.com/argilla-io/argilla-server/pull/153)) - Added `POST /api/v1/workspaces/:workspace_id/users` endpoind to add a user to a workspace. ([#156](https://github.com/argilla-io/argilla-server/pull/156)) - Added `DELETE /api/v1/workspaces/:workspace_id/users/:user_id` endpoint to remove a user from a workspace. ([#158](https://github.com/argilla-io/argilla-server/pull/158)) -- Added `GET /api/v1/version` endpoin to get the current Argilla version. ([#162](https://github.com/argilla-io/argilla-server/pull/162)) +- Added `GET /api/v1/version` endpoint to get the current Argilla version. ([#162](https://github.com/argilla-io/argilla-server/pull/162)) +- Added `GET /api/v1/status` endpoint to get Argilla service status. ([#165](https://github.com/argilla-io/argilla-server/pull/165)) ## [Unreleased]() diff --git a/src/argilla_server/apis/v1/handlers/info.py b/src/argilla_server/apis/v1/handlers/info.py index d1d4c890..7a1f1d27 100644 --- a/src/argilla_server/apis/v1/handlers/info.py +++ b/src/argilla_server/apis/v1/handlers/info.py @@ -12,14 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -from fastapi import APIRouter +from fastapi import APIRouter, Depends -from argilla_server._version import __version__ -from argilla_server.schemas.v1.info import Version +from argilla_server.contexts import info +from argilla_server.schemas.v1.info import Status, Version +from argilla_server.search_engine import SearchEngine, get_search_engine router = APIRouter(tags=["info"]) @router.get("/version", response_model=Version) async def get_version(): - return Version(version=__version__) + return Version(version=info.argilla_version()) + + +@router.get("/status", response_model=Status) +async def get_status(search_engine: SearchEngine = Depends(get_search_engine)): + return Status( + version=info.argilla_version(), + search_engine=await search_engine.info(), + memory=info.memory_status(), + ) diff --git a/src/argilla_server/contexts/info.py b/src/argilla_server/contexts/info.py new file mode 100644 index 00000000..704cc929 --- /dev/null +++ b/src/argilla_server/contexts/info.py @@ -0,0 +1,55 @@ +# Copyright 2021-present, the Recognai S.L. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import psutil + +from argilla_server._version import __version__ + + +def argilla_version() -> str: + return __version__ + + +def memory_status() -> dict: + process = psutil.Process(os.getpid()) + + return {k: _memory_size(v) for k, v in process.memory_info()._asdict().items()} + + +def _memory_size(bytes) -> str: + system = [ + (1024**5, "P"), + (1024**4, "T"), + (1024**3, "G"), + (1024**2, "M"), + (1024**1, "K"), + (1024**0, "B"), + ] + + factor, suffix = None, None + for factor, suffix in system: + if bytes >= factor: + break + + amount = int(bytes / factor) + if isinstance(suffix, tuple): + singular, multiple = suffix + if amount == 1: + suffix = singular + else: + suffix = multiple + + return str(amount) + suffix diff --git a/src/argilla_server/schemas/v1/info.py b/src/argilla_server/schemas/v1/info.py index dd528fc5..8ec7da13 100644 --- a/src/argilla_server/schemas/v1/info.py +++ b/src/argilla_server/schemas/v1/info.py @@ -17,3 +17,9 @@ class Version(BaseModel): version: str + + +class Status(BaseModel): + version: str + search_engine: dict + memory: dict diff --git a/src/argilla_server/search_engine/base.py b/src/argilla_server/search_engine/base.py index 2e0f402e..7396c3da 100644 --- a/src/argilla_server/search_engine/base.py +++ b/src/argilla_server/search_engine/base.py @@ -264,6 +264,10 @@ async def new_instance(cls) -> "SearchEngine": async def close(self): pass + @abstractmethod + async def info(self) -> dict: + pass + @classmethod def register(cls, engine_name: str): def decorator(engine_class): diff --git a/src/argilla_server/search_engine/elasticsearch.py b/src/argilla_server/search_engine/elasticsearch.py index c945e471..44b2d5c4 100644 --- a/src/argilla_server/search_engine/elasticsearch.py +++ b/src/argilla_server/search_engine/elasticsearch.py @@ -64,6 +64,9 @@ async def new_instance(cls) -> "ElasticSearchEngine": async def close(self): await self.client.close() + async def info(self) -> dict: + return await self.client.info() + def _mapping_for_vector_settings(self, vector_settings: VectorSettings) -> dict: return { es_field_for_vector_settings(vector_settings): { diff --git a/src/argilla_server/search_engine/opensearch.py b/src/argilla_server/search_engine/opensearch.py index 078ae051..6333dd64 100644 --- a/src/argilla_server/search_engine/opensearch.py +++ b/src/argilla_server/search_engine/opensearch.py @@ -56,6 +56,9 @@ async def new_instance(cls) -> "OpenSearchEngine": async def close(self): await self.client.close() + async def info(self) -> dict: + return await self.client.info() + def _configure_index_settings(self): base_settings = super()._configure_index_settings() return {**base_settings, "index.knn": False} diff --git a/src/argilla_server/services/info.py b/src/argilla_server/services/info.py index 496f9a75..79fcd4fc 100644 --- a/src/argilla_server/services/info.py +++ b/src/argilla_server/services/info.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +# TODO: Once we delete all API v0 endpoints, we can remove this file. + import os from typing import Any, Dict diff --git a/tests/unit/api/v1/info/test_get_status.py b/tests/unit/api/v1/info/test_get_status.py new file mode 100644 index 00000000..e337bf20 --- /dev/null +++ b/tests/unit/api/v1/info/test_get_status.py @@ -0,0 +1,36 @@ +# Copyright 2021-present, the Recognai S.L. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from argilla_server._version import __version__ +from argilla_server.search_engine import SearchEngine +from httpx import AsyncClient + + +@pytest.mark.asyncio +class TestGetStatus: + def url(self) -> str: + return "/api/v1/status" + + async def test_get_status(self, async_client: AsyncClient, mock_search_engine: SearchEngine): + mock_search_engine.info.return_value = {} + + response = await async_client.get(self.url()) + + assert response.status_code == 200 + + response_json = response.json() + assert response_json["version"] == __version__ + assert "search_engine" in response_json + assert "memory" in response_json