Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 16 additions & 16 deletions demo/demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
11 changes: 2 additions & 9 deletions gpu_queue/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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):

Expand Down Expand Up @@ -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))
74 changes: 74 additions & 0 deletions gpu_queue/static/index.html
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
<!DOCTYPE html>
<html>
<head>
<title>GPU Queue</title>

<style>
#progress {
height: 200px;
overflow-y: scroll;
}
</style>
</head>
<body>
<!-- <h1>Progress</h1>-->
<!-- <div id="progress"></div>-->

<h1>Update GPUs</h1>
<form id="gpu-form">
<label for="gpu-list">Enter GPU IDs (comma separated):</label><br>
<textarea id="gpu-list" name="gpu-list" rows="1" cols="30"></textarea><br><br>
<button type="button" onclick="updateGPUs()">Update</button>
</form>

<script>
const updateGPUs = async () => {
const form = document.getElementById("gpu-form");
const gpuList = form.elements["gpu-list"].value;

const response = await fetch("/update_gpus", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ "gpus": gpuList.split(",") }),
});

form.reset();

if (!response.ok) {
console.error("Error submitting form:", response.statusText);
return false;
}

return true;
};

// const handleMessage = (message) => {
// handleProgress(message);
// };

// const handleProgress = (message) => {
// if (!("progress" in message)) return;
// const { progress } = message;
// const progressDiv = document.getElementById("progress");
// const newLine = document.createElement("p");
// newLine.textContent = progress;
// progressDiv.appendChild(newLine);
// progressDiv.scrollTop = progressDiv.scrollHeight;
// };
//
// const host = location.host;
// const port = Number(host.split(":")[1])
// const ws = new WebSocket(`ws://localhost:${port}/ws`);
//
// ws.onmessage = (event) => {
// try {
// const message = JSON.parse(event.data);
// handleMessage(message);
// console.log(message);
// } catch (error) {
// console.error(error);
// }
// };
</script>
</body>
</html>
72 changes: 48 additions & 24 deletions gpu_queue/web.py
Original file line number Diff line number Diff line change
@@ -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)
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
"uvicorn",
"loguru",
"fastapi",
"websockets",
"pydantic",
"typing-extensions",
],
Expand Down