refactor(queue): Improve reliability of handling many requests

app.py

@@ -1,11 +1,10 @@
-from fastapi import FastAPI
-from fastapi.staticfiles import StaticFiles
-from fastapi.responses import FileResponse
-
-
 from time import time
 from statistics import mean
 
+from fastapi import BackgroundTasks, FastAPI
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+
 from modules.details import rand_details
 from modules.inference import generate_image
 
@@ -13,31 +12,59 @@ app = FastAPI()
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-
 tasks = {}
 
 
 def get_place_in_queue(task_id):
+    queued_tasks = list(task for task in tasks.values()
+                        if task["status"] == "queued" or task["status"] == "processing")
+
+    queued_tasks.sort(key=lambda task: task["created_at"])
 
-    pending_tasks = list(task for task in tasks.values()
-                         if task["status"] == "pending")
+    queued_task_ids = list(task["task_id"] for task in queued_tasks)
 
     try:
-        return pending_tasks.index(task_id) + 1
+        return queued_task_ids.index(task_id) + 1
     except:
         return 0
 
 
 def calculate_eta(task_id):
-    total_durations = list(task["completed_at"] - task["created_at"]
+    total_durations = list(task["completed_at"] - task["started_at"]
                            for task in tasks.values() if "completed_at" in task)
 
-    place = tasks[task_id]["initial_place_in_queue"] or 1
+    initial_place_in_queue = tasks[task_id]["initial_place_in_queue"]
 
     if len(total_durations):
-        return mean(total_durations) * place
+        eta = initial_place_in_queue * mean(total_durations)
+    else:
+        eta = initial_place_in_queue * 40
+
+    return round(eta, 1)
+
+
+def process_task(task_id):
+    if 'processing' in list(task['status'] for task in tasks.values()):
+        return
+
+    tasks[task_id]["status"] = "processing"
+    tasks[task_id]["started_at"] = time()
+
+    try:
+        tasks[task_id]["value"] = generate_image(tasks[task_id]["prompt"])
+    except Exception as ex:
+        tasks[task_id]["status"] = "failed"
+        tasks[task_id]["error"] = repr(ex)
     else:
-        return 40 * place
+        tasks[task_id]["status"] = "completed"
+    finally:
+        tasks[task_id]["completed_at"] = time()
+
+        queued_tasks = list(task for task in tasks.values() if task["status"] == "queued")
+
+        if queued_tasks:
+            print(f"Tasks remaining: {len(queued_tasks)}")
+            process_task(queued_tasks[0]["task_id"])
 
 
 @app.get('/')
@@ -45,8 +72,13 @@ def index():
     return FileResponse(path="static/index.html", media_type="text/html")
 
 
+@app.get('/details')
+def generate_details():
+    return rand_details()
+
+
 @app.get('/task/create')
-def create_task(prompt: str = "покемон"):
+def create_task(background_tasks: BackgroundTasks, prompt: str = "покемон"):
     created_at = time()
 
     task_id = f"{str(created_at)}_{prompt}"
@@ -55,75 +87,21 @@ def create_task(prompt: str = "покемон"):
         "task_id": task_id,
         "created_at": created_at,
         "prompt": prompt,
-        "initial_place_in_queue": get_place_in_queue(task_id),
-        "status": "pending",
+        "status": "queued",
         "poll_count": 0,
     }
 
-    print("Place in queue: ", get_place_in_queue(task_id))
-    print("ETA: ", calculate_eta(task_id))
-
-    return tasks[task_id]
-
+    tasks[task_id]["initial_place_in_queue"] = get_place_in_queue(task_id)
 
-@app.get('/task/queue')
-def queue_task(task_id: str):
-    try:
-        tasks[task_id]["value"] = generate_image(tasks[task_id]["prompt"])
-    except Exception as ex:
-        tasks[task_id]["status"] = "failed"
-        tasks[task_id]["error"] = repr(ex)
-    else:
-        tasks[task_id]["status"] = "completed"
-    finally:
-        tasks[task_id]["completed_at"] = time()
+    background_tasks.add_task(process_task, task_id)
 
     return tasks[task_id]
 
 
 @app.get('/task/poll')
 def poll_task(task_id: str):
-    pending_tasks = []
-    completed_durations = []
-
-    for task in tasks.values():
-        if task["status"] == "pending":
-            pending_tasks.append(task["task_id"])
-        elif task["status"] == "completed":
-            completed_durations.append(
-                task["completed_at"] - task["created_at"])
-
-    try:
-        place_in_queue = pending_tasks.index(task_id) + 1
-    except:
-        place_in_queue = 0
-
-    if (len(completed_durations)):
-        eta = sum(completed_durations) / \
-            len(completed_durations) * (place_in_queue or 1)
-    else:
-        eta = 40 * (place_in_queue or 1)
-
-    tasks[task_id]["place_in_queue"] = place_in_queue
-    tasks[task_id]["eta"] = round(eta, 1)
+    tasks[task_id]["place_in_queue"] = get_place_in_queue(task_id)
+    tasks[task_id]["eta"] = calculate_eta(task_id)
     tasks[task_id]["poll_count"] += 1
 
     return tasks[task_id]
-
-
-# @app.route('/details')
-@app.get('/details')
-async def generate_details():
-    return rand_details()
-
-
-@app.get('/duck/quack')
-async def test(query: str = "quack"):
-    print(query)
-    return {"duck": query}
-
-
-@app.get('/test')
-async def test(query: str = "test"):
-    print(query)
-    return {"query": query}

modules/inference.py

@@ -1,4 +1,6 @@
-print("Preparing for inference...")  # noqa
+from time import gmtime, strftime
+
+print(f'{strftime("%Y-%m-%d %H:%M:%S", gmtime())} Preparing for inference...')  # noqa
 
 from rudalle.pipelines import generate_images
 from rudalle import get_rudalle_model, get_tokenizer, get_vae
@@ -26,7 +28,7 @@ model.load_state_dict(torch.load(f"{file_dir}/{file_name}", map_location=f"{'cud
 vae = get_vae().to(device)
 tokenizer = get_tokenizer()
 
-print("Ready for inference")
+print(f'{strftime("%Y-%m-%d %H:%M:%S", gmtime())} Ready for inference')
 
 
 def english_to_russian(english_string):
@@ -48,7 +50,8 @@ def english_to_russian(english_string):
 
 
 def generate_image(prompt):
-    if prompt.lower() in ['grass', 'fire', 'water', 'lightning', 'fighting', 'psychic', 'colorless', 'darkness', 'metal', 'dragon', 'fairy']:
+    if prompt.lower() in ['grass', 'fire', 'water', 'lightning', 'fighting', 'psychic', 'colorless', 'darkness',
+                          'metal', 'dragon', 'fairy']:
         prompt = english_to_russian(prompt)
 
     result, _ = generate_images(prompt, tokenizer, model, vae, top_k=2048, images_num=1, top_p=0.995)

static/js/index.js

@@ -1,4 +1,4 @@
-import { generateDetails, createTask, queueTask, longPollTask } from './network.js';
+import { generateDetails, createTask, longPollTask } from './network.js';
 import { cardHTML } from './card-html.js';
 import {
   durationTimer,
@@ -30,11 +30,8 @@ const generate = async () => {
     scene.removeEventListener('mousemove', mousemoveHandlerForPreviousCard, true);
     cardSlot.innerHTML = '';
     generating = true;
-
     setOutput('booster', 'generating');
 
-    await new Promise((resolve) => setTimeout(resolve, 9999));
-
     const details = await generateDetails();
     pokeName = details.name;
     const task = await createTask(details.energy_type);
@@ -44,8 +41,7 @@ const generate = async () => {
     const timer = durationTimer(durationDisplay);
     const timerCleanup = timer().cleanup;
 
-    const longPromises = [queueTask(task.task_id), longPollTask(task)];
-    const completedTask = await Promise.any(longPromises);
+    const completedTask = await longPollTask(task);
 
     generating = false;
     timerCleanup(completedTask);

static/js/network.js

@@ -19,11 +19,6 @@ const createTask = async (prompt) => {
   return task;
 };
 
-const queueTask = async (task_id) => {
-  const queueResponse = await fetch(pathFor(`task/queue?task_id=${task_id}`));
-  return queueResponse.json();
-};
-
 const pollTask = async (task) => {
   const taskResponse = await fetch(pathFor(`task/poll?task_id=${task.task_id}`));
 
@@ -50,4 +45,4 @@ const longPollTask = async (task, interval = 10_000, max) => {
   return await longPollTask(task, interval, max);
 };
 
-export { generateDetails, createTask, queueTask, longPollTask };
+export { generateDetails, createTask, longPollTask };