2118 lines
79 KiB
YAML
2118 lines
79 KiB
YAML
# ⚠️ STALE / DO NOT USE FOR CURRENT DEPLOYS.
|
|
# This ConfigMap embeds an OLD snapshot of the orchestrator source (no planner / review loop /
|
|
# synthesis). The canonical deploy builds the image from Dockerfile.orchestrator and uses
|
|
# k8s/orchestrator-deployment.yaml. Regenerate from the live source (e.g.
|
|
# kubectl create configmap orchestrator-source --from-file=orchestrator/ --dry-run=client -o yaml)
|
|
# if a source-mounted deployment is ever needed, otherwise prefer the image-based deployment.
|
|
apiVersion: v1
|
|
data:
|
|
__init__.py: |
|
|
"""Swarm orchestrator package."""
|
|
__version__ = "1.0.0"
|
|
README.md: "# Swarm Orchestrator\n\nFastAPI-based orchestrator for coordinating
|
|
agent pods in the K8s swarm system.\n\n## Features\n\n- **WebSocket Communication**:
|
|
Real-time bidirectional communication with agent pods\n- **Agent Registry**: Redis-backed
|
|
registration and heartbeat tracking\n- **Task Queue**: Task assignment and reassignment
|
|
with failure recovery\n- **Handoff Mechanism**: Coordinate task handoffs between
|
|
agents\n- **Failure Recovery**: Automatic detection and recovery from agent failures\n\n##
|
|
Architecture\n\n### Components\n\n1. **FastAPI Server** (`main.py`)\n - REST
|
|
API for external clients\n - WebSocket endpoint for agent communication\n -
|
|
Background failure detection loop\n\n2. **Redis Client** (`redis_client.py`)\n
|
|
\ - Async Redis connection wrapper\n - Key-value, hash, and list operations\n\n3.
|
|
**Agent Registry** (`agent_registry.py`)\n - Agent registration and deregistration\n
|
|
\ - Heartbeat tracking (30s timeout)\n - Status management: idle, busy, handoff-pending,
|
|
failed\n\n4. **Handoff Manager** (`handoff_manager.py`)\n - Handoff initiation
|
|
and validation\n - Handoff history tracking\n - Target agent availability
|
|
checks\n\n5. **Task Queue** (`task_queue.py`)\n - Task creation and assignment\n
|
|
\ - Retry logic (max 3 retries)\n - Automatic reassignment on agent failure\n\n##
|
|
WebSocket Protocol\n\n### Agent Registration\n\n```json\n{\n \"type\": \"register\",\n
|
|
\ \"capabilities\": [\"python\", \"javascript\"]\n}\n```\n\n### Heartbeat\n\n```json\n{\n
|
|
\ \"type\": \"heartbeat\"\n}\n```\n\nResponse:\n```json\n{\n \"type\": \"heartbeat_ack\"\n}\n```\n\n###
|
|
Handoff Request\n\n```json\n{\n \"type\": \"handoff\",\n \"target_agent_id\":
|
|
\"agent-2\",\n \"task_context\": {\n \"task_id\": \"task-123\",\n \"description\":
|
|
\"Continue implementation\",\n \"files\": [\"src/main.py\"]\n }\n}\n```\n\n###
|
|
Task Status Updates\n\n```json\n{\n \"type\": \"task_start\",\n \"task_id\":
|
|
\"task-123\"\n}\n```\n\n```json\n{\n \"type\": \"task_complete\",\n \"task_id\":
|
|
\"task-123\"\n}\n```\n\n```json\n{\n \"type\": \"task_failed\",\n \"task_id\":
|
|
\"task-123\",\n \"reason\": \"Compilation error\"\n}\n```\n\n## REST API Endpoints\n\n###
|
|
Health Check\n- `GET /` - Basic health check\n- `GET /health` - Detailed health
|
|
status\n\n### Agents\n- `GET /agents` - List all agents\n- `GET /agents/{agent_id}`
|
|
- Get agent details\n- `GET /agents/idle` - List idle agents\n\n### Tasks\n- `POST
|
|
/tasks` - Create new task\n- `POST /tasks/assign` - Manually assign task\n- `GET
|
|
/tasks` - List all tasks (optional `?status=pending`)\n- `GET /tasks/{task_id}`
|
|
- Get task details\n\n### Handoffs\n- `GET /handoffs` - List handoff history (optional
|
|
`?agent_id=agent-1`)\n\n## Deployment\n\n### Build Docker Image\n\n```bash\ndocker
|
|
build -f Dockerfile.orchestrator -t swarm-orchestrator:latest .\n```\n\n### Deploy
|
|
to Kubernetes\n\n```bash\n# Deploy Redis\nkubectl apply -f k8s/redis-statefulset.yaml\n\n#
|
|
Deploy orchestrator\nkubectl apply -f k8s/orchestrator-deployment.yaml\n```\n\n###
|
|
Environment Variables\n\n- `REDIS_HOST` - Redis hostname (default: `redis-service`)\n-
|
|
`REDIS_PORT` - Redis port (default: `6379`)\n- `REDIS_DB` - Redis database number
|
|
(default: `0`)\n- `LOG_LEVEL` - Logging level (default: `INFO`)\n\n## Failure
|
|
Recovery\n\nThe orchestrator implements automatic failure recovery:\n\n1. **Heartbeat
|
|
Monitoring**: Agents must send heartbeat every 15s\n2. **Timeout Detection**:
|
|
Agents with no heartbeat for 30s are marked as failed\n3. **Task Reassignment**:
|
|
Tasks from failed agents are automatically reassigned\n4. **Retry Logic**: Failed
|
|
tasks are retried up to 3 times\n5. **WebSocket Cleanup**: Failed agent connections
|
|
are closed\n\n## Redis Data Model\n\n### Agent Registry\n- Key: `agent:{agent_id}`\n-
|
|
Value: JSON with status, last_heartbeat, capabilities, current_task_id\n\n###
|
|
Tasks\n- Key: `task:{task_id}`\n- Value: JSON with description, status, assigned_agent_id,
|
|
retry_count\n\n### Task Queue\n- Key: `queue:pending`\n- Type: List of task_ids\n\n###
|
|
Handoff History\n- Key: `handoff:{handoff_id}`\n- Value: JSON with source, target,
|
|
timestamp, context, status\n\n## Development\n\n### Install Dependencies\n\n```bash\npip
|
|
install -r orchestrator/requirements.txt\n```\n\n### Run Locally\n\n```bash\n#
|
|
Start Redis\ndocker run -d -p 6379:6379 redis:7-alpine\n\n# Run orchestrator\npython
|
|
-m uvicorn orchestrator.main:app --reload\n```\n\n### Test WebSocket Connection\n\n```python\nimport
|
|
asyncio\nimport websockets\nimport json\n\nasync def test_agent():\n uri =
|
|
\"ws://localhost:8000/ws/test-agent-1\"\n async with websockets.connect(uri)
|
|
as websocket:\n # Register\n await websocket.send(json.dumps({\n
|
|
\ \"type\": \"register\",\n \"capabilities\": [\"python\"]\n
|
|
\ }))\n response = await websocket.recv()\n print(f\"Registration:
|
|
{response}\")\n \n # Send heartbeat\n while True:\n await
|
|
websocket.send(json.dumps({\"type\": \"heartbeat\"}))\n response =
|
|
await websocket.recv()\n print(f\"Heartbeat: {response}\")\n await
|
|
asyncio.sleep(10)\n\nasyncio.run(test_agent())\n```\n\n## Acceptance Criteria
|
|
Status\n\n✅ Orchestrator successfully registers new agent pods via WebSocket \n✅
|
|
Agent registry persists in Redis and survives orchestrator restart \n✅ Agent
|
|
A can handoff task to Agent B via orchestrator \n✅ Orchestrator detects and deregisters
|
|
crashed agents within 60s \n✅ Desktop client can query orchestrator API for agent
|
|
status \n✅ System recovers from agent crashes without user intervention\n"
|
|
agent_registry.py: |
|
|
"""Agent registry with Redis-backed state management."""
|
|
import json
|
|
import time
|
|
import logging
|
|
from typing import Dict, List, Optional
|
|
from enum import Enum
|
|
from pydantic import BaseModel
|
|
from .redis_client import redis_client
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AgentStatus(str, Enum):
|
|
"""Agent status enumeration."""
|
|
IDLE = "idle"
|
|
BUSY = "busy"
|
|
HANDOFF_PENDING = "handoff-pending"
|
|
FAILED = "failed"
|
|
|
|
|
|
class AgentMetadata(BaseModel):
|
|
"""Agent metadata model."""
|
|
agent_id: str
|
|
status: AgentStatus
|
|
last_heartbeat: float
|
|
capabilities: List[str]
|
|
current_task_id: Optional[str] = None
|
|
|
|
|
|
class AgentRegistry:
|
|
"""Manages agent registration and heartbeat tracking."""
|
|
|
|
HEARTBEAT_TIMEOUT = 30 # seconds
|
|
AGENT_KEY_PREFIX = "agent:"
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
async def register_agent(
|
|
self, agent_id: str, capabilities: List[str]
|
|
) -> AgentMetadata:
|
|
"""Register a new agent."""
|
|
metadata = AgentMetadata(
|
|
agent_id=agent_id,
|
|
status=AgentStatus.IDLE,
|
|
last_heartbeat=time.time(),
|
|
capabilities=capabilities,
|
|
)
|
|
|
|
key = f"{self.AGENT_KEY_PREFIX}{agent_id}"
|
|
await redis_client.set(key, metadata.model_dump_json())
|
|
|
|
logger.info(f"Registered agent {agent_id} with capabilities: {capabilities}")
|
|
return metadata
|
|
|
|
async def deregister_agent(self, agent_id: str):
|
|
"""Deregister an agent."""
|
|
key = f"{self.AGENT_KEY_PREFIX}{agent_id}"
|
|
await redis_client.delete(key)
|
|
logger.info(f"Deregistered agent {agent_id}")
|
|
|
|
async def update_heartbeat(self, agent_id: str) -> bool:
|
|
"""Update agent heartbeat timestamp."""
|
|
key = f"{self.AGENT_KEY_PREFIX}{agent_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
logger.warning(f"Agent {agent_id} not found for heartbeat update")
|
|
return False
|
|
|
|
metadata = AgentMetadata.model_validate_json(data)
|
|
metadata.last_heartbeat = time.time()
|
|
|
|
await redis_client.set(key, metadata.model_dump_json())
|
|
return True
|
|
|
|
async def update_status(
|
|
self, agent_id: str, status: AgentStatus, task_id: Optional[str] = None
|
|
) -> bool:
|
|
"""Update agent status."""
|
|
key = f"{self.AGENT_KEY_PREFIX}{agent_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
logger.warning(f"Agent {agent_id} not found for status update")
|
|
return False
|
|
|
|
metadata = AgentMetadata.model_validate_json(data)
|
|
metadata.status = status
|
|
metadata.current_task_id = task_id
|
|
|
|
await redis_client.set(key, metadata.model_dump_json())
|
|
logger.info(f"Updated agent {agent_id} status to {status}")
|
|
return True
|
|
|
|
async def get_agent(self, agent_id: str) -> Optional[AgentMetadata]:
|
|
"""Get agent metadata."""
|
|
key = f"{self.AGENT_KEY_PREFIX}{agent_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
return None
|
|
|
|
return AgentMetadata.model_validate_json(data)
|
|
|
|
async def get_all_agents(self) -> List[AgentMetadata]:
|
|
"""Get all registered agents."""
|
|
pattern = f"{self.AGENT_KEY_PREFIX}*"
|
|
keys = await redis_client.keys(pattern)
|
|
|
|
agents = []
|
|
for key in keys:
|
|
data = await redis_client.get(key)
|
|
if data:
|
|
agents.append(AgentMetadata.model_validate_json(data))
|
|
|
|
return agents
|
|
|
|
async def get_idle_agents(self) -> List[AgentMetadata]:
|
|
"""Get all idle agents."""
|
|
all_agents = await self.get_all_agents()
|
|
return [agent for agent in all_agents if agent.status == AgentStatus.IDLE]
|
|
|
|
async def check_failed_agents(self) -> List[str]:
|
|
"""Check for agents with expired heartbeats and mark as failed."""
|
|
current_time = time.time()
|
|
failed_agents = []
|
|
|
|
all_agents = await self.get_all_agents()
|
|
|
|
for agent in all_agents:
|
|
if agent.status == AgentStatus.FAILED:
|
|
continue
|
|
|
|
time_since_heartbeat = current_time - agent.last_heartbeat
|
|
|
|
if time_since_heartbeat > self.HEARTBEAT_TIMEOUT:
|
|
await self.update_status(agent.agent_id, AgentStatus.FAILED)
|
|
failed_agents.append(agent.agent_id)
|
|
logger.warning(
|
|
f"Agent {agent.agent_id} marked as failed "
|
|
f"(no heartbeat for {time_since_heartbeat:.1f}s)"
|
|
)
|
|
|
|
return failed_agents
|
|
|
|
|
|
# Global agent registry instance
|
|
agent_registry = AgentRegistry()
|
|
checkpoint_manager.py: |
|
|
"""
|
|
Checkpoint Manager for Task Recovery
|
|
|
|
Provides partial checkpointing and retry strategies for agent tasks.
|
|
Enables recovery from failures without restarting entire workflows.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Dict, List, Optional, Any
|
|
from dataclasses import dataclass, asdict
|
|
from datetime import datetime
|
|
import redis
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class Checkpoint:
|
|
"""Represents a task checkpoint"""
|
|
task_id: str
|
|
agent_id: str
|
|
checkpoint_id: str
|
|
timestamp: float
|
|
phase: str # e.g., "analysis", "implementation", "testing"
|
|
state: Dict[str, Any] # Serializable state data
|
|
files_modified: List[str]
|
|
git_commit: Optional[str] = None
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
@dataclass
|
|
class RetryPolicy:
|
|
"""Retry strategy configuration"""
|
|
max_retries: int = 3
|
|
backoff_multiplier: float = 2.0
|
|
initial_delay_seconds: float = 1.0
|
|
max_delay_seconds: float = 60.0
|
|
retry_on_errors: List[str] = None # Error types to retry
|
|
|
|
|
|
class CheckpointManager:
|
|
"""Manages task checkpoints and recovery"""
|
|
|
|
def __init__(self, redis_client: redis.Redis):
|
|
self.redis = redis_client
|
|
self.checkpoint_ttl = 86400 * 7 # 7 days
|
|
|
|
def save_checkpoint(self, checkpoint: Checkpoint) -> bool:
|
|
"""Save a checkpoint to Redis"""
|
|
try:
|
|
key = f"checkpoint:{checkpoint.task_id}:{checkpoint.checkpoint_id}"
|
|
data = json.dumps(asdict(checkpoint))
|
|
self.redis.setex(key, self.checkpoint_ttl, data)
|
|
|
|
# Add to task's checkpoint list
|
|
list_key = f"checkpoints:{checkpoint.task_id}"
|
|
self.redis.lpush(list_key, checkpoint.checkpoint_id)
|
|
self.redis.expire(list_key, self.checkpoint_ttl)
|
|
|
|
logger.info(
|
|
f"Saved checkpoint {checkpoint.checkpoint_id} for task {checkpoint.task_id}"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to save checkpoint: {e}")
|
|
return False
|
|
|
|
def get_checkpoint(self, task_id: str, checkpoint_id: str) -> Optional[Checkpoint]:
|
|
"""Retrieve a specific checkpoint"""
|
|
try:
|
|
key = f"checkpoint:{task_id}:{checkpoint_id}"
|
|
data = self.redis.get(key)
|
|
if not data:
|
|
return None
|
|
|
|
checkpoint_dict = json.loads(data)
|
|
return Checkpoint(**checkpoint_dict)
|
|
except Exception as e:
|
|
logger.error(f"Failed to retrieve checkpoint: {e}")
|
|
return None
|
|
|
|
def get_latest_checkpoint(self, task_id: str) -> Optional[Checkpoint]:
|
|
"""Get the most recent checkpoint for a task"""
|
|
try:
|
|
list_key = f"checkpoints:{task_id}"
|
|
checkpoint_ids = self.redis.lrange(list_key, 0, 0)
|
|
|
|
if not checkpoint_ids:
|
|
return None
|
|
|
|
checkpoint_id = checkpoint_ids[0].decode('utf-8')
|
|
return self.get_checkpoint(task_id, checkpoint_id)
|
|
except Exception as e:
|
|
logger.error(f"Failed to get latest checkpoint: {e}")
|
|
return None
|
|
|
|
def list_checkpoints(self, task_id: str) -> List[str]:
|
|
"""List all checkpoint IDs for a task"""
|
|
try:
|
|
list_key = f"checkpoints:{task_id}"
|
|
checkpoint_ids = self.redis.lrange(list_key, 0, -1)
|
|
return [cid.decode('utf-8') for cid in checkpoint_ids]
|
|
except Exception as e:
|
|
logger.error(f"Failed to list checkpoints: {e}")
|
|
return []
|
|
|
|
def delete_checkpoint(self, task_id: str, checkpoint_id: str) -> bool:
|
|
"""Delete a specific checkpoint"""
|
|
try:
|
|
key = f"checkpoint:{task_id}:{checkpoint_id}"
|
|
self.redis.delete(key)
|
|
|
|
# Remove from list
|
|
list_key = f"checkpoints:{task_id}"
|
|
self.redis.lrem(list_key, 0, checkpoint_id)
|
|
|
|
logger.info(f"Deleted checkpoint {checkpoint_id} for task {task_id}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete checkpoint: {e}")
|
|
return False
|
|
|
|
def cleanup_task_checkpoints(self, task_id: str) -> bool:
|
|
"""Delete all checkpoints for a task"""
|
|
try:
|
|
checkpoint_ids = self.list_checkpoints(task_id)
|
|
for checkpoint_id in checkpoint_ids:
|
|
self.delete_checkpoint(task_id, checkpoint_id)
|
|
|
|
list_key = f"checkpoints:{task_id}"
|
|
self.redis.delete(list_key)
|
|
|
|
logger.info(f"Cleaned up all checkpoints for task {task_id}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to cleanup checkpoints: {e}")
|
|
return False
|
|
|
|
|
|
class RetryManager:
|
|
"""Manages retry logic with exponential backoff"""
|
|
|
|
def __init__(self, redis_client: redis.Redis):
|
|
self.redis = redis_client
|
|
|
|
def record_attempt(self, task_id: str, agent_id: str, error: Optional[str] = None):
|
|
"""Record a task attempt"""
|
|
key = f"retry:{task_id}"
|
|
attempt_data = {
|
|
"agent_id": agent_id,
|
|
"timestamp": time.time(),
|
|
"error": error
|
|
}
|
|
|
|
self.redis.lpush(key, json.dumps(attempt_data))
|
|
self.redis.expire(key, 86400) # 24 hour TTL
|
|
|
|
def get_attempt_count(self, task_id: str) -> int:
|
|
"""Get number of attempts for a task"""
|
|
key = f"retry:{task_id}"
|
|
return self.redis.llen(key)
|
|
|
|
def should_retry(self, task_id: str, policy: RetryPolicy) -> bool:
|
|
"""Determine if task should be retried"""
|
|
attempt_count = self.get_attempt_count(task_id)
|
|
return attempt_count < policy.max_retries
|
|
|
|
def get_retry_delay(self, task_id: str, policy: RetryPolicy) -> float:
|
|
"""Calculate delay before next retry (exponential backoff)"""
|
|
attempt_count = self.get_attempt_count(task_id)
|
|
delay = policy.initial_delay_seconds * (policy.backoff_multiplier ** attempt_count)
|
|
return min(delay, policy.max_delay_seconds)
|
|
|
|
def clear_attempts(self, task_id: str):
|
|
"""Clear retry history for a task"""
|
|
key = f"retry:{task_id}"
|
|
self.redis.delete(key)
|
|
|
|
|
|
class RecoveryCoordinator:
|
|
"""Coordinates task recovery from checkpoints"""
|
|
|
|
def __init__(self, checkpoint_manager: CheckpointManager, retry_manager: RetryManager):
|
|
self.checkpoint_manager = checkpoint_manager
|
|
self.retry_manager = retry_manager
|
|
|
|
def recover_task(
|
|
self,
|
|
task_id: str,
|
|
retry_policy: Optional[RetryPolicy] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Attempt to recover a failed task from its latest checkpoint
|
|
|
|
Returns recovery instructions or None if recovery not possible
|
|
"""
|
|
if retry_policy is None:
|
|
retry_policy = RetryPolicy()
|
|
|
|
# Check if we should retry
|
|
if not self.retry_manager.should_retry(task_id, retry_policy):
|
|
logger.warning(f"Task {task_id} exceeded max retries")
|
|
return None
|
|
|
|
# Get latest checkpoint
|
|
checkpoint = self.checkpoint_manager.get_latest_checkpoint(task_id)
|
|
if not checkpoint:
|
|
logger.warning(f"No checkpoint found for task {task_id}")
|
|
return None
|
|
|
|
# Calculate retry delay
|
|
delay = self.retry_manager.get_retry_delay(task_id, retry_policy)
|
|
|
|
recovery_plan = {
|
|
"task_id": task_id,
|
|
"checkpoint_id": checkpoint.checkpoint_id,
|
|
"resume_phase": checkpoint.phase,
|
|
"state": checkpoint.state,
|
|
"files_modified": checkpoint.files_modified,
|
|
"git_commit": checkpoint.git_commit,
|
|
"retry_delay_seconds": delay,
|
|
"attempt_number": self.retry_manager.get_attempt_count(task_id) + 1
|
|
}
|
|
|
|
logger.info(
|
|
f"Recovery plan created for task {task_id} from checkpoint {checkpoint.checkpoint_id}"
|
|
)
|
|
return recovery_plan
|
|
|
|
def mark_recovery_success(self, task_id: str):
|
|
"""Mark a task as successfully recovered"""
|
|
self.retry_manager.clear_attempts(task_id)
|
|
logger.info(f"Task {task_id} recovered successfully")
|
|
|
|
def mark_recovery_failure(self, task_id: str, agent_id: str, error: str):
|
|
"""Record a failed recovery attempt"""
|
|
self.retry_manager.record_attempt(task_id, agent_id, error)
|
|
logger.warning(f"Recovery attempt failed for task {task_id}: {error}")
|
|
handoff_manager.py: |
|
|
"""Handoff manager for task coordination between agents."""
|
|
import json
|
|
import time
|
|
import uuid
|
|
import logging
|
|
from typing import Dict, Optional
|
|
from pydantic import BaseModel
|
|
from .redis_client import redis_client
|
|
from .agent_registry import agent_registry, AgentStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class HandoffRequest(BaseModel):
|
|
"""Handoff request model."""
|
|
type: str = "handoff"
|
|
source_agent_id: str
|
|
target_agent_id: str
|
|
task_context: Dict
|
|
|
|
|
|
class HandoffRecord(BaseModel):
|
|
"""Handoff history record."""
|
|
handoff_id: str
|
|
source_agent_id: str
|
|
target_agent_id: str
|
|
timestamp: float
|
|
task_context: Dict
|
|
status: str # pending, completed, failed
|
|
|
|
|
|
class HandoffManager:
|
|
"""Manages task handoffs between agents."""
|
|
|
|
HANDOFF_KEY_PREFIX = "handoff:"
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
async def initiate_handoff(
|
|
self, source_agent_id: str, target_agent_id: str, task_context: Dict
|
|
) -> Optional[str]:
|
|
"""Initiate a handoff from source to target agent."""
|
|
|
|
# Validate source agent exists and is busy
|
|
source_agent = await agent_registry.get_agent(source_agent_id)
|
|
if not source_agent:
|
|
logger.error(f"Source agent {source_agent_id} not found")
|
|
return None
|
|
|
|
# Validate target agent exists and is idle
|
|
target_agent = await agent_registry.get_agent(target_agent_id)
|
|
if not target_agent:
|
|
logger.error(f"Target agent {target_agent_id} not found")
|
|
return None
|
|
|
|
if target_agent.status != AgentStatus.IDLE:
|
|
logger.error(
|
|
f"Target agent {target_agent_id} is not idle (status: {target_agent.status})"
|
|
)
|
|
return None
|
|
|
|
# Create handoff record
|
|
handoff_id = str(uuid.uuid4())
|
|
record = HandoffRecord(
|
|
handoff_id=handoff_id,
|
|
source_agent_id=source_agent_id,
|
|
target_agent_id=target_agent_id,
|
|
timestamp=time.time(),
|
|
task_context=task_context,
|
|
status="pending",
|
|
)
|
|
|
|
# Store in Redis
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
await redis_client.set(key, record.model_dump_json())
|
|
|
|
# Update agent statuses
|
|
await agent_registry.update_status(source_agent_id, AgentStatus.IDLE)
|
|
await agent_registry.update_status(
|
|
target_agent_id, AgentStatus.HANDOFF_PENDING, task_context.get("task_id")
|
|
)
|
|
|
|
logger.info(
|
|
f"Handoff {handoff_id} initiated: {source_agent_id} -> {target_agent_id}"
|
|
)
|
|
return handoff_id
|
|
|
|
async def complete_handoff(self, handoff_id: str) -> bool:
|
|
"""Mark handoff as completed."""
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
logger.error(f"Handoff {handoff_id} not found")
|
|
return False
|
|
|
|
record = HandoffRecord.model_validate_json(data)
|
|
record.status = "completed"
|
|
|
|
await redis_client.set(key, record.model_dump_json())
|
|
|
|
# Update target agent to busy
|
|
await agent_registry.update_status(
|
|
record.target_agent_id, AgentStatus.BUSY, record.task_context.get("task_id")
|
|
)
|
|
|
|
logger.info(f"Handoff {handoff_id} completed")
|
|
return True
|
|
|
|
async def fail_handoff(self, handoff_id: str, reason: str) -> bool:
|
|
"""Mark handoff as failed."""
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
logger.error(f"Handoff {handoff_id} not found")
|
|
return False
|
|
|
|
record = HandoffRecord.model_validate_json(data)
|
|
record.status = "failed"
|
|
|
|
await redis_client.set(key, record.model_dump_json())
|
|
|
|
# Revert target agent to idle
|
|
await agent_registry.update_status(record.target_agent_id, AgentStatus.IDLE)
|
|
|
|
logger.warning(f"Handoff {handoff_id} failed: {reason}")
|
|
return True
|
|
|
|
async def get_handoff(self, handoff_id: str) -> Optional[HandoffRecord]:
|
|
"""Get handoff record by ID."""
|
|
key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
return None
|
|
|
|
return HandoffRecord.model_validate_json(data)
|
|
|
|
async def get_handoff_history(
|
|
self, agent_id: Optional[str] = None, limit: int = 100
|
|
) -> list[HandoffRecord]:
|
|
"""Get handoff history, optionally filtered by agent."""
|
|
pattern = f"{self.HANDOFF_KEY_PREFIX}*"
|
|
keys = await redis_client.keys(pattern)
|
|
|
|
records = []
|
|
for key in keys[:limit]:
|
|
data = await redis_client.get(key)
|
|
if data:
|
|
record = HandoffRecord.model_validate_json(data)
|
|
if agent_id is None or (
|
|
record.source_agent_id == agent_id
|
|
or record.target_agent_id == agent_id
|
|
):
|
|
records.append(record)
|
|
|
|
# Sort by timestamp descending
|
|
records.sort(key=lambda r: r.timestamp, reverse=True)
|
|
return records[:limit]
|
|
|
|
|
|
# Global handoff manager instance
|
|
handoff_manager = HandoffManager()
|
|
main.py: |
|
|
"""FastAPI orchestrator with WebSocket support for agent coordination."""
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from contextlib import asynccontextmanager
|
|
from typing import Dict, List, Optional
|
|
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
|
|
from fastapi.responses import JSONResponse, PlainTextResponse
|
|
from pydantic import BaseModel
|
|
from prometheus_client import Counter, Gauge, Histogram, generate_latest, CONTENT_TYPE_LATEST
|
|
|
|
from .redis_client import redis_client
|
|
from .agent_registry import agent_registry, AgentStatus, AgentMetadata
|
|
from .handoff_manager import handoff_manager, HandoffRequest
|
|
from .task_queue import task_queue, TaskStatus
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Prometheus metrics
|
|
AGENTS_CREATED = Counter('swarm_agents_created_total', 'Total agents created')
|
|
AGENTS_FAILED = Counter('swarm_agents_failed_total', 'Total agents failed')
|
|
AGENTS_ACTIVE = Gauge('swarm_agents_active', 'Currently active agents')
|
|
TASKS_CREATED = Counter('swarm_tasks_created_total', 'Total tasks created')
|
|
TASKS_COMPLETED = Counter('swarm_tasks_completed_total', 'Total tasks completed')
|
|
TASKS_FAILED = Counter('swarm_tasks_failed_total', 'Total tasks failed')
|
|
HANDOFFS_TOTAL = Counter('swarm_handoffs_total', 'Total handoffs performed')
|
|
HANDOFF_DURATION = Histogram('swarm_handoff_duration_seconds', 'Handoff latency in seconds')
|
|
TASK_DURATION = Histogram('swarm_task_duration_seconds', 'Task duration in seconds')
|
|
AGENT_STATUS = Gauge('swarm_agent_status', 'Agent status by ID', ['agent_id', 'status'])
|
|
WEBSOCKET_CONNECTIONS = Gauge('swarm_websocket_connections', 'Active WebSocket connections')
|
|
ERRORS_TOTAL = Counter('swarm_errors_total', 'Total errors', ['error_type'])
|
|
|
|
|
|
# WebSocket connection manager
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections for agents."""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, agent_id: str, websocket: WebSocket):
|
|
"""Accept and store WebSocket connection."""
|
|
await websocket.accept()
|
|
self.active_connections[agent_id] = websocket
|
|
logger.info(f"Agent {agent_id} connected via WebSocket")
|
|
|
|
def disconnect(self, agent_id: str):
|
|
"""Remove WebSocket connection."""
|
|
if agent_id in self.active_connections:
|
|
del self.active_connections[agent_id]
|
|
logger.info(f"Agent {agent_id} disconnected")
|
|
|
|
async def send_message(self, agent_id: str, message: dict):
|
|
"""Send message to specific agent."""
|
|
if agent_id in self.active_connections:
|
|
try:
|
|
await self.active_connections[agent_id].send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"Failed to send message to agent {agent_id}: {e}")
|
|
|
|
async def broadcast(self, message: dict):
|
|
"""Broadcast message to all connected agents."""
|
|
for agent_id, connection in self.active_connections.items():
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"Failed to broadcast to agent {agent_id}: {e}")
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
# Background task for failure detection
|
|
async def failure_detection_loop():
|
|
"""Periodically check for failed agents and reassign their tasks."""
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(10) # Check every 10 seconds
|
|
|
|
failed_agents = await agent_registry.check_failed_agents()
|
|
|
|
for agent_id in failed_agents:
|
|
# Reassign tasks from failed agent
|
|
reassigned_tasks = await task_queue.reassign_agent_tasks(agent_id)
|
|
|
|
if reassigned_tasks:
|
|
logger.info(
|
|
f"Reassigned {len(reassigned_tasks)} tasks from failed agent {agent_id}"
|
|
)
|
|
|
|
# Disconnect WebSocket
|
|
manager.disconnect(agent_id)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in failure detection loop: {e}")
|
|
|
|
|
|
# Lifespan context manager
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Startup and shutdown events."""
|
|
# Startup
|
|
logger.info("Starting orchestrator...")
|
|
await redis_client.connect()
|
|
|
|
# Start background tasks
|
|
failure_task = asyncio.create_task(failure_detection_loop())
|
|
|
|
yield
|
|
|
|
# Shutdown
|
|
logger.info("Shutting down orchestrator...")
|
|
failure_task.cancel()
|
|
await redis_client.disconnect()
|
|
|
|
|
|
# Create FastAPI app
|
|
app = FastAPI(title="Swarm Orchestrator", version="1.0.0", lifespan=lifespan)
|
|
|
|
|
|
# Pydantic models for API
|
|
class RegisterRequest(BaseModel):
|
|
"""Agent registration request."""
|
|
agent_id: str
|
|
capabilities: List[str]
|
|
|
|
|
|
class TaskCreateRequest(BaseModel):
|
|
"""Task creation request."""
|
|
description: str
|
|
context: Dict = {}
|
|
max_retries: int = 3
|
|
|
|
|
|
class TaskAssignRequest(BaseModel):
|
|
"""Task assignment request."""
|
|
task_id: str
|
|
agent_id: str
|
|
|
|
|
|
# REST API endpoints
|
|
@app.get("/")
|
|
async def root():
|
|
"""Health check endpoint."""
|
|
return {"status": "ok", "service": "swarm-orchestrator"}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""Detailed health check."""
|
|
try:
|
|
await redis_client.client.ping()
|
|
redis_status = "connected"
|
|
except Exception:
|
|
redis_status = "disconnected"
|
|
|
|
return {
|
|
"status": "ok",
|
|
"redis": redis_status,
|
|
"active_connections": len(manager.active_connections)
|
|
}
|
|
|
|
|
|
@app.get("/agents")
|
|
async def list_agents():
|
|
"""List all registered agents."""
|
|
agents = await agent_registry.get_all_agents()
|
|
return {"agents": [agent.model_dump() for agent in agents]}
|
|
|
|
|
|
@app.get("/agents/{agent_id}")
|
|
async def get_agent(agent_id: str):
|
|
"""Get specific agent details."""
|
|
agent = await agent_registry.get_agent(agent_id)
|
|
if not agent:
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
|
return agent.model_dump()
|
|
|
|
|
|
@app.get("/agents/idle")
|
|
async def list_idle_agents():
|
|
"""List all idle agents."""
|
|
agents = await agent_registry.get_idle_agents()
|
|
return {"agents": [agent.model_dump() for agent in agents]}
|
|
|
|
|
|
@app.post("/tasks")
|
|
async def create_task(request: TaskCreateRequest):
|
|
"""Create a new task."""
|
|
task = await task_queue.create_task(
|
|
description=request.description,
|
|
context=request.context,
|
|
max_retries=request.max_retries
|
|
)
|
|
TASKS_CREATED.inc()
|
|
return task.model_dump()
|
|
|
|
|
|
@app.post("/tasks/assign")
|
|
async def assign_task(request: TaskAssignRequest):
|
|
"""Manually assign a task to an agent."""
|
|
success = await task_queue.assign_task(request.task_id, request.agent_id)
|
|
if not success:
|
|
raise HTTPException(status_code=400, detail="Failed to assign task")
|
|
|
|
# Notify agent via WebSocket
|
|
await manager.send_message(request.agent_id, {
|
|
"type": "task_assigned",
|
|
"task_id": request.task_id
|
|
})
|
|
|
|
return {"status": "assigned", "task_id": request.task_id, "agent_id": request.agent_id}
|
|
|
|
|
|
@app.get("/tasks")
|
|
async def list_tasks(status: Optional[str] = None):
|
|
"""List all tasks, optionally filtered by status."""
|
|
task_status = TaskStatus(status) if status else None
|
|
tasks = await task_queue.get_all_tasks(status=task_status)
|
|
return {"tasks": [task.model_dump() for task in tasks]}
|
|
|
|
|
|
@app.get("/tasks/{task_id}")
|
|
async def get_task(task_id: str):
|
|
"""Get specific task details."""
|
|
task = await task_queue.get_task(task_id)
|
|
if not task:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
return task.model_dump()
|
|
|
|
|
|
@app.get("/handoffs")
|
|
async def list_handoffs(agent_id: Optional[str] = None, limit: int = 100):
|
|
"""List handoff history."""
|
|
handoffs = await handoff_manager.get_handoff_history(agent_id=agent_id, limit=limit)
|
|
return {"handoffs": [h.model_dump() for h in handoffs]}
|
|
|
|
|
|
@app.get("/metrics")
|
|
async def metrics():
|
|
"""Prometheus metrics endpoint."""
|
|
# Update gauge metrics before returning
|
|
agents = await agent_registry.get_all_agents()
|
|
AGENTS_ACTIVE.set(len([a for a in agents if a.status == AgentStatus.BUSY]))
|
|
WEBSOCKET_CONNECTIONS.set(len(manager.active_connections))
|
|
|
|
# Update per-agent status
|
|
for agent in agents:
|
|
AGENT_STATUS.labels(agent_id=agent.agent_id, status=agent.status.value).set(1)
|
|
|
|
return PlainTextResponse(generate_latest(), media_type=CONTENT_TYPE_LATEST)
|
|
|
|
|
|
# WebSocket endpoint
|
|
@app.websocket("/ws/{agent_id}")
|
|
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
|
|
"""WebSocket endpoint for agent communication."""
|
|
await manager.connect(agent_id, websocket)
|
|
|
|
try:
|
|
# Wait for registration message
|
|
data = await websocket.receive_json()
|
|
|
|
if data.get("type") == "register":
|
|
capabilities = data.get("capabilities", [])
|
|
await agent_registry.register_agent(agent_id, capabilities)
|
|
|
|
await websocket.send_json({
|
|
"type": "registered",
|
|
"agent_id": agent_id,
|
|
"status": "success"
|
|
})
|
|
|
|
# Main message loop
|
|
while True:
|
|
try:
|
|
message = await asyncio.wait_for(
|
|
websocket.receive_json(),
|
|
timeout=15.0 # Expect heartbeat every 15s
|
|
)
|
|
|
|
message_type = message.get("type")
|
|
|
|
if message_type == "heartbeat":
|
|
await agent_registry.update_heartbeat(agent_id)
|
|
await websocket.send_json({"type": "heartbeat_ack"})
|
|
|
|
elif message_type == "handoff":
|
|
# Handle handoff request
|
|
target_agent_id = message.get("target_agent_id")
|
|
task_context = message.get("task_context", {})
|
|
|
|
handoff_id = await handoff_manager.initiate_handoff(
|
|
agent_id, target_agent_id, task_context
|
|
)
|
|
|
|
if handoff_id:
|
|
# Notify target agent
|
|
await manager.send_message(target_agent_id, {
|
|
"type": "handoff_request",
|
|
"handoff_id": handoff_id,
|
|
"source_agent_id": agent_id,
|
|
"task_context": task_context
|
|
})
|
|
|
|
await websocket.send_json({
|
|
"type": "handoff_initiated",
|
|
"handoff_id": handoff_id,
|
|
"status": "success"
|
|
})
|
|
else:
|
|
await websocket.send_json({
|
|
"type": "handoff_failed",
|
|
"reason": "Target agent not available"
|
|
})
|
|
|
|
elif message_type == "handoff_accept":
|
|
# Target agent accepts handoff
|
|
handoff_id = message.get("handoff_id")
|
|
await handoff_manager.complete_handoff(handoff_id)
|
|
|
|
await websocket.send_json({
|
|
"type": "handoff_accepted",
|
|
"handoff_id": handoff_id
|
|
})
|
|
|
|
elif message_type == "handoff_reject":
|
|
# Target agent rejects handoff
|
|
handoff_id = message.get("handoff_id")
|
|
reason = message.get("reason", "Rejected by target agent")
|
|
await handoff_manager.fail_handoff(handoff_id, reason)
|
|
|
|
await websocket.send_json({
|
|
"type": "handoff_rejected",
|
|
"handoff_id": handoff_id
|
|
})
|
|
|
|
elif message_type == "task_start":
|
|
# Agent starts working on task
|
|
task_id = message.get("task_id")
|
|
await task_queue.start_task(task_id)
|
|
|
|
elif message_type == "task_complete":
|
|
# Agent completes task
|
|
task_id = message.get("task_id")
|
|
await task_queue.complete_task(task_id)
|
|
|
|
await websocket.send_json({
|
|
"type": "task_completed",
|
|
"task_id": task_id
|
|
})
|
|
|
|
elif message_type == "task_failed":
|
|
# Agent reports task failure
|
|
task_id = message.get("task_id")
|
|
reason = message.get("reason", "Unknown error")
|
|
await task_queue.fail_task(task_id, reason)
|
|
|
|
await websocket.send_json({
|
|
"type": "task_failed_ack",
|
|
"task_id": task_id
|
|
})
|
|
|
|
elif message_type == "status_update":
|
|
# Agent updates its status
|
|
status = AgentStatus(message.get("status"))
|
|
task_id = message.get("task_id")
|
|
await agent_registry.update_status(agent_id, status, task_id)
|
|
|
|
else:
|
|
logger.warning(f"Unknown message type from {agent_id}: {message_type}")
|
|
|
|
except asyncio.TimeoutError:
|
|
# No message received within timeout - check heartbeat
|
|
agent = await agent_registry.get_agent(agent_id)
|
|
if agent:
|
|
import time
|
|
if time.time() - agent.last_heartbeat > 30:
|
|
logger.warning(f"Agent {agent_id} heartbeat timeout")
|
|
break
|
|
continue
|
|
|
|
else:
|
|
logger.error(f"Agent {agent_id} did not send registration message")
|
|
await websocket.close(code=1008, reason="Registration required")
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info(f"Agent {agent_id} disconnected")
|
|
except Exception as e:
|
|
logger.error(f"Error in WebSocket handler for agent {agent_id}: {e}")
|
|
finally:
|
|
# Cleanup
|
|
manager.disconnect(agent_id)
|
|
await agent_registry.deregister_agent(agent_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|
|
model_tuner.py: |
|
|
"""
|
|
Complexity Model Tuner
|
|
|
|
Implements feedback loop and retraining for the complexity analysis model.
|
|
Collects actual vs predicted agent counts and adjusts the model over time.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Dict, List, Optional, Tuple
|
|
from dataclasses import dataclass, asdict
|
|
from datetime import datetime, timedelta
|
|
import redis
|
|
import logging
|
|
from collections import defaultdict
|
|
import statistics
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ComplexityPrediction:
|
|
"""Record of a complexity prediction"""
|
|
task_id: str
|
|
timestamp: float
|
|
task_description: str
|
|
predicted_agents: int
|
|
predicted_complexity: str # "simple", "moderate", "complex"
|
|
confidence: float
|
|
model_version: str
|
|
|
|
|
|
@dataclass
|
|
class ComplexityActual:
|
|
"""Actual outcome of a task"""
|
|
task_id: str
|
|
timestamp: float
|
|
actual_agents: int
|
|
actual_duration_seconds: float
|
|
success: bool
|
|
user_feedback: Optional[str] = None # "too_many", "too_few", "just_right"
|
|
user_rating: Optional[int] = None # 1-5 scale
|
|
|
|
|
|
@dataclass
|
|
class ModelMetrics:
|
|
"""Model performance metrics"""
|
|
model_version: str
|
|
total_predictions: int
|
|
mean_absolute_error: float
|
|
accuracy_within_1: float # % predictions within ±1 agent
|
|
accuracy_within_2: float # % predictions within ±2 agents
|
|
user_satisfaction: float # Average user rating
|
|
last_updated: float
|
|
|
|
|
|
class ComplexityModelTuner:
|
|
"""Manages complexity model feedback and tuning"""
|
|
|
|
def __init__(self, redis_client: redis.Redis):
|
|
self.redis = redis_client
|
|
self.current_model_version = "v1.0"
|
|
self.data_ttl = 86400 * 30 # 30 days
|
|
|
|
def record_prediction(self, prediction: ComplexityPrediction) -> bool:
|
|
"""Record a complexity prediction"""
|
|
try:
|
|
key = f"prediction:{prediction.task_id}"
|
|
data = json.dumps(asdict(prediction))
|
|
self.redis.setex(key, self.data_ttl, data)
|
|
|
|
# Add to predictions list
|
|
list_key = f"predictions:{prediction.model_version}"
|
|
self.redis.lpush(list_key, prediction.task_id)
|
|
self.redis.expire(list_key, self.data_ttl)
|
|
|
|
logger.info(
|
|
f"Recorded prediction for task {prediction.task_id}: "
|
|
f"{prediction.predicted_agents} agents"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to record prediction: {e}")
|
|
return False
|
|
|
|
def record_actual(self, actual: ComplexityActual) -> bool:
|
|
"""Record actual task outcome"""
|
|
try:
|
|
key = f"actual:{actual.task_id}"
|
|
data = json.dumps(asdict(actual))
|
|
self.redis.setex(key, self.data_ttl, data)
|
|
|
|
# Add to actuals list
|
|
list_key = "actuals:all"
|
|
self.redis.lpush(list_key, actual.task_id)
|
|
self.redis.expire(list_key, self.data_ttl)
|
|
|
|
logger.info(
|
|
f"Recorded actual for task {actual.task_id}: "
|
|
f"{actual.actual_agents} agents, feedback: {actual.user_feedback}"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to record actual: {e}")
|
|
return False
|
|
|
|
def get_prediction(self, task_id: str) -> Optional[ComplexityPrediction]:
|
|
"""Retrieve a prediction"""
|
|
try:
|
|
key = f"prediction:{task_id}"
|
|
data = self.redis.get(key)
|
|
if not data:
|
|
return None
|
|
return ComplexityPrediction(**json.loads(data))
|
|
except Exception as e:
|
|
logger.error(f"Failed to get prediction: {e}")
|
|
return None
|
|
|
|
def get_actual(self, task_id: str) -> Optional[ComplexityActual]:
|
|
"""Retrieve actual outcome"""
|
|
try:
|
|
key = f"actual:{task_id}"
|
|
data = self.redis.get(key)
|
|
if not data:
|
|
return None
|
|
return ComplexityActual(**json.loads(data))
|
|
except Exception as e:
|
|
logger.error(f"Failed to get actual: {e}")
|
|
return None
|
|
|
|
def calculate_metrics(
|
|
self,
|
|
model_version: Optional[str] = None,
|
|
limit: int = 1000
|
|
) -> Optional[ModelMetrics]:
|
|
"""Calculate model performance metrics"""
|
|
if model_version is None:
|
|
model_version = self.current_model_version
|
|
|
|
try:
|
|
# Get predictions for this model version
|
|
list_key = f"predictions:{model_version}"
|
|
task_ids = self.redis.lrange(list_key, 0, limit - 1)
|
|
|
|
if not task_ids:
|
|
logger.warning(f"No predictions found for model {model_version}")
|
|
return None
|
|
|
|
errors = []
|
|
within_1 = 0
|
|
within_2 = 0
|
|
ratings = []
|
|
|
|
for task_id_bytes in task_ids:
|
|
task_id = task_id_bytes.decode('utf-8')
|
|
|
|
prediction = self.get_prediction(task_id)
|
|
actual = self.get_actual(task_id)
|
|
|
|
if not prediction or not actual:
|
|
continue
|
|
|
|
# Calculate error
|
|
error = abs(prediction.predicted_agents - actual.actual_agents)
|
|
errors.append(error)
|
|
|
|
# Check accuracy thresholds
|
|
if error <= 1:
|
|
within_1 += 1
|
|
if error <= 2:
|
|
within_2 += 1
|
|
|
|
# Collect user ratings
|
|
if actual.user_rating:
|
|
ratings.append(actual.user_rating)
|
|
|
|
if not errors:
|
|
logger.warning(f"No matched predictions/actuals for model {model_version}")
|
|
return None
|
|
|
|
total = len(errors)
|
|
mae = statistics.mean(errors)
|
|
acc_1 = (within_1 / total) * 100
|
|
acc_2 = (within_2 / total) * 100
|
|
avg_rating = statistics.mean(ratings) if ratings else 0.0
|
|
|
|
metrics = ModelMetrics(
|
|
model_version=model_version,
|
|
total_predictions=total,
|
|
mean_absolute_error=mae,
|
|
accuracy_within_1=acc_1,
|
|
accuracy_within_2=acc_2,
|
|
user_satisfaction=avg_rating,
|
|
last_updated=time.time()
|
|
)
|
|
|
|
# Cache metrics
|
|
metrics_key = f"metrics:{model_version}"
|
|
self.redis.setex(metrics_key, 3600, json.dumps(asdict(metrics)))
|
|
|
|
logger.info(
|
|
f"Model {model_version} metrics: MAE={mae:.2f}, "
|
|
f"Acc±1={acc_1:.1f}%, Acc±2={acc_2:.1f}%, "
|
|
f"Satisfaction={avg_rating:.2f}/5"
|
|
)
|
|
|
|
return metrics
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to calculate metrics: {e}")
|
|
return None
|
|
|
|
def get_error_patterns(self, limit: int = 100) -> Dict[str, List[Tuple[str, int, int]]]:
|
|
"""
|
|
Analyze error patterns to identify systematic biases
|
|
|
|
Returns dict with categories:
|
|
- overestimated: tasks where we predicted too many agents
|
|
- underestimated: tasks where we predicted too few agents
|
|
- accurate: tasks where prediction was close
|
|
"""
|
|
patterns = {
|
|
"overestimated": [],
|
|
"underestimated": [],
|
|
"accurate": []
|
|
}
|
|
|
|
try:
|
|
list_key = "actuals:all"
|
|
task_ids = self.redis.lrange(list_key, 0, limit - 1)
|
|
|
|
for task_id_bytes in task_ids:
|
|
task_id = task_id_bytes.decode('utf-8')
|
|
|
|
prediction = self.get_prediction(task_id)
|
|
actual = self.get_actual(task_id)
|
|
|
|
if not prediction or not actual:
|
|
continue
|
|
|
|
error = prediction.predicted_agents - actual.actual_agents
|
|
|
|
entry = (
|
|
task_id,
|
|
prediction.predicted_agents,
|
|
actual.actual_agents
|
|
)
|
|
|
|
if error > 1:
|
|
patterns["overestimated"].append(entry)
|
|
elif error < -1:
|
|
patterns["underestimated"].append(entry)
|
|
else:
|
|
patterns["accurate"].append(entry)
|
|
|
|
return patterns
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to analyze error patterns: {e}")
|
|
return patterns
|
|
|
|
def get_feedback_summary(self, limit: int = 100) -> Dict[str, int]:
|
|
"""Summarize user feedback"""
|
|
feedback_counts = defaultdict(int)
|
|
|
|
try:
|
|
list_key = "actuals:all"
|
|
task_ids = self.redis.lrange(list_key, 0, limit - 1)
|
|
|
|
for task_id_bytes in task_ids:
|
|
task_id = task_id_bytes.decode('utf-8')
|
|
actual = self.get_actual(task_id)
|
|
|
|
if actual and actual.user_feedback:
|
|
feedback_counts[actual.user_feedback] += 1
|
|
|
|
return dict(feedback_counts)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to get feedback summary: {e}")
|
|
return {}
|
|
|
|
def generate_tuning_recommendations(self) -> List[str]:
|
|
"""Generate recommendations for model tuning based on data"""
|
|
recommendations = []
|
|
|
|
try:
|
|
# Get current metrics
|
|
metrics = self.calculate_metrics()
|
|
if not metrics:
|
|
return ["Insufficient data for recommendations"]
|
|
|
|
# Check accuracy
|
|
if metrics.accuracy_within_1 < 60:
|
|
recommendations.append(
|
|
f"Low accuracy ({metrics.accuracy_within_1:.1f}%). "
|
|
"Consider retraining with more diverse examples."
|
|
)
|
|
|
|
# Check error patterns
|
|
patterns = self.get_error_patterns()
|
|
overestimated = len(patterns["overestimated"])
|
|
underestimated = len(patterns["underestimated"])
|
|
total = overestimated + underestimated + len(patterns["accurate"])
|
|
|
|
if total > 0:
|
|
over_pct = (overestimated / total) * 100
|
|
under_pct = (underestimated / total) * 100
|
|
|
|
if over_pct > 40:
|
|
recommendations.append(
|
|
f"Model overestimates in {over_pct:.1f}% of cases. "
|
|
"Consider reducing base agent count or adjusting complexity thresholds."
|
|
)
|
|
|
|
if under_pct > 40:
|
|
recommendations.append(
|
|
f"Model underestimates in {under_pct:.1f}% of cases. "
|
|
"Consider increasing base agent count or lowering complexity thresholds."
|
|
)
|
|
|
|
# Check user satisfaction
|
|
if metrics.user_satisfaction < 3.5:
|
|
recommendations.append(
|
|
f"Low user satisfaction ({metrics.user_satisfaction:.1f}/5). "
|
|
"Review user feedback and adjust model accordingly."
|
|
)
|
|
|
|
# Check feedback
|
|
feedback = self.get_feedback_summary()
|
|
if feedback.get("too_many", 0) > feedback.get("too_few", 0) * 2:
|
|
recommendations.append(
|
|
"Users frequently report 'too many agents'. "
|
|
"Consider reducing default agent counts."
|
|
)
|
|
elif feedback.get("too_few", 0) > feedback.get("too_many", 0) * 2:
|
|
recommendations.append(
|
|
"Users frequently report 'too few agents'. "
|
|
"Consider increasing default agent counts."
|
|
)
|
|
|
|
if not recommendations:
|
|
recommendations.append(
|
|
f"Model performing well (MAE={metrics.mean_absolute_error:.2f}, "
|
|
f"Acc±1={metrics.accuracy_within_1:.1f}%). Continue monitoring."
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to generate recommendations: {e}")
|
|
recommendations.append(f"Error generating recommendations: {e}")
|
|
|
|
return recommendations
|
|
|
|
def export_training_data(self, limit: int = 1000) -> List[Dict]:
|
|
"""Export prediction/actual pairs for model retraining"""
|
|
training_data = []
|
|
|
|
try:
|
|
list_key = "actuals:all"
|
|
task_ids = self.redis.lrange(list_key, 0, limit - 1)
|
|
|
|
for task_id_bytes in task_ids:
|
|
task_id = task_id_bytes.decode('utf-8')
|
|
|
|
prediction = self.get_prediction(task_id)
|
|
actual = self.get_actual(task_id)
|
|
|
|
if not prediction or not actual:
|
|
continue
|
|
|
|
training_data.append({
|
|
"task_description": prediction.task_description,
|
|
"predicted_agents": prediction.predicted_agents,
|
|
"actual_agents": actual.actual_agents,
|
|
"duration_seconds": actual.actual_duration_seconds,
|
|
"success": actual.success,
|
|
"user_feedback": actual.user_feedback,
|
|
"user_rating": actual.user_rating
|
|
})
|
|
|
|
logger.info(f"Exported {len(training_data)} training examples")
|
|
return training_data
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to export training data: {e}")
|
|
return []
|
|
redis_client.py: |
|
|
"""Redis client for orchestrator state management."""
|
|
import os
|
|
import redis.asyncio as redis
|
|
from typing import Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RedisClient:
|
|
"""Async Redis client wrapper for orchestrator operations."""
|
|
|
|
def __init__(self):
|
|
self.client: Optional[redis.Redis] = None
|
|
self.host = os.getenv("REDIS_HOST", "redis-service")
|
|
self.port = int(os.getenv("REDIS_PORT", "6379"))
|
|
self.db = int(os.getenv("REDIS_DB", "0"))
|
|
|
|
async def connect(self):
|
|
"""Establish Redis connection."""
|
|
try:
|
|
self.client = await redis.Redis(
|
|
host=self.host,
|
|
port=self.port,
|
|
db=self.db,
|
|
decode_responses=True,
|
|
socket_connect_timeout=5,
|
|
socket_keepalive=True,
|
|
)
|
|
await self.client.ping()
|
|
logger.info(f"Connected to Redis at {self.host}:{self.port}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to Redis: {e}")
|
|
raise
|
|
|
|
async def disconnect(self):
|
|
"""Close Redis connection."""
|
|
if self.client:
|
|
await self.client.close()
|
|
logger.info("Disconnected from Redis")
|
|
|
|
async def set(self, key: str, value: str, ex: Optional[int] = None):
|
|
"""Set key-value pair with optional expiration."""
|
|
await self.client.set(key, value, ex=ex)
|
|
|
|
async def get(self, key: str) -> Optional[str]:
|
|
"""Get value by key."""
|
|
return await self.client.get(key)
|
|
|
|
async def delete(self, key: str):
|
|
"""Delete key."""
|
|
await self.client.delete(key)
|
|
|
|
async def exists(self, key: str) -> bool:
|
|
"""Check if key exists."""
|
|
return await self.client.exists(key) > 0
|
|
|
|
async def hset(self, name: str, key: str, value: str):
|
|
"""Set hash field."""
|
|
await self.client.hset(name, key, value)
|
|
|
|
async def hget(self, name: str, key: str) -> Optional[str]:
|
|
"""Get hash field."""
|
|
return await self.client.hget(name, key)
|
|
|
|
async def hgetall(self, name: str) -> dict:
|
|
"""Get all hash fields."""
|
|
return await self.client.hgetall(name)
|
|
|
|
async def hdel(self, name: str, *keys: str):
|
|
"""Delete hash fields."""
|
|
await self.client.hdel(name, *keys)
|
|
|
|
async def keys(self, pattern: str) -> list:
|
|
"""Get keys matching pattern."""
|
|
return await self.client.keys(pattern)
|
|
|
|
async def lpush(self, key: str, *values: str):
|
|
"""Push values to list head."""
|
|
await self.client.lpush(key, *values)
|
|
|
|
async def rpop(self, key: str) -> Optional[str]:
|
|
"""Pop value from list tail."""
|
|
return await self.client.rpop(key)
|
|
|
|
async def llen(self, key: str) -> int:
|
|
"""Get list length."""
|
|
return await self.client.llen(key)
|
|
|
|
async def lrange(self, key: str, start: int, end: int) -> list:
|
|
"""Get list range."""
|
|
return await self.client.lrange(key, start, end)
|
|
|
|
|
|
# Global Redis client instance
|
|
redis_client = RedisClient()
|
|
requirements.txt: |
|
|
fastapi==0.115.0
|
|
uvicorn[standard]==0.32.0
|
|
websockets==13.1
|
|
redis==5.2.0
|
|
pydantic==2.9.2
|
|
python-dotenv==1.0.1
|
|
prometheus-client==0.20.0
|
|
opentelemetry-api==1.24.0
|
|
opentelemetry-sdk==1.24.0
|
|
opentelemetry-exporter-otlp==1.24.0
|
|
opentelemetry-instrumentation-redis==0.45b0
|
|
opentelemetry-instrumentation-requests==0.45b0
|
|
opentelemetry-instrumentation-logging==0.45b0
|
|
task_queue.py: |
|
|
"""Task queue management with failure recovery."""
|
|
import json
|
|
import time
|
|
import uuid
|
|
import logging
|
|
from typing import Dict, List, Optional
|
|
from enum import Enum
|
|
from pydantic import BaseModel
|
|
from .redis_client import redis_client
|
|
from .agent_registry import agent_registry, AgentStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TaskStatus(str, Enum):
|
|
"""Task status enumeration."""
|
|
PENDING = "pending"
|
|
ASSIGNED = "assigned"
|
|
IN_PROGRESS = "in_progress"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
|
|
class Task(BaseModel):
|
|
"""Task model."""
|
|
task_id: str
|
|
description: str
|
|
status: TaskStatus
|
|
assigned_agent_id: Optional[str] = None
|
|
created_at: float
|
|
started_at: Optional[float] = None
|
|
completed_at: Optional[float] = None
|
|
retry_count: int = 0
|
|
max_retries: int = 3
|
|
context: Dict = {}
|
|
|
|
|
|
class TaskQueue:
|
|
"""Manages task assignment and reassignment with failure recovery."""
|
|
|
|
TASK_KEY_PREFIX = "task:"
|
|
PENDING_QUEUE_KEY = "queue:pending"
|
|
AGENT_TASK_KEY_PREFIX = "agent_task:"
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
async def create_task(
|
|
self, description: str, context: Dict = {}, max_retries: int = 3
|
|
) -> Task:
|
|
"""Create a new task and add to pending queue."""
|
|
task = Task(
|
|
task_id=str(uuid.uuid4()),
|
|
description=description,
|
|
status=TaskStatus.PENDING,
|
|
created_at=time.time(),
|
|
max_retries=max_retries,
|
|
context=context,
|
|
)
|
|
|
|
# Store task
|
|
key = f"{self.TASK_KEY_PREFIX}{task.task_id}"
|
|
await redis_client.set(key, task.model_dump_json())
|
|
|
|
# Add to pending queue
|
|
await redis_client.lpush(self.PENDING_QUEUE_KEY, task.task_id)
|
|
|
|
logger.info(f"Created task {task.task_id}: {description}")
|
|
return task
|
|
|
|
async def assign_task(self, task_id: str, agent_id: str) -> bool:
|
|
"""Assign a task to an agent."""
|
|
# Get task
|
|
task = await self.get_task(task_id)
|
|
if not task:
|
|
logger.error(f"Task {task_id} not found")
|
|
return False
|
|
|
|
if task.status not in [TaskStatus.PENDING, TaskStatus.FAILED]:
|
|
logger.error(f"Task {task_id} cannot be assigned (status: {task.status})")
|
|
return False
|
|
|
|
# Verify agent is idle
|
|
agent = await agent_registry.get_agent(agent_id)
|
|
if not agent or agent.status != AgentStatus.IDLE:
|
|
logger.error(f"Agent {agent_id} is not available for task assignment")
|
|
return False
|
|
|
|
# Update task
|
|
task.status = TaskStatus.ASSIGNED
|
|
task.assigned_agent_id = agent_id
|
|
task.started_at = time.time()
|
|
|
|
key = f"{self.TASK_KEY_PREFIX}{task_id}"
|
|
await redis_client.set(key, task.model_dump_json())
|
|
|
|
# Track agent's current task
|
|
agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{agent_id}"
|
|
await redis_client.set(agent_task_key, task_id)
|
|
|
|
# Update agent status
|
|
await agent_registry.update_status(agent_id, AgentStatus.BUSY, task_id)
|
|
|
|
logger.info(f"Assigned task {task_id} to agent {agent_id}")
|
|
return True
|
|
|
|
async def start_task(self, task_id: str) -> bool:
|
|
"""Mark task as in progress."""
|
|
task = await self.get_task(task_id)
|
|
if not task:
|
|
return False
|
|
|
|
task.status = TaskStatus.IN_PROGRESS
|
|
key = f"{self.TASK_KEY_PREFIX}{task_id}"
|
|
await redis_client.set(key, task.model_dump_json())
|
|
|
|
logger.info(f"Task {task_id} started")
|
|
return True
|
|
|
|
async def complete_task(self, task_id: str) -> bool:
|
|
"""Mark task as completed."""
|
|
task = await self.get_task(task_id)
|
|
if not task:
|
|
return False
|
|
|
|
task.status = TaskStatus.COMPLETED
|
|
task.completed_at = time.time()
|
|
|
|
key = f"{self.TASK_KEY_PREFIX}{task_id}"
|
|
await redis_client.set(key, task.model_dump_json())
|
|
|
|
# Clear agent's current task
|
|
if task.assigned_agent_id:
|
|
agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{task.assigned_agent_id}"
|
|
await redis_client.delete(agent_task_key)
|
|
|
|
# Update agent to idle
|
|
await agent_registry.update_status(task.assigned_agent_id, AgentStatus.IDLE)
|
|
|
|
logger.info(f"Task {task_id} completed")
|
|
return True
|
|
|
|
async def fail_task(self, task_id: str, reason: str = "") -> bool:
|
|
"""Mark task as failed and handle retry logic."""
|
|
task = await self.get_task(task_id)
|
|
if not task:
|
|
return False
|
|
|
|
task.retry_count += 1
|
|
|
|
# Clear agent's current task
|
|
if task.assigned_agent_id:
|
|
agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{task.assigned_agent_id}"
|
|
await redis_client.delete(agent_task_key)
|
|
|
|
# Update agent to idle
|
|
await agent_registry.update_status(task.assigned_agent_id, AgentStatus.IDLE)
|
|
|
|
# Check if we should retry
|
|
if task.retry_count < task.max_retries:
|
|
task.status = TaskStatus.PENDING
|
|
task.assigned_agent_id = None
|
|
task.started_at = None
|
|
|
|
# Re-add to pending queue
|
|
await redis_client.lpush(self.PENDING_QUEUE_KEY, task_id)
|
|
|
|
logger.warning(
|
|
f"Task {task_id} failed (retry {task.retry_count}/{task.max_retries}): {reason}"
|
|
)
|
|
else:
|
|
task.status = TaskStatus.FAILED
|
|
task.completed_at = time.time()
|
|
|
|
logger.error(
|
|
f"Task {task_id} permanently failed after {task.retry_count} retries: {reason}"
|
|
)
|
|
|
|
key = f"{self.TASK_KEY_PREFIX}{task_id}"
|
|
await redis_client.set(key, task.model_dump_json())
|
|
|
|
return True
|
|
|
|
async def reassign_agent_tasks(self, failed_agent_id: str) -> List[str]:
|
|
"""Reassign all tasks from a failed agent back to pending queue."""
|
|
# Get agent's current task
|
|
agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{failed_agent_id}"
|
|
task_id = await redis_client.get(agent_task_key)
|
|
|
|
reassigned_tasks = []
|
|
|
|
if task_id:
|
|
await self.fail_task(task_id, f"Agent {failed_agent_id} failed")
|
|
reassigned_tasks.append(task_id)
|
|
|
|
logger.info(
|
|
f"Reassigned {len(reassigned_tasks)} tasks from failed agent {failed_agent_id}"
|
|
)
|
|
return reassigned_tasks
|
|
|
|
async def get_task(self, task_id: str) -> Optional[Task]:
|
|
"""Get task by ID."""
|
|
key = f"{self.TASK_KEY_PREFIX}{task_id}"
|
|
data = await redis_client.get(key)
|
|
|
|
if not data:
|
|
return None
|
|
|
|
return Task.model_validate_json(data)
|
|
|
|
async def get_next_pending_task(self) -> Optional[Task]:
|
|
"""Get next pending task from queue."""
|
|
task_id = await redis_client.rpop(self.PENDING_QUEUE_KEY)
|
|
|
|
if not task_id:
|
|
return None
|
|
|
|
return await self.get_task(task_id)
|
|
|
|
async def get_pending_count(self) -> int:
|
|
"""Get count of pending tasks."""
|
|
return await redis_client.llen(self.PENDING_QUEUE_KEY)
|
|
|
|
async def get_all_tasks(self, status: Optional[TaskStatus] = None) -> List[Task]:
|
|
"""Get all tasks, optionally filtered by status."""
|
|
pattern = f"{self.TASK_KEY_PREFIX}*"
|
|
keys = await redis_client.keys(pattern)
|
|
|
|
tasks = []
|
|
for key in keys:
|
|
data = await redis_client.get(key)
|
|
if data:
|
|
task = Task.model_validate_json(data)
|
|
if status is None or task.status == status:
|
|
tasks.append(task)
|
|
|
|
return tasks
|
|
|
|
|
|
# Global task queue instance
|
|
task_queue = TaskQueue()
|
|
tracing.py: |
|
|
"""
|
|
OpenTelemetry Distributed Tracing Setup
|
|
|
|
Provides distributed tracing for agent handoffs and task flows.
|
|
Enables end-to-end visibility across the swarm system.
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
from typing import Optional, Dict, Any
|
|
from contextlib import contextmanager
|
|
|
|
from opentelemetry import trace
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
|
|
from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION
|
|
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
|
|
from opentelemetry.instrumentation.redis import RedisInstrumentor
|
|
from opentelemetry.instrumentation.requests import RequestsInstrumentor
|
|
from opentelemetry.instrumentation.logging import LoggingInstrumentor
|
|
from opentelemetry.trace import Status, StatusCode, SpanKind
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SwarmTracer:
|
|
"""Manages distributed tracing for the swarm system"""
|
|
|
|
def __init__(
|
|
self,
|
|
service_name: str,
|
|
service_version: str = "1.0.0",
|
|
otlp_endpoint: Optional[str] = None,
|
|
enable_console: bool = False
|
|
):
|
|
self.service_name = service_name
|
|
self.service_version = service_version
|
|
self.otlp_endpoint = otlp_endpoint or os.getenv(
|
|
"OTEL_EXPORTER_OTLP_ENDPOINT",
|
|
"http://localhost:4317"
|
|
)
|
|
self.enable_console = enable_console
|
|
|
|
self._setup_tracing()
|
|
|
|
def _setup_tracing(self):
|
|
"""Initialize OpenTelemetry tracing"""
|
|
# Create resource with service information
|
|
resource = Resource.create({
|
|
SERVICE_NAME: self.service_name,
|
|
SERVICE_VERSION: self.service_version,
|
|
"deployment.environment": os.getenv("ENVIRONMENT", "production"),
|
|
"k8s.namespace": os.getenv("K8S_NAMESPACE", "swarm-system"),
|
|
"k8s.pod.name": os.getenv("HOSTNAME", "unknown"),
|
|
})
|
|
|
|
# Create tracer provider
|
|
provider = TracerProvider(resource=resource)
|
|
|
|
# Add OTLP exporter
|
|
try:
|
|
otlp_exporter = OTLPSpanExporter(endpoint=self.otlp_endpoint)
|
|
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
|
|
logger.info(f"OTLP exporter configured: {self.otlp_endpoint}")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to configure OTLP exporter: {e}")
|
|
|
|
# Add console exporter for debugging
|
|
if self.enable_console:
|
|
console_exporter = ConsoleSpanExporter()
|
|
provider.add_span_processor(BatchSpanProcessor(console_exporter))
|
|
|
|
# Set global tracer provider
|
|
trace.set_tracer_provider(provider)
|
|
|
|
# Auto-instrument libraries
|
|
self._instrument_libraries()
|
|
|
|
self.tracer = trace.get_tracer(__name__)
|
|
logger.info(f"Tracing initialized for service: {self.service_name}")
|
|
|
|
def _instrument_libraries(self):
|
|
"""Auto-instrument common libraries"""
|
|
try:
|
|
RedisInstrumentor().instrument()
|
|
RequestsInstrumentor().instrument()
|
|
LoggingInstrumentor().instrument()
|
|
logger.info("Auto-instrumentation enabled")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to auto-instrument libraries: {e}")
|
|
|
|
@contextmanager
|
|
def trace_operation(
|
|
self,
|
|
operation_name: str,
|
|
attributes: Optional[Dict[str, Any]] = None,
|
|
kind: SpanKind = SpanKind.INTERNAL
|
|
):
|
|
"""
|
|
Context manager for tracing an operation
|
|
|
|
Usage:
|
|
with tracer.trace_operation("process_task", {"task_id": "123"}):
|
|
# do work
|
|
pass
|
|
"""
|
|
with self.tracer.start_as_current_span(
|
|
operation_name,
|
|
kind=kind,
|
|
attributes=attributes or {}
|
|
) as span:
|
|
try:
|
|
yield span
|
|
except Exception as e:
|
|
span.set_status(Status(StatusCode.ERROR, str(e)))
|
|
span.record_exception(e)
|
|
raise
|
|
|
|
def trace_task_submission(self, task_id: str, task_description: str):
|
|
"""Trace task submission"""
|
|
with self.trace_operation(
|
|
"task.submit",
|
|
{
|
|
"task.id": task_id,
|
|
"task.description": task_description[:100] # Truncate
|
|
},
|
|
kind=SpanKind.PRODUCER
|
|
) as span:
|
|
span.add_event("task_submitted")
|
|
return span
|
|
|
|
def trace_agent_creation(self, agent_id: str, task_id: str, pod_name: str):
|
|
"""Trace agent pod creation"""
|
|
with self.trace_operation(
|
|
"agent.create",
|
|
{
|
|
"agent.id": agent_id,
|
|
"task.id": task_id,
|
|
"k8s.pod.name": pod_name
|
|
}
|
|
) as span:
|
|
span.add_event("agent_pod_created")
|
|
return span
|
|
|
|
def trace_agent_execution(self, agent_id: str, task_id: str):
|
|
"""Trace agent task execution"""
|
|
with self.trace_operation(
|
|
"agent.execute",
|
|
{
|
|
"agent.id": agent_id,
|
|
"task.id": task_id
|
|
}
|
|
) as span:
|
|
span.add_event("agent_started")
|
|
return span
|
|
|
|
def trace_handoff(
|
|
self,
|
|
from_agent_id: str,
|
|
to_agent_id: str,
|
|
task_id: str,
|
|
handoff_reason: str
|
|
):
|
|
"""Trace agent handoff"""
|
|
with self.trace_operation(
|
|
"agent.handoff",
|
|
{
|
|
"handoff.from_agent": from_agent_id,
|
|
"handoff.to_agent": to_agent_id,
|
|
"task.id": task_id,
|
|
"handoff.reason": handoff_reason
|
|
},
|
|
kind=SpanKind.CLIENT
|
|
) as span:
|
|
span.add_event("handoff_initiated")
|
|
return span
|
|
|
|
def trace_result_aggregation(self, task_id: str, agent_count: int):
|
|
"""Trace result aggregation"""
|
|
with self.trace_operation(
|
|
"result.aggregate",
|
|
{
|
|
"task.id": task_id,
|
|
"agent.count": agent_count
|
|
}
|
|
) as span:
|
|
span.add_event("aggregation_started")
|
|
return span
|
|
|
|
def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None):
|
|
"""Add an event to the current span"""
|
|
span = trace.get_current_span()
|
|
if span:
|
|
span.add_event(name, attributes or {})
|
|
|
|
def set_attribute(self, key: str, value: Any):
|
|
"""Set an attribute on the current span"""
|
|
span = trace.get_current_span()
|
|
if span:
|
|
span.set_attribute(key, value)
|
|
|
|
def record_error(self, error: Exception):
|
|
"""Record an error in the current span"""
|
|
span = trace.get_current_span()
|
|
if span:
|
|
span.set_status(Status(StatusCode.ERROR, str(error)))
|
|
span.record_exception(error)
|
|
|
|
|
|
# Singleton instance
|
|
_tracer_instance: Optional[SwarmTracer] = None
|
|
|
|
|
|
def initialize_tracing(
|
|
service_name: str,
|
|
service_version: str = "1.0.0",
|
|
otlp_endpoint: Optional[str] = None,
|
|
enable_console: bool = False
|
|
) -> SwarmTracer:
|
|
"""Initialize global tracing instance"""
|
|
global _tracer_instance
|
|
_tracer_instance = SwarmTracer(
|
|
service_name=service_name,
|
|
service_version=service_version,
|
|
otlp_endpoint=otlp_endpoint,
|
|
enable_console=enable_console
|
|
)
|
|
return _tracer_instance
|
|
|
|
|
|
def get_tracer() -> Optional[SwarmTracer]:
|
|
"""Get the global tracer instance"""
|
|
return _tracer_instance
|
|
|
|
|
|
# Decorator for tracing functions
|
|
def traced(operation_name: Optional[str] = None, **span_attributes):
|
|
"""
|
|
Decorator to automatically trace a function
|
|
|
|
Usage:
|
|
@traced("my_operation", task_id="123")
|
|
def my_function():
|
|
pass
|
|
"""
|
|
def decorator(func):
|
|
def wrapper(*args, **kwargs):
|
|
tracer = get_tracer()
|
|
if not tracer:
|
|
return func(*args, **kwargs)
|
|
|
|
op_name = operation_name or f"{func.__module__}.{func.__name__}"
|
|
with tracer.trace_operation(op_name, span_attributes):
|
|
return func(*args, **kwargs)
|
|
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
# Context propagation helpers
|
|
def inject_trace_context(headers: Dict[str, str]) -> Dict[str, str]:
|
|
"""
|
|
Inject trace context into HTTP headers for propagation
|
|
|
|
Usage:
|
|
headers = inject_trace_context({})
|
|
requests.post(url, headers=headers)
|
|
"""
|
|
from opentelemetry.propagate import inject
|
|
inject(headers)
|
|
return headers
|
|
|
|
|
|
def extract_trace_context(headers: Dict[str, str]):
|
|
"""
|
|
Extract trace context from HTTP headers
|
|
|
|
Usage:
|
|
extract_trace_context(request.headers)
|
|
"""
|
|
from opentelemetry.propagate import extract
|
|
return extract(headers)
|
|
|
|
|
|
# Example usage patterns
|
|
"""
|
|
# In orchestrator/main.py:
|
|
from orchestrator.tracing import initialize_tracing, get_tracer
|
|
|
|
tracer = initialize_tracing(
|
|
service_name="swarm-orchestrator",
|
|
service_version="1.0.0",
|
|
otlp_endpoint="http://otel-collector:4317"
|
|
)
|
|
|
|
# Trace task submission
|
|
with tracer.trace_task_submission(task_id, description):
|
|
# Submit task logic
|
|
pass
|
|
|
|
# In agent/main.py:
|
|
from orchestrator.tracing import initialize_tracing, get_tracer
|
|
|
|
tracer = initialize_tracing(
|
|
service_name="swarm-agent",
|
|
service_version="1.0.0"
|
|
)
|
|
|
|
# Trace agent execution
|
|
with tracer.trace_agent_execution(agent_id, task_id):
|
|
# Execute task
|
|
pass
|
|
|
|
# Trace handoff
|
|
with tracer.trace_handoff(from_agent, to_agent, task_id, reason):
|
|
# Perform handoff
|
|
pass
|
|
|
|
# Using decorator
|
|
@traced("process_subtask", task_id="123")
|
|
def process_subtask():
|
|
pass
|
|
|
|
# Manual span management
|
|
tracer = get_tracer()
|
|
with tracer.trace_operation("custom_operation", {"key": "value"}):
|
|
tracer.add_event("checkpoint_reached")
|
|
tracer.set_attribute("result_count", 42)
|
|
"""
|
|
kind: ConfigMap
|
|
metadata:
|
|
creationTimestamp: null
|
|
name: orchestrator-source
|
|
namespace: swarm-system
|