File size: 6,819 Bytes
0af083b | 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 asyncio
from datetime import datetime, timezone
from typing import Dict, List, Optional
from .models import Job, JobStatus, Pipeline, PipelineStep
from src.agents import executor, registry, AgentContext
from src.agents.base import AgentStatus
from src.config import settings
from src.telemetry import collector
# Per-agent-type concurrency weight: fraction of max_concurrent_agents slots reserved
# Values > 1.0 mean the type can use more than its fair share; < 1.0 means restricted
AGENT_TYPE_WEIGHTS: Dict[str, float] = {
"echo": 1.0,
"compute": 2.0, # compute-heavy, more slots
"sleep": 0.5, # low priority, fewer slots
"assistant": 1.5,
"router": 1.5,
}
_PRIORITY_BOOST: Dict[str, float] = {
"vip": settings.priority_boost_vip,
"high": settings.priority_boost_high,
"normal": settings.priority_boost_normal,
"low": settings.priority_boost_low,
}
_DECAY: float = settings.weight_decay
def _effective_weight(job: Job, attempt: int = 0) -> float:
"""Compute scheduling weight. Higher = executed sooner."""
boost = _PRIORITY_BOOST.get(job.priority, 1.0)
type_weight = AGENT_TYPE_WEIGHTS.get(job.agent_type, 1.0)
decay = _DECAY ** attempt
return job.weight * boost * type_weight * decay
class JobQueue:
def __init__(self):
self._jobs: Dict[str, Job] = {}
self._queue: asyncio.PriorityQueue = asyncio.PriorityQueue()
self._pipelines: Dict[str, Pipeline] = {}
self._worker_task: Optional[asyncio.Task] = None
# ── Job management ──────────────────────────────────────────────
def enqueue(self, job: Job) -> Job:
job.weight = _effective_weight(job)
self._jobs[job.job_id] = job
priority_score = -job.weight # negate so higher weight = dequeued first
self._queue.put_nowait((priority_score, job.created_at.timestamp(), job.job_id))
collector.emit_raw("job.queued", {"job_id": job.job_id, "weight": job.weight,
"priority": job.priority, "agent_type": job.agent_type})
return job
def get_job(self, job_id: str) -> Optional[Job]:
return self._jobs.get(job_id)
def list_jobs(self, status: Optional[JobStatus] = None, limit: int = 100) -> List[Job]:
jobs = list(self._jobs.values())
if status:
jobs = [j for j in jobs if j.status == status]
return sorted(jobs, key=lambda j: j.created_at, reverse=True)[:limit]
def cancel_job(self, job_id: str) -> bool:
job = self._jobs.get(job_id)
if job and job.status == JobStatus.QUEUED:
job.status = JobStatus.CANCELLED
return True
return False
# ── Pipeline management ─────────────────────────────────────────
def enqueue_pipeline(self, pipeline: Pipeline) -> Pipeline:
self._pipelines[pipeline.pipeline_id] = pipeline
asyncio.create_task(self._run_pipeline(pipeline))
return pipeline
def get_pipeline(self, pipeline_id: str) -> Optional[Pipeline]:
return self._pipelines.get(pipeline_id)
def list_pipelines(self, limit: int = 50) -> List[Pipeline]:
return sorted(self._pipelines.values(), key=lambda p: p.created_at, reverse=True)[:limit]
# ── Worker ──────────────────────────────────────────────────────
async def start(self) -> None:
if self._worker_task is None or self._worker_task.done():
self._worker_task = asyncio.create_task(self._worker_loop())
async def stop(self) -> None:
if self._worker_task:
self._worker_task.cancel()
async def _worker_loop(self) -> None:
while True:
_, _, job_id = await self._queue.get()
job = self._jobs.get(job_id)
if not job or job.status == JobStatus.CANCELLED:
self._queue.task_done()
continue
asyncio.create_task(self._execute_job(job))
self._queue.task_done()
async def _execute_job(self, job: Job) -> None:
job.status = JobStatus.RUNNING
job.started_at = datetime.now(timezone.utc)
try:
agent = registry.create(job.agent_type)
ctx = AgentContext(agent_id=agent.agent_id)
ctx.metadata["job_id"] = job.job_id
if job.retry:
result = await executor.run_with_retry(agent, ctx,
retries=job.retries, **job.kwargs)
else:
result = await executor.run(agent, ctx, **job.kwargs)
job.result = result.output
job.error = result.error
job.duration_ms = result.duration_ms
job.status = (JobStatus.COMPLETED if result.status == AgentStatus.COMPLETED
else JobStatus.FAILED)
except Exception as exc:
job.status = JobStatus.FAILED
job.error = str(exc)
finally:
job.completed_at = datetime.now(timezone.utc)
collector.emit_raw(f"job.{job.status.value}",
{"job_id": job.job_id, "duration_ms": job.duration_ms})
async def _run_pipeline(self, pipeline: Pipeline) -> None:
pipeline.status = JobStatus.RUNNING
previous_output = None
for i, step in enumerate(pipeline.steps):
try:
agent = registry.create(step.agent_type)
ctx = AgentContext(agent_id=agent.agent_id)
ctx.metadata["pipeline_id"] = pipeline.pipeline_id
ctx.metadata["step"] = i
kwargs = dict(step.kwargs)
if step.pass_output and previous_output is not None:
kwargs["previous_output"] = previous_output
result = await executor.run(agent, ctx, **kwargs)
pipeline.results.append(result.output)
previous_output = result.output
if result.status != AgentStatus.COMPLETED:
pipeline.status = JobStatus.FAILED
pipeline.error = f"Step {i} ({step.agent_type}) failed: {result.error}"
return
except Exception as exc:
pipeline.status = JobStatus.FAILED
pipeline.error = f"Step {i} ({step.agent_type}) exception: {exc}"
return
pipeline.status = JobStatus.COMPLETED
pipeline.completed_at = datetime.now(timezone.utc)
job_queue = JobQueue()
|