increase workers and fix async issues

This commit is contained in:
2026-08-01 22:24:29 -04:00
parent 73b2b7972c
commit d487fad8fa
5 changed files with 340 additions and 23 deletions
+138 -13
View File
@@ -16,6 +16,11 @@ import valkey
from dotenv import load_dotenv
from flask import Flask, Request, Response, request, stream_with_context
if os.getenv("GUNICORN_WORKER_ID"):
from gevent import monkey
monkey.patch_all()
lxml_html: Any = importlib.import_module("lxml.html")
logging.basicConfig(
@@ -88,6 +93,24 @@ body {
color: #6366f1;
font-variant-numeric: tabular-nums;
}
.retry-btn {
margin-top: 20px;
padding: 10px 24px;
background: rgba(99, 102, 241, 0.2);
color: #e0e0e0;
border: 1px solid rgba(99, 102, 241, 0.5);
border-radius: 8px;
font-size: 0.95rem;
cursor: pointer;
font-family: system-ui, -apple-system, sans-serif;
display: none;
transition: all 0.15s ease;
}
.retry-btn:hover {
background: rgba(99, 102, 241, 0.4);
border-color: #6366f1;
color: #fff;
}
.particles {
position: fixed;
inset: 0;
@@ -116,9 +139,11 @@ body {
<div class="title">Crafting your page...</div>
<div class="subtitle">The AI is designing something unique for you</div>
<div class="stats" id="stats">Characters: 0</div>
<button class="retry-btn" id="retryBtn">Retry</button>
<script>
const stats = document.getElementById('stats');
const titleEl = document.querySelector('.title');
const retryBtn = document.getElementById('retryBtn');
const particles = document.getElementById('particles');
for (let i = 0; i < 30; i++) {
const p = document.createElement('div');
@@ -128,6 +153,7 @@ for (let i = 0; i < 30; i++) {
p.style.animationDuration = (2 + Math.random() * 2) + 's';
particles.appendChild(p);
}
retryBtn.onclick = () => { window.location.href = '/'; };
const start = Date.now();
const es = new EventSource('/stream' + window.location.search + window.location.hash);
let chars = 0;
@@ -139,11 +165,19 @@ const interval = setInterval(() => {
stats.textContent = 'Characters: ' + chars.toLocaleString() + ' | ' + ((Date.now() - start) / 1000).toFixed(1) + 's';
}
}, 100);
const timeout = setTimeout(() => {
es.close();
clearInterval(interval);
titleEl.textContent = 'Request timed out';
document.querySelector('.subtitle').textContent = 'Generation took too long. Please try again.';
retryBtn.style.display = 'inline-block';
}, 120000);
window.__wotf_html = '';
es.onmessage = (e) => {
if (e.data === '[DONE]') {
titleEl.textContent = 'Validating HTML...';
} else if (e.data === '[RENDER]') {
clearTimeout(timeout);
clearInterval(interval);
es.close();
titleEl.textContent = 'Rendering...';
@@ -164,10 +198,12 @@ es.onmessage = (e) => {
mode = 'status';
}
} else if (e.data === '[ERROR]') {
clearTimeout(timeout);
clearInterval(interval);
es.close();
titleEl.textContent = 'Error generating page';
document.querySelector('.subtitle').textContent = 'Please refresh and try again';
retryBtn.style.display = 'inline-block';
} else {
mode = 'generating';
const bytes = Uint8Array.from(atob(e.data), c => c.charCodeAt(0));
@@ -178,10 +214,12 @@ es.onmessage = (e) => {
}
};
es.onerror = () => {
clearTimeout(timeout);
clearInterval(interval);
es.close();
titleEl.textContent = 'Connection error';
document.querySelector('.subtitle').textContent = 'Please refresh and try again';
retryBtn.style.display = 'inline-block';
};
function injectMenu() {
@@ -299,6 +337,10 @@ function injectMenu() {
</html>"""
def make_httpx_client() -> httpx.Client:
return httpx.Client(timeout=300)
def call_llm(prompt: str, system: str = "You are a helpful assistant.") -> str:
headers = {"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"}
payload: dict[str, Any] = {
@@ -311,7 +353,7 @@ def call_llm(prompt: str, system: str = "You are a helpful assistant.") -> str:
"stream": False,
}
try:
with httpx.Client(timeout=300) as client:
with make_httpx_client() as client:
resp = client.post(API_URL, headers=headers, json=payload)
resp.raise_for_status()
data: dict[str, Any] = resp.json()
@@ -509,23 +551,41 @@ def get_generated_html(request_id: str) -> str:
def worker_loop() -> None:
worker_id = f"worker-{os.getpid()}-{threading.current_thread().name}"
hb_key = "wotf:processing:hb"
while worker_running:
request_id: str | None = None
try:
current = r.getset("wotf:processing", "locked") # type: ignore[union-attr]
if current is not None:
lock_holder = r.get("wotf:processing") # type: ignore[union-attr]
if lock_holder is not None:
hb_ttl_raw = r.ttl(hb_key) # type: ignore[union-attr]
hb_ttl: int = int(hb_ttl_raw) if isinstance(hb_ttl_raw, (int, float)) else -2
if hb_ttl == -2:
logger.info(f"[worker] Reclaiming stale lock from {lock_holder}")
r.delete("wotf:processing") # type: ignore[union-attr]
r.delete(hb_key) # type: ignore[union-attr]
else:
time.sleep(0.5)
continue
acquired = r.set("wotf:processing", worker_id, ex=120, nx=True) # type: ignore[union-attr]
if not acquired:
time.sleep(0.5)
continue
r.set(hb_key, worker_id, ex=15) # type: ignore[union-attr]
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]
r.delete(hb_key) # 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]
r.delete(hb_key) # type: ignore[union-attr]
time.sleep(0.5)
continue
@@ -544,9 +604,60 @@ def worker_loop() -> None:
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]
hb_start = time.time()
html_buffer: list[str] = []
try:
for chunk in stream_llm(ctx):
if time.time() - hb_start >= 5:
r.set(hb_key, worker_id, ex=15) # type: ignore[union-attr]
hb_start = time.time()
html_buffer.append(chunk)
encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
r.rpush(f"wotf:stream:{request_id}:chunks", encoded) # type: ignore[union-attr]
except BaseException as e: # noqa: BLE001
elapsed = time.time() - start
logger.error(f"[worker] stream_llm failed for {request_id} in {elapsed:.1f}s: {e}")
retries_raw = r.incr(f"wotf:req:{request_id}:retries") # type: ignore[union-attr]
retries: int = int(retries_raw) if isinstance(retries_raw, (int, float)) else 1
if retries <= 2:
logger.info(f"[worker] Requeuing {request_id} (attempt {retries}/2)")
r.set(f"wotf:req:{request_id}:status", "pending") # type: ignore[union-attr]
r.delete(f"wotf:stream:{request_id}:chunks") # type: ignore[union-attr]
r.rpush("wotf:queue", request_id) # type: ignore[union-attr]
else:
logger.error(f"[worker] {request_id} failed after {retries} attempts")
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]
r.delete(hb_key) # type: ignore[union-attr]
time.sleep(1)
continue
html = "".join(html_buffer)
r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validating HTML structure...]") # type: ignore[union-attr]
errors = validate_html(html)
if errors:
short = summarize_errors(errors)
r.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...]") # type: ignore[union-attr]
html = fix_html(html, errors)
errors = validate_html(html)
if errors:
short = summarize_errors(errors)
r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Still invalid. Fixing (attempt 2/2)...]") # type: ignore[union-attr]
html = fix_html(html, errors)
errors = validate_html(html)
if errors:
short = summarize_errors(errors)
r.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Still has {len(errors)} issue(s): {short}. Serving as-is.]") # type: ignore[union-attr]
else:
r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Fixed successfully.]") # type: ignore[union-attr]
r.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]") # type: ignore[union-attr]
for i in range(0, len(html), 1000):
chunk = html[i : i + 1000]
encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
r.rpush(f"wotf:stream:{request_id}:chunks", encoded) # type: ignore[union-attr]
else:
r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validation passed.]") # type: ignore[union-attr]
r.rpush(f"wotf:stream:{request_id}:chunks", "[RENDER]") # type: ignore[union-attr]
elapsed = time.time() - start
r.set(f"wotf:req:{request_id}:status", "done") # type: ignore[union-attr]
@@ -561,6 +672,7 @@ def worker_loop() -> None:
r.set("wotf:stats:gen_count", str(n + 1)) # type: ignore[union-attr]
r.delete("wotf:processing") # type: ignore[union-attr]
r.delete(hb_key) # 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
@@ -569,6 +681,7 @@ def worker_loop() -> None:
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]
r.delete(hb_key) # type: ignore[union-attr]
time.sleep(1)
@@ -630,7 +743,7 @@ def stream_llm(context: dict[str, Any]) -> Generator[str]:
logger.info(f"[{req_id}] Streaming LLM call for IP={context['ip']}, lang={context['language']}")
try:
with httpx.Client(timeout=300) as client, client.stream("POST", API_URL, headers=headers, json=payload) as resp:
with make_httpx_client() as client, client.stream("POST", API_URL, headers=headers, json=payload) as resp:
elapsed = time.time() - start
resp.raise_for_status()
buffer = ""
@@ -767,6 +880,9 @@ def stream() -> Response:
chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else [] # type: ignore[arg-type]
if chunks:
for chunk in chunks:
if chunk == "[RENDER]":
yield f"data: {chunk}\n\n"
return
yield f"data: {chunk}\n\n"
last_idx += len(chunks)
time.sleep(0.05)
@@ -775,11 +891,14 @@ def stream() -> Response:
status = status_raw if isinstance(status_raw, str) else "pending"
if status == "done":
break
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
time.sleep(0.2)
yield "data: [DONE]\n\n"
yield "data: [RENDER]\n\n"
return Response(
stream_with_context(generate()),
mimetype="text/event-stream",
@@ -798,9 +917,15 @@ def start_worker_if_needed() -> None:
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")
if os.getenv("GUNICORN_WORKER_ID"):
from gevent import spawn
worker_thread = spawn(worker_loop) # type: ignore[assignment]
logger.info("[startup] Worker greenlet started in this process")
else:
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")