increase workers and fix async issues
This commit is contained in:
@@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user