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

monocle_apptrace/exporters/base_exporter.py

@@ -0,0 +1,48 @@
+import time
+import random
+import logging
+from abc import ABC, abstractmethod
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from typing import Sequence
+import asyncio
+
+logger = logging.getLogger(__name__)
+
+class SpanExporterBase(ABC):
+    def __init__(self):
+        self.backoff_factor = 2
+        self.max_retries = 10
+        self.export_queue = []
+        self.last_export_time = time.time()
+
+    @abstractmethod
+    async def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        pass
+
+    @abstractmethod
+    async def force_flush(self, timeout_millis: int = 30000) -> bool:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    @staticmethod
+    def retry_with_backoff(retries=3, backoff_in_seconds=1, max_backoff_in_seconds=32, exceptions=(Exception,)):
+        def decorator(func):
+            def wrapper(*args, **kwargs):
+                attempt = 0
+                while attempt < retries:
+                    try:
+                        return func(*args, **kwargs)
+                    except exceptions as e:
+                        attempt += 1
+                        sleep_time = min(max_backoff_in_seconds, backoff_in_seconds * (2 ** (attempt - 1)))
+                        sleep_time = sleep_time * (1 + random.uniform(-0.1, 0.1))  # Add jitter
+                        logger.warning(f"Network connectivity error, Attempt {attempt} failed: {e}. Retrying in {sleep_time:.2f} seconds...")
+                        time.sleep(sleep_time)
+                raise Exception(f"Failed after {retries} attempts")
+
+            return wrapper
+
+        return decorator

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/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

@@ -0,0 +1,55 @@
+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
+
+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"},
+    "memory": {"module": "opentelemetry.sdk.trace.export.in_memory_span_exporter", "class": "InMemorySpanExporter"},
+    "console": {"module": "opentelemetry.sdk.trace.export", "class": "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

@@ -0,0 +1,117 @@
+import json
+import logging
+import os
+from typing import Callable, Optional, Sequence
+import requests
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult, ConsoleSpanExporter
+from requests.exceptions import ReadTimeout
+
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
+
+REQUESTS_SUCCESS_STATUS_CODES = (200, 202)
+OKAHU_PROD_INGEST_ENDPOINT = "https://ingest.okahu.co/api/v1/trace/ingest"
+
+logger = logging.getLogger(__name__)
+
+
+class OkahuSpanExporter(SpanExporter):
+    def __init__(
+            self,
+            endpoint: Optional[str] = None,
+            timeout: Optional[int] = None,
+            session: Optional[requests.Session] = None,
+            task_processor: ExportTaskProcessor = None
+    ):
+        """Okahu exporter."""
+        okahu_endpoint: str = os.environ.get("OKAHU_INGESTION_ENDPOINT", OKAHU_PROD_INGEST_ENDPOINT)
+        self.endpoint = endpoint or okahu_endpoint
+        api_key: str = os.environ.get("OKAHU_API_KEY")
+        self._closed = False
+        if not api_key:
+            raise ValueError("OKAHU_API_KEY not set.")
+        self.timeout = timeout or 15
+        self.session = session or requests.Session()
+        self.session.headers.update(
+            {"Content-Type": "application/json", "x-api-key": api_key}
+        )
+
+        self.task_processor = task_processor or None
+        if task_processor is not None:
+            task_processor.start()
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        # After the call to Shutdown subsequent calls to Export are
+        # not allowed and should return a Failure result
+        if not hasattr(self, 'session'):
+            return self.exporter.export(spans)
+
+        if self._closed:
+            logger.warning("Exporter already shutdown, ignoring batch")
+            return SpanExportResult.FAILURE
+        if len(spans) == 0:
+            return
+
+        span_list = {
+            "batch": []
+        }
+
+        # append the batch object with all the spans object
+        for span in spans:
+            # create a object from serialized span
+            obj = json.loads(span.to_json())
+            if obj["parent_id"] is None:
+                obj["parent_id"] = "None"
+            else:
+                obj["parent_id"] = remove_0x_from_start(obj["parent_id"])
+            if obj["context"] is not None:
+                obj["context"]["trace_id"] = remove_0x_from_start(obj["context"]["trace_id"])
+                obj["context"]["span_id"] = remove_0x_from_start(obj["context"]["span_id"])
+            span_list["batch"].append(obj)
+
+        # 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,
+                    data=json.dumps(span_list_local),
+                    timeout=self.timeout,
+                )
+                if result.status_code not in REQUESTS_SUCCESS_STATUS_CODES:
+                    logger.error(
+                        "Traces cannot be uploaded; status code: %s, message %s",
+                        result.status_code,
+                        result.text,
+                    )
+                    return SpanExportResult.FAILURE
+                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, is_root_span)
+            return SpanExportResult.SUCCESS
+        return send_spans_to_okahu(span_list, is_root_span)
+
+    def shutdown(self) -> None:
+        if self._closed:
+            logger.warning("Exporter already shutdown, ignoring call")
+            return
+        if hasattr(self, 'session'):
+            self.session.close()
+        self._closed = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+
+# only removes the first occurrence of 0x from the string
+def remove_0x_from_start(my_str: str):
+    if my_str.startswith("0x"):
+        return my_str.replace("0x", "", 1)
+    return my_str

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"