330 lines
9.8 KiB
Python
330 lines
9.8 KiB
Python
"""
|
|
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)
|
|
"""
|