From 0bc9cdbf91b40ab6dfadc5b05ccaa0baa2224cfa Mon Sep 17 00:00:00 2001 From: bgoelTT Date: Mon, 30 Dec 2024 13:41:15 -0500 Subject: [PATCH 1/3] Add JWT API key authentication --- tt-metal-yolov4/README.md | 10 ++++ tt-metal-yolov4/requirements.txt | 1 + tt-metal-yolov4/server/fast_api_yolov4.py | 59 +++++++++++++++++++++- tt-metal-yolov4/server/scripts/jwt_util.py | 52 +++++++++++++++++++ 4 files changed, 120 insertions(+), 2 deletions(-) create mode 100644 tt-metal-yolov4/server/scripts/jwt_util.py diff --git a/tt-metal-yolov4/README.md b/tt-metal-yolov4/README.md index 646ffcb..4ea7045 100644 --- a/tt-metal-yolov4/README.md +++ b/tt-metal-yolov4/README.md @@ -5,6 +5,7 @@ This implementation supports YOLOv4 execution on Grayskull and Worhmole. ## Table of Contents - [Run server](#run-server) +- [JWT_TOKEN Authorization](#jwt_token-authorization) - [Development](#development) - [Tests](#tests) @@ -19,6 +20,15 @@ docker compose --env-file tt-metal-yolov4/.env.default -f tt-metal-yolov4/docker This will start the default Docker container with the entrypoint command set to `server/run_uvicorn.sh`. The next section describes how to override the container's default command with an interractive shell via `bash`. +### JWT_TOKEN Authorization + +To authenticate requests use the header `Authorization`. The JWT token can be computed using the script `jwt_util.py`. This is an example: +```bash +export JWT_SECRET= +export AUTHORIZATION="Bearer $(python scripts/jwt_util.py --secret ${JWT_SECRET?ERROR env var JWT_SECRET must be set} encode '{"team_id": "tenstorrent", "token_id":"debug-test"}')" +``` + + ## Development Inside the container you can then start the server with: ```bash diff --git a/tt-metal-yolov4/requirements.txt b/tt-metal-yolov4/requirements.txt index 4ed5e13..bf32dc8 100644 --- a/tt-metal-yolov4/requirements.txt +++ b/tt-metal-yolov4/requirements.txt @@ -1,6 +1,7 @@ # inference server requirements fastapi==0.85.1 uvicorn==0.19.0 +pyjwt==2.7.0 python-multipart==0.0.5 -f https://download.pytorch.org/whl/cpu/torch_stable.html diff --git a/tt-metal-yolov4/server/fast_api_yolov4.py b/tt-metal-yolov4/server/fast_api_yolov4.py index 0262d49..1e759ae 100644 --- a/tt-metal-yolov4/server/fast_api_yolov4.py +++ b/tt-metal-yolov4/server/fast_api_yolov4.py @@ -3,11 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 import os import logging -from fastapi import FastAPI, File, UploadFile +from fastapi import FastAPI, File, HTTPException, Request, status, UploadFile +from functools import wraps from io import BytesIO +import jwt from PIL import Image from models.demos.yolov4.tests.yolov4_perfomant_webdemo import Yolov4Trace2CQ import ttnn +from typing import Optional import numpy as np import torch @@ -199,8 +202,60 @@ def nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False): return np.array(keep) +def normalize_token(token) -> [str, str]: + """ + Note that scheme is case insensitive for the authorization header. + See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Authorization#directives + """ # noqa: E501 + one_space = " " + words = token.split(one_space) + scheme = words[0].lower() + return [scheme, " ".join(words[1:])] + + +def read_authorization( + headers, +) -> Optional[dict]: + authorization = headers.get("authorization") + if not authorization: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Must provide Authorization header.", + ) + [scheme, parameters] = normalize_token(authorization) + if scheme != "bearer": + user_error_msg = f"Authorization scheme was '{scheme}' instead of bearer" + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, detail=user_error_msg + ) + try: + payload = jwt.decode(parameters, os.getenv("JWT_SECRET"), algorithms=["HS256"]) + if not payload: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + return payload + except jwt.InvalidTokenError as exc: + user_error_msg = f"JWT payload decode error: {exc}" + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=user_error_msg + ) + + +def api_key_required(f): + """Decorates an endpoint to require API key validation""" + + @wraps(f) + async def wrapper(*args, **kwargs): + request: Request = kwargs.get("request") + _ = read_authorization(request.headers) + + return await f(*args, **kwargs) + + return wrapper + + @app.post("/objdetection_v2") -async def objdetection_v2(file: UploadFile = File(...)): +@api_key_required +async def objdetection_v2(request: Request, file: UploadFile = File(...)): contents = await file.read() # Load and convert the image to RGB image = Image.open(BytesIO(contents)).convert("RGB") diff --git a/tt-metal-yolov4/server/scripts/jwt_util.py b/tt-metal-yolov4/server/scripts/jwt_util.py new file mode 100644 index 0000000..74fa466 --- /dev/null +++ b/tt-metal-yolov4/server/scripts/jwt_util.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC + +#!/usr/bin/env python3 +import argparse +import json +import os + +import sys + +import jwt + + +def main(): + parser = argparse.ArgumentParser(description="Generate signed JWT payload") + parser.add_argument("mode", type=str, help="'encode' or 'decode'") + parser.add_argument( + "payload", type=str, help="JSON string if 'encode', token if 'decode'" + ) + parser.add_argument( + "--secret", + type=str, + dest="secret", + help="JWT secret if not provided as environment variable JWT_SECRET", + ) + args = parser.parse_args() + + try: + jwt_secret = os.environ.get("JWT_SECRET", args.secret) + except KeyError: + print("ERROR: Expected JWT_SECRET environment variable to be provided") + sys.exit(1) + + try: + if args.mode == "encode": + json_payload = json.loads(args.payload) + encoded_jwt = jwt.encode(json_payload, jwt_secret, algorithm="HS256") + print(encoded_jwt) + elif args.mode == "decode": + decoded_jwt = jwt.decode(args.payload, jwt_secret, algorithms="HS256") + print(decoded_jwt) + else: + print("ERROR: Expected mode to be 'encode' or 'decode'") + sys.exit(1) + except json.decoder.JSONDecodeError: + print("ERROR: Expected payload to be a valid JSON string") + sys.exit(1) + + +if __name__ == "__main__": + main() From 297cdec8e447bde040626f903c78a094c8795220 Mon Sep 17 00:00:00 2001 From: bgoelTT Date: Mon, 30 Dec 2024 20:21:41 -0500 Subject: [PATCH 2/3] Add authentication header to locust testing and add API demonstration test --- tt-metal-yolov4/requirements-test.txt | 1 + tt-metal-yolov4/tests/locustfile.py | 22 ++++--------- tt-metal-yolov4/tests/test_inference_api.py | 36 +++++++++++++++++++++ tt-metal-yolov4/tests/utils.py | 34 +++++++++++++++++++ 4 files changed, 77 insertions(+), 16 deletions(-) create mode 100644 tt-metal-yolov4/tests/test_inference_api.py create mode 100644 tt-metal-yolov4/tests/utils.py diff --git a/tt-metal-yolov4/requirements-test.txt b/tt-metal-yolov4/requirements-test.txt index 01b78ec..4db54de 100644 --- a/tt-metal-yolov4/requirements-test.txt +++ b/tt-metal-yolov4/requirements-test.txt @@ -1,2 +1,3 @@ pillow==10.3.0 locust==2.25.0 +pytest==7.2.2 diff --git a/tt-metal-yolov4/tests/locustfile.py b/tt-metal-yolov4/tests/locustfile.py index 6ff8e7c..2b05872 100644 --- a/tt-metal-yolov4/tests/locustfile.py +++ b/tt-metal-yolov4/tests/locustfile.py @@ -2,26 +2,16 @@ # # SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC -import io -import requests -from PIL import Image from locust import HttpUser, task +from utils import get_auth_header, sample_file -# Save image as JPEG in-memory for load testing -# Load sample image -url = "http://images.cocodataset.org/val2017/000000039769.jpg" -pil_image = Image.open(requests.get(url, stream=True).raw) -pil_image = pil_image.resize((320, 320)) # Resize to target dimensions -buf = io.BytesIO() -pil_image.save( - buf, - format="JPEG", -) -byte_im = buf.getvalue() -file = {"file": byte_im} + +# load sample file in memory +file = sample_file() class HelloWorldUser(HttpUser): @task def hello_world(self): - self.client.post("/objdetection_v2", files=file) + headers = get_auth_header() + self.client.post("/objdetection_v2", files=file, headers=headers) diff --git a/tt-metal-yolov4/tests/test_inference_api.py b/tt-metal-yolov4/tests/test_inference_api.py new file mode 100644 index 0000000..d4f630a --- /dev/null +++ b/tt-metal-yolov4/tests/test_inference_api.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC + +from http import HTTPStatus +import os +import pytest +import requests +from utils import get_auth_header, sample_file + + +DEPLOY_URL = "http://127.0.0.1" +SERVICE_PORT = int(os.getenv("SERVICE_PORT", 7000)) +API_BASE_URL = f"{DEPLOY_URL}:{SERVICE_PORT}" +API_URL = f"{API_BASE_URL}/objdetection_v2" +HEALTH_URL = f"{API_BASE_URL}/health" + + +def test_valid_api_call(): + # get sample image file + file = sample_file() + # make request with auth headers + headers = get_auth_header() + response = requests.post(API_URL, files=file, headers=headers) + # perform status and value checking + assert response.status_code == HTTPStatus.OK + assert isinstance(response.json(), list) + + +@pytest.mark.skip( + reason="Not implemented, see https://github.com/tenstorrent/tt-inference-server/issues/63" +) +def test_get_health(): + headers = {} + response = requests.get(HEALTH_URL, headers=headers, timeout=35) + assert response.status_code == 200 diff --git a/tt-metal-yolov4/tests/utils.py b/tt-metal-yolov4/tests/utils.py new file mode 100644 index 0000000..e884c07 --- /dev/null +++ b/tt-metal-yolov4/tests/utils.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC + +import io +import os +from PIL import Image +import requests + + +def get_auth_header(): + if authorization_header := os.getenv("AUTHORIZATION", None): + headers = {"Authorization": authorization_header} + return headers + else: + raise RuntimeError("AUTHORIZATION environment variable is undefined.") + + +# save image as JPEG in-memory +def sample_file(): + # load sample image + url = "http://images.cocodataset.org/val2017/000000039769.jpg" + pil_image = Image.open(requests.get(url, stream=True).raw) + pil_image = pil_image.resize((320, 320)) # Resize to target dimensions + # convert to bytes + buf = io.BytesIO() + # format as JPEG + pil_image.save( + buf, + format="JPEG", + ) + byte_im = buf.getvalue() + file = {"file": byte_im} + return file From f12988d55417b36f78e999cd4eb26936d0413c8c Mon Sep 17 00:00:00 2001 From: bgoelTT Date: Mon, 30 Dec 2024 20:24:44 -0500 Subject: [PATCH 3/3] Add invalid API key test --- tt-metal-yolov4/tests/test_inference_api.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tt-metal-yolov4/tests/test_inference_api.py b/tt-metal-yolov4/tests/test_inference_api.py index d4f630a..d38602f 100644 --- a/tt-metal-yolov4/tests/test_inference_api.py +++ b/tt-metal-yolov4/tests/test_inference_api.py @@ -27,6 +27,17 @@ def test_valid_api_call(): assert isinstance(response.json(), list) +def test_invalid_api_call(): + # get sample image file + file = sample_file() + # make request with INVALID auth header + headers = get_auth_header() + headers.update(Authorization="INVALID API KEY") + response = requests.post(API_URL, files=file, headers=headers) + # assert request was unauthorized + assert response.status_code == HTTPStatus.UNAUTHORIZED + + @pytest.mark.skip( reason="Not implemented, see https://github.com/tenstorrent/tt-inference-server/issues/63" )