From 15a426ca6917c33c9bb3f5c580f6f14ccc591e1f Mon Sep 17 00:00:00 2001 From: PE39806 <185931318+PE39806@users.noreply.github.com> Date: Mon, 18 Nov 2024 13:46:56 +0000 Subject: [PATCH] BAI-1502 apply PR suggestions --- docker-compose-prod.yml | 2 +- docker-compose.yml | 2 +- lib/modelscan_api/.pre-commit-config.yaml | 2 + lib/modelscan_api/Dockerfile | 2 +- lib/modelscan_api/Dockerfile.dev | 2 +- lib/modelscan_api/bailo_modelscan_api/main.py | 38 +++++++++++++------ .../bailo_modelscan_api/test_dependencies.py | 1 - .../bailo_modelscan_api/test_main.py | 17 +++++++-- 8 files changed, 47 insertions(+), 19 deletions(-) diff --git a/docker-compose-prod.yml b/docker-compose-prod.yml index b822489f8..a9097ec08 100644 --- a/docker-compose-prod.yml +++ b/docker-compose-prod.yml @@ -41,7 +41,7 @@ services: ports: - 3311:80 healthcheck: - test: ['CMD-SHELL', 'curl --fail http://localhost:80/health || exit 1'] + test: ['CMD-SHELL', 'curl --fail http://localhost:80/info || exit 1'] interval: 30s timeout: 10s retries: 5 diff --git a/docker-compose.yml b/docker-compose.yml index 6cf653de6..3bc163cd1 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -44,7 +44,7 @@ services: ports: - 3311:80 healthcheck: - test: ['CMD-SHELL', 'curl --fail http://localhost:80/health || exit 1'] + test: ['CMD-SHELL', 'curl --fail http://localhost:80/info || exit 1'] interval: 30s timeout: 10s retries: 5 diff --git a/lib/modelscan_api/.pre-commit-config.yaml b/lib/modelscan_api/.pre-commit-config.yaml index d957f567f..d39d72cbd 100644 --- a/lib/modelscan_api/.pre-commit-config.yaml +++ b/lib/modelscan_api/.pre-commit-config.yaml @@ -2,6 +2,8 @@ ci: autoupdate_commit_msg: 'chore: update pre-commit hooks' autofix_commit_msg: 'style: pre-commit fixes' +files: ^lib/modelscan_api/ + repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v4.5.0 diff --git a/lib/modelscan_api/Dockerfile b/lib/modelscan_api/Dockerfile index 4d8d35fa4..cb7b231dd 100644 --- a/lib/modelscan_api/Dockerfile +++ b/lib/modelscan_api/Dockerfile @@ -1,4 +1,4 @@ -FROM python:3.12-slim +FROM python:3.12-slim-bullseye # Prevents Python from writing pyc files. ENV PYTHONDONTWRITEBYTECODE=1 diff --git a/lib/modelscan_api/Dockerfile.dev b/lib/modelscan_api/Dockerfile.dev index 4ef55254d..30c655495 100644 --- a/lib/modelscan_api/Dockerfile.dev +++ b/lib/modelscan_api/Dockerfile.dev @@ -1,4 +1,4 @@ -FROM python:3.12-slim +FROM python:3.12-slim-bullseye # Prevents Python from writing pyc files. ENV PYTHONDONTWRITEBYTECODE=1 diff --git a/lib/modelscan_api/bailo_modelscan_api/main.py b/lib/modelscan_api/bailo_modelscan_api/main.py index f9537163a..ac56b4262 100644 --- a/lib/modelscan_api/bailo_modelscan_api/main.py +++ b/lib/modelscan_api/bailo_modelscan_api/main.py @@ -3,17 +3,19 @@ from __future__ import annotations -from contextlib import nullcontext import logging +from contextlib import nullcontext from functools import lru_cache from http import HTTPStatus from pathlib import Path from tempfile import TemporaryDirectory from typing import Annotated, Any +import modelscan import uvicorn from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException, UploadFile from modelscan.modelscan import ModelScan +from pydantic import BaseModel from bailo_modelscan_api.config import Settings from bailo_modelscan_api.dependencies import safe_join @@ -40,22 +42,34 @@ def get_settings() -> Settings: ) # Instantiating ModelScan -modelscan = ModelScan(settings=get_settings().modelscan_settings) +modelscan_model = ModelScan(settings=get_settings().modelscan_settings) + + +class ApiInformation(BaseModel): + apiName: str + apiVersion: str + scannerName: str + modelscanVersion: str @app.get( - "/health", - summary="Simple health check endpoint", - description="Utility to check the operational status of the API.", + "/info", + summary="Simple information endpoint", + description="Utility to get the key information about the API.", status_code=HTTPStatus.OK, - response_description='Always `{"status": "healthy"}`', + response_description="A populated ApiInformation object", ) -def health_check() -> dict[str, str]: - """Minimal health check for the API endpoint. +def info(settings: Annotated[Settings, Depends(get_settings)]) -> ApiInformation: + """Information about the API. - :return: always `{"status": "healthy"}` + :return: a JSON representable object with keys from ApiInformation """ - return {"status": "healthy"} + return ApiInformation( + apiName=settings.app_name, + apiVersion=settings.app_version, + scannerName=modelscan.__name__, + modelscanVersion=modelscan.__version__, + ) @app.post( @@ -110,7 +124,9 @@ def scan_file( ) from exception # Scan the uploaded file. - result = modelscan.scan(pathlib_path) + logger.info("Initiating ModelScan scan of %s", pathlib_path) + result = modelscan_model.scan(pathlib_path) + logger.info("ModelScan result: %s", result) # Finally, return the result. return result diff --git a/lib/modelscan_api/bailo_modelscan_api/test_dependencies.py b/lib/modelscan_api/bailo_modelscan_api/test_dependencies.py index 4f4c20d5c..fe27ce221 100644 --- a/lib/modelscan_api/bailo_modelscan_api/test_dependencies.py +++ b/lib/modelscan_api/bailo_modelscan_api/test_dependencies.py @@ -11,7 +11,6 @@ from .dependencies import safe_join - # Helpers diff --git a/lib/modelscan_api/bailo_modelscan_api/test_main.py b/lib/modelscan_api/bailo_modelscan_api/test_main.py index 05492e8aa..456b163e6 100644 --- a/lib/modelscan_api/bailo_modelscan_api/test_main.py +++ b/lib/modelscan_api/bailo_modelscan_api/test_main.py @@ -1,8 +1,13 @@ """Test for the main.py file. """ +from __future__ import annotations + +from functools import lru_cache from pathlib import Path from unittest.mock import Mock, patch + +import modelscan from fastapi.testclient import TestClient from .config import Settings @@ -12,6 +17,7 @@ client = TestClient(app) +@lru_cache def get_settings_override(): return Settings(download_dir=".") @@ -19,11 +25,16 @@ def get_settings_override(): app.dependency_overrides[get_settings] = get_settings_override -def test_health(): - response = client.get("/health") +def test_info(): + response = client.get("/info") assert response.status_code == 200 - assert response.json() == {"status": "healthy"} + assert response.json() == { + "apiName": get_settings_override().app_name, + "apiVersion": get_settings_override().app_version, + "scannerName": modelscan.__name__, + "modelscanVersion": modelscan.__version__, + } @patch("modelscan.modelscan.ModelScan.scan")