Merge pull request #49 from quantmind/ls-extra

allow to set port

Author: lsbardel

fluid/tools_fastapi/backdoor.py

@@ -1,11 +1,12 @@
 from functools import partial
+from typing import Any
 
 from fastapi import FastAPI
 
 from fluid.utils.backdoor import ConsoleManager
 
 
-def setup(app: FastAPI) -> None:
-    console = ConsoleManager()
+def setup(app: FastAPI, **kwargs: Any) -> None:
+    console = ConsoleManager(**kwargs)
     app.add_event_handler("startup", partial(console.on_startup, app))
     app.add_event_handler("shutdown", partial(console.on_cleanup, app))

fluid/utils/backdoor.py

@@ -1,5 +1,6 @@
 import logging
 import sys
+from dataclasses import dataclass
 from functools import partial
 
 import aioconsole
@@ -22,16 +23,11 @@ def get_default_banner(self):
         return f"Python {sys.version} on {sys.platform}\n{cprt}\n{CONSOLE_MESSAGE}"
 
 
+@dataclass
 class ConsoleManager:
     aio_console = None
-
-    def __init__(
-        self,
-        port: int = settings.AIO_BACKDOOR_PORT,
-        host: str = "0.0.0.0",
-    ) -> None:
-        self.port = port
-        self.host = host
+    port: int = settings.AIO_BACKDOOR_PORT
+    host: str = "0.0.0.0"
 
     async def on_startup(self, app) -> None:
         self.aio_console = await aioconsole.start_interactive_server(

tests/scheduler/conftest.py

@@ -9,11 +9,13 @@
 from fluid.scheduler import TaskManager, TaskScheduler
 from fluid.scheduler.broker import RedisTaskBroker
 from fluid.scheduler.endpoints import get_task_manger, setup_fastapi
+from fluid.tools_fastapi import backdoor
 from tests.scheduler.tasks import TaskClient, task_application
 
 
 @asynccontextmanager
 async def start_fastapi(app: FastAPI) -> AsyncIterator:
+    backdoor.setup(app, port=0)
     async with app.router.lifespan_context(app):
         yield app