Migrate backend jobs to Redis Streams
This commit is contained in:
“ydy0615”
@@ -0,0 +1,704 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncIterator, Awaitable, Callable, Optional
|
||||
|
||||
logger = logging.getLogger("job_system")
|
||||
|
||||
try: # pragma: no cover - optional dependency in tests
|
||||
from redis import asyncio as redis_asyncio # type: ignore
|
||||
except Exception: # pragma: no cover - optional dependency in tests
|
||||
redis_asyncio = None
|
||||
|
||||
|
||||
TERMINAL_STATUSES = {"completed", "failed", "cancelled"}
|
||||
TERMINAL_EVENTS = {"done", "error", "cancelled"}
|
||||
|
||||
JOB_TYPES = (
|
||||
"completion",
|
||||
"pro_completion",
|
||||
"compress",
|
||||
"ocr",
|
||||
"convert",
|
||||
"tts",
|
||||
"asr",
|
||||
)
|
||||
|
||||
DEFAULT_CONCURRENCY = {
|
||||
"completion": 2,
|
||||
"pro_completion": 1,
|
||||
"compress": 1,
|
||||
"ocr": 1,
|
||||
"convert": 1,
|
||||
"tts": 1,
|
||||
"asr": 1,
|
||||
}
|
||||
|
||||
DEFAULT_QUEUE_SIZE = {
|
||||
"completion": 16,
|
||||
"pro_completion": 8,
|
||||
"compress": 8,
|
||||
"ocr": 8,
|
||||
"convert": 8,
|
||||
"tts": 4,
|
||||
"asr": 4,
|
||||
}
|
||||
|
||||
|
||||
class JobSystemError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class QueueFullError(JobSystemError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueConfig:
|
||||
job_type: str
|
||||
concurrency: int
|
||||
max_queue: int
|
||||
|
||||
|
||||
Handler = Callable[[dict[str, Any], Callable[[str, dict[str, Any]], Awaitable[None]], Callable[[], bool]], Awaitable[dict[str, Any]]]
|
||||
|
||||
|
||||
def _bool_env(name: str, default: bool) -> bool:
|
||||
value = os.getenv(name)
|
||||
if value is None:
|
||||
return default
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _int_env(name: str, default: int) -> int:
|
||||
try:
|
||||
return max(1, int(os.getenv(name, str(default))))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _float_env(name: str, default: float) -> float:
|
||||
try:
|
||||
return float(os.getenv(name, str(default)))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
|
||||
|
||||
def _busy_level(ratio: float) -> str:
|
||||
if ratio >= _float_env("JOB_BUSY_FULL_THRESHOLD", 1.0):
|
||||
return "full"
|
||||
if ratio >= _float_env("JOB_BUSY_HIGH_THRESHOLD", 0.75):
|
||||
return "busy"
|
||||
if ratio >= _float_env("JOB_BUSY_NORMAL_THRESHOLD", 0.25):
|
||||
return "normal"
|
||||
return "idle"
|
||||
|
||||
|
||||
def _json_dumps(value: Any) -> str:
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
|
||||
|
||||
def _json_loads(value: str | bytes | None, default: Any = None) -> Any:
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode("utf-8")
|
||||
if not value:
|
||||
return default
|
||||
return json.loads(value)
|
||||
|
||||
|
||||
def _queue_config(job_type: str) -> QueueConfig:
|
||||
upper = job_type.upper()
|
||||
concurrency = _int_env(f"JOB_{upper}_CONCURRENCY", DEFAULT_CONCURRENCY[job_type])
|
||||
max_queue = _int_env(f"JOB_{upper}_MAX_QUEUE", DEFAULT_QUEUE_SIZE[job_type])
|
||||
return QueueConfig(job_type=job_type, concurrency=concurrency, max_queue=max_queue)
|
||||
|
||||
|
||||
def get_job_backend_name() -> str:
|
||||
value = (os.getenv("JOB_BACKEND") or "").strip().lower()
|
||||
if value:
|
||||
return value
|
||||
if redis_asyncio is not None:
|
||||
return "redis"
|
||||
return "memory"
|
||||
|
||||
|
||||
def _shared_temp_dir() -> Path:
|
||||
path = Path(os.getenv("JOB_SHARED_TEMP_DIR", tempfile.gettempdir()) or tempfile.gettempdir())
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def persist_temp_input(raw_bytes: bytes, suffix: str) -> str:
|
||||
directory = _shared_temp_dir()
|
||||
fd, path = tempfile.mkstemp(prefix="job-input-", suffix=suffix, dir=directory)
|
||||
os.close(fd)
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(raw_bytes)
|
||||
return path
|
||||
|
||||
|
||||
async def _maybe_await(value: Any) -> Any:
|
||||
if inspect.isawaitable(value):
|
||||
return await value
|
||||
return value
|
||||
|
||||
|
||||
class BaseJobManager:
|
||||
def __init__(self) -> None:
|
||||
self.handlers: dict[str, Handler] = {}
|
||||
|
||||
def register_handler(self, job_type: str, handler: Handler) -> None:
|
||||
self.handlers[job_type] = handler
|
||||
|
||||
async def submit(self, job_type: str, payload: dict[str, Any], request_id: str | None = None) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
async def cancel(self, job_id: str, reason: str = "abort") -> dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def get_status(self, job_id: str) -> dict[str, Any] | None:
|
||||
raise NotImplementedError
|
||||
|
||||
async def stream_events(self, job_id: str) -> AsyncIterator[dict[str, Any]]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class InMemoryJobManager(BaseJobManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.jobs: dict[str, dict[str, Any]] = {}
|
||||
self.event_history: dict[str, list[dict[str, Any]]] = {}
|
||||
self.subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
self.queues = {job_type: asyncio.Queue() for job_type in JOB_TYPES}
|
||||
self.semaphores = {job_type: asyncio.Semaphore(_queue_config(job_type).concurrency) for job_type in JOB_TYPES}
|
||||
self.queue_counts = {job_type: 0 for job_type in JOB_TYPES}
|
||||
self.running_counts = {job_type: 0 for job_type in JOB_TYPES}
|
||||
self.running_tasks: dict[str, asyncio.Task] = {}
|
||||
self.worker_tasks: list[asyncio.Task] = []
|
||||
self.started = False
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def _ensure_started(self) -> None:
|
||||
if self.started:
|
||||
return
|
||||
self.started = True
|
||||
for job_type in JOB_TYPES:
|
||||
self.worker_tasks.append(asyncio.create_task(self._worker_loop(job_type)))
|
||||
|
||||
def _metrics(self, job_type: str) -> dict[str, Any]:
|
||||
config = _queue_config(job_type)
|
||||
queued = self.queue_counts[job_type]
|
||||
running = self.running_counts[job_type]
|
||||
capacity = max(config.max_queue + config.concurrency, 1)
|
||||
ratio = min((queued + running) / capacity, 1.0)
|
||||
return {
|
||||
"queue_position": queued if queued > 0 else 0,
|
||||
"queued_count": queued,
|
||||
"running_count": running,
|
||||
"concurrency_limit": config.concurrency,
|
||||
"max_queue": config.max_queue,
|
||||
"busy_ratio": round(ratio, 4),
|
||||
"busy_level": _busy_level(ratio),
|
||||
}
|
||||
|
||||
async def _publish(self, job_id: str, event: str, data: dict[str, Any]) -> None:
|
||||
event_payload = {"event": event, **data}
|
||||
self.event_history.setdefault(job_id, []).append(event_payload)
|
||||
for queue in self.subscribers.get(job_id, []):
|
||||
await queue.put(event_payload)
|
||||
|
||||
async def submit(self, job_type: str, payload: dict[str, Any], request_id: str | None = None) -> str:
|
||||
await self._ensure_started()
|
||||
if job_type not in self.handlers:
|
||||
raise JobSystemError(f"missing handler for job type: {job_type}")
|
||||
|
||||
async with self.lock:
|
||||
config = _queue_config(job_type)
|
||||
if self.queue_counts[job_type] >= config.max_queue:
|
||||
raise QueueFullError(f"{job_type} queue is full")
|
||||
job_id = request_id or str(uuid.uuid4())
|
||||
self.jobs[job_id] = {
|
||||
"job_id": job_id,
|
||||
"request_id": job_id,
|
||||
"job_type": job_type,
|
||||
"status": "queued",
|
||||
"payload": payload,
|
||||
"result": None,
|
||||
"error": "",
|
||||
"cancel_requested": False,
|
||||
"created_at": _now_ms(),
|
||||
"updated_at": _now_ms(),
|
||||
}
|
||||
self.event_history[job_id] = []
|
||||
self.queue_counts[job_type] += 1
|
||||
metrics = self._metrics(job_type)
|
||||
|
||||
await self._publish(job_id, "queued", {"job_id": job_id, "type": job_type, "status": "queued", **metrics})
|
||||
await self.queues[job_type].put(job_id)
|
||||
return job_id
|
||||
|
||||
async def cancel(self, job_id: str, reason: str = "abort") -> dict[str, Any]:
|
||||
async with self.lock:
|
||||
job = self.jobs.get(job_id)
|
||||
if not job:
|
||||
return {"cancelled": False, "status": "not_found"}
|
||||
if job["status"] in TERMINAL_STATUSES:
|
||||
return {"cancelled": False, "status": job["status"]}
|
||||
job["cancel_requested"] = True
|
||||
job["updated_at"] = _now_ms()
|
||||
task = self.running_tasks.get(job_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
if job["status"] == "queued":
|
||||
job["status"] = "cancelled"
|
||||
self.queue_counts[job["job_type"]] = max(0, self.queue_counts[job["job_type"]] - 1)
|
||||
metrics = self._metrics(job["job_type"])
|
||||
else:
|
||||
job["status"] = "cancelled"
|
||||
metrics = self._metrics(job["job_type"])
|
||||
|
||||
await self._publish(job_id, "cancelled", {"job_id": job_id, "type": job["job_type"], "status": "cancelled", "reason": reason, **metrics})
|
||||
return {"cancelled": True, "status": "ok"}
|
||||
|
||||
async def get_status(self, job_id: str) -> dict[str, Any] | None:
|
||||
job = self.jobs.get(job_id)
|
||||
if not job:
|
||||
return None
|
||||
metrics = self._metrics(job["job_type"])
|
||||
return {
|
||||
"job_id": job_id,
|
||||
"request_id": job["request_id"],
|
||||
"type": job["job_type"],
|
||||
"status": job["status"],
|
||||
"result": job["result"],
|
||||
"error": job["error"],
|
||||
**metrics,
|
||||
}
|
||||
|
||||
async def stream_events(self, job_id: str) -> AsyncIterator[dict[str, Any]]:
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
history = list(self.event_history.get(job_id, []))
|
||||
for item in history:
|
||||
yield item
|
||||
self.subscribers.setdefault(job_id, []).append(queue)
|
||||
try:
|
||||
while True:
|
||||
event = await queue.get()
|
||||
yield event
|
||||
if event["event"] in TERMINAL_EVENTS:
|
||||
break
|
||||
finally:
|
||||
with suppress(ValueError):
|
||||
self.subscribers.get(job_id, []).remove(queue)
|
||||
|
||||
async def _worker_loop(self, job_type: str) -> None:
|
||||
queue = self.queues[job_type]
|
||||
sem = self.semaphores[job_type]
|
||||
while True:
|
||||
job_id = await queue.get()
|
||||
async with self.lock:
|
||||
job = self.jobs.get(job_id)
|
||||
if not job or job["status"] == "cancelled":
|
||||
continue
|
||||
await sem.acquire()
|
||||
task = asyncio.create_task(self._run_job(job_id))
|
||||
self.running_tasks[job_id] = task
|
||||
|
||||
async def _run_job(self, job_id: str) -> None:
|
||||
job = self.jobs[job_id]
|
||||
job_type = job["job_type"]
|
||||
try:
|
||||
async with self.lock:
|
||||
if job["status"] == "cancelled":
|
||||
return
|
||||
self.queue_counts[job_type] = max(0, self.queue_counts[job_type] - 1)
|
||||
self.running_counts[job_type] += 1
|
||||
job["status"] = "running"
|
||||
job["updated_at"] = _now_ms()
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "started", {"job_id": job_id, "type": job_type, "status": "running", **metrics})
|
||||
|
||||
async def emit(event: str, data: dict[str, Any]) -> None:
|
||||
metrics_now = self._metrics(job_type)
|
||||
await self._publish(job_id, event, {"job_id": job_id, "type": job_type, "status": job["status"], **metrics_now, **data})
|
||||
|
||||
def is_cancelled() -> bool:
|
||||
return bool(job.get("cancel_requested"))
|
||||
|
||||
result = await self.handlers[job_type](job["payload"], emit, is_cancelled)
|
||||
async with self.lock:
|
||||
if job["cancel_requested"]:
|
||||
job["status"] = "cancelled"
|
||||
metrics = self._metrics(job_type)
|
||||
await emit("cancelled", {"reason": "abort"})
|
||||
return
|
||||
job["status"] = "completed"
|
||||
job["result"] = result
|
||||
job["updated_at"] = _now_ms()
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "done", {"job_id": job_id, "type": job_type, "status": "completed", "result": result, **metrics})
|
||||
except asyncio.CancelledError:
|
||||
async with self.lock:
|
||||
job["status"] = "cancelled"
|
||||
job["cancel_requested"] = True
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "cancelled", {"job_id": job_id, "type": job_type, "status": "cancelled", **metrics})
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("in-memory job failed job_id=%s type=%s", job_id, job_type)
|
||||
async with self.lock:
|
||||
job["status"] = "failed"
|
||||
job["error"] = str(exc)
|
||||
job["updated_at"] = _now_ms()
|
||||
metrics = self._metrics(job_type)
|
||||
await self._publish(job_id, "error", {"job_id": job_id, "type": job_type, "status": "failed", "error": str(exc), **metrics})
|
||||
finally:
|
||||
async with self.lock:
|
||||
self.running_counts[job_type] = max(0, self.running_counts[job_type] - 1)
|
||||
self.running_tasks.pop(job_id, None)
|
||||
self.semaphores[job_type].release()
|
||||
|
||||
async def close(self) -> None:
|
||||
for task in self.worker_tasks:
|
||||
task.cancel()
|
||||
for task in self.running_tasks.values():
|
||||
task.cancel()
|
||||
for task in self.worker_tasks:
|
||||
with suppress(asyncio.CancelledError):
|
||||
await task
|
||||
self.worker_tasks.clear()
|
||||
self.running_tasks.clear()
|
||||
self.started = False
|
||||
|
||||
|
||||
class RedisJobManager(BaseJobManager):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
if redis_asyncio is None:
|
||||
raise JobSystemError("redis package is not installed")
|
||||
self.redis = redis_asyncio.from_url(
|
||||
os.getenv("REDIS_URL", "redis://localhost:6379/0"),
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
)
|
||||
self.prefix = (os.getenv("JOB_REDIS_PREFIX") or "llmtext:jobs").strip() or "llmtext:jobs"
|
||||
self.state_ttl = _int_env("JOB_STATE_TTL_SECONDS", 600)
|
||||
self.event_ttl = _int_env("JOB_EVENT_TTL_SECONDS", 600)
|
||||
|
||||
def _queue_key(self, job_type: str) -> str:
|
||||
return f"{self.prefix}:queue:{job_type}"
|
||||
|
||||
def _event_key(self, job_id: str) -> str:
|
||||
return f"{self.prefix}:events:{job_id}"
|
||||
|
||||
def _state_key(self, job_id: str) -> str:
|
||||
return f"{self.prefix}:state:{job_id}"
|
||||
|
||||
def _metrics_key(self, job_type: str) -> str:
|
||||
return f"{self.prefix}:metrics:{job_type}"
|
||||
|
||||
def _group_name(self, job_type: str) -> str:
|
||||
return f"{self.prefix}:group:{job_type}"
|
||||
|
||||
async def ensure_groups(self) -> None:
|
||||
for job_type in JOB_TYPES:
|
||||
stream = self._queue_key(job_type)
|
||||
group = self._group_name(job_type)
|
||||
try:
|
||||
await self.redis.xgroup_create(stream, group, id="0-0", mkstream=True)
|
||||
except Exception as exc: # pragma: no cover - redis-specific
|
||||
if "BUSYGROUP" not in str(exc):
|
||||
raise
|
||||
|
||||
async def close(self) -> None:
|
||||
await self.redis.aclose()
|
||||
|
||||
async def _metrics(self, job_type: str) -> dict[str, Any]:
|
||||
raw = await self.redis.hgetall(self._metrics_key(job_type))
|
||||
queued = int(raw.get("queued_count", "0") or 0)
|
||||
running = int(raw.get("running_count", "0") or 0)
|
||||
config = _queue_config(job_type)
|
||||
capacity = max(config.max_queue + config.concurrency, 1)
|
||||
ratio = min((queued + running) / capacity, 1.0)
|
||||
return {
|
||||
"queued_count": queued,
|
||||
"running_count": running,
|
||||
"concurrency_limit": config.concurrency,
|
||||
"max_queue": config.max_queue,
|
||||
"busy_ratio": round(ratio, 4),
|
||||
"busy_level": _busy_level(ratio),
|
||||
}
|
||||
|
||||
async def _emit_event(self, job_id: str, event: str, data: dict[str, Any]) -> None:
|
||||
key = self._event_key(job_id)
|
||||
payload = {k: _json_dumps(v) if not isinstance(v, str) else v for k, v in data.items()}
|
||||
payload["event"] = event
|
||||
await self.redis.xadd(key, payload, maxlen=_int_env("JOB_EVENT_STREAM_MAXLEN", 512), approximate=True)
|
||||
await self.redis.expire(key, self.event_ttl)
|
||||
|
||||
async def _set_state(self, job_id: str, state: dict[str, Any]) -> None:
|
||||
serializable = {k: _json_dumps(v) if isinstance(v, (dict, list)) else str(v) for k, v in state.items()}
|
||||
await self.redis.hset(self._state_key(job_id), mapping=serializable)
|
||||
await self.redis.expire(self._state_key(job_id), self.state_ttl)
|
||||
|
||||
async def submit(self, job_type: str, payload: dict[str, Any], request_id: str | None = None) -> str:
|
||||
config = _queue_config(job_type)
|
||||
metrics = await self._metrics(job_type)
|
||||
if metrics["queued_count"] >= config.max_queue:
|
||||
raise QueueFullError(f"{job_type} queue is full")
|
||||
|
||||
job_id = request_id or str(uuid.uuid4())
|
||||
created_at = _now_ms()
|
||||
state = {
|
||||
"job_id": job_id,
|
||||
"request_id": job_id,
|
||||
"type": job_type,
|
||||
"status": "queued",
|
||||
"error": "",
|
||||
"created_at": created_at,
|
||||
"updated_at": created_at,
|
||||
"cancel_requested": "0",
|
||||
}
|
||||
await self._set_state(job_id, state)
|
||||
await self.redis.hincrby(self._metrics_key(job_type), "queued_count", 1)
|
||||
await self.redis.expire(self._metrics_key(job_type), self.state_ttl)
|
||||
metrics = await self._metrics(job_type)
|
||||
await self._emit_event(job_id, "queued", {"job_id": job_id, "type": job_type, "status": "queued", **metrics})
|
||||
await self.redis.xadd(self._queue_key(job_type), {"job_id": job_id, "payload": _json_dumps(payload), "request_id": job_id})
|
||||
return job_id
|
||||
|
||||
async def cancel(self, job_id: str, reason: str = "abort") -> dict[str, Any]:
|
||||
state = await self.get_status(job_id)
|
||||
if not state:
|
||||
return {"cancelled": False, "status": "not_found"}
|
||||
if state["status"] in TERMINAL_STATUSES:
|
||||
return {"cancelled": False, "status": state["status"]}
|
||||
await self.redis.hset(self._state_key(job_id), mapping={"cancel_requested": "1", "status": "cancelled", "updated_at": _now_ms(), "cancel_reason": reason})
|
||||
metrics = await self._metrics(state["type"])
|
||||
await self._emit_event(job_id, "cancelled", {"job_id": job_id, "type": state["type"], "status": "cancelled", "reason": reason, **metrics})
|
||||
return {"cancelled": True, "status": "ok"}
|
||||
|
||||
async def get_status(self, job_id: str) -> dict[str, Any] | None:
|
||||
state = await self.redis.hgetall(self._state_key(job_id))
|
||||
if not state:
|
||||
return None
|
||||
job_type = state.get("type", "")
|
||||
metrics = await self._metrics(job_type) if job_type else {}
|
||||
result = state.get("result")
|
||||
error = state.get("error", "")
|
||||
return {
|
||||
"job_id": state.get("job_id", job_id),
|
||||
"request_id": state.get("request_id", job_id),
|
||||
"type": job_type,
|
||||
"status": state.get("status", "queued"),
|
||||
"error": error,
|
||||
"result": _json_loads(result, result),
|
||||
"cancel_requested": state.get("cancel_requested") == "1",
|
||||
**metrics,
|
||||
}
|
||||
|
||||
async def stream_events(self, job_id: str) -> AsyncIterator[dict[str, Any]]:
|
||||
stream = self._event_key(job_id)
|
||||
last_id = "0-0"
|
||||
while True:
|
||||
events = await self.redis.xread({stream: last_id}, block=1000, count=20)
|
||||
if not events:
|
||||
state = await self.get_status(job_id)
|
||||
if state and state["status"] in TERMINAL_STATUSES:
|
||||
break
|
||||
continue
|
||||
for _, entries in events:
|
||||
for entry_id, fields in entries:
|
||||
last_id = entry_id
|
||||
event_payload: dict[str, Any] = {}
|
||||
for key, value in fields.items():
|
||||
if key == "event":
|
||||
event_payload[key] = value
|
||||
continue
|
||||
try:
|
||||
event_payload[key] = json.loads(value)
|
||||
except Exception:
|
||||
event_payload[key] = value
|
||||
yield event_payload
|
||||
if event_payload.get("event") in TERMINAL_EVENTS:
|
||||
return
|
||||
|
||||
|
||||
class RedisWorker:
|
||||
def __init__(self, manager: RedisJobManager) -> None:
|
||||
self.manager = manager
|
||||
self.running_tasks: dict[str, asyncio.Task] = {}
|
||||
self.queue_semaphores = {job_type: asyncio.Semaphore(_queue_config(job_type).concurrency) for job_type in JOB_TYPES}
|
||||
self.poll_interval = _float_env("JOB_CANCEL_POLL_SECONDS", 0.5)
|
||||
self.consumer_name = (os.getenv("JOB_CONSUMER_NAME") or f"worker-{uuid.uuid4().hex[:8]}").strip()
|
||||
|
||||
async def run_forever(self) -> None:
|
||||
await self.manager.ensure_groups()
|
||||
cancel_task = asyncio.create_task(self._cancel_watch_loop())
|
||||
consumers = [asyncio.create_task(self._consume_loop(job_type)) for job_type in JOB_TYPES]
|
||||
try:
|
||||
await asyncio.gather(*consumers)
|
||||
finally:
|
||||
cancel_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await cancel_task
|
||||
|
||||
async def _cancel_watch_loop(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
for job_id, task in list(self.running_tasks.items()):
|
||||
state = await self.manager.get_status(job_id)
|
||||
if state and state.get("cancel_requested") and not task.done():
|
||||
task.cancel()
|
||||
|
||||
async def _consume_loop(self, job_type: str) -> None:
|
||||
queue_key = self.manager._queue_key(job_type)
|
||||
group = self.manager._group_name(job_type)
|
||||
semaphore = self.queue_semaphores[job_type]
|
||||
while True:
|
||||
streams = await self.manager.redis.xreadgroup(group, self.consumer_name, {queue_key: ">"}, count=1, block=1000)
|
||||
if not streams:
|
||||
continue
|
||||
for _, messages in streams:
|
||||
for message_id, fields in messages:
|
||||
await semaphore.acquire()
|
||||
task = asyncio.create_task(self._run_message(job_type, queue_key, group, message_id, fields, semaphore))
|
||||
self.running_tasks[fields["job_id"]] = task
|
||||
|
||||
async def _run_message(
|
||||
self,
|
||||
job_type: str,
|
||||
queue_key: str,
|
||||
group: str,
|
||||
message_id: str,
|
||||
fields: dict[str, str],
|
||||
semaphore: asyncio.Semaphore,
|
||||
) -> None:
|
||||
job_id = fields["job_id"]
|
||||
try:
|
||||
state = await self.manager.get_status(job_id)
|
||||
if not state or state["status"] == "cancelled":
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
return
|
||||
|
||||
await self.manager.redis.hincrby(self.manager._metrics_key(job_type), "queued_count", -1)
|
||||
await self.manager.redis.hincrby(self.manager._metrics_key(job_type), "running_count", 1)
|
||||
await self.manager._set_state(job_id, {
|
||||
"job_id": job_id,
|
||||
"request_id": state["request_id"],
|
||||
"type": job_type,
|
||||
"status": "running",
|
||||
"updated_at": _now_ms(),
|
||||
"created_at": state.get("created_at", _now_ms()),
|
||||
"cancel_requested": "1" if state.get("cancel_requested") else "0",
|
||||
"error": "",
|
||||
})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "started", {"job_id": job_id, "type": job_type, "status": "running", **metrics})
|
||||
|
||||
payload = _json_loads(fields["payload"], {})
|
||||
|
||||
async def emit(event: str, data: dict[str, Any]) -> None:
|
||||
live_state = await self.manager.get_status(job_id) or {"status": "running"}
|
||||
live_metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, event, {"job_id": job_id, "type": job_type, "status": live_state["status"], **live_metrics, **data})
|
||||
|
||||
def is_cancelled() -> bool:
|
||||
task = self.running_tasks.get(job_id)
|
||||
return bool(task and task.cancelled())
|
||||
|
||||
result = await self.manager.handlers[job_type](payload, emit, is_cancelled)
|
||||
current = await self.manager.get_status(job_id)
|
||||
if current and current["status"] == "cancelled":
|
||||
return
|
||||
|
||||
await self.manager._set_state(job_id, {
|
||||
"job_id": job_id,
|
||||
"request_id": state["request_id"],
|
||||
"type": job_type,
|
||||
"status": "completed",
|
||||
"updated_at": _now_ms(),
|
||||
"created_at": state.get("created_at", _now_ms()),
|
||||
"cancel_requested": "0",
|
||||
"error": "",
|
||||
"result": _json_dumps(result),
|
||||
})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "done", {"job_id": job_id, "type": job_type, "status": "completed", "result": result, **metrics})
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
except asyncio.CancelledError:
|
||||
await self.manager.redis.hset(self.manager._state_key(job_id), mapping={"status": "cancelled", "cancel_requested": "1", "updated_at": _now_ms()})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "cancelled", {"job_id": job_id, "type": job_type, "status": "cancelled", **metrics})
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("redis worker failed job_id=%s type=%s", job_id, job_type)
|
||||
state = await self.manager.get_status(job_id)
|
||||
request_id = state["request_id"] if state else job_id
|
||||
await self.manager._set_state(job_id, {
|
||||
"job_id": job_id,
|
||||
"request_id": request_id,
|
||||
"type": job_type,
|
||||
"status": "failed",
|
||||
"updated_at": _now_ms(),
|
||||
"created_at": state.get("created_at", _now_ms()) if state else _now_ms(),
|
||||
"cancel_requested": "0",
|
||||
"error": str(exc),
|
||||
})
|
||||
metrics = await self.manager._metrics(job_type)
|
||||
await self.manager._emit_event(job_id, "error", {"job_id": job_id, "type": job_type, "status": "failed", "error": str(exc), **metrics})
|
||||
await self.manager.redis.xack(queue_key, group, message_id)
|
||||
finally:
|
||||
self.running_tasks.pop(job_id, None)
|
||||
await self.manager.redis.hincrby(self.manager._metrics_key(job_type), "running_count", -1)
|
||||
semaphore.release()
|
||||
|
||||
|
||||
_job_manager: BaseJobManager | None = None
|
||||
|
||||
|
||||
def get_job_manager() -> BaseJobManager:
|
||||
global _job_manager
|
||||
if _job_manager is None:
|
||||
backend = get_job_backend_name()
|
||||
if backend == "redis":
|
||||
_job_manager = RedisJobManager()
|
||||
else:
|
||||
_job_manager = InMemoryJobManager()
|
||||
return _job_manager
|
||||
|
||||
|
||||
def reset_job_manager() -> None:
|
||||
global _job_manager
|
||||
manager = _job_manager
|
||||
if manager is not None:
|
||||
close = getattr(manager, "close", None)
|
||||
if close is not None:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
try:
|
||||
asyncio.run(close())
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
loop.create_task(close())
|
||||
_job_manager = None
|
||||
Reference in New Issue
Block a user