monocle-apptrace 0.3.0b2__py3-none-any.whl → 0.3.0b3__py3-none-any.whl

monocle_apptrace/exporters/aws/s3_exporter.py

@@ -25,7 +25,7 @@ class S3SpanExporter(SpanExporterBase):
         super().__init__()
         # Use environment variables if credentials are not provided
         DEFAULT_FILE_PREFIX = "monocle_trace_"
-        DEFAULT_TIME_FORMAT = "%Y-%m-%d_%H.%M.%S"
+        DEFAULT_TIME_FORMAT = "%Y-%m-%d__%H.%M.%S"
         self.max_batch_size = 500
         self.export_interval = 1
         self.s3_client = boto3.client(

monocle_apptrace/exporters/aws/s3_exporter_opendal.py

@@ -0,0 +1,126 @@
+import os
+import time
+import datetime
+import logging
+import asyncio
+from typing import Sequence
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+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):
+        super().__init__()
+        DEFAULT_FILE_PREFIX = "monocle_trace_"
+        DEFAULT_TIME_FORMAT = "%Y-%m-%d__%H.%M.%S"
+        self.max_batch_size = 500
+        self.export_interval = 1
+        self.file_prefix = DEFAULT_FILE_PREFIX
+        self.time_format = DEFAULT_TIME_FORMAT
+        self.export_queue = []
+        self.last_export_time = time.time()
+        self.bucket_name = bucket_name or os.getenv("MONOCLE_S3_BUCKET_NAME", "default-bucket")
+
+        # Initialize OpenDAL S3 operator
+        self.op = Operator(
+            "s3",
+            root = "/",
+            region=os.getenv("AWS_REGION", region_name),
+            bucket=self.bucket_name,
+            access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
+            secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
+        )
+
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Synchronous export method that internally handles async logic."""
+        try:
+            # Run the asynchronous export logic in an event loop
+            asyncio.run(self.__export_async(spans))
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Error exporting spans: {e}")
+            return SpanExportResult.FAILURE
+
+    async def __export_async(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        try:
+            # Add spans to the export queue
+            for span in spans:
+                self.export_queue.append(span)
+                if len(self.export_queue) >= self.max_batch_size:
+                    await self.__export_spans()
+
+            # Check if it's time to force a flush
+            current_time = time.time()
+            if current_time - self.last_export_time >= self.export_interval:
+                await self.__export_spans()
+                self.last_export_time = current_time
+
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Error exporting spans: {e}")
+            return SpanExportResult.FAILURE
+
+    def __serialize_spans(self, spans: Sequence[ReadableSpan]) -> str:
+        try:
+            # Serialize spans to JSON or any other format you prefer
+            valid_json_list = []
+            for span in spans:
+                try:
+                    valid_json_list.append(span.to_json(indent=0).replace("\n", ""))
+                except json.JSONDecodeError as e:
+                    logger.warning(f"Invalid JSON format in span data: {span.context.span_id}. Error: {e}")
+                    continue
+            return "\n".join(valid_json_list) + "\n"
+        except Exception as e:
+            logger.warning(f"Error serializing spans: {e}")
+
+    async def __export_spans(self):
+        if not self.export_queue:
+            return
+        # Take a batch of spans from the queue
+        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}")
+
+    @SpanExporterBase.retry_with_backoff(exceptions=(Unexpected))
+    def __upload_to_s3(self, span_data_batch: str):
+
+        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}.")
+
+        except PermissionDenied as e:
+            # S3 bucket is forbidden.
+            logger.error(f"Access to bucket {self.bucket_name} is forbidden (403).")
+            raise PermissionError(f"Access to bucket {self.bucket_name} is forbidden.")
+
+        except ConfigInvalid as e:
+            # Bucket does not exist.
+            if "404" in str(e):
+                logger.error("Bucket does not exist. Please check the bucket name and region.")
+                raise Exception(f"Bucket does not exist. Error: {e}")
+            else:
+                logger.error(f"Unexpected error when accessing bucket {self.bucket_name}: {e}")
+                raise e
+
+
+    async def force_flush(self, timeout_millis: int = 30000) -> bool:
+        await self.__export_spans()
+        return True
+
+    def shutdown(self) -> None:
+        logger.info("S3SpanExporter has been shut down.")

monocle_apptrace/exporters/azure/blob_exporter_opendal.py

@@ -0,0 +1,147 @@
+import os
+import time
+import datetime
+import logging
+import asyncio
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from typing import Sequence
+from opendal import Operator
+from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+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):
+        super().__init__()
+        DEFAULT_FILE_PREFIX = "monocle_trace_"
+        DEFAULT_TIME_FORMAT = "%Y-%m-%d_%H.%M.%S"
+        self.max_batch_size = 500
+        self.export_interval = 1
+        self.container_name = container_name
+
+        # Default values
+        self.file_prefix = DEFAULT_FILE_PREFIX
+        self.time_format = DEFAULT_TIME_FORMAT
+
+        # Validate input
+        if not connection_string:
+            connection_string = os.getenv('MONOCLE_BLOB_CONNECTION_STRING')
+            if not connection_string:
+                raise ValueError("Azure Storage connection string is not provided or set in environment variables.")
+
+        if not container_name:
+            container_name = os.getenv('MONOCLE_BLOB_CONTAINER_NAME', 'default-container')
+        endpoint, account_name , account_key = self.parse_connection_string(connection_string)
+
+        if not account_name or not account_key:
+            raise ValueError("AccountName or AccountKey missing in the connection string.")
+
+        try:
+            # Initialize OpenDAL operator with explicit credentials
+            self.operator = Operator(
+                "azblob",
+                endpoint=endpoint,
+                account_name=account_name,
+                account_key=account_key,
+                container=container_name
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize OpenDAL operator: {e}")
+
+
+    def parse_connection_string(self,connection_string):
+        connection_params = dict(item.split('=', 1) for item in connection_string.split(';') if '=' in item)
+
+        account_name = connection_params.get('AccountName')
+        account_key = connection_params.get('AccountKey')
+        endpoint_suffix = connection_params.get('EndpointSuffix')
+
+        if not all([account_name, account_key, endpoint_suffix]):
+            raise ValueError("Invalid connection string. Ensure it contains AccountName, AccountKey, and EndpointSuffix.")
+
+        endpoint = f"https://{account_name}.blob.{endpoint_suffix}"
+        return endpoint, account_name, account_key
+
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Synchronous export method that internally handles async logic."""
+        try:
+            # Run the asynchronous export logic in an event loop
+            asyncio.run(self._export_async(spans))
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Error exporting spans: {e}")
+            return SpanExportResult.FAILURE
+
+    async def _export_async(self, spans: Sequence[ReadableSpan]):
+        """The actual async export logic is run here."""
+        # Add spans to the export queue
+        for span in spans:
+            self.export_queue.append(span)
+            if len(self.export_queue) >= self.max_batch_size:
+                await self.__export_spans()
+
+        # Force a flush if the interval has passed
+        current_time = time.time()
+        if current_time - self.last_export_time >= self.export_interval:
+            await self.__export_spans()
+            self.last_export_time = current_time
+
+    def __serialize_spans(self, spans: Sequence[ReadableSpan]) -> str:
+        try:
+            valid_json_list = []
+            for span in spans:
+                try:
+                    valid_json_list.append(span.to_json(indent=0).replace("\n", ""))
+                except json.JSONDecodeError as e:
+                    logger.warning(f"Invalid JSON format in span data: {span.context.span_id}. Error: {e}")
+                    continue
+
+            ndjson_data = "\n".join(valid_json_list) + "\n"
+            return ndjson_data
+        except Exception as e:
+            logger.warning(f"Error serializing spans: {e}")
+
+    async def __export_spans(self):
+        if len(self.export_queue) == 0:
+            return
+
+        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}")
+
+    @SpanExporterBase.retry_with_backoff(exceptions=(Unexpected,))
+    def __upload_to_opendal(self, span_data_batch: str):
+        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}.")
+        except PermissionDenied as e:
+            # Azure Container is forbidden.
+            logger.error(f"Access to container {self.container_name} is forbidden (403).")
+            raise PermissionError(f"Access to container {self.container_name} is forbidden.")
+
+        except NotFound as e:
+            # Container does not exist.
+            if "404" in str(e):
+                logger.error("Container does not exist. Please check the container name.")
+                raise Exception(f"Container does not exist. Error: {e}")
+            else:
+                logger.error(f"Unexpected NotFound error when accessing container {self.container_name}: {e}")
+                raise e
+
+    async def force_flush(self, timeout_millis: int = 30000) -> bool:
+        await self.__export_spans()
+        return True
+
+    def shutdown(self) -> None:
+        logger.info("OpenDALAzureExporter has been shut down.")

monocle_apptrace/exporters/monocle_exporters.py

@@ -1,27 +1,45 @@
-from typing import Dict, Any
-import os, warnings
+from typing import Dict, Any, List
+import os
+import warnings
 from importlib import import_module
 from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
 from monocle_apptrace.exporters.file_exporter import FileSpanExporter
 
-monocle_exporters:Dict[str, Any] = {
+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 = []
+
+    for exporter_name in exporter_names:
+        exporter_name = exporter_name.strip()
+        try:
+            exporter_class_path = monocle_exporters[exporter_name]
+        except KeyError:
+            warnings.warn(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"])
+            exporters.append(exporter_class())
+        except Exception as ex:
+            warnings.warn(
+                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:
+        warnings.warn("No valid Monocle span exporters configured. Defaulting to FileSpanExporter.")
+        exporters.append(FileSpanExporter())
+
+    return exporters

monocle_apptrace/{constants.py → instrumentation/common/constants.py}

@@ -9,6 +9,8 @@ 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"
@@ -34,3 +36,14 @@ service_name_map = {
     AWS_LAMBDA_SERVICE_NAME: AWS_LAMBDA_FUNCTION_IDENTIFIER_ENV_NAME,
     GITHUB_CODESPACE_SERVICE_NAME: GITHUB_CODESPACE_IDENTIFIER_ENV_NAME
 }
+
+WORKFLOW_TYPE_KEY = "workflow_type"
+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_apptrace/instrumentation/common/instrumentor.py

@@ -0,0 +1,208 @@
+import logging
+from typing import Collection, Dict, List, Union
+import random
+import uuid
+from opentelemetry import trace
+from opentelemetry.context import attach, get_value, set_value, get_current, detach
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.trace import SpanContext
+from opentelemetry.sdk.trace import TracerProvider, Span, id_generator
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import Span, TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanProcessor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+from opentelemetry.trace.propagation import set_span_in_context
+from monocle_apptrace.exporters.monocle_exporters import get_monocle_exporter
+from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
+from monocle_apptrace.instrumentation.common.wrapper_method import (
+    DEFAULT_METHODS_LIST,
+    WrapperMethod,
+)
+
+logger = logging.getLogger(__name__)
+
+SESSION_PROPERTIES_KEY = "session"
+
+_instruments = ()
+
+monocle_tracer_provider: TracerProvider = None
+
+MONOCLE_INSTRUMENTOR = "monocle_apptrace"
+
+class MonocleInstrumentor(BaseInstrumentor):
+    workflow_name: str = ""
+    user_wrapper_methods: list[Union[dict,WrapperMethod]] = []
+    instrumented_method_list: list[object] = []
+    handlers:Dict[str,SpanHandler] = {} # dict of handlers
+    union_with_default_methods: bool = False
+
+    def __init__(
+            self,
+            handlers,
+            user_wrapper_methods: list[Union[dict,WrapperMethod]] = None,
+            union_with_default_methods: bool = True
+            ) -> None:
+        self.user_wrapper_methods = user_wrapper_methods or []
+        self.handlers = handlers or {'default':SpanHandler()}
+        self.union_with_default_methods = union_with_default_methods
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        tracer_provider: TracerProvider = kwargs.get("tracer_provider")
+        global monocle_tracer_provider
+        monocle_tracer_provider = tracer_provider
+        tracer = get_tracer(instrumenting_module_name=MONOCLE_INSTRUMENTOR, tracer_provider=tracer_provider)
+
+        final_method_list = []
+        if self.union_with_default_methods is True:
+            final_method_list= final_method_list + DEFAULT_METHODS_LIST
+
+        for method in self.user_wrapper_methods:
+            if isinstance(method, dict):
+                final_method_list.append(method)
+            elif isinstance(method, WrapperMethod):
+                final_method_list.append(method.to_dict())
+        
+        for method_config in final_method_list:
+            target_package = method_config.get("package", None)
+            target_object = method_config.get("object", None)
+            target_method = method_config.get("method", None)
+            wrapped_by = method_config.get("wrapper_method", None)
+            #get the requisite handler or default one
+            handler_key = method_config.get("span_handler",'default')
+            try:
+                handler =  self.handlers.get(handler_key)
+                wrap_function_wrapper(
+                    target_package,
+                    f"{target_object}.{target_method}" if target_object else target_method,
+                    wrapped_by(tracer, handler, method_config),
+                )
+                self.instrumented_method_list.append(method_config)
+            except ModuleNotFoundError as e:
+                logger.debug(f"ignoring module {e.name}")
+
+            except Exception as ex:
+                logger.error(f"""_instrument wrap exception: {str(ex)}
+                            for package: {target_package},
+                            object:{target_object},
+                            method:{target_method}""")
+
+    def _uninstrument(self, **kwargs):
+        for wrapped_method in self.instrumented_method_list:
+            try:
+                wrap_package = wrapped_method.get("package")
+                wrap_object = wrapped_method.get("object")
+                wrap_method = wrapped_method.get("method")
+                unwrap(
+                    f"{wrap_package}.{wrap_object}" if wrap_object else wrap_package,
+                    wrap_method,
+                )
+            except Exception as ex:
+                logger.error(f"""_instrument unwrap exception: {str(ex)}
+                             for package: {wrap_package},
+                             object:{wrap_object},
+                             method:{wrap_method}""")
+
+def setup_monocle_telemetry(
+        workflow_name: str,
+        span_processors: List[SpanProcessor] = None,
+        span_handlers: Dict[str,SpanHandler] = None,
+        wrapper_methods: List[Union[dict,WrapperMethod]] = None,
+        union_with_default_methods: bool = True) -> None:
+    resource = Resource(attributes={
+        SERVICE_NAME: workflow_name
+    })
+    exporters = get_monocle_exporter()
+    span_processors = span_processors or [BatchSpanProcessor(exporter) for exporter in exporters]
+    trace_provider = TracerProvider(resource=resource)
+    attach(set_value("workflow_name", workflow_name))
+    tracer_provider_default = trace.get_tracer_provider()
+    provider_type = type(tracer_provider_default).__name__
+    is_proxy_provider = "Proxy" in provider_type
+    for processor in span_processors:
+        processor.on_start = on_processor_start
+        if not is_proxy_provider:
+            tracer_provider_default.add_span_processor(processor)
+        else:
+            trace_provider.add_span_processor(processor)
+    if is_proxy_provider:
+        trace.set_tracer_provider(trace_provider)
+    instrumentor = MonocleInstrumentor(user_wrapper_methods=wrapper_methods or [], 
+                                       handlers=span_handlers, union_with_default_methods = union_with_default_methods)
+    # instrumentor.app_name = workflow_name
+    if not instrumentor.is_instrumented_by_opentelemetry:
+        instrumentor.instrument(trace_provider=trace_provider)
+
+    return instrumentor
+
+def on_processor_start(span: Span, parent_context):
+    context_properties = get_value(SESSION_PROPERTIES_KEY)
+    if context_properties is not None:
+        for key, value in context_properties.items():
+            span.set_attribute(
+                f"{SESSION_PROPERTIES_KEY}.{key}", value
+            )
+
+def set_context_properties(properties: dict) -> None:
+    attach(set_value(SESSION_PROPERTIES_KEY, properties))
+
+
+def propagate_trace_id(traceId = "", use_trace_context = False):
+    try:
+        if traceId.startswith("0x"):
+            traceId = traceId.lstrip("0x")
+        tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= monocle_tracer_provider)
+        initial_id_generator = tracer.id_generator
+        _parent_span_context = get_current() if use_trace_context else None
+        if traceId and is_valid_trace_id_uuid(traceId):
+            tracer.id_generator = FixedIdGenerator(uuid.UUID(traceId).int)
+
+        span = tracer.start_span(name = "parent_placeholder_span", context= _parent_span_context)
+        updated_span_context = set_span_in_context(span=span, context= _parent_span_context)
+        updated_span_context = set_value("root_span_id", span.get_span_context().span_id, updated_span_context)
+        token = attach(updated_span_context)
+
+        span.end()
+        tracer.id_generator = initial_id_generator
+        return token
+    except:
+        logger.warning("Failed to propagate trace id")
+    return
+
+
+def propagate_trace_id_from_traceparent():
+    propagate_trace_id(use_trace_context = True)
+
+
+def stop_propagate_trace_id(token) -> None:
+    try:
+        detach(token)
+    except:
+        logger.warning("Failed to stop propagating trace id")
+
+
+def is_valid_trace_id_uuid(traceId: str) -> bool:
+    try:
+        uuid.UUID(traceId)
+        return True
+    except:
+        pass
+    return False
+
+
+class FixedIdGenerator(id_generator.IdGenerator):
+    def __init__(
+            self,
+            trace_id: int) -> None:
+        self.trace_id = trace_id
+
+    def generate_span_id(self) -> int:
+        return random.getrandbits(64)
+
+    def generate_trace_id(self) -> int:
+        return self.trace_id