monocle-apptrace 0.4.2__py3-none-any.whl → 0.5.0b1__py3-none-any.whl

monocle_apptrace/__main__.py

@@ -1,6 +1,6 @@
 import sys, os
 import runpy
-from monocle_apptrace.instrumentation.common.instrumentor import setup_monocle_telemetry
+from monocle_apptrace import setup_monocle_telemetry
 
 def main():
     if len(sys.argv) < 2 or not sys.argv[1].endswith(".py"):

monocle_apptrace/exporters/file_exporter.py

@@ -3,7 +3,7 @@
 from os import linesep, path
 from io import TextIOWrapper
 from datetime import datetime
-from typing import Optional, Callable, Sequence
+from typing import Optional, Callable, Sequence, Dict, Tuple
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from opentelemetry.sdk.resources import SERVICE_NAME
@@ -12,11 +12,9 @@ from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
 
 DEFAULT_FILE_PREFIX:str = "monocle_trace_"
 DEFAULT_TIME_FORMAT:str = "%Y-%m-%d_%H.%M.%S"
+HANDLE_TIMEOUT_SECONDS: int = 60  # 1 minute timeout
 
 class FileSpanExporter(SpanExporterBase):
-    current_trace_id: int = None
-    current_file_path: str = None
-
     def __init__(
         self,
         service_name: Optional[str] = None,
@@ -30,13 +28,15 @@ class FileSpanExporter(SpanExporterBase):
         task_processor: Optional[ExportTaskProcessor] = None
     ):
         super().__init__()
-        self.out_handle:TextIOWrapper = None
+        # Dictionary to store file handles: {trace_id: (file_handle, file_path, creation_time, first_span)}
+        self.file_handles: Dict[int, Tuple[TextIOWrapper, str, datetime, bool]] = {}
         self.formatter = formatter
         self.service_name = service_name
         self.output_path = out_path
         self.file_prefix = file_prefix
         self.time_format = time_format
         self.task_processor = task_processor
+        self.is_first_span_in_file = True  # Track if this is the first span in the current file
         if self.task_processor is not None:
             self.task_processor.start()
 
@@ -49,46 +49,133 @@ class FileSpanExporter(SpanExporterBase):
         else:
             return self._process_spans(spans, is_root_span=is_root_span)
 
+    def _cleanup_expired_handles(self) -> None:
+        """Close and remove file handles that have exceeded the timeout."""
+        current_time = datetime.now()
+        expired_trace_ids = []
+        
+        for trace_id, (handle, file_path, creation_time, _) in self.file_handles.items():
+            if (current_time - creation_time).total_seconds() > HANDLE_TIMEOUT_SECONDS:
+                expired_trace_ids.append(trace_id)
+        
+        for trace_id in expired_trace_ids:
+            self._close_trace_handle(trace_id)
+
+    def _get_or_create_handle(self, trace_id: int, service_name: str) -> Tuple[TextIOWrapper, str, bool]:
+        """Get existing handle or create new one for the trace_id."""
+        self._cleanup_expired_handles()
+        
+        if trace_id in self.file_handles:
+            handle, file_path, creation_time, first_span = self.file_handles[trace_id]
+            return handle, file_path, first_span
+        
+        # Create new handle
+        file_path = path.join(self.output_path,
+                             self.file_prefix + service_name + "_" + hex(trace_id) + "_"
+                             + datetime.now().strftime(self.time_format) + ".json")
+        
+        try:
+            handle = open(file_path, "w", encoding='UTF-8')
+            handle.write("[")
+            self.file_handles[trace_id] = (handle, file_path, datetime.now(), True)
+            return handle, file_path, True
+        except Exception as e:
+            print(f"Error creating file {file_path}: {e}")
+            return None, file_path, True
+
+    def _close_trace_handle(self, trace_id: int) -> None:
+        """Close and remove a specific trace handle."""
+        if trace_id in self.file_handles:
+            handle, file_path, creation_time, _ = self.file_handles[trace_id]
+            try:
+                if handle is not None:
+                    handle.write("]")
+                    handle.close()
+            except Exception as e:
+                print(f"Error closing file {file_path}: {e}")
+            finally:
+                del self.file_handles[trace_id]
+
+    def _mark_span_written(self, trace_id: int) -> None:
+        """Mark that a span has been written for this trace (no longer first span)."""
+        if trace_id in self.file_handles:
+            handle, file_path, creation_time, _ = self.file_handles[trace_id]
+            self.file_handles[trace_id] = (handle, file_path, creation_time, False)
+
     def _process_spans(self, spans: Sequence[ReadableSpan], is_root_span: bool = False) -> SpanExportResult:
-        first_span:bool = True
+        # Group spans by trace_id for efficient processing
+        spans_by_trace = {}
+        root_span_traces = set()
+        
         for span in spans:
             if self.skip_export(span):
                 continue
-            if span.context.trace_id != self.current_trace_id:
-                self.rotate_file(span.resource.attributes[SERVICE_NAME],
-                                span.context.trace_id)
-                first_span = True
-            if first_span:
-                first_span = False
-            else:
-                self.out_handle.write(",")
-            self.out_handle.write(self.formatter(span))
-        if is_root_span:
-            self.reset_handle()
-        else:
-            self.out_handle.flush()
+            
+            trace_id = span.context.trace_id
+            if trace_id not in spans_by_trace:
+                spans_by_trace[trace_id] = []
+            spans_by_trace[trace_id].append(span)
+            
+            # Check if this span is a root span
+            if not span.parent:
+                root_span_traces.add(trace_id)
+        
+        # Process spans for each trace
+        for trace_id, trace_spans in spans_by_trace.items():
+            service_name = trace_spans[0].resource.attributes.get(SERVICE_NAME, "unknown")
+            handle, file_path, is_first_span = self._get_or_create_handle(trace_id, service_name)
+            
+            if handle is None:
+                continue
+            
+            for span in trace_spans:
+                if not is_first_span:
+                    try:
+                        handle.write(",")
+                    except Exception as e:
+                        print(f"Error writing comma to file {file_path} for span {span.context.span_id}: {e}")
+                        continue
+                
+                try:
+                    handle.write(self.formatter(span))
+                    if is_first_span:
+                        self._mark_span_written(trace_id)
+                        is_first_span = False
+                except Exception as e:
+                    print(f"Error formatting span {span.context.span_id}: {e}")
+                    continue
+        
+        # Close handles for traces with root spans
+        for trace_id in root_span_traces:
+            self._close_trace_handle(trace_id)
+        
+        # Flush remaining handles
+        for trace_id, (handle, file_path, _, _) in self.file_handles.items():
+            if trace_id not in root_span_traces:
+                try:
+                    if handle is not None:
+                        handle.flush()
+                except Exception as e:
+                    print(f"Error flushing file {file_path}: {e}")
+        
         return SpanExportResult.SUCCESS
 
-    def rotate_file(self, trace_name:str, trace_id:int) -> None:
-        self.reset_handle()
-        self.current_file_path = path.join(self.output_path,
-                        self.file_prefix + trace_name + "_" + hex(trace_id) + "_"
-                        + datetime.now().strftime(self.time_format) + ".json")
-        self.out_handle = open(self.current_file_path, "w", encoding='UTF-8')
-        self.out_handle.write("[")
-        self.current_trace_id = trace_id
-
     def force_flush(self, timeout_millis: int = 30000) -> bool:
-        self.out_handle.flush()
+        """Flush all open file handles."""
+        for trace_id, (handle, file_path, _, _) in self.file_handles.items():
+            try:
+                if handle is not None:
+                    handle.flush()
+            except Exception as e:
+                print(f"Error flushing file {file_path}: {e}")
         return True
 
-    def reset_handle(self) -> None:
-        if self.out_handle is not None:
-            self.out_handle.write("]")
-            self.out_handle.close()
-            self.out_handle = None
-
     def shutdown(self) -> None:
+        """Close all file handles and stop task processor."""
         if hasattr(self, 'task_processor') and self.task_processor is not None:
             self.task_processor.stop()
-        self.reset_handle()
+        
+        # Close all remaining file handles
+        trace_ids_to_close = list(self.file_handles.keys())
+        for trace_id in trace_ids_to_close:
+            self._close_trace_handle(trace_id)

monocle_apptrace/instrumentation/common/__init__.py

@@ -1,2 +1,17 @@
-from .instrumentor import setup_monocle_telemetry, start_trace, stop_trace, start_scope, stop_scope, http_route_handler, monocle_trace_scope, monocle_trace_scope_method, monocle_trace
+from .instrumentor import (
+    setup_monocle_telemetry, 
+    start_trace, 
+    stop_trace, 
+    start_scope, 
+    stop_scope, 
+    http_route_handler, 
+    monocle_trace_scope, 
+    amonocle_trace_scope,
+    monocle_trace_scope_method, 
+    monocle_trace, 
+    amonocle_trace,
+    monocle_trace_method,
+    monocle_trace_http_route,
+    is_valid_trace_id_uuid
+)
 from .utils import MonocleSpanException

monocle_apptrace/instrumentation/common/constants.py

@@ -57,6 +57,10 @@ llm_type_map = {
     "chatcompletionsclient": "azure_ai_inference",
     "embeddingsclient": "azure_ai_inference",
     "imageembeddingsclient": "azure_ai_inference",
+    "chatbedrockconverse": "aws_bedrock",
+    "googleaigeminichatgenerator": "gemini",
+    "gemini": "gemini",
+    "chatgooglegenerativeai": "gemini",
 }
 
 MONOCLE_INSTRUMENTOR = "monocle_apptrace"
@@ -81,4 +85,5 @@ WORKFLOW_TYPE_GENERIC = "workflow.generic"
 MONOCLE_SDK_VERSION = "monocle_apptrace.version"
 MONOCLE_SDK_LANGUAGE = "monocle_apptrace.language"
 MONOCLE_DETECTED_SPAN_ERROR = "monocle_apptrace.detected_span_error"
-HTTP_SUCCESS_CODES = ('200', '201', '202', '204', '205', '206')
+HTTP_SUCCESS_CODES = ('200', '201', '202', '204', '205', '206')
+CHILD_ERROR_CODE = "child.error.code"

monocle_apptrace/instrumentation/common/instrumentor.py

@@ -1,23 +1,19 @@
 import logging
 import inspect
 from typing import Collection, Dict, List, Union
-import random
 import uuid
 import inspect
 from opentelemetry import trace
-from contextlib import contextmanager
-from opentelemetry.context import attach, get_value, set_value, get_current, detach
+from opentelemetry.context import attach, get_value, set_value
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
-from opentelemetry.trace import SpanContext
-from opentelemetry.sdk.trace import TracerProvider, Span, id_generator
+from opentelemetry.sdk.trace import TracerProvider, Span
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from opentelemetry.sdk.trace import Span, TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanProcessor
 from opentelemetry.sdk.trace.export import SpanExporter
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper
-from opentelemetry.trace.propagation import set_span_in_context, _SPAN_KEY
 from monocle_apptrace.exporters.monocle_exporters import get_monocle_exporter
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler, NonFrameworkSpanHandler
 from monocle_apptrace.instrumentation.common.wrapper_method import (
@@ -27,10 +23,11 @@ from monocle_apptrace.instrumentation.common.wrapper_method import (
 )
 from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper, ascope_wrapper, monocle_wrapper, amonocle_wrapper
 from monocle_apptrace.instrumentation.common.utils import (
-    set_scope, remove_scope, http_route_handler, load_scopes, http_async_route_handler
+    load_scopes
 )
-from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR, WORKFLOW_TYPE_GENERIC
+from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR
 from functools import wraps
+
 logger = logging.getLogger(__name__)
 
 SESSION_PROPERTIES_KEY = "session"
@@ -229,49 +226,6 @@ def on_processor_start(span: Span, parent_context):
 def set_context_properties(properties: dict) -> None:
     attach(set_value(SESSION_PROPERTIES_KEY, properties))
 
-def start_trace():
-    """
-    Starts a new trace. All the spans created after this call will be part of the same trace. 
-    Returns:
-        Token: A token representing the attached context for the workflow span.
-                      This token is to be used later to stop the current trace.
-                      Returns None if tracing fails.
-    
-    Raises:
-        Exception: The function catches all exceptions internally and logs a warning.
-    """
-    try:
-        tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= get_tracer_provider())
-        span = tracer.start_span(name = "workflow")
-        updated_span_context = set_span_in_context(span=span)
-        SpanHandler.set_default_monocle_attributes(span)
-        SpanHandler.set_workflow_properties(span)
-        token = attach(updated_span_context)
-        return token
-    except Exception as e:
-        logger.warning("Failed to start trace {e}")
-        return None
-
-def stop_trace(token) -> None:
-    """
-    Stop the active trace and detach workflow type if token is provided. All the spans created after this will not be part of the trace.
-    Args:
-        token: The token that was returned when the trace was started. Used to detach 
-               workflow type. Can be None in which case only the span is ended.
-    Returns:
-        None
-    """
-    try:
-        _parent_span_context = get_current()
-        if _parent_span_context is not None:
-            parent_span: Span = _parent_span_context.get(_SPAN_KEY, None)
-            if parent_span is not None:
-                parent_span.end()
-        if token is not None:
-            detach(token)
-    except:
-        logger.warning("Failed to stop trace")
-
 def is_valid_trace_id_uuid(traceId: str) -> bool:
     try:
         uuid.UUID(traceId)
@@ -280,105 +234,18 @@ def is_valid_trace_id_uuid(traceId: str) -> bool:
         pass
     return False
 
-def start_scope(scope_name: str, scope_value:str = None) -> object:
-    """
-    Start a new scope with the given name and and optional value. If no value is provided, a random UUID will be generated.
-    All the spans, across traces created after this call will have the scope attached until the scope is stopped.
-    Args:
-        scope_name: The name of the scope.
-        scope_value: Optional value of the scope. If None, a random UUID will be generated.
-    Returns:
-        Token: A token representing the attached context for the scope. This token is to be used later to stop the current scope.
-    """
-    return set_scope(scope_name, scope_value)
-
-def stop_scope(token:object) -> None:
-    """
-    Stop the active scope. All the spans created after this will not have the scope attached.
-    Args:
-        token: The token that was returned when the scope was started.
-    Returns:
-        None
-    """
-    remove_scope(token)
-    return
-
-@contextmanager
-def monocle_trace():
-    """
-    Context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have same trace ID
-    """
-    token = start_trace()
-    try:
-        yield
-    finally:
-        stop_trace(token)
-
-@contextmanager
-def monocle_trace_scope(scope_name: str, scope_value:str = None):
-    """
-    Context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have the scope attached.
-    Args:
-        scope_name: The name of the scope.
-        scope_value: Optional value of the scope. If None, a random UUID will be generated."""
-    token = start_scope(scope_name, scope_value)
-    try:
-        yield
-    finally:
-        stop_scope(token)
-    
-def monocle_trace_scope_method(scope_name: str, scope_value:str=None):
-    """
-    Decorator to start and stop a scope for a method. All the spans, across traces created in the method will have the scope attached.
-    """
-    def decorator(func):
-        if inspect.iscoroutinefunction(func):
-            @wraps(func)
-            async def wrapper(*args, **kwargs):
-                token = start_scope(scope_name, scope_value)
-                try:
-                    result = await func(*args, **kwargs)
-                    return result
-                finally:
-                    stop_scope(token)
-            return wrapper
-        else:
-            @wraps(func)
-            def wrapper(*args, **kwargs):
-                token = start_scope(scope_name, scope_value)
-                try:
-                    result = func(*args, **kwargs)
-                    return result
-                finally:
-                    stop_scope(token)
-            return wrapper
-    return decorator
-
-def monocle_trace_http_route(func):
-    """
-    Decorator to start and stop a continue traces and scope for a http route. It will also initiate new scopes from the http headers if configured in ``monocle_scopes.json``
-    All the spans, across traces created in the route will have the scope attached.
-    """
-    if inspect.iscoroutinefunction(func):
-        @wraps(func)
-        async def wrapper(*args, **kwargs):
-            return await http_async_route_handler(func, *args, **kwargs)
-        return wrapper
-    else:
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            return http_route_handler(func, *args, **kwargs)
-        return wrapper
-
-class FixedIdGenerator(id_generator.IdGenerator):
-    def __init__(
-            self,
-            trace_id: int) -> None:
-        self.trace_id = trace_id
-
-    def generate_span_id(self) -> int:
-        return random.getrandbits(64)
-
-    def generate_trace_id(self) -> int:
-        return self.trace_id
+from monocle_apptrace.instrumentation.common.method_wrappers import (
+    monocle_trace_scope_method,
+    amonocle_trace_scope,
+    monocle_trace,
+    amonocle_trace,
+    monocle_trace_method,
+    monocle_trace_http_route,
+    start_scope,
+    stop_scope,
+    start_trace,
+    stop_trace,
+    http_route_handler,
+    monocle_trace_scope,
+)