Merge pull request #2230 from FedML-AI/inference_runner_custom_host_port

Parameterizing deploy host, port.
FedML-AI · Nov 6, 2024 · 9fc5b4d · 9fc5b4d
2 parents c6d4daf + 55ff447
commit 9fc5b4d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/fedml/serving/fedml_inference_runner.py b/python/fedml/serving/fedml_inference_runner.py
@@ -6,8 +6,10 @@
 
 
 class FedMLInferenceRunner(ABC):
-    def __init__(self, client_predictor):
+    def __init__(self, client_predictor, host="0.0.0.0", port=2345):
         self.client_predictor = client_predictor
+        self.host = host
+        self.port = port
 
     def run(self):
         api = FastAPI()
@@ -43,5 +45,4 @@ async def ready():
                 return Response(status_code=status.HTTP_202_ACCEPTED)
 
         import uvicorn
-        port = 2345
-        uvicorn.run(api, host="0.0.0.0", port=port)
+        uvicorn.run(api, host=self.host, port=self.port)