Spaces:
Running
Running
File size: 4,230 Bytes
99a1943 68c5d6f 99a1943 68c5d6f 99a1943 68c5d6f 99a1943 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """SSE progress streaming with tqdm integration.
Usage in a FastAPI endpoint:
from shared.progress import StreamingTask
from fastapi.responses import StreamingResponse
@app.post("/some-endpoint")
async def my_endpoint(...):
task = StreamingTask()
def work(p):
p.update(10, "Loading data...")
# ... heavy work ...
p.update(80, "Computing...")
# ...
p.finish(result={"key": "value"})
return StreamingResponse(task.stream(work), media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
"""
from __future__ import annotations
import asyncio
import json
import math
import threading
from tqdm import tqdm
def _sanitize(obj):
"""Recursively replace NaN/Inf with None and convert numpy scalar types to Python natives."""
if isinstance(obj, float):
return None if (math.isnan(obj) or math.isinf(obj)) else obj
if isinstance(obj, dict):
return {k: _sanitize(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_sanitize(v) for v in obj]
# coerce numpy int/float types that slipped through
try:
import numpy as _np # noqa: PLC0415
if isinstance(obj, (_np.integer,)):
return int(obj)
if isinstance(obj, (_np.floating,)):
v = float(obj)
return None if (math.isnan(v) or math.isinf(v)) else v
except ImportError:
pass
return obj
class _Progress:
"""Passed to the worker function; wraps tqdm and SSE put."""
def __init__(self, put_fn, total: int = 100):
self._put = put_fn
self._pct = 0
self._bar = tqdm(total=total, unit="%", ncols=70, colour="green",
bar_format="{l_bar}{bar}| {n}/{total}%")
def update(self, target_pct: int, msg: str = "") -> None:
"""Advance to target_pct (0-100) and send an SSE event."""
delta = max(0, target_pct - self._pct)
if delta:
self._bar.update(delta)
self._bar.set_description(msg[:40] if msg else "")
self._pct = max(self._pct, target_pct)
self._put({"pct": self._pct, "msg": msg})
def finish(self, result: dict | None = None, error: str | None = None) -> None:
"""Send the final event (with result payload or error)."""
self._bar.close()
payload: dict = {"done": True}
if error:
payload["error"] = error
payload["pct"] = -1
else:
payload["pct"] = 100
payload["msg"] = "Done"
if result is not None:
payload["result"] = result
self._put(payload)
class StreamingTask:
"""Runs a synchronous worker in a thread and yields SSE events."""
def __init__(self) -> None:
self._loop: asyncio.AbstractEventLoop | None = None
self._queue: asyncio.Queue | None = None
def _put(self, data: dict) -> None:
if self._loop and self._queue:
self._loop.call_soon_threadsafe(self._queue.put_nowait, data)
def stream(self, worker_fn):
"""
Returns an async generator of SSE-formatted strings.
worker_fn(p: _Progress) — synchronous; call p.update() and p.finish().
"""
self._loop = asyncio.get_event_loop()
self._queue = asyncio.Queue()
put = self._put
def _run():
p = _Progress(put)
try:
worker_fn(p)
except Exception as exc:
p.finish(error=str(exc))
async def _generate():
t = threading.Thread(target=_run, daemon=True)
t.start()
while True:
try:
item = await asyncio.wait_for(self._queue.get(), timeout=2.0)
except asyncio.TimeoutError:
yield ": keepalive\n\n" # SSE comment — ignored by browser, prevents proxy timeout
continue
yield f"data: {json.dumps(_sanitize(item))}\n\n"
if item.get("done"):
break
t.join()
return _generate()
|