From cd480430b343fb6c0663b0978832e46ea93383d1 Mon Sep 17 00:00:00 2001 From: ducoterra Date: Sat, 1 Aug 2026 21:42:38 -0400 Subject: [PATCH] add queue system --- .env.example | 1 + Containerfile | 3 +- app.py | 192 ++++++++++++++++++++++++++++++++++++++++++++++- compose.yaml | 10 +++ gunicorn.conf.py | 11 +++ pyproject.toml | 1 + uv.lock | 11 +++ 7 files changed, 224 insertions(+), 5 deletions(-) create mode 100644 compose.yaml create mode 100644 gunicorn.conf.py diff --git a/.env.example b/.env.example index 85c8561..ffad7a0 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,4 @@ OPENAI_API_KEY=api-key-here OPENAI_API_URL=https://aipi.reeseapps.com/v1/chat/completions MODEL=turbo +VALKEY_URL=redis://localhost:6379/0 diff --git a/Containerfile b/Containerfile index fea262f..bc57ed0 100644 --- a/Containerfile +++ b/Containerfile @@ -9,7 +9,8 @@ COPY pyproject.toml uv.lock ./ RUN uv sync --frozen --no-dev COPY app.py ./ +COPY gunicorn.conf.py ./ EXPOSE 5000 -CMD ["uv", "run", "--no-sync", "gunicorn", "app:app", "--bind", "0.0.0.0:5000", "--workers", "2", "--timeout", "300"] +CMD ["uv", "run", "--no-sync", "gunicorn", "-c", "gunicorn.conf.py", "app:app"] diff --git a/app.py b/app.py index b4180d2..2026da4 100644 --- a/app.py +++ b/app.py @@ -5,12 +5,14 @@ import logging import os import random import re +import threading import time import uuid from collections.abc import Generator from typing import Any import httpx +import valkey from dotenv import load_dotenv from flask import Flask, Request, Response, request, stream_with_context @@ -30,6 +32,12 @@ API_URL = os.getenv("OPENAI_API_URL", "https://api.openai.com/v1/chat/completion API_KEY = os.getenv("OPENAI_API_KEY", "") MODEL = os.getenv("MODEL", "gpt-4o-mini") +VALKEY_URL = os.getenv("VALKEY_URL", "redis://localhost:6379/0") +r = valkey.Valkey.from_url(VALKEY_URL, decode_responses=True) # type: ignore[assignment] + +worker_running = False +worker_thread: threading.Thread | None = None + logger.info(f"Using model: {MODEL}") logger.info(f"API URL: {API_URL}") logger.info(f"API key present: {bool(API_KEY)}") @@ -125,8 +133,11 @@ const es = new EventSource('/stream' + window.location.search + window.location. let chars = 0; let htmlChunks = []; let htmlSize = 0; +let mode = 'idle'; const interval = setInterval(() => { - stats.textContent = 'Characters: ' + chars.toLocaleString() + ' | ' + ((Date.now() - start) / 1000).toFixed(1) + 's'; + if (mode === 'generating') { + stats.textContent = 'Characters: ' + chars.toLocaleString() + ' | ' + ((Date.now() - start) / 1000).toFixed(1) + 's'; + } }, 100); window.__wotf_html = ''; es.onmessage = (e) => { @@ -147,14 +158,18 @@ es.onmessage = (e) => { htmlSize = 0; chars = 0; } else if (e.data.startsWith('[STATUS:')) { - const msg = e.data.slice(9); - stats.textContent = msg; + const msg = e.data.slice(8); + if (mode !== 'generating') { + stats.textContent = msg; + mode = 'status'; + } } else if (e.data === '[ERROR]') { clearInterval(interval); es.close(); titleEl.textContent = 'Error generating page'; document.querySelector('.subtitle').textContent = 'Please refresh and try again'; } else { + mode = 'generating'; const bytes = Uint8Array.from(atob(e.data), c => c.charCodeAt(0)); const decoded = new TextDecoder().decode(bytes); htmlChunks.push(decoded); @@ -463,6 +478,100 @@ THEMES = [ ] +def enqueue_request(context: dict[str, Any]) -> tuple[str, int]: + request_id = uuid.uuid4().hex[:12] + ctx_key = f"wotf:req:{request_id}:ctx" + r.hset(ctx_key, mapping={k: json.dumps(v) for k, v in context.items()}) # type: ignore[union-attr] + r.set(f"wotf:req:{request_id}:status", "pending") # type: ignore[union-attr] + r.rpush("wotf:queue", request_id) # type: ignore[union-attr] + position_raw = r.llen("wotf:queue") # type: ignore[union-attr] + position: int = int(position_raw) if isinstance(position_raw, (int, float)) else 0 + return (request_id, position) + + +def get_queue_info(request_id: str) -> dict[str, Any]: + queue_items_raw = r.lrange("wotf:queue", 0, -1) # type: ignore[union-attr] + queue_items: list[str] = list(queue_items_raw) if isinstance(queue_items_raw, list) else [] # type: ignore[arg-type] + if request_id in queue_items: + position = queue_items.index(request_id) + 1 + else: + position = 0 + avg_raw = r.get("wotf:stats:avg_time") # type: ignore[union-attr] + avg_str: str = str(avg_raw) if avg_raw is not None else "15" + avg = float(avg_str) + eta = position * avg + return {"position": position, "avg_time": avg, "eta": eta} + + +def get_generated_html(request_id: str) -> str: + result = r.get(f"wotf:req:{request_id}:html") # type: ignore[union-attr] + return result if isinstance(result, str) else "" + + +def worker_loop() -> None: + while worker_running: + request_id: str | None = None + try: + current = r.getset("wotf:processing", "locked") # type: ignore[union-attr] + if current is not None: + time.sleep(0.5) + continue + + request_id_raw = r.lpop("wotf:queue") # type: ignore[union-attr] + if request_id_raw is None: + r.delete("wotf:processing") # type: ignore[union-attr] + time.sleep(0.5) + continue + + request_id = request_id_raw if isinstance(request_id_raw, str) else None + if request_id is None: + r.delete("wotf:processing") # type: ignore[union-attr] + time.sleep(0.5) + continue + + r.set(f"wotf:req:{request_id}:status", "processing") # type: ignore[union-attr] + r.delete(f"wotf:stream:{request_id}:chunks") # type: ignore[union-attr] + + ctx_raw = r.hgetall(f"wotf:req:{request_id}:ctx") # type: ignore[union-attr] + ctx: dict[str, Any] = {} + if isinstance(ctx_raw, dict): + for k_str, v_raw in ctx_raw.items(): # type: ignore[union-attr] + k: str = str(k_str) # type: ignore[arg-type] + v: Any = v_raw # type: ignore[assignment] + try: + ctx[k] = json.loads(v) if isinstance(v, str) else v + except json.JSONDecodeError: + ctx[k] = v + + start = time.time() + for chunk in stream_llm(ctx): + encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii") + r.rpush(f"wotf:stream:{request_id}:chunks", encoded) # type: ignore[union-attr] + elapsed = time.time() - start + r.set(f"wotf:req:{request_id}:status", "done") # type: ignore[union-attr] + + n_raw = r.get("wotf:stats:gen_count") # type: ignore[union-attr] + n_str: str = str(n_raw) if n_raw is not None else "0" + n = int(n_str) + avg_raw = r.get("wotf:stats:avg_time") # type: ignore[union-attr] + avg_str: str = str(avg_raw) if avg_raw is not None else "15" + avg = float(avg_str) + new_avg = (avg * n + elapsed) / (n + 1) if n > 0 else elapsed + r.set("wotf:stats:avg_time", str(new_avg)) # type: ignore[union-attr] + r.set("wotf:stats:gen_count", str(n + 1)) # type: ignore[union-attr] + + r.delete("wotf:processing") # type: ignore[union-attr] + logger.info(f"[worker] Generated {request_id} in {elapsed:.1f}s (avg={new_avg:.1f}s)") + + except BaseException as e: # noqa: BLE001 + logger.error(f"[worker] Error: {e}") + if request_id: + r.set(f"wotf:req:{request_id}:status", "error") # type: ignore[union-attr] + r.set(f"wotf:req:{request_id}:error", str(e)) # type: ignore[union-attr] + r.delete("wotf:processing") # type: ignore[union-attr] + time.sleep(1) + + def build_prompt(context: dict[str, Any]) -> str: site_type, theme = random.choice(THEMES) mood = random.choice(["playful", "mysterious", "cozy", "epic", "dreamy", "energetic", "melancholic", "whimsical"]) @@ -614,8 +723,62 @@ def stream() -> Response: context = get_user_context(request) logger.info(f"Stream request from {request.remote_addr}") + request_id, position = enqueue_request(context) + logger.info(f"[{request_id}] Enqueued at position {position}") + def generate() -> Generator[str]: - yield from sse_stream(context) + yield "data: [CLEAR]\n\n" + + while True: + status_raw = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr] + status = status_raw if isinstance(status_raw, str) else "pending" + + if status == "error": + error_raw = r.get(f"wotf:req:{request_id}:error") # type: ignore[union-attr] + error_msg = error_raw if isinstance(error_raw, str) else "Unknown error" + yield f"data: [STATUS:Error: {error_msg}\n\n" + yield "data: [ERROR]\n\n" + return + + if status == "done": + break + + if status == "processing": + chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", 0, -1) # type: ignore[union-attr] + chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type] + if chunks: + break + yield "data: [STATUS:Generating your page...\n\n" + time.sleep(0.5) + continue + + info = get_queue_info(request_id) + pos = info["position"] + eta = info["eta"] + if pos == 0: + yield "data: [STATUS:Your turn is next...\n\n" + else: + yield f"data: [STATUS:Position {pos} in line | Est. wait: ~{eta:.0f}s\n\n" + time.sleep(1) + + last_idx = 0 + while True: + chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1) # type: ignore[union-attr] + chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type] + if chunks: + for chunk in chunks: + yield f"data: {chunk}\n\n" + last_idx += len(chunks) + time.sleep(0.05) + else: + status_raw = r.get(f"wotf:req:{request_id}:status") # type: ignore[union-attr] + status = status_raw if isinstance(status_raw, str) else "pending" + if status == "done": + break + time.sleep(0.2) + + yield "data: [DONE]\n\n" + yield "data: [RENDER]\n\n" return Response( stream_with_context(generate()), @@ -628,5 +791,26 @@ def stream() -> Response: ) +def start_worker_if_needed() -> None: + global worker_running, worker_thread + if worker_running: + return + claimed = r.set("wotf:worker_lock", "1", nx=True, ex=300) # type: ignore[union-attr] + if claimed: + worker_running = True + worker_thread = threading.Thread(target=worker_loop, daemon=True) + worker_thread.start() + logger.info("[startup] Worker thread started in this process") + else: + logger.info("[startup] Worker thread already running in another process") + + +def gunicorn_post_fork(server: Any, worker: Any) -> None: + start_worker_if_needed() + + if __name__ == "__main__": + worker_running = True + worker_thread = threading.Thread(target=worker_loop, daemon=True) + worker_thread.start() app.run(host="0.0.0.0", port=5000, debug=True) diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 0000000..5908ffd --- /dev/null +++ b/compose.yaml @@ -0,0 +1,10 @@ +services: + valkey: + image: valkey/valkey:latest + ports: + - "6379:6379" + volumes: + - valkey_data:/data + +volumes: + valkey_data: diff --git a/gunicorn.conf.py b/gunicorn.conf.py new file mode 100644 index 0000000..fe8f8bc --- /dev/null +++ b/gunicorn.conf.py @@ -0,0 +1,11 @@ +from typing import Any + +import app + +bind = "0.0.0.0:5000" +workers = 2 +timeout = 300 + + +def post_fork(server: Any, worker: Any) -> None: + app.gunicorn_post_fork(server, worker) diff --git a/pyproject.toml b/pyproject.toml index f198f74..73711e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "httpx>=0.28.1", "lxml>=5.0.0", "python-dotenv>=1.2.2", + "valkey>=6.0.0", ] [project.scripts] diff --git a/uv.lock b/uv.lock index 959b284..487b341 100644 --- a/uv.lock +++ b/uv.lock @@ -337,6 +337,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, ] +[[package]] +name = "valkey" +version = "6.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/ee/7fd930fc712275084722ddd464a0ea296abdb997d2da396320507968daeb/valkey-6.1.1.tar.gz", hash = "sha256:5880792990c6c2b5eb604a5ed5f98f300880b6dd92d123819b66ed54bb259731", size = 4601372, upload-time = "2025-08-11T06:41:10.63Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/a2/252afa4da08c714460f49e943070f86a02931f99f886182765194002fe33/valkey-6.1.1-py3-none-any.whl", hash = "sha256:e2691541c6e1503b53c714ad9a35551ac9b7c0bbac93865f063dbc859a46de92", size = 259474, upload-time = "2025-08-11T06:41:08.769Z" }, +] + [[package]] name = "web-on-the-fly" version = "0.1.0" @@ -347,6 +356,7 @@ dependencies = [ { name = "httpx" }, { name = "lxml" }, { name = "python-dotenv" }, + { name = "valkey" }, ] [package.dev-dependencies] @@ -362,6 +372,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.28.1" }, { name = "lxml", specifier = ">=5.0.0" }, { name = "python-dotenv", specifier = ">=1.2.2" }, + { name = "valkey", specifier = ">=6.0.0" }, ] [package.metadata.requires-dev]