From 8574f9b682d3b7fd727a369894fc8181338b7027 Mon Sep 17 00:00:00 2001 From: peelstnac Date: Sun, 31 Mar 2024 17:40:39 -0400 Subject: [PATCH] minimal ui --- demo/demo.py | 32 ++++++++-------- gpu_queue/main.py | 11 +----- gpu_queue/static/index.html | 74 +++++++++++++++++++++++++++++++++++++ gpu_queue/web.py | 72 ++++++++++++++++++++++++------------ setup.py | 1 + 5 files changed, 141 insertions(+), 49 deletions(-) create mode 100644 gpu_queue/static/index.html diff --git a/demo/demo.py b/demo/demo.py index aa3c582..3663b92 100644 --- a/demo/demo.py +++ b/demo/demo.py @@ -3,24 +3,24 @@ from gpu_queue import JobSubmitter job_array = [ - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.1)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(2.3)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.5)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(0.5)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.6)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.1)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.5)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.123)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(2.23)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(0.15)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(0.325)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.12123)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.123)\'', - 'python -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.1123232)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.1)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(2.3)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.5)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(0.5)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.6)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.1)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.5)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.123)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(2.23)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(0.15)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(0.325)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.12123)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(1.123)\'', + 'python3 -c \'import os, time;print("GPU num utilized",os.environ["CUDA_VISIBLE_DEVICES"]);time.sleep(3.1123232)\'', ] * 20 -J = JobSubmitter(job_array, [0, 1, 2], wait_second=0, first_time_wait_second=0) +J = JobSubmitter(job_array, [0], wait_second=0, first_time_wait_second=0) worker = J.submit(block=False) time.sleep(10) -# launch_server(8080) +J.app.launch_server(8080) worker.join() diff --git a/gpu_queue/main.py b/gpu_queue/main.py index 673e64e..87e9f1b 100644 --- a/gpu_queue/main.py +++ b/gpu_queue/main.py @@ -6,17 +6,14 @@ from queue import Queue, Empty from subprocess import run from threading import Lock, Thread - -import uvicorn from loguru import logger - from gpu_queue.utils import wait_thread, threaded, _SingletonMeta try: from stdout_writer import log_writer except ModuleNotFoundError: from .stdout_writer import log_writer -from gpu_queue.web import app +from gpu_queue.web import App locker = Lock() @@ -83,7 +80,7 @@ def __init__( self.first_job_wait_second = first_time_wait_second or wait_second self.cur_job = 0 - app.submitter = self + self.app = App(self) def __submit_jobs(self): @@ -177,7 +174,3 @@ def _print(result_dict): k = " ".join(re.split(" +|\n+", k)).strip() print(f"Job:\n{k}") print("result_code", v) - - -def launch_server(port: int = 8080): - uvicorn.run(app, host="0.0.0.0", port=int(port)) diff --git a/gpu_queue/static/index.html b/gpu_queue/static/index.html new file mode 100644 index 0000000..cf7b15c --- /dev/null +++ b/gpu_queue/static/index.html @@ -0,0 +1,74 @@ + + + + GPU Queue + + + + + + + +

Update GPUs

+
+
+

+ +
+ + + + diff --git a/gpu_queue/web.py b/gpu_queue/web.py index b39cda7..bae5988 100644 --- a/gpu_queue/web.py +++ b/gpu_queue/web.py @@ -1,39 +1,63 @@ -import typing +import typing as t +import uvicorn +from pathlib import Path +from fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnect +from fastapi.staticfiles import StaticFiles +from pydantic import BaseModel -from fastapi import FastAPI -if typing.TYPE_CHECKING: - from gpu_queue.main import JobSubmitter +class GPUList(BaseModel): + gpus: t.List[str | int] -app = FastAPI() -app.submitter: "JobSubmitter" +class ConnectionManager: + def __init__(self): + self.active_connections: t.List[WebSocket] = [] -@app.get("/") -async def root(): - return f"{app.submitter.cur_job}/{len(app.submitter.job_array)}" + async def connect(self, ws: WebSocket): + await ws.accept() + self.active_connections.append(ws) + def disconnect(self, ws: WebSocket): + self.active_connections.remove(ws) -@app.get("/liveness/", status_code=200) -def liveness_check(): - return "Liveness check succeeded." + async def send(self, message: dict): + for connection in self.active_connections: + await connection.send_json(message) -@app.get("/readiness/", status_code=200) -def readiness_check(): - return "Readiness check succeeded." +class App: + def __init__( + self, + job_submitter: "JobSubmitter", + ): + self.app = FastAPI() + self.connection_manager = ConnectionManager() + self.job_submitter = job_submitter + self.register_routes() -@app.get("/startup/", status_code=200) -def startup_check(): - return "Startup check succeeded." + def register_routes(self): + @self.app.post("/update_gpus") + async def update_gpus(req: GPUList): + self.job_submitter.update_available_gpus(req.gpus) + @self.app.websocket("/ws") + async def websocket(ws: WebSocket): + await self.connection_manager.connect(ws) -@app.get("/update_gpu/{gpus}") -async def update_gpu(gpus): - if any([x for x in gpus if not x.isdigit()]): - return "Invalid GPU number" + try: + while True: + _ = await websocket.receive_text() + except WebSocketDisconnect: + self.connection_manager.disconnect(ws) - app.submitter.update_available_gpus(list(gpus)) + self.app.mount("/", StaticFiles(directory=Path(__file__).parent / "static", html=True), name="static") - return f"okay, updated to {gpus}" + async def update_progress(self, job, gpu): + await self.connection_manager.send({ + "progress": job, + }) + + def launch_server(self, port: int): + uvicorn.run(self.app, host="0.0.0.0", port=port) diff --git a/setup.py b/setup.py index 59c5070..2efa252 100644 --- a/setup.py +++ b/setup.py @@ -18,6 +18,7 @@ "uvicorn", "loguru", "fastapi", + "websockets", "pydantic", "typing-extensions", ],