lava_callback: migrate to fastapi

nuclearcat · nuclearcat · commit 67aab85fae4b · 2024-07-10T09:06:02.000Z
It will be easier to maintain API and Pipeline, as
both will be powered by FastAPI framework.

Signed-off-by: Denys Fedoryshchenko &lt;denys.f@collabora.com&gt;
diff --git a/docker-compose-production.yaml b/docker-compose-production.yaml
@@ -9,15 +9,16 @@ services:
 
   lava-callback:
     # With uWSGI in socket mode, to use with reverse proxy e.g. Nginx and SSL
-    command:
-      - '/usr/local/bin/uwsgi'
-      - '--master'
-      - '--socket=:8000'
-      - '--buffer-size=32768'
-      - '-p${NPROC:-4}'
-      - '--enable-threads'
-      - '--wsgi-file=/home/kernelci/pipeline/lava_callback.py'
-      - '--callable=app'
+    command: uvicorn lava_callback:app --host 0.0.0.0 --port 8000 --app-dir /home/kernelci/pipeline/
+    #command:
+    #  - '/usr/local/bin/uwsgi'
+    #  - '--master'
+    #  - '--socket=:8000'
+    #  - '--buffer-size=32768'
+    #  - '-p${NPROC:-4}'
+    #  - '--enable-threads'
+    #  - '--wsgi-file=/home/kernelci/pipeline/lava_callback.py'
+    #  - '--callable=app'
 
     # With uWSGI HTTP server, suitable for a public instance but no SSL
     # command:
diff --git a/docker/lava-callback/requirements.txt b/docker/lava-callback/requirements.txt
@@ -1,2 +1,3 @@
-flask==2.3.2
 uwsgi==2.0.22
+uvicorn==0.30.1
+fastapi==0.111.0
diff --git a/src/lava_callback.py b/src/lava_callback.py
@@ -9,9 +9,10 @@
 import gzip
 import json
 import requests
-from flask import Flask, request
 import toml
 import threading
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
 
 import kernelci.api.helper
 import kernelci.config
@@ -26,7 +27,7 @@
 )
 SETTINGS_PREFIX = 'runtime'
 
-app = Flask(__name__)
+app = FastAPI()
 executor = ThreadPoolExecutor(max_workers=16)
 
 
@@ -86,15 +87,20 @@ def _upload_log(log_parser, job_node, storage):
         return _upload_file(storage, job_node, src, 'log.txt.gz')
 
 
-@app.errorhandler(requests.exceptions.HTTPError)
-def handle_http_error(ex):
-    detail = ex.response.json().get('detail') or str(ex)
-    return detail, ex.response.status_code
-
-
-@app.route('/')
-def hello():
-    return "KernelCI API & Pipeline LAVA callback handler"
+@app.get('/')
+async def read_root():
+    page = '''
+    <html>
+    <head>
+    <title>KernelCI Pipeline Callback</title>
+    </head>
+    <body>
+    <h1>KernelCI Pipeline Callback</h1>
+    <p>This is a callback endpoint for the KernelCI pipeline.</p>
+    </body>
+    </html>
+    '''
+    return page
 
 
 def async_job_submit(api_helper, node_id, job_callback):
@@ -146,8 +152,9 @@ def submit_job(api_helper, node_id, job_callback):
     executor.submit(async_job_submit, api_helper, node_id, job_callback)
 
 
-@app.post('/node/<node_id>')
-def callback(node_id):
+# POST /node/<node_id>
+@app.post('/node/{node_id}')
+async def callback(node_id: str, request: Request):
     tokens = SETTINGS.get(SETTINGS_PREFIX)
     if not tokens:
         return 'Unauthorized', 401
@@ -169,7 +176,7 @@ def callback(node_id):
     if not lab_name:
         return 'Unauthorized', 401
 
-    data = request.get_json()
+    data = await request.json()
     job_callback = kernelci.runtime.lava.Callback(data)
     api_config_name = job_callback.get_meta('api_config_name')
     api_token = os.getenv('KCI_API_TOKEN')
@@ -185,4 +192,4 @@ def callback(node_id):
     tokens = SETTINGS.get(SETTINGS_PREFIX)
     if not tokens:
         print('No tokens configured in toml file')
-    app.run(host='0.0.0.0', port=8000)
+    uvicorn.run(app, host='0.0.0.0', port=8000)
