From 4b742240e5056403f2b71bc118ef1eade9af27c7 Mon Sep 17 00:00:00 2001
From: nidhal baccouri <nidhalbacc@gmail.com>
Date: Fri, 18 Jun 2021 22:30:47 +0200
Subject: [PATCH] added uvicorn params

---
 igel/cli.py                    | 20 ++++++++++++++++----
 igel/servers/fastapi_server.py |  2 +-
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/igel/cli.py b/igel/cli.py
index d74021d..9842357 100644
--- a/igel/cli.py
+++ b/igel/cli.py
@@ -32,6 +32,9 @@ class CLI:
         "model": "model_name",
         "type": "model_type",
         "tg": "target",
+        # host and port for serving the model
+        "h": "host",
+        "p": "port",
     }
 
     def __init__(self):
@@ -570,10 +573,19 @@ def serve(self):
         """
         expose a REST endpoint in order to use the trained ML model
         """
-        os.environ[Constants.model_results_path] = self.dict_args[
-            "model_results_dir"
-        ]
-        fastapi_server.run()
+        try:
+            os.environ[Constants.model_results_path] = self.dict_args[
+                "model_results_dir"
+            ]
+            uvicorn_params = {}
+            if "host" in self.dict_args and "port" in self.dict_args:
+                uvicorn_params["host"] = self.dict_args.get("host")
+                uvicorn_params["port"] = int(self.dict_args.get("port"))
+
+            fastapi_server.run(**uvicorn_params)
+
+        except Exception as ex:
+            logger.exception(ex)
 
     def _tableize(self, df):
         """
diff --git a/igel/servers/fastapi_server.py b/igel/servers/fastapi_server.py
index 3c210db..ef86929 100644
--- a/igel/servers/fastapi_server.py
+++ b/igel/servers/fastapi_server.py
@@ -83,7 +83,7 @@ async def predict(data: dict = Body(...)):
 
 def run(**kwargs):
 
-    uvicorn.run(app)
+    uvicorn.run(app, **kwargs)
 
 
 if __name__ == "__main__":