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", "web_search", "compress", "ocr", "convert", "tts", "asr", ) DEFAULT_CONCURRENCY = { "completion": 2, "pro_completion": 1, "web_search": 1, "compress": 1, "ocr": 1, "convert": 1, "tts": 1, "asr": 1, } DEFAULT_QUEUE_SIZE = { "completion": 16, "pro_completion": 8, "web_search": 4, "compress": 8, "ocr": 8, "convert": 8, "tts": 4, "asr": 4, } class JobSystemError(RuntimeError): """任务系统内部错误""" def __init__(self, message: str = "任务系统错误", error_code: str = "job_system_error") -> None: super().__init__(message) self.message = message self.error_code = error_code def __str__(self) -> str: return f"{self.error_code}: {self.message}" class QueueFullError(JobSystemError): """任务队列已满""" def __init__(self, job_type: str, max_queue: int) -> None: super().__init__(f"{job_type} 队列已满 (当前: {max_queue}/{max_queue})", "queue_full") self.job_type = job_type self.max_queue = max_queue @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() self.subscribers.setdefault(job_id, []).append(queue) last_index = 0 try: while True: history = list(self.event_history.get(job_id, [])) while last_index < len(history): item = history[last_index] last_index += 1 yield item if item.get("event") in TERMINAL_EVENTS: return event = await queue.get() last_index = len(self.event_history.get(job_id, [])) 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: try: streams = await self.manager.redis.xreadgroup(group, self.consumer_name, {queue_key: ">"}, count=1, block=1000) except asyncio.CancelledError: raise except Exception as exc: logger.warning("redis worker consume loop retrying type=%s error=%s", job_type, exc) await asyncio.sleep(1) continue 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