monocle-apptrace 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

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

@@ -1,27 +1,55 @@
-from typing import Dict, Any
-import os, warnings
+from typing import Dict, Any, List
+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
 
-monocle_exporters:Dict[str, Any] = {
+logger = logging.getLogger(__name__)
+
+monocle_exporters: Dict[str, Any] = {
     "s3": {"module": "monocle_apptrace.exporters.aws.s3_exporter", "class": "S3SpanExporter"},
-    "blob": {"module":"monocle_apptrace.exporters.azure.blob_exporter", "class": "AzureBlobSpanExporter"},
-    "okahu": {"module":"monocle_apptrace.exporters.okahu.okahu_exporter", "class": "OkahuSpanExporter"},
-    "file": {"module":"monocle_apptrace.exporters.file_exporter", "class": "FileSpanExporter"}
+    "blob": {"module": "monocle_apptrace.exporters.azure.blob_exporter", "class": "AzureBlobSpanExporter"},
+    "okahu": {"module": "monocle_apptrace.exporters.okahu.okahu_exporter", "class": "OkahuSpanExporter"},
+    "file": {"module": "monocle_apptrace.exporters.file_exporter", "class": "FileSpanExporter"},
+    "memory": {"module": "opentelemetry.sdk.trace.export.in_memory_span_exporter", "class": "InMemorySpanExporter"},
+    "console": {"module": "opentelemetry.sdk.trace.export", "class": "ConsoleSpanExporter"}
 }
 
-def get_monocle_exporter() -> SpanExporter:
-    exporter_name = os.environ.get("MONOCLE_EXPORTER", "file")
-    try:
-        exporter_class_path  = monocle_exporters[exporter_name]
-    except Exception as ex:
-        warnings.warn(f"Unsupported Monocle span exporter setting {exporter_name}, using default FileSpanExporter.")
-        return FileSpanExporter()
-    try:
-        exporter_module = import_module(exporter_class_path.get("module"))
-        exporter_class = getattr(exporter_module, exporter_class_path.get("class"))
-        return exporter_class()
-    except Exception as ex:
-        warnings.warn(f"Unable to set Monocle span exporter to {exporter_name}, error {ex}. Using ConsoleSpanExporter")
-        return ConsoleSpanExporter()
+
+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()
+        try:
+            exporter_class_path = monocle_exporters[exporter_name]
+        except KeyError:
+            logger.debug(f"Unsupported Monocle span exporter '{exporter_name}', skipping.")
+            continue
+        try:
+            exporter_module = import_module(exporter_class_path["module"])
+            exporter_class = getattr(exporter_module, exporter_class_path["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.")
+            exporters.append(ConsoleSpanExporter())
+            continue
+
+    # If no exporters were created, default to FileSpanExporter
+    if not exporters:
+        logger.debug("No valid Monocle span exporters configured. Defaulting to FileSpanExporter.")
+        exporters.append(FileSpanExporter())
+
+    return exporters

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

@@ -0,0 +1,70 @@
+# Azure environment constants
+AZURE_ML_ENDPOINT_ENV_NAME = "AZUREML_ENTRY_SCRIPT"
+AZURE_FUNCTION_WORKER_ENV_NAME = "FUNCTIONS_WORKER_RUNTIME"
+AZURE_APP_SERVICE_ENV_NAME = "WEBSITE_SITE_NAME"
+AWS_LAMBDA_ENV_NAME = "AWS_LAMBDA_RUNTIME_API"
+GITHUB_CODESPACE_ENV_NAME = "CODESPACES"
+
+AWS_LAMBDA_FUNCTION_IDENTIFIER_ENV_NAME = "AWS_LAMBDA_FUNCTION_NAME"
+AZURE_FUNCTION_IDENTIFIER_ENV_NAME = "WEBSITE_SITE_NAME"
+AZURE_APP_SERVICE_IDENTIFIER_ENV_NAME = "WEBSITE_DEPLOYMENT_ID"
+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
+AZURE_FUNCTION_NAME = "azure.func"
+AZURE_APP_SERVICE_NAME = "azure.asp"
+AZURE_ML_SERVICE_NAME = "azure.mlw"
+AWS_LAMBDA_SERVICE_NAME = "aws.lambda"
+GITHUB_CODESPACE_SERVICE_NAME = "github_codespace"
+
+# Env variables to identify infra service type
+service_type_map = {
+    AZURE_ML_ENDPOINT_ENV_NAME: AZURE_ML_SERVICE_NAME,
+    AZURE_APP_SERVICE_ENV_NAME: AZURE_APP_SERVICE_NAME,
+    AZURE_FUNCTION_WORKER_ENV_NAME: AZURE_FUNCTION_NAME,
+    AWS_LAMBDA_ENV_NAME: AWS_LAMBDA_SERVICE_NAME,
+    GITHUB_CODESPACE_ENV_NAME: GITHUB_CODESPACE_SERVICE_NAME
+}
+
+# Env variables to identify infra service name
+service_name_map = {
+    AZURE_APP_SERVICE_NAME: AZURE_APP_SERVICE_IDENTIFIER_ENV_NAME,
+    AZURE_FUNCTION_NAME: AZURE_FUNCTION_IDENTIFIER_ENV_NAME,
+    AZURE_ML_SERVICE_NAME: AZURE_ML_ENDPOINT_ENV_NAME,
+    AWS_LAMBDA_SERVICE_NAME: AWS_LAMBDA_FUNCTION_IDENTIFIER_ENV_NAME,
+    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"
+DATA_INPUT_KEY = "data.input"
+DATA_OUTPUT_KEY = "data.output"
+PROMPT_INPUT_KEY = "data.input"
+PROMPT_OUTPUT_KEY = "data.output"
+QUERY = "input"
+RESPONSE = "response"
+SESSION_PROPERTIES_KEY = "session"
+INFRA_SERVICE_KEY = "infra_service_name"
+META_DATA = 'metadata'
+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"