Spaces:
Running
Running
File size: 5,994 Bytes
df6cd5e | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """
core/task_queue.py — lightweight persistent task queue.
This is intentionally small and dependency-free: it gives the app an async
execution path without requiring Celery/Redis just to package a deployable
ZIP. Queue state is persisted to JSON so the UI/API can poll task progress.
"""
from __future__ import annotations
import copy
import json
import os
import queue
import tempfile
import threading
import time
import uuid
from typing import Any, Callable, Dict, Optional
class InMemoryTaskQueue:
def __init__(self, workspace: str, max_workers: int = 1) -> None:
self.workspace = os.path.abspath(workspace)
self.max_workers = max(1, int(max_workers))
self.tasks_dir = os.path.join(self.workspace, "queue", "tasks")
os.makedirs(self.tasks_dir, exist_ok=True)
self._queue: "queue.Queue[tuple[str, Callable[..., Any], Dict[str, Any]]]" = queue.Queue()
self._lock = threading.RLock()
self._started = False
self._workers: list[threading.Thread] = []
def _task_path(self, task_id: str) -> str:
return os.path.join(self.tasks_dir, f"{task_id}.json")
def _atomic_write(self, path: str, payload: Dict[str, Any]) -> None:
fd, tmp = tempfile.mkstemp(prefix=".task_", suffix=".tmp", dir=os.path.dirname(path))
try:
with os.fdopen(fd, "w", encoding="utf-8") as fh:
json.dump(payload, fh, indent=2, ensure_ascii=False, default=str)
fh.flush()
try:
os.fsync(fh.fileno())
except OSError:
pass
os.replace(tmp, path)
finally:
try:
if os.path.exists(tmp):
os.unlink(tmp)
except OSError:
pass
def _load_task(self, task_id: str) -> Optional[Dict[str, Any]]:
path = self._task_path(task_id)
if not os.path.exists(path):
return None
try:
with open(path, "r", encoding="utf-8") as fh:
payload = json.load(fh)
return payload if isinstance(payload, dict) else None
except Exception:
return None
def _save_task(self, task_id: str, payload: Dict[str, Any]) -> None:
self._atomic_write(self._task_path(task_id), payload)
def _ensure_started(self) -> None:
if self._started:
return
with self._lock:
if self._started:
return
for idx in range(self.max_workers):
t = threading.Thread(target=self._worker, name=f"devai-queue-{idx}", daemon=True)
t.start()
self._workers.append(t)
self._started = True
def enqueue(
self,
func: Callable[..., Any],
*,
kwargs: Optional[Dict[str, Any]] = None,
tenant_id: str = "default",
label: str = "pipeline",
) -> str:
self._ensure_started()
task_id = f"task-{uuid.uuid4().hex[:16]}"
now = time.time()
payload = {
"task_id": task_id,
"tenant_id": tenant_id,
"label": label,
"status": "queued",
"created_at": now,
"started_at": None,
"finished_at": None,
"progress_events": [],
"result": None,
"error": None,
}
self._save_task(task_id, payload)
self._queue.put((task_id, func, dict(kwargs or {})))
return task_id
def _worker(self) -> None:
while True:
task_id, func, kwargs = self._queue.get()
try:
payload = self._load_task(task_id) or {"task_id": task_id}
payload.update({"status": "running", "started_at": time.time(), "error": None})
self._save_task(task_id, payload)
def _progress(event: Dict[str, Any]) -> None:
current = self._load_task(task_id) or payload
events = list(current.get("progress_events") or [])
events.append(copy.deepcopy(event))
current["progress_events"] = events[-500:]
self._save_task(task_id, current)
call_kwargs = dict(kwargs)
call_kwargs.setdefault("_progress_callback", _progress)
result = func(**call_kwargs)
current = self._load_task(task_id) or payload
current.update({
"status": "completed",
"finished_at": time.time(),
"result": result,
})
self._save_task(task_id, current)
except Exception as exc:
current = self._load_task(task_id) or {"task_id": task_id}
current.update({
"status": "failed",
"finished_at": time.time(),
"error": str(exc),
})
self._save_task(task_id, current)
finally:
self._queue.task_done()
def get_status(self, task_id: str) -> Optional[Dict[str, Any]]:
payload = self._load_task(task_id)
if payload is None:
return None
return {
"task_id": payload.get("task_id"),
"tenant_id": payload.get("tenant_id"),
"label": payload.get("label"),
"status": payload.get("status"),
"created_at": payload.get("created_at"),
"started_at": payload.get("started_at"),
"finished_at": payload.get("finished_at"),
"progress_events": list(payload.get("progress_events") or []),
"error": payload.get("error"),
}
def get_result(self, task_id: str) -> Optional[Dict[str, Any]]:
payload = self._load_task(task_id)
if payload is None:
return None
return copy.deepcopy(payload)
__all__ = ["InMemoryTaskQueue"]
|