Upload 3 files

aworld/trace/opentelemetry/__init__.py

@@ -0,0 +1,2 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.

aworld/trace/opentelemetry/memory_storage.py

@@ -0,0 +1,238 @@
+import os
+import json
+import time
+import threading
+from datetime import datetime
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from pydantic import BaseModel
+from typing import Optional, Dict, Any, Union
+from opentelemetry.sdk.trace import Span, SpanContext
+from opentelemetry.sdk.trace.export import SpanExporter
+from aworld.logs.util import logger
+from aworld.trace.constants import ATTRIBUTES_MESSAGE_RUN_TYPE_KEY, RunType
+
+
+class SpanStatus(BaseModel):
+    code: str = "UNSET"
+    description: Optional[str] = None
+
+
+class SpanModel(BaseModel):
+    trace_id: str
+    span_id: str
+    name: str
+    start_time: str
+    end_time: str
+    duration_ms: float
+    attributes: Dict[str, Any]
+    status: SpanStatus
+    parent_id: Optional[str]
+    children: list['SpanModel'] = []
+    run_type: Optional[str] = RunType.OTHER.value
+    is_event: bool = False
+
+    @staticmethod
+    def from_span(span):
+        start_timestamp = span.start_time / 1e9
+        end_timestamp = span.end_time / 1e9
+        start_ms = int((span.start_time % 1e9) / 1e6)
+        end_ms = int((span.end_time % 1e9) / 1e6)
+
+        return SpanModel(
+            trace_id=f"{span.get_span_context().trace_id:032x}",
+            span_id=SpanModel.get_span_id(span),
+            name=span.name,
+            start_time=time.strftime(
+                '%Y-%m-%d %H:%M:%S', time.localtime(start_timestamp)) + f'.{start_ms:03d}',
+            end_time=time.strftime(
+                '%Y-%m-%d %H:%M:%S', time.localtime(end_timestamp)) + f'.{end_ms:03d}',
+            duration_ms=(span.end_time - span.start_time)/1e6,
+            attributes={k: v for k, v in span.attributes.items()},
+            status=SpanStatus(
+                code=str(
+                    span.status.status_code) if span.status.status_code else "UNSET",
+                description=span.status.description or None
+            ),
+            parent_id=SpanModel.get_span_id(
+                span.parent) if span.parent else None,
+            run_type=span.attributes.get(
+                ATTRIBUTES_MESSAGE_RUN_TYPE_KEY, RunType.OTHER.value),
+            is_event=(span.attributes.get("event.id") is not None)
+        )
+
+    @staticmethod
+    def get_span_id(span: Union[Span, SpanContext]):
+        if isinstance(span, SpanContext):
+            return f"{span.span_id:016x}"
+        return f"{span.get_span_context().span_id:016x}"
+
+
+class TraceStorage(ABC):
+    """
+    Storage for traces.
+    """
+    @abstractmethod
+    def add_span(self, span: Span) -> None:
+        """
+        Add a span to the storage.
+        """
+
+    @abstractmethod
+    def get_all_traces(self) -> list[str]:
+        """
+        Get all trace ids.
+        """
+
+    @abstractmethod
+    def get_all_spans(self, trace_id) -> list[SpanModel]:
+        """
+        Get all spans of a trace.
+        """
+
+
+class InMemoryStorage(TraceStorage):
+    """
+    In-memory storage for spans.
+    """
+
+    def __init__(self,  max_traces=1000):
+        self._traces = defaultdict(list)
+        self._trace_order = []
+        self.max_traces = max_traces
+
+    def add_span(self, span: Span):
+        trace_id = f"{span.get_span_context().trace_id:032x}"
+        if trace_id not in self._traces:
+            self._trace_order.append(trace_id)
+            if len(self._trace_order) > self.max_traces:
+                oldest_trace = self._trace_order.pop(0)
+                del self._traces[oldest_trace]
+        self._traces[trace_id].append(SpanModel.from_span(span))
+
+    def get_all_traces(self):
+        return list(self._traces.keys())
+
+    def get_all_spans(self, trace_id):
+        return self._traces.get(trace_id, [])
+
+
+class InMemoryWithPersistStorage(TraceStorage):
+    """
+    In-memory storage for spans with optimized disk persistence.
+    """
+
+    def __init__(self, storage_dir: str = "./trace_data"):
+        self._traces = defaultdict(list)
+        self._pending_spans = []
+        self.storage_dir = os.path.abspath(storage_dir)
+        os.makedirs(self.storage_dir, exist_ok=True)
+        self._lock = threading.Lock()
+        self._persist_thread = None
+        self._load_today_traces()
+        self.current_filename = None
+
+    def _get_today_filename(self):
+        if not self.current_filename:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            self.current_filename = f"trace_{timestamp}.json"
+        return self.current_filename
+
+    def _load_today_traces(self):
+        today = datetime.now().strftime("%Y%m%d")
+        for filename in os.listdir(self.storage_dir):
+            if filename.startswith(f"trace_{today}") and filename.endswith(".json"):
+                filepath = os.path.join(self.storage_dir, filename)
+                try:
+                    with self._lock, open(filepath, 'r') as f:
+                        data = json.load(f)
+                        for span_data in data:
+                            trace_id = span_data.get("trace_id")
+                            span_json = span_data.get("span")
+                            self._traces[trace_id].append(
+                                SpanModel.parse_raw(span_json))
+                except Exception as e:
+                    logger.error(
+                        f"Error loading trace file {filename}: {str(e)}")
+
+    def _start_persist_thread(self):
+        if self._persist_thread is None:
+            self._persist_thread = threading.Thread(
+                target=self._persist_worker, daemon=True)
+            self._persist_thread.start()
+
+    def _persist_worker(self):
+        while True:
+            time.sleep(5)
+            self._persist()
+
+    def _persist(self):
+        if not self._pending_spans:
+            return
+
+        temp_filepath = os.path.join(
+            self.storage_dir, f"temp_{time.time_ns()}.json")
+        final_filepath = os.path.join(
+            self.storage_dir, self._get_today_filename())
+
+        try:
+            spans_to_persist = []
+            with self._lock:
+                spans_to_persist = self._pending_spans.copy()
+                self._pending_spans.clear()
+
+            if spans_to_persist:
+                existing_data = []
+                if os.path.exists(final_filepath):
+                    try:
+                        with open(final_filepath, 'r') as f:
+                            existing_data = json.load(f)
+                    except Exception as e:
+                        logger.error(
+                            f"Error reading existing trace file: {str(e)}")
+
+                merged_spans = existing_data + spans_to_persist
+
+                with open(temp_filepath, 'w') as f:
+                    json.dump(merged_spans, f, default=str)
+                os.replace(temp_filepath, final_filepath)
+        except Exception as e:
+            logger.error(f"Error persisting traces: {str(e)}")
+            try:
+                os.unlink(temp_filepath)
+            except:
+                pass
+
+    def add_span(self, span: Span):
+        span_model = SpanModel.from_span(span)
+        with self._lock:
+            self._traces[span_model.trace_id].append(span_model)
+            self._pending_spans.append({
+                "trace_id": span_model.trace_id,
+                "span": span_model.json()
+            })
+        self._start_persist_thread()
+
+    def get_all_traces(self):
+        with self._lock:
+            return list(self._traces.keys())
+
+    def get_all_spans(self, trace_id):
+        with self._lock:
+            return self._traces.get(trace_id, [])
+
+
+class InMemorySpanExporter(SpanExporter):
+    """
+    Span exporter that stores spans in memory.
+    """
+
+    def __init__(self, storage: TraceStorage):
+        self._storage = storage
+
+    def export(self, spans):
+        for span in spans:
+            self._storage.add_span(span)
+
+    def shutdown(self):
+        pass

aworld/trace/opentelemetry/opentelemetry_adapter.py

@@ -0,0 +1,436 @@
+import sys
+import os
+import traceback
+import time
+import datetime
+import requests
+from threading import Lock
+from typing import Any, Iterator, Sequence, Optional, TYPE_CHECKING
+from contextvars import Token
+from urllib.parse import urljoin
+import opentelemetry.context as otlp_context_api
+from opentelemetry.trace import (
+    SpanKind,
+    set_span_in_context,
+    get_current_span as get_current_otlp_span,
+    NonRecordingSpan,
+    SpanContext,
+    TraceFlags
+)
+from opentelemetry.trace.status import StatusCode
+from opentelemetry.sdk.trace import (
+    ReadableSpan,
+    SynchronousMultiSpanProcessor,
+    Tracer as SDKTracer,
+    Span as SDKSpan,
+    TracerProvider as SDKTracerProvider
+)
+from opentelemetry.context import Context as OTLPContext
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
+
+from aworld.trace.base import (
+    AttributeValueType,
+    NoOpTracer,
+    SpanType,
+    TraceProvider,
+    Tracer,
+    Span,
+    TraceContext,
+    set_tracer_provider
+)
+from aworld.trace.span_cosumer import SpanConsumer
+from aworld.trace.propagator import get_global_trace_context
+from aworld.trace.baggage.sofa_tracer import SofaSpanHelper
+from aworld.logs.util import logger
+from aworld.utils.common import get_local_ip
+from .memory_storage import InMemorySpanExporter, InMemoryStorage
+from ..constants import ATTRIBUTES_MESSAGE_KEY
+from .export import FileSpanExporter, NoOpSpanExporter, SpanConsumerExporter
+from ..server import set_trace_server
+
+
+class OTLPTraceProvider(TraceProvider):
+    """A TraceProvider that wraps an existing `SDKTracerProvider`.
+    This class provides a way to use a `SDKTracerProvider` as a `TraceProvider`.
+    When the context manager is entered, it returns the `SDKTracerProvider` itself.
+    When the context manager is exited, it calls `shutdown` on the `SDKTracerProvider`.
+    Args:
+        provider: The internal provider to wrap.
+    """
+
+    def __init__(self, provider: SDKTracerProvider, suppressed_scopes: Optional[set[str]] = None):
+        self._provider: SDKTracerProvider = provider
+        self._suppressed_scopes = set()
+        if suppressed_scopes:
+            self._suppressed_scopes.update(suppressed_scopes)
+        self._lock: Lock = Lock()
+
+    def get_tracer(
+            self,
+            name: str,
+            version: Optional[str] = None
+    ):
+        with self._lock:
+            if name in self._suppressed_scopes:
+                return NoOpTracer()
+            else:
+                tracer = self._provider.get_tracer(instrumenting_module_name=name,
+                                                   instrumenting_library_version=version)
+                return OTLPTracer(tracer)
+
+    def shutdown(self) -> None:
+        with self._lock:
+            if isinstance(self._provider, SDKTracerProvider):
+                self._provider.shutdown()
+
+    def force_flush(self, timeout: Optional[float] = None) -> bool:
+        with self._lock:
+            if isinstance(self._provider, SDKTracerProvider):
+                return self._provider.force_flush(timeout)
+            else:
+                return False
+
+    def get_current_span(self) -> Optional["Span"]:
+        otlp_span = get_current_otlp_span()
+        return OTLPSpan(otlp_span, is_new_span=False)
+
+
+class OTLPTracer(Tracer):
+    """A Tracer represents a collection of Spans.
+    Args:
+        tracer: The internal tracer to wrap.
+    """
+
+    def __init__(self, tracer: SDKTracer):
+        self._tracer = tracer
+
+    def start_span(
+            self,
+            name: str,
+            span_type: SpanType = SpanType.INTERNAL,
+            attributes: dict[str, AttributeValueType] = None,
+            start_time: Optional[int] = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+            trace_context: Optional[TraceContext] = None
+    ) -> "Span":
+        otel_context = None
+        trace_context = trace_context or get_global_trace_context().get_and_clear()
+        if trace_context:
+            otel_context = self._get_otel_context_from_trace_context(
+                trace_context)
+        start_time = start_time or time.time_ns()
+        attributes = {**(attributes or {})}
+        attributes.setdefault(ATTRIBUTES_MESSAGE_KEY, name)
+        SofaSpanHelper.set_sofa_context_to_attr(attributes)
+        attributes = {k: v for k, v in attributes.items(
+        ) if is_valid_attribute_value(k, v)}
+
+        span_kind = self._convert_to_span_kind(
+            span_type) if span_type else SpanKind.INTERNAL
+        span = self._tracer.start_span(name=name,
+                                       kind=span_kind,
+                                       context=otel_context,
+                                       attributes=attributes,
+                                       start_time=start_time,
+                                       record_exception=record_exception,
+                                       set_status_on_exception=set_status_on_exception)
+        return OTLPSpan(span)
+
+    def start_as_current_span(
+            self,
+            name: str,
+            span_type: SpanType = SpanType.INTERNAL,
+            attributes: dict[str, AttributeValueType] = None,
+            start_time: Optional[int] = None,
+            record_exception: bool = True,
+            set_status_on_exception: bool = True,
+            end_on_exit: bool = True,
+            trace_context: Optional[TraceContext] = None
+    ) -> Iterator["Span"]:
+
+        start_time = start_time or time.time_ns()
+        attributes = {**(attributes or {})}
+        attributes.setdefault(ATTRIBUTES_MESSAGE_KEY, name)
+        SofaSpanHelper.set_sofa_context_to_attr(attributes)
+        attributes = {k: v for k, v in attributes.items(
+        ) if is_valid_attribute_value(k, v)}
+
+        span_kind = self._convert_to_span_kind(
+            span_type) if span_type else SpanKind.INTERNAL
+        otel_context = None
+        trace_context = trace_context or get_global_trace_context().get_and_clear()
+        if trace_context:
+            otel_context = self._get_otel_context_from_trace_context(
+                trace_context)
+
+        class _OTLPSpanContextManager:
+            def __init__(self, tracer: SDKTracer):
+                self._span_cm = None
+                self._tracer = tracer
+
+            def __enter__(self):
+                self._span_cm = self._tracer.start_as_current_span(
+                    name=name,
+                    kind=span_kind,
+                    context=otel_context,
+                    attributes=attributes,
+                    start_time=start_time,
+                    record_exception=record_exception,
+                    set_status_on_exception=set_status_on_exception,
+                    end_on_exit=end_on_exit
+                )
+                inner_span = self._span_cm.__enter__()
+                return OTLPSpan(inner_span)
+
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                return self._span_cm.__exit__(exc_type, exc_val, exc_tb)
+
+        return _OTLPSpanContextManager(self._tracer)
+
+    def _convert_to_span_kind(self, span_type: SpanType) -> str:
+        if span_type == SpanType.INTERNAL:
+            return SpanKind.INTERNAL
+        elif span_type == SpanType.CLIENT:
+            return SpanKind.CLIENT
+        elif span_type == SpanType.SERVER:
+            return SpanKind.SERVER
+        elif span_type == SpanType.PRODUCER:
+            return SpanKind.PRODUCER
+        elif span_type == SpanType.CONSUMER:
+            return SpanKind.CONSUMER
+        else:
+            return SpanKind.INTERNAL
+
+    def _get_otel_context_from_trace_context(self, trace_context: TraceContext) -> OTLPContext:
+        trace_flags = None
+        if trace_context.trace_flags:
+            trace_flags = TraceFlags(int(trace_context.trace_flags, 16))
+        otel_context = otlp_context_api.Context()
+        return set_span_in_context(
+            NonRecordingSpan(
+                SpanContext(
+                    trace_id=int(trace_context.trace_id, 16),
+                    span_id=int(trace_context.span_id, 16),
+                    is_remote=True,
+                    trace_flags=trace_flags
+                )
+            ),
+            otel_context,
+        )
+
+
+class OTLPSpan(Span, ReadableSpan):
+    """A Span represents a single operation within a trace.
+    """
+
+    def __init__(self, span: SDKSpan, is_new_span=True):
+        self._span = span
+        self._token: Optional[Token[OTLPContext]] = None
+        if is_new_span:
+            self._attach()
+            self._add_to_open_spans()
+
+    if not TYPE_CHECKING:  # pragma: no branch
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self._span, name)
+
+    def end(self, end_time: Optional[int] = None) -> None:
+        self._remove_from_open_spans()
+        end_time = end_time or time.time_ns()
+        if not self._span._status or self._span._status.status_code == StatusCode.UNSET:
+            self._span.set_status(
+                status=StatusCode.OK,
+                description="",
+            )
+        self._span.end(end_time=end_time)
+        self._detach()
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        if not is_valid_attribute_value(key, value):
+            return
+        self._span.set_attribute(key=key, value=value)
+
+    def set_attributes(self, attributes: dict[str, Any]) -> None:
+        attributes = {k: v for k, v in attributes.items(
+        ) if is_valid_attribute_value(k, v)}
+        self._span.set_attributes(attributes=attributes)
+
+    def is_recording(self) -> bool:
+        return self._span.is_recording()
+
+    def record_exception(
+            self,
+            exception: BaseException,
+            attributes: dict[str, Any] = None,
+            timestamp: Optional[int] = None,
+            escaped: bool = False,
+    ) -> None:
+        timestamp = timestamp or time.time_ns()
+        attributes = {**(attributes or {})}
+
+        stacktrace = ''.join(traceback.format_exception(
+            type(exception), exception, exception.__traceback__))
+        self._span.set_attributes({
+            SpanAttributes.EXCEPTION_STACKTRACE: stacktrace,
+            SpanAttributes.EXCEPTION_TYPE: type(exception).__name__,
+            SpanAttributes.EXCEPTION_MESSAGE: str(exception),
+            SpanAttributes.EXCEPTION_ESCAPED: escaped
+        })
+        if exception is not sys.exc_info()[1]:
+            attributes[SpanAttributes.EXCEPTION_STACKTRACE] = stacktrace
+
+        self._span.record_exception(exception=exception,
+                                    attributes=attributes,
+                                    timestamp=timestamp,
+                                    escaped=escaped)
+        self._span.set_status(
+            status=StatusCode.ERROR,
+            description=str(exception),
+        )
+
+    def get_trace_id(self) -> str:
+        """Get the trace ID of the span.
+        Returns:
+            The trace ID of the span.
+        """
+        if not self._span or not self._span.get_span_context() or not self.is_recording():
+            return None
+        return f"{self._span.get_span_context().trace_id:032x}"
+
+    def get_span_id(self) -> str:
+        """Get the span ID of the span.
+        Returns:
+            The span ID of the span.
+        """
+        if not self._span or not self._span.get_span_context() or not self.is_recording():
+            return None
+        return f"{self._span.get_span_context().span_id:016x}"
+
+    def _attach(self):
+        if self._token is not None:
+            return
+        self._token = otlp_context_api.attach(set_span_in_context(self._span))
+
+    def _detach(self):
+        if self._token is None:
+            return
+        try:
+            otlp_context_api.detach(self._token)
+        except ValueError as e:
+            logger.warning(f"Failed to detach context: {e}")
+        finally:
+            self._token = None
+
+
+def configure_otlp_provider(
+        backends: Sequence[str] = None,
+        base_url: str = None,
+        write_token: str = None,
+        span_consumers: Optional[Sequence[SpanConsumer]] = None,
+        **kwargs
+) -> None:
+    """Configure the OTLP provider.
+    Args:
+        backend: The backend to use.
+        write_token: The write token to use.
+        **kwargs: Additional keyword arguments to pass to the provider.
+    """
+    from aworld.metrics.opentelemetry.opentelemetry_adapter import build_otel_resource
+    backends = backends or ["logfire"]
+    processor = SynchronousMultiSpanProcessor()
+    processor.add_span_processor(BatchSpanProcessor(
+        SpanConsumerExporter(span_consumers)))
+    for backend in backends:
+        if backend == "logfire":
+            span_exporter = _configure_logfire_exporter(
+                write_token=write_token, base_url=base_url, **kwargs)
+            processor.add_span_processor(BatchSpanProcessor(span_exporter))
+        elif backend == "console":
+            from opentelemetry.sdk.trace.export import ConsoleSpanExporter
+            processor.add_span_processor(
+                BatchSpanProcessor(ConsoleSpanExporter()))
+        elif backend == "file":
+            timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+            file_path = kwargs.get("file_path", f"traces_{timestamp}.json")
+            processor.add_span_processor(
+                BatchSpanProcessor(FileSpanExporter(file_path)))
+        elif backend == "memory":
+            logger.info("Using in-memory storage for traces.")
+            storage = kwargs.get(
+                "storage", InMemoryStorage())
+            processor.add_span_processor(
+                SimpleSpanProcessor(InMemorySpanExporter(storage=storage)))
+            server_enabled = kwargs.get("server_enabled") or os.getenv(
+                "START_TRACE_SERVER") or "true"
+            server_port = kwargs.get("server_port") or 7079
+            if (server_enabled.lower() == "true"):
+                logger.info(f"Starting trace server on port {server_port}.")
+                set_trace_server(storage=storage, port=int(
+                    server_port), start_server=True)
+            else:
+                logger.info("Trace server is not started.")
+                set_trace_server(storage=storage, port=int(
+                    server_port), start_server=False)
+        else:
+            span_exporter = _configure_otlp_exporter(
+                base_url=base_url, **kwargs)
+            processor.add_span_processor(BatchSpanProcessor(span_exporter))
+
+    set_tracer_provider(OTLPTraceProvider(SDKTracerProvider(active_span_processor=processor,
+                                                            resource=build_otel_resource())))
+
+
+def _configure_logfire_exporter(write_token: str, base_url: str = None) -> None:
+    """Configure the Logfire exporter.
+    Args:
+        write_token: The write token to use.
+        base_url: The base URL to use.
+        **kwargs: Additional keyword arguments to pass to the exporter.
+    """
+    from opentelemetry.exporter.otlp.proto.http import Compression
+    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
+    base_url = base_url or "https://logfire-us.pydantic.dev"
+    headers = {'User-Agent': f'logfire/3.14.0', 'Authorization': write_token}
+    session = requests.Session()
+    session.headers.update(headers)
+    return OTLPSpanExporter(
+        endpoint=urljoin(base_url, '/v1/traces'),
+        session=session,
+        compression=Compression.Gzip,
+    )
+
+
+def _configure_otlp_exporter(base_url: str = None,  **kwargs) -> None:
+    """Configure the OTLP exporter.
+    Args:
+        write_token: The write token to use.
+        base_url: The base URL to use.
+        **kwargs: Additional keyword arguments to pass to the exporter.
+    """
+    import requests
+    from opentelemetry.exporter.otlp.proto.http import Compression
+    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
+    otlp_traces_endpoint = os.getenv("OTLP_TRACES_ENDPOINT")
+    base_url = base_url or otlp_traces_endpoint
+    session = requests.Session()
+    return OTLPSpanExporter(
+        endpoint=base_url,
+        session=session,
+        compression=Compression.Gzip,
+    )
+
+
+def is_valid_attribute_value(k, v):
+    valid = True
+    if not v:
+        valid = False
+    valid = isinstance(v, (str, bool, int, float)) or \
+        (isinstance(v, Sequence) and
+            all(isinstance(i, (str, bool, int, float)) for i in v))
+    if not valid:
+        logger.warning(f"value of attribute[{k}] is invalid: {v}")
+    return valid