monocle-apptrace 0.3.0b5__py3-none-any.whl → 0.3.0b7__py3-none-any.whl

monocle_apptrace/__init__.py

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

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/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/exporters/azure/blob_exporter.py

@@ -8,13 +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"
@@ -43,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)
@@ -72,6 +78,12 @@ class AzureBlobSpanExporter(SpanExporterBase):
         """The actual async export logic is run here."""
         # Add spans to the export queue
         for span in spans:
+            # Azure blob library has a check to generate it's own span if OpenTelemetry is loaded and Azure trace package is installed (just pip install azure-trace-opentelemetry)
+            # With Monocle,OpenTelemetry is always loaded. If the Azure trace package is installed, then it triggers the blob trace generation on every blob operation.
+            # Thus, the Monocle span write ends up generating a blob span which again comes back to the exporter .. and would result in an infinite loop.
+            # To avoid this, we check if the span has the Monocle SDK version attribute and skip it if it doesn't. That way the blob span genearted by Azure library are not exported.
+            if not span.attributes.get(MONOCLE_SDK_VERSION):
+                continue # TODO: All exporters to use same base class and check it there
             self.export_queue.append(span)
             if len(self.export_queue) >= self.max_batch_size:
                 await self.__export_spans()
@@ -104,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/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/exporters/exporter_processor.py

@@ -1,11 +1,18 @@
 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
@@ -15,5 +22,123 @@ class ExportTaskProcessor(ABC):
         return
 
     @abstractmethod
-    def queue_task(self, async_task: Callable[[Callable, any], any] = None, args: any = None):
-        return
+    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/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/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/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.warning("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/instrumentation/__init__.py

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

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/instrumentation/common/constants.py

@@ -37,6 +37,19 @@ service_name_map = {
     GITHUB_CODESPACE_SERVICE_NAME: GITHUB_CODESPACE_IDENTIFIER_ENV_NAME
 }
 
+
+llm_type_map = {
+    "sagemakerendpoint": "aws_sagemaker",
+    "azureopenai": "azure_openai",
+    "openai": "openai",
+    "chatopenai": "openai",
+    "azurechatopenai": "azure_openai",
+    "bedrock": "aws_bedrock",
+    "sagemakerllm": "aws_sagemaker",
+    "chatbedrock": "aws_bedrock",
+    "openaigenerator": "openai",
+}
+
 MONOCLE_INSTRUMENTOR = "monocle_apptrace"
 WORKFLOW_TYPE_KEY = "workflow_type"
 DATA_INPUT_KEY = "data.input"
@@ -52,3 +65,7 @@ MONOCLE_SCOPE_NAME_PREFIX = "monocle.scope."
 SCOPE_METHOD_LIST = 'MONOCLE_SCOPE_METHODS'
 SCOPE_METHOD_FILE = 'monocle_scopes.json'
 SCOPE_CONFIG_PATH = 'MONOCLE_SCOPE_CONFIG_PATH'
+TRACE_PROPOGATION_URLS = "MONOCLE_TRACE_PROPAGATATION_URLS"
+WORKFLOW_TYPE_KEY = "monocle.workflow_type"
+WORKFLOW_TYPE_GENERIC = "workflow.generic"
+MONOCLE_SDK_VERSION = "monocle_apptrace.version"