From 927aadc5ef3561152dbcdc0558ae00d44eaeb2a7 Mon Sep 17 00:00:00 2001 From: bgoelTT Date: Thu, 2 Jan 2025 13:57:16 -0500 Subject: [PATCH] Add API authentication to YOLOv4 backend --- app/api/model_control/views.py | 12 +++++++++++- app/docker-compose.yml | 6 +++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/app/api/model_control/views.py b/app/api/model_control/views.py index bc73bf8b..b4b4e1af 100644 --- a/app/api/model_control/views.py +++ b/app/api/model_control/views.py @@ -15,6 +15,7 @@ from .serializers import InferenceSerializer, ModelWeightsSerializer from model_control.model_utils import ( + encoded_jwt, get_deploy_cache, stream_response_from_external_api, health_check, @@ -123,7 +124,16 @@ def post(self, request, *args, **kwargs): ) byte_im = buf.getvalue() file = {"file": byte_im} - inference_data = requests.post(internal_url, files=file, timeout=5) + try: + headers = {"Authorization": f"Bearer {encoded_jwt}"} + inference_data = requests.post(internal_url, files=file, headers=headers, timeout=5) + inference_data.raise_for_status() + except requests.exceptions.HTTPError as http_err: + if inference_data.status_code == status.HTTP_401_UNAUTHORIZED: + return Response(status=status.HTTP_401_UNAUTHORIZED) + else: + return Response(status=status.HTTP_500_INTERNAL_SERVER_ERROR) + return Response(inference_data.json(), status=status.HTTP_200_OK) else: return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST) diff --git a/app/docker-compose.yml b/app/docker-compose.yml index bd7e85bc..76895274 100644 --- a/app/docker-compose.yml +++ b/app/docker-compose.yml @@ -21,7 +21,7 @@ services: - "8000:8000" # command: bash # dev server can be used for breakpoint debugging, does not support streaming - # command: ./manage.py runserver 0.0.0.0:8000 + # command: python ./manage.py runserver 0.0.0.0:8000 # gunicorn is used from production, supports streaming command: gunicorn --workers 3 --bind 0.0.0.0:8000 --preload --timeout 1200 api.wsgi:application @@ -52,7 +52,7 @@ services: # On first application load resources for transformers/etc # are downloaded. The UI should not start until these resources # have been downloaded. Adjust timeout if on a very slow connection - test: ["CMD", "curl", "-f", "http://localhost:8000/up/"] + test: [ "CMD", "curl", "-f", "http://localhost:8000/up/" ] timeout: 120s interval: 10s retries: 5 @@ -93,7 +93,7 @@ services: - "8111:8111" healthcheck: # Adjust below to match your container port - test: ["CMD", "curl", "-f", "http://localhost:8111/api/v1/heartbeat"] + test: [ "CMD", "curl", "-f", "http://localhost:8111/api/v1/heartbeat" ] interval: 10s timeout: 10s retries: 3