File size: 4,685 Bytes
345855e | 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 | import uuid
import asyncio
from datetime import datetime
from utils.logger import logger
from utils.job_queue import create_job, get_job
# =====================================================
# VALIDATION
# =====================================================
def normalize_payload(payload: dict | None):
payload = payload or {}
items = payload.get("items", [])
if not isinstance(items, list):
raise ValueError("batch items must be a list")
normalized = []
for i, item in enumerate(items):
if not isinstance(item, dict):
raise ValueError(f"batch item {i} must be an object")
normalized.append({
"id": item.get("id", str(uuid.uuid4())),
"task": item.get("task"),
"video_path": item.get("video_path"),
"source": item.get("source"),
"payload": item.get("payload", {}),
"webhook": item.get("webhook"),
})
if not normalized:
raise ValueError("batch cannot be empty")
return normalized
# =====================================================
# SAFE TASK EXECUTION WRAPPER
# =====================================================
async def execute_single(task_executor, item, index: int):
try:
logger.info(f"[BATCH] Executing item {index} → {item['task']}")
result = await task_executor(
item["task"],
{
**(item.get("payload") or {}),
"video_path": item.get("video_path"),
"source": item.get("source"),
},
item.get("webhook"),
)
return {
"index": index,
"id": item["id"],
"task": item["task"],
"status": "success",
"result": result,
}
except Exception as e:
logger.exception(f"[BATCH ERROR] item {index}")
return {
"index": index,
"id": item["id"],
"task": item.get("task"),
"status": "failed",
"error": str(e),
}
# =====================================================
# CONCURRENCY CONTROLLER
# =====================================================
async def run_concurrent(tasks, executor, max_concurrency: int = 3):
semaphore = asyncio.Semaphore(max_concurrency)
async def bound(item, index):
async with semaphore:
return await execute_single(executor, item, index)
return await asyncio.gather(
*[bound(item, i) for i, item in enumerate(tasks)]
)
# =====================================================
# MAIN ENTRYPOINT (REGISTRY COMPATIBLE)
# =====================================================
async def run(payload: dict | None = None, context: dict | None = None):
batch_id = str(uuid.uuid4())
started_at = datetime.utcnow().isoformat()
try:
items = normalize_payload(payload)
logger.info(f"[BATCH] Starting batch_id={batch_id}, items={len(items)}")
# -------------------------------------------------
# EXECUTOR HOOK (inject from registry context)
# -------------------------------------------------
def executor(task_name, task_payload, webhook=None):
"""
This is intentionally abstract so it can plug into:
- registry executor
- legacy executor
- FastAPI layer
"""
from core.execution.executor import execute_task
return execute_task(task_name, task_payload, webhook)
# -------------------------------------------------
# RUN BATCH
# -------------------------------------------------
results = await run_concurrent(items, executor)
success_count = len([r for r in results if r["status"] == "success"])
failed_count = len(results) - success_count
# -------------------------------------------------
# OUTPUT CONTRACT
# -------------------------------------------------
return {
"status": "completed",
"task": "batch_runner",
"batch_id": batch_id,
"started_at": started_at,
"completed_at": datetime.utcnow().isoformat(),
"total": len(items),
"success": success_count,
"failed": failed_count,
"results": results,
}
except Exception as e:
logger.exception("[BATCH FATAL ERROR]")
return {
"status": "error",
"task": "batch_runner",
"batch_id": batch_id,
"message": str(e),
"stage": "batch_execution_failed"
} |