“ydy0615”
873 lines
34 KiB
Python
873 lines
34 KiB
Python
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",
|
|
)
|
|
|
|
|
|
def _int_env(name: str, default: int) -> int:
|
|
try:
|
|
return max(1, int(os.getenv(name, str(default))))
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
DEFAULT_CONCURRENCY = {
|
|
"completion": 2,
|
|
"pro_completion": 1,
|
|
"web_search": 1,
|
|
"compress": 1,
|
|
"ocr": 1,
|
|
"convert": 1,
|
|
"tts": _int_env("JOB_TTS_CONCURRENCY", 2),
|
|
"asr": _int_env("JOB_ASR_CONCURRENCY", 1),
|
|
}
|
|
|
|
DEFAULT_QUEUE_SIZE = {
|
|
"completion": 16,
|
|
"pro_completion": 8,
|
|
"web_search": 4,
|
|
"compress": 8,
|
|
"ocr": 8,
|
|
"convert": 8,
|
|
"tts": _int_env("JOB_TTS_MAX_QUEUE", 8),
|
|
"asr": _int_env("JOB_ASR_MAX_QUEUE", 8),
|
|
}
|
|
|
|
|
|
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 _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(job_type, config.max_queue)
|
|
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(),
|
|
"started_at": 0,
|
|
"completed_at": 0,
|
|
"queue_ms": 0,
|
|
"run_ms": 0,
|
|
"total_ms": 0,
|
|
}
|
|
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"],
|
|
"created_at": job.get("created_at", 0),
|
|
"started_at": job.get("started_at", 0),
|
|
"completed_at": job.get("completed_at", 0),
|
|
"queue_ms": job.get("queue_ms", 0),
|
|
"run_ms": job.get("run_ms", 0),
|
|
"total_ms": job.get("total_ms", 0),
|
|
**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"
|
|
started_at = _now_ms()
|
|
job["updated_at"] = started_at
|
|
job["started_at"] = started_at
|
|
job["queue_ms"] = max(0, started_at - int(job.get("created_at", started_at)))
|
|
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"))
|
|
|
|
job_payload = dict(job["payload"])
|
|
job_payload["job_context"] = {
|
|
"job_id": job_id,
|
|
"created_at": int(job.get("created_at", 0) or 0),
|
|
"started_at": int(job.get("started_at", 0) or 0),
|
|
"queue_ms": int(job.get("queue_ms", 0) or 0),
|
|
}
|
|
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
|
|
completed_at = _now_ms()
|
|
job["updated_at"] = completed_at
|
|
job["completed_at"] = completed_at
|
|
job["run_ms"] = max(0, completed_at - int(job.get("started_at", completed_at)))
|
|
job["total_ms"] = max(0, completed_at - int(job.get("created_at", completed_at)))
|
|
metrics = self._metrics(job_type)
|
|
await self._publish(
|
|
job_id,
|
|
"done",
|
|
{
|
|
"job_id": job_id,
|
|
"type": job_type,
|
|
"status": "completed",
|
|
"result": result,
|
|
"queue_ms": job.get("queue_ms", 0),
|
|
"run_ms": job.get("run_ms", 0),
|
|
"total_ms": job.get("total_ms", 0),
|
|
**metrics,
|
|
},
|
|
)
|
|
except asyncio.CancelledError:
|
|
async with self.lock:
|
|
job["status"] = "cancelled"
|
|
job["cancel_requested"] = True
|
|
completed_at = _now_ms()
|
|
job["updated_at"] = completed_at
|
|
job["completed_at"] = completed_at
|
|
job["run_ms"] = max(0, completed_at - int(job.get("started_at", completed_at)))
|
|
job["total_ms"] = max(0, completed_at - int(job.get("created_at", completed_at)))
|
|
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)
|
|
completed_at = _now_ms()
|
|
job["updated_at"] = completed_at
|
|
job["completed_at"] = completed_at
|
|
job["run_ms"] = max(0, completed_at - int(job.get("started_at", completed_at)))
|
|
job["total_ms"] = max(0, completed_at - int(job.get("created_at", completed_at)))
|
|
metrics = self._metrics(job_type)
|
|
await self._publish(
|
|
job_id,
|
|
"error",
|
|
{
|
|
"job_id": job_id,
|
|
"type": job_type,
|
|
"status": "failed",
|
|
"error": str(exc),
|
|
"queue_ms": job.get("queue_ms", 0),
|
|
"run_ms": job.get("run_ms", 0),
|
|
"total_ms": job.get("total_ms", 0),
|
|
**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(job_type, config.max_queue)
|
|
|
|
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,
|
|
"started_at": 0,
|
|
"completed_at": 0,
|
|
"queue_ms": 0,
|
|
"run_ms": 0,
|
|
"total_ms": 0,
|
|
"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",
|
|
"created_at": int(state.get("created_at", "0") or 0),
|
|
"started_at": int(state.get("started_at", "0") or 0),
|
|
"completed_at": int(state.get("completed_at", "0") or 0),
|
|
"queue_ms": int(state.get("queue_ms", "0") or 0),
|
|
"run_ms": int(state.get("run_ms", "0") or 0),
|
|
"total_ms": int(state.get("total_ms", "0") or 0),
|
|
**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"]
|
|
started_at = 0
|
|
created_at = 0
|
|
queue_ms = 0
|
|
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)
|
|
started_at = _now_ms()
|
|
created_at = int(state.get("created_at", 0) or 0)
|
|
queue_ms = max(0, started_at - created_at)
|
|
await self.manager._set_state(job_id, {
|
|
"job_id": job_id,
|
|
"request_id": state["request_id"],
|
|
"type": job_type,
|
|
"status": "running",
|
|
"updated_at": started_at,
|
|
"created_at": created_at or started_at,
|
|
"started_at": started_at,
|
|
"completed_at": 0,
|
|
"queue_ms": queue_ms,
|
|
"run_ms": 0,
|
|
"total_ms": 0,
|
|
"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"], {})
|
|
payload["job_context"] = {
|
|
"job_id": job_id,
|
|
"created_at": created_at,
|
|
"started_at": started_at,
|
|
"queue_ms": queue_ms,
|
|
}
|
|
|
|
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":
|
|
await self.manager.redis.xack(queue_key, group, message_id)
|
|
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": created_at or started_at,
|
|
"started_at": started_at,
|
|
"completed_at": _now_ms(),
|
|
"queue_ms": queue_ms,
|
|
"run_ms": max(0, _now_ms() - started_at),
|
|
"total_ms": max(0, _now_ms() - (created_at or started_at)),
|
|
"cancel_requested": "0",
|
|
"error": "",
|
|
"result": _json_dumps(result),
|
|
})
|
|
metrics = await self.manager._metrics(job_type)
|
|
final_state = await self.manager.get_status(job_id) or {}
|
|
await self.manager._emit_event(
|
|
job_id,
|
|
"done",
|
|
{
|
|
"job_id": job_id,
|
|
"type": job_type,
|
|
"status": "completed",
|
|
"result": result,
|
|
"queue_ms": final_state.get("queue_ms", queue_ms),
|
|
"run_ms": final_state.get("run_ms", 0),
|
|
"total_ms": final_state.get("total_ms", 0),
|
|
**metrics,
|
|
},
|
|
)
|
|
await self.manager.redis.xack(queue_key, group, message_id)
|
|
except asyncio.CancelledError:
|
|
cancelled_at = _now_ms()
|
|
await self.manager.redis.hset(
|
|
self.manager._state_key(job_id),
|
|
mapping={
|
|
"status": "cancelled",
|
|
"cancel_requested": "1",
|
|
"updated_at": cancelled_at,
|
|
"completed_at": cancelled_at,
|
|
"run_ms": max(0, cancelled_at - started_at),
|
|
"total_ms": max(0, cancelled_at - (created_at or started_at)),
|
|
},
|
|
)
|
|
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": created_at or (_now_ms() if state else _now_ms()),
|
|
"started_at": started_at,
|
|
"completed_at": _now_ms(),
|
|
"queue_ms": queue_ms,
|
|
"run_ms": max(0, _now_ms() - started_at),
|
|
"total_ms": max(0, _now_ms() - (created_at or started_at)),
|
|
"cancel_requested": "0",
|
|
"error": str(exc),
|
|
})
|
|
metrics = await self.manager._metrics(job_type)
|
|
final_state = await self.manager.get_status(job_id) or {}
|
|
await self.manager._emit_event(
|
|
job_id,
|
|
"error",
|
|
{
|
|
"job_id": job_id,
|
|
"type": job_type,
|
|
"status": "failed",
|
|
"error": str(exc),
|
|
"queue_ms": final_state.get("queue_ms", queue_ms),
|
|
"run_ms": final_state.get("run_ms", 0),
|
|
"total_ms": final_state.get("total_ms", 0),
|
|
**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
|