monocle-apptrace 0.3.0b6__py3-none-any.whl → 0.3.1b1__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,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/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.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/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

@@ -13,6 +13,7 @@ GITHUB_CODESPACE_IDENTIFIER_ENV_NAME = "GITHUB_REPOSITORY"
 
 # Azure naming reference can be found here
 # https://learn.microsoft.com/en-us/azure/cloud-adoption-framework/ready/azure-best-practices/resource-abbreviations
+# https://docs.aws.amazon.com/resource-explorer/latest/userguide/supported-resource-types.html#services-lookoutmetrics
 AZURE_FUNCTION_NAME = "azure.func"
 AZURE_APP_SERVICE_NAME = "azure.asp"
 AZURE_ML_SERVICE_NAME = "azure.mlw"
@@ -48,6 +49,8 @@ llm_type_map = {
     "sagemakerllm": "aws_sagemaker",
     "chatbedrock": "aws_bedrock",
     "openaigenerator": "openai",
+    "bedrockruntime":"aws_bedrock",
+    "sagemakerruntime":"aws_sagemaker",
 }
 
 MONOCLE_INSTRUMENTOR = "monocle_apptrace"