add json logging support

.gitignore

@@ -1,2 +1,3 @@
 maestro_worker_python.egg-info
-__pycache__
+__pycache__
+venv

maestro_worker_python/cli.py

@@ -1,4 +1,5 @@
 from importlib.machinery import SourceFileLoader
+import logging
 import sys
 
 def main():
@@ -10,6 +11,6 @@ def main():
             key, value = arg[2:].split("=",1)
             params[key] = value
 
-    print("Running with", params)
+    logging.info("Running with", params)
     data = model.inference(params)
-    print(data)
+    logging.info(data)

maestro_worker_python/convert_files.py

@@ -5,7 +5,6 @@
 import requests
 import logging
 from subprocess import check_call
-logging.basicConfig(level=logging.INFO)
 
 
 @dataclass

maestro_worker_python/data/worker.py

@@ -9,9 +9,9 @@ def your_model(input):
 
 class MoisesWorker(object):
     def __init__(self):
-        print("Loading model...")
+        logging.info("Loading model...")
         self.model = your_model
-        print("Model loaded")
+        logging.info("Model loaded")
 
     def inference(self, input_data):
         try:

maestro_worker_python/download_file.py

@@ -1,10 +1,10 @@
 
 import logging
 import urllib.request
-logging.basicConfig(level=logging.INFO)
+
 
 def download_file(signed_url: str):
     logging.info(f"Downloading input")
     file_name, headers = urllib.request.urlretrieve(signed_url)
     logging.info(f"Downloaded input")
-    return file_name
+    return file_name

maestro_worker_python/get_duration.py

@@ -1,7 +1,6 @@
 
 import logging
 from subprocess import check_output
-logging.basicConfig(level=logging.INFO)
 
 
 def get_duration(local_file_path: str) -> int:

maestro_worker_python/init.py

@@ -2,6 +2,7 @@
 import argparse
 import os
 import pathlib
+import logging
 from xml.etree.ElementInclude import include
 
 def main():
@@ -22,7 +23,7 @@ def main():
 \_|  |_/\__,_|\___||___/\__|_|  \___/ 
                                                                           
 """)
-    print(f"Initializing Maestro worker on folder: {args.get('folder')}")
+    logging.info(f"Initializing Maestro worker on folder: {args.get('folder')}")
     shutil.copy(f"{dir}/data/worker.py", f"{args.get('folder')}/worker.py")
     shutil.copy(f"{dir}/data/requirements.txt", f"{args.get('folder')}/requirements.txt")
     shutil.copy(f"{dir}/data/docker-compose.yaml", f"{args.get('folder')}/docker-compose.yaml")

maestro_worker_python/serve.py

@@ -2,15 +2,25 @@
 import os
 import socket
 import datetime
+import logging
+import json_logging
 from starlette.concurrency import run_in_threadpool
 from importlib.machinery import SourceFileLoader
 
 
 MODEL_PATH = os.environ.get("MODEL_PATH", "./worker.py")
 
+app = FastAPI()
+
+logging.basicConfig(level=os.environ.get('LOG_LEVEL', 'INFO').upper())
+
+json_logging.init_fastapi(enable_json=True)
+json_logging.init_request_instrument(app)
+json_logging.config_root_logger()
+
 worker = SourceFileLoader("worker",MODEL_PATH).load_module()
 model = worker.MoisesWorker()
-app = FastAPI()
+
 
 @app.post("/inference")
 async def inference(request: Request, response: Response):

maestro_worker_python/server.py

@@ -1,6 +1,8 @@
 import uvicorn
 import argparse
 import os
+import logging
+
 
 def main():
     parser = argparse.ArgumentParser(
@@ -14,7 +16,29 @@ def main():
                         help="Reload the server on code changes")
 
     args = parser.parse_args().__dict__
-    print(f"Running maestro server with {str(args)}")
+    logging.info(f"Running maestro server with {str(args)}")
     os.environ["MODEL_PATH"] = args.get("worker")
-    uvicorn.run("maestro_worker_python.serve:app", host='0.0.0.0', port=int(args.get(
-        "port")), reload=args.get("reload"))
+
+    log_level = os.environ.get('LOG_LEVEL', 'INFO').upper()
+    logging_config = {
+        'version': 1,
+        'disable_existing_loggers': False,
+        'handlers': {
+            'default_handler': {
+                'class': 'logging.StreamHandler',
+                'level': log_level,
+            },
+        },
+        'loggers': {
+            '': {
+                '#handlers': ['default_handler'],
+            },
+            'root': {
+                '#handlers': ['default_handler'],
+            }
+        }
+    }
+    uvicorn.run(
+        "maestro_worker_python.serve:app", host='0.0.0.0', port=int(args.get("port")), reload=args.get("reload"),
+        log_level=log_level.lower(), log_config=logging_config,
+    )

maestro_worker_python/upload_files.py

@@ -4,7 +4,6 @@
 from typing import List
 import requests
 import logging
-logging.basicConfig(level=logging.INFO)
 
 
 @dataclass
@@ -41,4 +40,4 @@ def _upload(upload_file: UploadFile, did_raise_exception):
     except Exception as e:
         did_raise_exception.set()
         logging.exception(e)
-        raise e
+        raise e

requirements.txt

@@ -1,3 +1,4 @@
 fastapi==0.85.0
 uvicorn==0.16.0
-starlette==0.20.4
+starlette==0.20.4
+json-logging~=1.3.0