fix: inference api with jwt token

Author: vancauwe

README.md

@@ -71,6 +71,37 @@ Then interactively run training by :
 
 (remember to set any env variables needed beforehand)
 
+#### Test Inference
+
+```
+runai workspace submit openpulse-inference \
+  -i ghcr.io/sdsc-ordes/open-pulse-graph-classifier-inference:latest \
+  --image-pull-policy Always \
+  --gpu-devices-request 1 --preemptible
+```
+
+Testing from inside the container:
+
+1. entering into the container: `runai workspace bash openpulse-inference`
+2. going to the right repository: `cd ../app/`
+3. install curl for all testing the api : `apt-get -y update; apt-get -y install curl`
+4. Try the test endpoint:
+
+```bash
+curl -X GET "http://localhost:8000/v1/test"
+```
+
+the output should be: `{"data":{"type":"test","id":"1","attributes":{"message":"Test endpoint is working!"}}}`
+
+5. Make a valid JWT token with the `utils/api_token.py` (if you have configured secret key etc.)
+6. Try the inference endpoint:
+
+```bash
+curl -X GET "http://localhost:8000/v1/inference/epfl/YOUR_DB" \
+  -H "accept: application/json" \
+  -H "Authorization: Bearer YOUR_JWT_TOKEN"
+```
+
 #### Inference
 
 To do: this needs update to new runai cli

src/inference_api.py

@@ -23,7 +23,8 @@ def verify_jwt(credentials: HTTPAuthorizationCredentials = Depends(security)):
         if payload.get("service") != "airflow":
             raise HTTPException(status_code=403, detail="Invalid service identity")
         return payload
-    except JWTError:
+    except JWTError as e:
+        print("JWT decode error:", e)
         raise HTTPException(status_code=403, detail="Invalid token")
 
 
@@ -40,12 +41,13 @@ async def test():
 
 @app.get("/v1/inference/epfl/{neo4j_database}")
 async def do_inference(neo4j_database: str, token_data: dict = Depends(verify_jwt)):
-    output = inference(neo4j_database)
-    # should the predictions be returned in the response or should it be saved to NEO4J or something?
+    inference(neo4j_database)
     return {
         "data": {
             "type": "inference",
             "id": neo4j_database,
-            "attributes": {"results": output},
+            "attributes": {
+                "results": "Predictions were uploaded to your Neo4J database under the node property 'predition'."
+            },
         }
     }

src/utils/api_token.py

@@ -2,6 +2,10 @@
 from dotenv import load_dotenv
 import os
 import requests
+from datetime import datetime, timedelta
+from fastapi.security import HTTPAuthorizationCredentials
+
+from inference_api import verify_jwt
 
 load_dotenv()
 JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY")
@@ -17,8 +21,20 @@
 NEO4J_DATABASE = os.getenv("NEO4J_DATABASE")
 
 # FOR RUNAI INFERENCE API
-payload = {"sub": AIRFLOW_SUB, "service": AIRFLOW_SERVICE}
+print("All configurations for token in place:")
+print(JWT_SECRET_KEY, JWT_ALGORITHM)
+payload = {
+    "sub": AIRFLOW_SUB,
+    "service": AIRFLOW_SERVICE,
+    "exp": datetime.utcnow() + timedelta(hours=1),
+    "iat": datetime.utcnow(),
+}
 token = jwt.encode(payload, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
+print("Validating token")
+credentials = HTTPAuthorizationCredentials(scheme="Bearer", credentials=token)
+print(credentials)
+payload = verify_jwt(credentials)
+print(payload)
 print(f"Generated token: {token}")
 
 # for AIRFLOW API