monocle-apptrace 0.3.0b6__tar.gz → 0.3.0b7__tar.gz

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/CHANGELOG.md

@@ -1,4 +1,4 @@
-## Version 0.3.0b6 (2024-12-10)
+## Version 0.3.0b7 (2024-12-10)
 
 - Add dev dependency for Mistral AI integration ([#81](https://github.com/monocle2ai/monocle/pull/81))
 - Add VectorStore deployment URL capture support ([#80](https://github.com/monocle2ai/monocle/pull/80))  

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: monocle_apptrace
-Version: 0.3.0b6
+Version: 0.3.0b7
 Summary: package with monocle genAI tracing
 Project-URL: Homepage, https://github.com/monocle2ai/monocle
 Project-URL: Issues, https://github.com/monocle2ai/monocle/issues

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "monocle_apptrace"
-version = "0.3.0b6"
+version = "0.3.0b7"
 authors = []
 description = "package with monocle genAI tracing"
 readme = "README.md"

monocle_apptrace-0.3.0b7/src/monocle_apptrace/__init__.py

@@ -0,0 +1 @@
+from .instrumentation import *

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/aws/s3_exporter.py

@@ -16,12 +16,13 @@ from botocore.exceptions import (
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from monocle_apptrace.exporters.base_exporter import SpanExporterBase
-from typing import Sequence
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
+from typing import Sequence, Optional
 import json
 logger = logging.getLogger(__name__)
 
 class S3SpanExporter(SpanExporterBase):
-    def __init__(self, bucket_name=None, region_name=None):
+    def __init__(self, bucket_name=None, region_name=None, task_processor: Optional[ExportTaskProcessor] = None):
         super().__init__()
         # Use environment variables if credentials are not provided
         DEFAULT_FILE_PREFIX = "monocle_trace_"
@@ -47,6 +48,9 @@ class S3SpanExporter(SpanExporterBase):
         self.time_format = DEFAULT_TIME_FORMAT
         self.export_queue = []
         self.last_export_time = time.time()
+        self.task_processor = task_processor
+        if self.task_processor is not None:
+            self.task_processor.start()
 
         # Check if bucket exists or create it
         if not self.__bucket_exists(self.bucket_name):
@@ -92,6 +96,7 @@ class S3SpanExporter(SpanExporterBase):
         """Synchronous export method that internally handles async logic."""
         try:
             # Run the asynchronous export logic in an event loop
+            logger.info(f"Exporting {len(spans)} spans to S3.")
             asyncio.run(self.__export_async(spans))
             return SpanExportResult.SUCCESS
         except Exception as e:
@@ -100,6 +105,7 @@ class S3SpanExporter(SpanExporterBase):
 
     async def __export_async(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         try:
+            logger.info(f"__export_async {len(spans)} spans to S3.")
             # Add spans to the export queue
             for span in spans:
                 self.export_queue.append(span)
@@ -142,10 +148,16 @@ class S3SpanExporter(SpanExporterBase):
         batch_to_export = self.export_queue[:self.max_batch_size]
         serialized_data = self.__serialize_spans(batch_to_export)
         self.export_queue = self.export_queue[self.max_batch_size:]
-        try:
-            self.__upload_to_s3(serialized_data)
-        except Exception as e:
-            logger.error(f"Failed to upload span batch: {e}")
+        # to calculate is_root_span loop over each span in batch_to_export and check if parent id is none or null
+        is_root_span = any(not span.parent for span in batch_to_export)
+        logger.info(f"Exporting {len(batch_to_export)} spans to S3 is_root_span : {is_root_span}.")
+        if self.task_processor is not None and callable(getattr(self.task_processor, 'queue_task', None)):
+            self.task_processor.queue_task(self.__upload_to_s3, serialized_data, is_root_span)
+        else:
+            try:
+                self.__upload_to_s3(serialized_data)
+            except Exception as e:
+                logger.error(f"Failed to upload span batch: {e}")
 
     @SpanExporterBase.retry_with_backoff(exceptions=(EndpointConnectionError, ConnectionClosedError, ReadTimeoutError, ConnectTimeoutError))
     def __upload_to_s3(self, span_data_batch: str):
@@ -164,4 +176,6 @@ class S3SpanExporter(SpanExporterBase):
         return True
 
     def shutdown(self) -> None:
+        if hasattr(self, 'task_processor') and self.task_processor is not None:
+            self.task_processor.stop()
         logger.info("S3SpanExporter has been shut down.")

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/aws/s3_exporter_opendal.py

@@ -3,19 +3,19 @@ import time
 import datetime
 import logging
 import asyncio
-from typing import Sequence
+from typing import Sequence, Optional
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
 from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
 from opendal import Operator
 from opendal.exceptions import PermissionDenied, ConfigInvalid, Unexpected
 
-
 import json
 
 logger = logging.getLogger(__name__)
 class OpenDALS3Exporter(SpanExporterBase):
-    def __init__(self, bucket_name=None, region_name=None):
+    def __init__(self, bucket_name=None, region_name=None, task_processor: Optional[ExportTaskProcessor] = None):
         super().__init__()
         DEFAULT_FILE_PREFIX = "monocle_trace_"
         DEFAULT_TIME_FORMAT = "%Y-%m-%d__%H.%M.%S"
@@ -36,7 +36,10 @@ class OpenDALS3Exporter(SpanExporterBase):
             access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
             secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
         )
-
+        
+        self.task_processor = task_processor
+        if self.task_processor is not None:
+            self.task_processor.start()
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         """Synchronous export method that internally handles async logic."""
@@ -88,20 +91,26 @@ class OpenDALS3Exporter(SpanExporterBase):
         batch_to_export = self.export_queue[:self.max_batch_size]
         serialized_data = self.__serialize_spans(batch_to_export)
         self.export_queue = self.export_queue[self.max_batch_size:]
-        try:
-            self.__upload_to_s3(serialized_data)
-        except Exception as e:
-            logger.error(f"Failed to upload span batch: {e}")
+        
+        # Calculate is_root_span by checking if any span has no parent
+        is_root_span = any(not span.parent for span in batch_to_export)
+        
+        if self.task_processor is not None and callable(getattr(self.task_processor, 'queue_task', None)):
+            self.task_processor.queue_task(self.__upload_to_s3, serialized_data, is_root_span)
+        else:
+            try:
+                self.__upload_to_s3(serialized_data, is_root_span)
+            except Exception as e:
+                logger.error(f"Failed to upload span batch: {e}")
 
     @SpanExporterBase.retry_with_backoff(exceptions=(Unexpected))
-    def __upload_to_s3(self, span_data_batch: str):
-
+    def __upload_to_s3(self, span_data_batch: str, is_root_span: bool = False):
         current_time = datetime.datetime.now().strftime(self.time_format)
         file_name = f"{self.file_prefix}{current_time}.ndjson"
         try:
             # Attempt to write the span data batch to S3
             self.op.write(file_name, span_data_batch.encode("utf-8"))
-            logger.info(f"Span batch uploaded to S3 as {file_name}.")
+            logger.info(f"Span batch uploaded to S3 as {file_name}. Is root span: {is_root_span}")
 
         except PermissionDenied as e:
             # S3 bucket is forbidden.
@@ -123,4 +132,6 @@ class OpenDALS3Exporter(SpanExporterBase):
         return True
 
     def shutdown(self) -> None:
+        if hasattr(self, 'task_processor') and self.task_processor is not None:
+            self.task_processor.stop()
         logger.info("S3SpanExporter has been shut down.")

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/azure/blob_exporter.py

@@ -8,14 +8,15 @@ from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient
 from azure.core.exceptions import ResourceNotFoundError, ClientAuthenticationError, ServiceRequestError
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
-from typing import Sequence
+from typing import Sequence, Optional
 from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
 import json
 from monocle_apptrace.instrumentation.common.constants import MONOCLE_SDK_VERSION
 logger = logging.getLogger(__name__)
 
 class AzureBlobSpanExporter(SpanExporterBase):
-    def __init__(self, connection_string=None, container_name=None):
+    def __init__(self, connection_string=None, container_name=None, task_processor: Optional[ExportTaskProcessor] = None):
         super().__init__()
         DEFAULT_FILE_PREFIX = "monocle_trace_"
         DEFAULT_TIME_FORMAT = "%Y-%m-%d_%H.%M.%S"
@@ -44,6 +45,10 @@ class AzureBlobSpanExporter(SpanExporterBase):
                 logger.error(f"Error creating container {container_name}: {e}")
                 raise e
 
+        self.task_processor = task_processor
+        if self.task_processor is not None:
+            self.task_processor.start()
+
     def __container_exists(self, container_name):
         try:
             container_client = self.blob_service_client.get_container_client(container_name)
@@ -111,22 +116,31 @@ class AzureBlobSpanExporter(SpanExporterBase):
         batch_to_export = self.export_queue[:self.max_batch_size]
         serialized_data = self.__serialize_spans(batch_to_export)
         self.export_queue = self.export_queue[self.max_batch_size:]
-        try:
-            self.__upload_to_blob(serialized_data)
-        except Exception as e:
-            logger.error(f"Failed to upload span batch: {e}")
+        
+        # Calculate is_root_span by checking if any span has no parent
+        is_root_span = any(not span.parent for span in batch_to_export)
+        
+        if self.task_processor is not None and callable(getattr(self.task_processor, 'queue_task', None)):
+            self.task_processor.queue_task(self.__upload_to_blob, serialized_data, is_root_span)
+        else:
+            try:
+                self.__upload_to_blob(serialized_data, is_root_span)
+            except Exception as e:
+                logger.error(f"Failed to upload span batch: {e}")
 
     @SpanExporterBase.retry_with_backoff(exceptions=(ResourceNotFoundError, ClientAuthenticationError, ServiceRequestError))
-    def __upload_to_blob(self, span_data_batch: str):
+    def __upload_to_blob(self, span_data_batch: str, is_root_span: bool = False):
         current_time = datetime.datetime.now().strftime(self.time_format)
         file_name = f"{self.file_prefix}{current_time}.ndjson"
         blob_client = self.blob_service_client.get_blob_client(container=self.container_name, blob=file_name)
         blob_client.upload_blob(span_data_batch, overwrite=True)
-        logger.info(f"Span batch uploaded to Azure Blob Storage as {file_name}.")
+        logger.info(f"Span batch uploaded to Azure Blob Storage as {file_name}. Is root span: {is_root_span}")
 
     async def force_flush(self, timeout_millis: int = 30000) -> bool:
         await self.__export_spans()
         return True
 
     def shutdown(self) -> None:
+        if hasattr(self, 'task_processor') and self.task_processor is not None:
+            self.task_processor.stop()
         logger.info("AzureBlobSpanExporter has been shut down.")

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/azure/blob_exporter_opendal.py

@@ -5,16 +5,17 @@ import logging
 import asyncio
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
-from typing import Sequence
+from typing import Sequence, Optional
 from opendal import Operator
 from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
 from opendal.exceptions import Unexpected, PermissionDenied, NotFound
 import json
 
 logger = logging.getLogger(__name__)
 
 class OpenDALAzureExporter(SpanExporterBase):
-    def __init__(self, connection_string=None, container_name=None):
+    def __init__(self, connection_string=None, container_name=None, task_processor: Optional[ExportTaskProcessor] = None):
         super().__init__()
         DEFAULT_FILE_PREFIX = "monocle_trace_"
         DEFAULT_TIME_FORMAT = "%Y-%m-%d_%H.%M.%S"
@@ -25,6 +26,8 @@ class OpenDALAzureExporter(SpanExporterBase):
         # Default values
         self.file_prefix = DEFAULT_FILE_PREFIX
         self.time_format = DEFAULT_TIME_FORMAT
+        self.export_queue = []  # Add this line to initialize export_queue
+        self.last_export_time = time.time()  # Add this line to initialize last_export_time
 
         # Validate input
         if not connection_string:
@@ -51,6 +54,9 @@ class OpenDALAzureExporter(SpanExporterBase):
         except Exception as e:
             raise RuntimeError(f"Failed to initialize OpenDAL operator: {e}")
 
+        self.task_processor = task_processor
+        if self.task_processor is not None:
+            self.task_processor.start()
 
     def parse_connection_string(self,connection_string):
         connection_params = dict(item.split('=', 1) for item in connection_string.split(';') if '=' in item)
@@ -112,19 +118,26 @@ class OpenDALAzureExporter(SpanExporterBase):
         batch_to_export = self.export_queue[:self.max_batch_size]
         serialized_data = self.__serialize_spans(batch_to_export)
         self.export_queue = self.export_queue[self.max_batch_size:]
-        try:
-            self.__upload_to_opendal(serialized_data)
-        except Exception as e:
-            logger.error(f"Failed to upload span batch: {e}")
+        
+        # Calculate is_root_span by checking if any span has no parent
+        is_root_span = any(not span.parent for span in batch_to_export)
+        
+        if self.task_processor is not None and callable(getattr(self.task_processor, 'queue_task', None)):
+            self.task_processor.queue_task(self.__upload_to_opendal, serialized_data, is_root_span)
+        else:
+            try:
+                self.__upload_to_opendal(serialized_data, is_root_span)
+            except Exception as e:
+                logger.error(f"Failed to upload span batch: {e}")
 
     @SpanExporterBase.retry_with_backoff(exceptions=(Unexpected,))
-    def __upload_to_opendal(self, span_data_batch: str):
+    def __upload_to_opendal(self, span_data_batch: str, is_root_span: bool = False):
         current_time = datetime.datetime.now().strftime(self.time_format)
         file_name = f"{self.file_prefix}{current_time}.ndjson"
 
         try:
             self.operator.write(file_name, span_data_batch.encode('utf-8'))
-            logger.info(f"Span batch uploaded to Azure Blob Storage as {file_name}.")
+            logger.info(f"Span batch uploaded to Azure Blob Storage as {file_name}. Is root span: {is_root_span}")
         except PermissionDenied as e:
             # Azure Container is forbidden.
             logger.error(f"Access to container {self.container_name} is forbidden (403).")
@@ -144,4 +157,6 @@ class OpenDALAzureExporter(SpanExporterBase):
         return True
 
     def shutdown(self) -> None:
+        if hasattr(self, 'task_processor') and self.task_processor is not None:
+            self.task_processor.stop()
         logger.info("OpenDALAzureExporter has been shut down.")

monocle_apptrace-0.3.0b7/src/monocle_apptrace/exporters/exporter_processor.py

@@ -0,0 +1,144 @@
+from abc import ABC, abstractmethod
+import logging
+import os
+import queue
+import threading
+import time
+from typing import Callable
+import requests
+from monocle_apptrace.instrumentation.common.constants import AWS_LAMBDA_ENV_NAME
+
+logger = logging.getLogger(__name__)
+LAMBDA_EXTENSION_NAME = "AsyncProcessorMonocle"
+
+class ExportTaskProcessor(ABC):
+    
+    @abstractmethod
+    def start(self):
+        return
+
+    @abstractmethod
+    def stop(self):
+        return
+
+    @abstractmethod
+    def queue_task(self, async_task: Callable[[Callable, any], any] = None, args: any = None, is_root_span: bool = False):
+        return
+
+class LambdaExportTaskProcessor(ExportTaskProcessor):
+    
+    def __init__(
+        self,
+        span_check_interval_seconds: int = 1,
+        max_time_allowed_seconds: int = 30):
+        # An internal queue used by the handler to notify the extension that it can
+        # start processing the async task.
+        self.async_tasks_queue = queue.Queue()
+        self.span_check_interval = span_check_interval_seconds
+        self.max_time_allowed = max_time_allowed_seconds
+
+    def start(self):
+        try:
+            self._start_async_processor()
+        except Exception as e:
+            logger.error(f"LambdaExportTaskProcessor| Failed to start. {e}")
+
+    def stop(self):
+        return
+
+    def queue_task(self, async_task=None, args=None, is_root_span=False):
+        self.async_tasks_queue.put((async_task, args, is_root_span))
+    
+    def set_sagemaker_model(self, endpoint_name: str, span: dict[str, dict[str, str]]):
+        try:
+            try:
+                import boto3
+            except ImportError:
+                logger.error("LambdaExportTaskProcessor| Failed to import boto3")
+                return
+            
+            client = boto3.client('sagemaker')
+            response = client.describe_endpoint(
+                EndpointName=endpoint_name
+            )
+            endpoint_config_name = response["EndpointConfigName"]
+            endpoint_config_response = client.describe_endpoint_config(
+                EndpointConfigName=endpoint_config_name
+            )
+            model_name = endpoint_config_response["ProductionVariants"][0]["ModelName"]
+            model_name_response = client.describe_model(ModelName = model_name)
+            model_name_id = ""
+            try:
+                model_name_id = model_name_response["PrimaryContainer"]["Environment"]["HF_MODEL_ID"]
+            except:
+                pass
+            span["attributes"]["model_name"] = model_name_id
+        except Exception as e:
+            logger.error(f"LambdaExportTaskProcessor| Failed to get sagemaker model. {e}")
+
+    def update_spans(self, export_args):
+        try:
+            if 'batch' in export_args:
+                for span in export_args["batch"]:
+                    try:
+                        if len(span["attributes"]["sagemaker_endpoint_name"]) > 0 :
+                            self.set_sagemaker_model(endpoint_name=span["attributes"]["sagemaker_endpoint_name"], span=span)
+                    except:
+                        pass
+        except Exception as e:
+            logger.error(f"LambdaExportTaskProcessor| Failed to update spans. {e}")
+
+    def _start_async_processor(self):
+        # Register internal extension
+        logger.debug(f"[{LAMBDA_EXTENSION_NAME}] Registering with Lambda service...")
+        response = requests.post(
+            url=f"http://{os.environ['AWS_LAMBDA_RUNTIME_API']}/2020-01-01/extension/register",
+            json={'events': ['INVOKE']},
+            headers={'Lambda-Extension-Name': LAMBDA_EXTENSION_NAME}
+        )
+        ext_id = response.headers['Lambda-Extension-Identifier']
+        logger.debug(f"[{LAMBDA_EXTENSION_NAME}] Registered with ID: {ext_id}")
+
+        def process_tasks():
+            while True:
+                # Call /next to get notified when there is a new invocation and let
+                # Lambda know that we are done processing the previous task.
+
+                logger.debug(f"[{LAMBDA_EXTENSION_NAME}] Waiting for invocation...")
+                response = requests.get(
+                    url=f"http://{os.environ['AWS_LAMBDA_RUNTIME_API']}/2020-01-01/extension/event/next",
+                    headers={'Lambda-Extension-Identifier': ext_id},
+                    timeout=None
+                )
+                root_span_found = False
+                # all values in seconds
+                total_time_elapsed = 0
+                while root_span_found is False and total_time_elapsed < self.max_time_allowed:
+                    logger.debug(response.json())
+                    # Get next task from internal queue
+                    logger.info(f"[{LAMBDA_EXTENSION_NAME}] Async thread running, waiting for task from handler")
+                    while self.async_tasks_queue.empty() is False :
+                        logger.info(f"[{LAMBDA_EXTENSION_NAME}] Processing task from handler")
+                        async_task, arg, is_root_span = self.async_tasks_queue.get()
+                        root_span_found = is_root_span
+                        # self.update_spans(export_args=arg)
+
+                        if async_task is None:
+                            # No task to run this invocation
+                            logger.debug(f"[{LAMBDA_EXTENSION_NAME}] Received null task. Ignoring.")
+                        else:
+                            # Invoke task
+                            logger.debug(f"[{LAMBDA_EXTENSION_NAME}] Received async task from handler. Starting task.")
+                            async_task(arg)
+                    total_time_elapsed+=self.span_check_interval
+                    logger.info(f"[{LAMBDA_EXTENSION_NAME}] Waiting for root span. total_time_elapsed: {total_time_elapsed}, root_span_found: {root_span_found}.")
+                    time.sleep(self.span_check_interval)
+                
+                logger.debug(f"[{LAMBDA_EXTENSION_NAME}] Finished processing task. total_time_elapsed: {total_time_elapsed}, root_span_found: {root_span_found}.")
+
+        # Start processing extension events in a separate thread
+        threading.Thread(target=process_tasks, daemon=True, name=LAMBDA_EXTENSION_NAME).start() 
+
+
+def is_aws_lambda_environment():
+    return AWS_LAMBDA_ENV_NAME in os.environ

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/file_exporter.py

@@ -7,6 +7,7 @@ from typing import Optional, Callable, Sequence
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from opentelemetry.sdk.resources import SERVICE_NAME
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
 
 DEFAULT_FILE_PREFIX:str = "monocle_trace_"
 DEFAULT_TIME_FORMAT:str = "%Y-%m-%d_%H.%M.%S"
@@ -25,6 +26,7 @@ class FileSpanExporter(SpanExporter):
             [ReadableSpan], str
         ] = lambda span: span.to_json()
         + linesep,
+        task_processor: Optional[ExportTaskProcessor] = None
     ):
         self.out_handle:TextIOWrapper = None
         self.formatter = formatter
@@ -32,8 +34,20 @@ class FileSpanExporter(SpanExporter):
         self.output_path = out_path
         self.file_prefix = file_prefix
         self.time_format = time_format
+        self.task_processor = task_processor
+        if self.task_processor is not None:
+            self.task_processor.start()
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        if self.task_processor is not None and callable(getattr(self.task_processor, 'queue_task', None)):
+            # Check if any span is a root span (no parent)
+            is_root_span = any(not span.parent for span in spans)
+            self.task_processor.queue_task(self._process_spans, spans, is_root_span)
+            return SpanExportResult.SUCCESS
+        else:
+            return self._process_spans(spans)
+
+    def _process_spans(self, spans: Sequence[ReadableSpan], is_root_span: bool = False) -> SpanExportResult:
         for span in spans:
             if span.context.trace_id != self.current_trace_id:
                 self.rotate_file(span.resource.attributes[SERVICE_NAME],
@@ -60,4 +74,6 @@ class FileSpanExporter(SpanExporter):
             self.out_handle = None
 
     def shutdown(self) -> None:
+        if hasattr(self, 'task_processor') and self.task_processor is not None:
+            self.task_processor.stop()
         self.reset_handle()

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/monocle_exporters.py

@@ -3,7 +3,9 @@ import os
 import logging
 from importlib import import_module
 from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
+from monocle_apptrace.exporters.exporter_processor import LambdaExportTaskProcessor, is_aws_lambda_environment
 from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+
 logger = logging.getLogger(__name__)
 
 monocle_exporters: Dict[str, Any] = {
@@ -20,6 +22,9 @@ def get_monocle_exporter() -> List[SpanExporter]:
     # Retrieve the MONOCLE_EXPORTER environment variable and split it into a list
     exporter_names = os.environ.get("MONOCLE_EXPORTER", "file").split(",")
     exporters = []
+    
+    # Create task processor for AWS Lambda environment
+    task_processor = LambdaExportTaskProcessor() if is_aws_lambda_environment() else None
 
     for exporter_name in exporter_names:
         exporter_name = exporter_name.strip()
@@ -31,7 +36,11 @@ def get_monocle_exporter() -> List[SpanExporter]:
         try:
             exporter_module = import_module(exporter_class_path["module"])
             exporter_class = getattr(exporter_module, exporter_class_path["class"])
-            exporters.append(exporter_class())
+            # Pass task_processor to all exporters when in AWS Lambda environment
+            if task_processor is not None and exporter_module.__name__.startswith("monocle_apptrace"):
+                exporters.append(exporter_class(task_processor=task_processor))
+            else:
+                exporters.append(exporter_class())
         except Exception as ex:
             logger.debug(
                 f"Unable to initialize Monocle span exporter '{exporter_name}', error: {ex}. Using ConsoleSpanExporter as a fallback.")

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/exporters/okahu/okahu_exporter.py

@@ -48,7 +48,7 @@ class OkahuSpanExporter(SpanExporter):
 
         if self._closed:
             logger.warning("Exporter already shutdown, ignoring batch")
-            return SpanExportResult.FAILUREencoder
+            return SpanExportResult.FAILURE
         if len(spans) == 0:
             return
 
@@ -69,7 +69,10 @@ class OkahuSpanExporter(SpanExporter):
                 obj["context"]["span_id"] = remove_0x_from_start(obj["context"]["span_id"])
             span_list["batch"].append(obj)
 
-        def send_spans_to_okahu(span_list_local=None):
+        # Calculate is_root_span by checking if any span has no parent
+        is_root_span = any(not span.parent for span in spans)
+
+        def send_spans_to_okahu(span_list_local=None, is_root=False):
             try:
                 result = self.session.post(
                     url=self.endpoint,
@@ -83,18 +86,17 @@ class OkahuSpanExporter(SpanExporter):
                         result.text,
                     )
                     return SpanExportResult.FAILURE
-                logger.debug("spans successfully exported to okahu")
+                logger.debug("spans successfully exported to okahu. Is root span: %s", is_root)
                 return SpanExportResult.SUCCESS
             except ReadTimeout as e:
                 logger.warning("Trace export timed out: %s", str(e))
                 return SpanExportResult.FAILURE
 
         # if async task function is present, then push the request to asnc task
-
         if self.task_processor is not None and callable(self.task_processor.queue_task):
-            self.task_processor.queue_task(send_spans_to_okahu, span_list)
+            self.task_processor.queue_task(send_spans_to_okahu, span_list, is_root_span)
             return SpanExportResult.SUCCESS
-        return send_spans_to_okahu(span_list)
+        return send_spans_to_okahu(span_list, is_root_span)
 
     def shutdown(self) -> None:
         if self._closed:

monocle_apptrace-0.3.0b7/src/monocle_apptrace/instrumentation/__init__.py

@@ -0,0 +1 @@
+from .common import *

monocle_apptrace-0.3.0b7/src/monocle_apptrace/instrumentation/common/__init__.py

@@ -0,0 +1,2 @@
+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 .utils import MonocleSpanException

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/instrumentation/common/instrumentor.py

@@ -17,13 +17,13 @@ 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
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler, NonFrameworkSpanHandler
 from monocle_apptrace.instrumentation.common.wrapper_method import (
     DEFAULT_METHODS_LIST,
     WrapperMethod,
     MONOCLE_SPAN_HANDLERS
 )
-from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper, ascope_wrapper
+from monocle_apptrace.instrumentation.common.wrapper import scope_wrapper, ascope_wrapper, wrapper_processor
 from monocle_apptrace.instrumentation.common.utils import (
     set_scope, remove_scope, http_route_handler, load_scopes, async_wrapper, http_async_route_handler
 )
@@ -65,13 +65,11 @@ class MonocleInstrumentor(BaseInstrumentor):
         def instrumented_endpoint_invoke(to_wrap,wrapped, span_name, instance,fn):
             @wraps(fn)
             def with_instrumentation(*args, **kwargs):
-                handler = SpanHandler()
-                with tracer.start_as_current_span(span_name) as span:
-                    response = fn(*args, **kwargs)
-                    handler.hydrate_span(to_wrap, wrapped=wrapped, instance=instance, args=args, kwargs=kwargs,
-                                         result=response, span=span)
-                    return response
-
+                async_task = inspect.iscoroutinefunction(fn)
+                boto_method_to_wrap = to_wrap.copy()
+                boto_method_to_wrap['skip_span'] = False
+                return wrapper_processor(async_task, tracer, NonFrameworkSpanHandler(),
+                            boto_method_to_wrap, fn, instance, args, kwargs)
             return with_instrumentation
         return instrumented_endpoint_invoke
 
@@ -158,6 +156,24 @@ def setup_monocle_telemetry(
         span_handlers: Dict[str,SpanHandler] = None,
         wrapper_methods: List[Union[dict,WrapperMethod]] = None,
         union_with_default_methods: bool = True) -> None:
+    """
+    Set up Monocle telemetry for the application.
+
+    Parameters
+    ----------
+    workflow_name : str
+        The name of the workflow to be used as the service name in telemetry.
+    span_processors : List[SpanProcessor], optional
+        Custom span processors to use instead of the default ones. If None, 
+        BatchSpanProcessors with Monocle exporters will be used.
+    span_handlers : Dict[str, SpanHandler], optional
+        Dictionary of span handlers to be used by the instrumentor, mapping handler names to handler objects.
+    wrapper_methods : List[Union[dict, WrapperMethod]], optional
+        Custom wrapper methods for instrumentation. If None, default methods will be used.
+    union_with_default_methods : bool, default=True
+        If True, combine the provided wrapper_methods with the default methods.
+        If False, only use the provided wrapper_methods.
+    """
     resource = Resource(attributes={
         SERVICE_NAME: workflow_name
     })
@@ -196,6 +212,16 @@ 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")
@@ -209,6 +235,14 @@ def start_trace():
         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:
@@ -229,32 +263,67 @@ def is_valid_trace_id_uuid(traceId: str) -> bool:
     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):
+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):
-                result = async_wrapper(func, scope_name, None, *args, **kwargs)
+                result = async_wrapper(func, scope_name, scope_value, None, *args, **kwargs)
                 return result
             return wrapper
         else:
             @wraps(func)
             def wrapper(*args, **kwargs):
-                token = start_scope(scope_name)
+                token = start_scope(scope_name, scope_value)
                 try:
                     result = func(*args, **kwargs)
                     return result
@@ -264,6 +333,10 @@ def monocle_trace_scope_method(scope_name: str):
     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):
@@ -286,3 +359,4 @@ class FixedIdGenerator(id_generator.IdGenerator):
 
     def generate_trace_id(self) -> int:
         return self.trace_id
+

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/instrumentation/common/span_handler.py

@@ -3,14 +3,14 @@ import os
 from importlib.metadata import version
 from opentelemetry.context import get_value, set_value, attach, detach
 from opentelemetry.sdk.trace import Span
-
+from opentelemetry.trace.status import Status, StatusCode
 from monocle_apptrace.instrumentation.common.constants import (
     QUERY,
     service_name_map,
     service_type_map,
     MONOCLE_SDK_VERSION
 )
-from monocle_apptrace.instrumentation.common.utils import set_attribute, get_scopes
+from monocle_apptrace.instrumentation.common.utils import set_attribute, get_scopes, MonocleSpanException
 from monocle_apptrace.instrumentation.common.constants import WORKFLOW_TYPE_KEY, WORKFLOW_TYPE_GENERIC
 
 logger = logging.getLogger(__name__)
@@ -64,9 +64,12 @@ class SpanHandler:
         """ Set attributes of workflow if this is a root span"""
         SpanHandler.set_workflow_attributes(to_wrap, span)
         SpanHandler.set_app_hosting_identifier_attribute(span)
+        span.set_status(StatusCode.OK)
 
-    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span):
-        pass
+
+    def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, span:Span):
+        if span.status.status_code == StatusCode.UNSET:
+            span.set_status(StatusCode.OK)
 
     def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span):
         self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span)
@@ -95,6 +98,8 @@ class SpanHandler:
                                 result = accessor(arguments)
                                 if result and isinstance(result, (str, list)):
                                     span.set_attribute(attribute_name, result)
+                            except MonocleSpanException as e:
+                                span.set_status(StatusCode.ERROR, e.message)
                             except Exception as e:
                                 logger.debug(f"Error processing accessor: {e}")
                         else:
@@ -131,6 +136,8 @@ class SpanHandler:
                                     event_attributes[attribute_key] = accessor(arguments)
                                 else:
                                     event_attributes.update(accessor(arguments))
+                            except MonocleSpanException as e:
+                                span.set_status(StatusCode.ERROR, e.message)
                             except Exception as e:
                                 logger.debug(f"Error evaluating accessor for attribute '{attribute_key}': {e}")
                     span.add_event(name=event_name, attributes=event_attributes)

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/instrumentation/common/utils.py

@@ -21,6 +21,21 @@ embedding_model_context = {}
 scope_id_generator = id_generator.RandomIdGenerator()
 http_scopes:dict[str:str] = {}
 
+class MonocleSpanException(Exception):
+    def __init__(self, err_message:str):
+        """
+        Monocle exeption to indicate error in span processing.        
+        Parameters:
+        - err_message (str): Error message.
+        - status (str): Status code
+        """
+        super().__init__(err_message)
+        self.message = err_message
+
+    def __str__(self):
+        """String representation of the exception."""
+        return f"[Monocle Span Error: {self.message} {self.status}"
+
 def set_tracer_provider(tracer_provider: TracerProvider):
     global monocle_tracer_provider
     monocle_tracer_provider = tracer_provider
@@ -252,35 +267,49 @@ async def http_async_route_handler(func, *args, **kwargs):
         headers = kwargs['req'].headers
     else:
         headers = None
-    return async_wrapper(func, None, headers, *args, **kwargs)
+    return async_wrapper(func, None, None, headers, *args, **kwargs)
 
-def run_async_with_scope(method, scope_name, headers, *args, **kwargs):
+def run_async_with_scope(method, current_context, exceptions, *args, **kwargs):
     token = None
-    if scope_name:
-        token = set_scope(scope_name)
-    elif headers:
-        token = extract_http_headers(headers)
     try:
+        if current_context:
+            token = attach(current_context)
         return asyncio.run(method(*args, **kwargs))
+    except Exception as e:
+        exceptions['exception'] = e
+        raise e
     finally:
         if token:
-            remove_scope(token)
+            detach(token)
 
-def async_wrapper(method, scope_name=None, headers=None, *args, **kwargs):
+def async_wrapper(method, scope_name=None, scope_value=None, headers=None, *args, **kwargs):
     try:
         run_loop = asyncio.get_running_loop()
     except RuntimeError:
         run_loop = None
 
-    if run_loop and run_loop.is_running():
-        results = []
-        thread = threading.Thread(target=lambda: results.append(run_async_with_scope(method, scope_name, headers, *args, **kwargs)))
-        thread.start()
-        thread.join()
-        return_value = results[0] if len(results) > 0 else None
-        return return_value
-    else:
-        return run_async_with_scope(method, scope_name, headers, *args, **kwargs)
+    token = None
+    exceptions = {}
+    if scope_name:
+        token = set_scope(scope_name, scope_value)
+    elif headers:
+        token = extract_http_headers(headers)
+    current_context = get_current()
+    try:
+        if run_loop and run_loop.is_running():
+            results = []
+            thread = threading.Thread(target=lambda: results.append(run_async_with_scope(method, current_context, exceptions, *args, **kwargs)))
+            thread.start()
+            thread.join()
+            if 'exception' in exceptions:
+                raise exceptions['exception']
+            return_value = results[0] if len(results) > 0 else None
+            return return_value
+        else:
+            return run_async_with_scope(method, None, exceptions, *args, **kwargs)
+    finally:
+        if token:
+            remove_scope(token)
 
 class Option(Generic[T]):
     def __init__(self, value: Optional[T]):

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/instrumentation/common/wrapper.py

@@ -31,10 +31,11 @@ def wrapper_processor(async_task: bool, tracer: Tracer, handler: SpanHandler, to
     try:
         handler.pre_tracing(to_wrap, wrapped, instance, args, kwargs)
         skip_scan:bool = to_wrap.get('skip_span') or handler.skip_span(to_wrap, wrapped, instance, args, kwargs)
-        token = SpanHandler.attach_workflow_type(to_wrap=to_wrap)
+        if not to_wrap.get('skip_span'):
+            token = SpanHandler.attach_workflow_type(to_wrap=to_wrap)
         if skip_scan:
             if async_task:
-                return_value = async_wrapper(wrapped, None, None, *args, **kwargs)
+                return_value = async_wrapper(wrapped, None, None, None, *args, **kwargs)
             else:
                 return_value = wrapped(*args, **kwargs)
         else:
@@ -58,7 +59,7 @@ def span_processor(name: str, async_task: bool, tracer: Tracer, handler: SpanHan
         else:
             handler.pre_task_processing(to_wrap, wrapped, instance, args, kwargs, span)
             if async_task:
-                return_value = async_wrapper(wrapped, None, None, *args, **kwargs)
+                return_value = async_wrapper(wrapped, None, None, None, *args, **kwargs)
             else:                    
                 return_value = wrapped(*args, **kwargs)
             handler.hydrate_span(to_wrap, wrapped, instance, args, kwargs, return_value, span)
@@ -86,5 +87,6 @@ def scope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instan
 @with_tracer_wrapper
 async def ascope_wrapper(tracer: Tracer, handler: SpanHandler, to_wrap, wrapped, instance, args, kwargs):
     scope_name = to_wrap.get('scope_name', None)
-    return_value = async_wrapper(wrapped, scope_name, None, *args, **kwargs)
+    scope_value = to_wrap.get('scope_value', None)
+    return_value = async_wrapper(wrapped, scope_name, scope_value, None, *args, **kwargs)
     return return_value

{monocle_apptrace-0.3.0b6 → monocle_apptrace-0.3.0b7}/src/monocle_apptrace/instrumentation/metamodel/botocore/handlers/botocore_span_handler.py

@@ -1,3 +1,4 @@
+from opentelemetry.context import get_value, set_value, attach, detach
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
 
 class BotoCoreSpanHandler(SpanHandler):
@@ -22,4 +23,4 @@ class BotoCoreSpanHandler(SpanHandler):
     def post_tracing(self, to_wrap, wrapped, instance, args, kwargs, return_value):
         self._botocore_processor(to_wrap=to_wrap, wrapped=wrapped, instance=instance, return_value=return_value, args=args,
                                  kwargs=kwargs)
-        return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)
+        return super().pre_tracing(to_wrap, wrapped, instance, args, kwargs)

monocle_apptrace-0.3.0b6/src/monocle_apptrace/exporters/exporter_processor.py

@@ -1,19 +0,0 @@
-from abc import ABC, abstractmethod
-import logging
-from typing import Callable
-
-logger = logging.getLogger(__name__)
-
-class ExportTaskProcessor(ABC):
-
-    @abstractmethod
-    def start(self):
-        return
-
-    @abstractmethod
-    def stop(self):
-        return
-
-    @abstractmethod
-    def queue_task(self, async_task: Callable[[Callable, any], any] = None, args: any = None):
-        return