mantisdk 0.1.0__py3-none-any.whl

mantisdk/emitter/annotation.py

@@ -0,0 +1,370 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helpers for emitting annotation/operation spans."""
+
+import asyncio
+import functools
+import inspect
+import logging
+from types import TracebackType
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from mantisdk.semconv import AGL_ANNOTATION, AGL_OPERATION, LightningSpanAttributes
+from mantisdk.tracer.base import get_active_tracer
+from mantisdk.tracer.dummy import DummyTracer
+from mantisdk.types import SpanCoreFields, SpanRecordingContext, TraceStatus
+from mantisdk.utils.otel import check_attributes_sanity, flatten_attributes, sanitize_attributes
+
+_FnType = TypeVar("_FnType", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def emit_annotation(annotation: Dict[str, Any], propagate: bool = True) -> SpanCoreFields:
+    """Emit a new annotation span.
+
+    This is the underlying implementation of [`emit_reward`][mantisdk.emit_reward].
+
+    Annotation spans are used to annotate a specific event or a part of rollout.
+    See [semconv][mantisdk.semconv] for conventional annotation keys in Mantisdk.
+
+    If annotations contain nested dicts, they will be flattened before emitting.
+    Complex objects will lead to emitting failures.
+
+    Args:
+        annotation: Dictionary containing annotation key-value pairs.
+            Representatives are rewards, tags, and metadata.
+        propagate: Whether to propagate the span to tracers automatically.
+    """
+    annotation_attributes = flatten_attributes(annotation, expand_leaf_lists=False)
+    check_attributes_sanity(annotation_attributes)
+    sanitized_attributes = sanitize_attributes(annotation_attributes)
+    logger.debug("Emitting annotation span with keys %s", sanitized_attributes.keys())
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit annotation span.")
+    else:
+        tracer = DummyTracer()
+
+    return tracer.create_span(
+        name=AGL_ANNOTATION,
+        attributes=sanitized_attributes,
+        status=TraceStatus(status_code="OK"),
+    )
+
+
+class OperationContext:
+    """Context manager and decorator for tracing operations.
+
+    This class manages a tracer-backed span for a logical unit of work. It can be
+    used either:
+
+    * As a decorator, in which case inputs and outputs are inferred
+      automatically from the wrapped function's signature.
+    * As a context manager, in which case inputs and outputs can be recorded
+      explicitly via [`set_input`][mantisdk.emitter.annotation.OperationContext.set_input]
+      and [`set_output`][mantisdk.emitter.annotation.OperationContext.set_output].
+
+    Attributes:
+        name: Human-readable span name.
+        initial_attributes: Attributes applied when the span is created.
+        tracer: Tracer implementation used to create spans.
+    """
+
+    def __init__(self, name: str, attributes: Dict[str, Any], propagate: bool = True) -> None:
+        """Initialize a new operation context.
+
+        Args:
+            name: Human-readable name of the span.
+            attributes: Initial attributes attached to the span. Values are
+                JSON-serialized where necessary.
+            propagate: Whether the span should be sent to active exporters.
+        """
+        self.name = name
+        self.initial_attributes = flatten_attributes(attributes, expand_leaf_lists=False)
+        self.propagate = propagate
+        if propagate:
+            tracer = get_active_tracer()
+            if tracer is None:
+                raise RuntimeError("No active tracer found. Cannot trace operation spans.")
+            self.tracer = tracer
+        else:
+            self.tracer = DummyTracer()
+        self._ctx_manager: Optional[ContextManager[SpanRecordingContext]] = None
+        self._recording_context: Optional[SpanRecordingContext] = None
+        self._span: Optional[SpanCoreFields] = None
+
+    def __enter__(self) -> "OperationContext":
+        """Enter the context manager and start a new span.
+
+        Returns:
+            The current :class:`OperationContext` instance with an active span.
+        """
+        sanitized_attrs = sanitize_attributes(self.initial_attributes)
+        self._ctx_manager = self.tracer.operation_context(self.name, attributes=sanitized_attrs)
+        recording_context = self._ctx_manager.__enter__()
+        self._recording_context = recording_context
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Exit the context manager and finish the span."""
+        if self._ctx_manager:
+            self._ctx_manager.__exit__(exc_type, exc_val, exc_tb)
+        if self._recording_context:
+            self._span = self._recording_context.get_recorded_span()
+        self._ctx_manager = None
+        self._recording_context = None
+
+    def span(self) -> SpanCoreFields:
+        """Get the span that was created by this context manager."""
+        if self._span is None:
+            raise RuntimeError("Span is not ready yet.")
+        return self._span
+
+    def set_input(self, *args: Any, **kwargs: Any) -> None:
+        """Record input arguments on the current span.
+
+        Positional arguments are stored under the `input.args.<index>` attributes,
+        and keyword arguments are stored under `input.<name>` attributes.
+
+        This is intended for use inside a `with operation(...) as op` block.
+
+        Args:
+            *args: Positional arguments to record.
+            **kwargs: Keyword arguments to record.
+        """
+        if not self._recording_context:
+            raise RuntimeError("No recording context found. Cannot set input.")
+
+        prefix = LightningSpanAttributes.OPERATION_INPUT.value
+        attributes: Dict[str, Any] = {}
+        if args:
+            for idx, value in enumerate(args):
+                flattened = flatten_attributes({str(idx): value})
+                for nested_key, nested_value in flattened.items():
+                    attributes[f"{prefix}.args.{nested_key}"] = nested_value
+        if kwargs:
+            for key, value in kwargs.items():
+                flattened = flatten_attributes({key: value})
+                for nested_key, nested_value in flattened.items():
+                    attributes[f"{prefix}.{nested_key}"] = nested_value
+        if attributes:
+            self._recording_context.record_attributes(sanitize_attributes(attributes))
+
+    def set_output(self, output: Any) -> None:
+        """Record the output value on the current span.
+
+        This is intended for use inside a `with operation(...) as op` block.
+
+        Args:
+            output: The output value to record.
+        """
+        if not self._recording_context:
+            raise RuntimeError("No recording context found. Cannot set output.")
+
+        flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: output})
+        self._recording_context.record_attributes(sanitize_attributes(flattened))
+
+    def __call__(self, fn: _FnType) -> _FnType:
+        """Wrap a callable so its execution is traced in a span.
+
+        When used as a decorator, a new span is created for each call to
+        the wrapped function. The bound arguments are recorded as input
+        attributes, the return value is recorded as an output attribute,
+        and any exception is recorded and marks the span as an error.
+
+        Args:
+            fn: The function or coroutine function to wrap.
+
+        Returns:
+            The wrapped callable.
+        """
+        function_name = fn.__name__
+
+        sig = inspect.signature(fn)
+
+        sanitized_init_attrs = sanitize_attributes(
+            {LightningSpanAttributes.OPERATION_NAME.value: function_name, **self.initial_attributes}
+        )
+
+        def _record_auto_inputs(
+            recording_ctx: SpanRecordingContext, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+        ) -> None:
+            """Bind arguments to signature and log them on the span."""
+            attributes: Dict[str, Any] = {}
+            try:
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+                for name, value in bound.arguments.items():
+                    parameter = sig.parameters.get(name)
+                    if parameter and parameter.kind is inspect.Parameter.VAR_POSITIONAL:
+                        attr_prefix = f"{LightningSpanAttributes.OPERATION_INPUT.value}.{name}"
+                        for idx, item in enumerate(value):
+                            flattened = flatten_attributes({str(idx): item})
+                            for nested_key, nested_value in flattened.items():
+                                attributes[f"{attr_prefix}.{nested_key}"] = nested_value
+                    else:
+                        flattened = flatten_attributes({name: value})
+                        for nested_key, nested_value in flattened.items():
+                            attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
+            except Exception:
+                if args:
+                    for idx, value in enumerate(args):
+                        flattened = flatten_attributes({str(idx): value})
+                        for nested_key, nested_value in flattened.items():
+                            attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.args.{nested_key}"] = (
+                                nested_value
+                            )
+                if kwargs:
+                    flattened = flatten_attributes({"kwargs": kwargs})
+                    for nested_key, nested_value in flattened.items():
+                        attributes[f"{LightningSpanAttributes.OPERATION_INPUT.value}.{nested_key}"] = nested_value
+            if attributes:
+                recording_ctx.record_attributes(sanitize_attributes(attributes))
+
+        def _record_auto_outputs(recording_ctx: SpanRecordingContext, result: Any) -> None:
+            """Record the output value on the span."""
+            flattened = flatten_attributes({LightningSpanAttributes.OPERATION_OUTPUT.value: result})
+            recording_ctx.record_attributes(sanitize_attributes(flattened))
+
+        if asyncio.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                """Async wrapper that traces the wrapped coroutine."""
+                with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
+                    _record_auto_inputs(recording_ctx, args, kwargs)
+                    result = await fn(*args, **kwargs)
+                    _record_auto_outputs(recording_ctx, result)
+                    return result
+
+            return cast(_FnType, async_wrapper)
+
+        else:
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                """Sync wrapper that traces the wrapped callable."""
+                with self.tracer.operation_context(self.name, attributes=sanitized_init_attrs) as recording_ctx:
+                    _record_auto_inputs(recording_ctx, args, kwargs)
+                    result = fn(*args, **kwargs)
+                    _record_auto_outputs(recording_ctx, result)
+                    return result
+
+            return cast(_FnType, sync_wrapper)
+
+
+@overload
+def operation(
+    fn: _FnType, *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
+) -> _FnType: ...
+
+
+@overload
+def operation(
+    *, propagate: bool = True, name: Optional[str] = None, **additional_attributes: Any
+) -> OperationContext: ...
+
+
+@overload
+def operation(fn: _FnType, *, name: Optional[str] = None, **additional_attributes: Any) -> _FnType: ...
+
+
+@overload
+def operation(*, name: Optional[str] = None, **additional_attributes: Any) -> OperationContext: ...
+
+
+@overload
+def operation(fn: _FnType, **additional_attributes: Any) -> _FnType: ...
+
+
+@overload
+def operation(**additional_attributes: Any) -> OperationContext: ...
+
+
+def operation(
+    fn: Optional[_FnType] = None,
+    *,
+    propagate: bool = True,
+    name: Optional[str] = None,
+    **additional_attributes: Any,
+) -> Union[_FnType, OperationContext]:
+    """Entry point for tracking operations.
+
+    This helper can be used either as a decorator or as a context manager.
+    The span name is fixed to [`AGL_OPERATION`][mantisdk.semconv.AGL_OPERATION];
+    custom span names are not supported. Any keyword arguments are recorded as span attributes.
+
+    Usage as a decorator:
+
+    ```python
+    @operation
+    def func(...):
+        ...
+
+    @operation(category="compute")
+    def func(...):
+        ...
+    ```
+
+    Usage as a context manager:
+
+    ```python
+    with operation(user_id=123) as op:
+        op.set_input(data=data)
+        # ... do work ...
+        op.set_output(result)
+    ```
+
+    Args:
+        fn: When used as `@operation`, this is the wrapped function.
+            When used as `operation(**attrs)`, this should be omitted (or
+            left as `None`) and only keyword attributes are provided.
+        propagate: Whether spans should use the active span processor. When False,
+            spans will stay local and not be exported.
+        name: Optional alias that populates
+            [`LightningSpanAttributes.OPERATION_NAME`][mantisdk.semconv.LightningSpanAttributes.OPERATION_NAME]
+            when `additional_attributes` does not already define it.
+        **additional_attributes: Additional span attributes to attach at
+            creation time.
+
+    Returns:
+        Either a wrapped callable (when used as a decorator) or an
+        [`OperationContext`][mantisdk.emitter.annotation.OperationContext]
+        (when used as a context manager factory).
+    """
+
+    if name is not None:
+        if LightningSpanAttributes.OPERATION_NAME.value in additional_attributes:
+            raise ValueError("Cannot specify both `name` and `additional_attributes.operation_name`.")
+        additional_attributes[LightningSpanAttributes.OPERATION_NAME.value] = name
+
+    # Case 1: Used as @operation (bare decorator or with attributes)
+    if callable(fn):
+        # Create context with fixed name, then immediately wrap the function
+        return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)(fn)
+
+    # Case 2: Used as operation(...) / with operation(...)
+    # Custom span names are intentionally not supported; use AGL_OPERATION.
+    if fn is not None:
+        raise ValueError("Custom span names are intentionally not supported when used as a context manager.")
+    return OperationContext(AGL_OPERATION, additional_attributes, propagate=propagate)

mantisdk/emitter/exception.py

@@ -0,0 +1,54 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import Any, Dict, Optional
+
+from mantisdk.semconv import AGL_EXCEPTION
+from mantisdk.tracer.base import get_active_tracer
+from mantisdk.tracer.dummy import DummyTracer
+from mantisdk.types import TraceStatus
+from mantisdk.utils.otel import flatten_attributes, format_exception_attributes, sanitize_attributes
+
+logger = logging.getLogger(__name__)
+
+
+def emit_exception(
+    exception: BaseException, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True
+) -> None:
+    """Record an exception with OpenTelemetry metadata.
+
+    Classic OpenTelemetry records exceptions in a dedicated logging service.
+    We simplify the model and use trace spans to record exceptions as well.
+
+    Args:
+        exception: Raised exception instance to serialize into telemetry attributes.
+        attributes: Additional attributes to attach to the exception span.
+        propagate: Whether to propagate the span to exporters automatically.
+
+    !!! note
+
+        The helper validates its input. If a non-exception value is provided,
+        a TypeError is raised to indicate a programming mistake.
+    """
+    if not isinstance(exception, BaseException):  # type: ignore
+        raise TypeError(f"Expected a BaseException instance, got: {type(exception)}.")
+    span_attributes = format_exception_attributes(exception)
+
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
+
+    logger.debug("Emitting exception span for %s", type(exception).__name__)
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit exception span.")
+    else:
+        tracer = DummyTracer()
+    tracer.create_span(
+        AGL_EXCEPTION,
+        attributes=span_attributes,
+        # The exception span is successful by itself.
+        status=TraceStatus(status_code="OK"),
+    )

mantisdk/emitter/message.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import Any, Dict, Optional
+
+from mantisdk.semconv import AGL_MESSAGE, LightningSpanAttributes
+from mantisdk.tracer.base import get_active_tracer
+from mantisdk.tracer.dummy import DummyTracer
+from mantisdk.types import Attributes, SpanLike
+from mantisdk.utils.otel import flatten_attributes, sanitize_attributes
+
+logger = logging.getLogger(__name__)
+
+
+def emit_message(message: str, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> None:
+    """Emit a textual message as an OpenTelemetry span.
+
+    Commonly used for sending debugging and logging messages.
+
+    Args:
+        message: Human readable message to attach as a span attribute.
+        attributes: Additional attributes to attach to the message span.
+        propagate: Whether to propagate the span to exporters automatically.
+
+    !!! note
+        OpenTelemetry distinguishes between logs and spans. Emitting the message as a
+        span keeps all Mantisdk telemetry in a single data store for analysis.
+    """
+    if not isinstance(message, str):  # type: ignore
+        raise TypeError(f"Message must be a string or list of strings, got: {type(message)}.")
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit message span.")
+    else:
+        tracer = DummyTracer()
+    span_attributes: Attributes = {LightningSpanAttributes.MESSAGE_BODY.value: message}
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
+    logger.debug("Emitting message span with message: %s", message)
+    tracer.create_span(
+        AGL_MESSAGE,
+        attributes=span_attributes,
+    )
+
+
+def get_message_value(span: SpanLike) -> Optional[str]:
+    """Extract the message string from a message span.
+
+    Args:
+        span: Span-like object to extract the message from.
+    """
+    span_attributes = span.attributes or {}
+    if LightningSpanAttributes.MESSAGE_BODY.value not in span_attributes:
+        return None
+    message = span_attributes[LightningSpanAttributes.MESSAGE_BODY.value]
+    if isinstance(message, str):
+        return message
+    raise TypeError(f"Message must be a string, got: {type(message)}.")

mantisdk/emitter/object.py

@@ -0,0 +1,117 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import base64
+import json
+import logging
+from typing import Any, Dict, Optional
+
+from mantisdk.semconv import AGL_OBJECT, LightningSpanAttributes
+from mantisdk.tracer.base import get_active_tracer
+from mantisdk.tracer.dummy import DummyTracer
+from mantisdk.types import SpanCoreFields, SpanLike, TraceStatus
+from mantisdk.utils.otel import flatten_attributes, full_qualified_name, sanitize_attributes
+
+logger = logging.getLogger(__name__)
+
+
+def emit_object(object: Any, attributes: Optional[Dict[str, Any]] = None, propagate: bool = True) -> SpanCoreFields:
+    """Emit an object's serialized representation as an OpenTelemetry span.
+
+    Args:
+        object: Data structure to encode as JSON and attach to the span payload.
+        attributes: Additional attributes to attach to the object span.
+        propagate: Whether to propagate the span to exporters automatically.
+
+    !!! note
+        The payload must be JSON serializable. Non-serializable objects will lead to a RuntimeError.
+    """
+    span_attributes = encode_object(object)
+    if attributes:
+        flattened = flatten_attributes(attributes, expand_leaf_lists=False)
+        span_attributes.update(sanitize_attributes(flattened))
+
+    attr_length = 0
+    if LightningSpanAttributes.OBJECT_JSON.value in span_attributes:
+        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_JSON.value])
+    elif LightningSpanAttributes.OBJECT_LITERAL.value in span_attributes:
+        attr_length = len(span_attributes[LightningSpanAttributes.OBJECT_LITERAL.value])
+    logger.debug("Emitting object span with payload size %d characters", attr_length)
+
+    if propagate:
+        tracer = get_active_tracer()
+        if tracer is None:
+            raise RuntimeError("No active tracer found. Cannot emit object span.")
+    else:
+        # Do not actually propagate to any store or tracer backend.
+        tracer = DummyTracer()
+
+    return tracer.create_span(
+        name=AGL_OBJECT,
+        attributes=span_attributes,
+        status=TraceStatus(status_code="OK"),
+    )
+
+
+def encode_object(object: Any) -> Dict[str, Any]:
+    """Encode an object as span attributes.
+
+    Args:
+        object: Data structure to encode as JSON.
+    """
+    span_attributes = {}
+    if isinstance(object, (str, int, float, bool)):
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: type(object).__name__,
+            LightningSpanAttributes.OBJECT_LITERAL.value: str(object),
+        }
+    elif isinstance(object, bytes):
+        b64_encoded = base64.b64encode(object).decode("utf-8")
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: "bytes",
+            LightningSpanAttributes.OBJECT_LITERAL.value: b64_encoded,
+        }
+    else:
+        try:
+            serialized = json.dumps(object)
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError(f"Object must be JSON serializable, got: {type(object)}.") from exc
+
+        span_attributes = {
+            LightningSpanAttributes.OBJECT_TYPE.value: full_qualified_name(type(object)),  # type: ignore
+            LightningSpanAttributes.OBJECT_JSON.value: serialized,
+        }
+
+    return span_attributes
+
+
+def get_object_value(span: SpanLike) -> Any:
+    """Extract the object payload from an object span.
+
+    Args:
+        span: Span object produced by Mantisdk emitters.
+    """
+    attributes = span.attributes or {}
+    if LightningSpanAttributes.OBJECT_JSON.value in attributes:
+        serialized = attributes[LightningSpanAttributes.OBJECT_JSON.value]
+        try:
+            return json.loads(serialized)  # type: ignore
+        except (TypeError, ValueError) as exc:
+            raise RuntimeError("Failed to deserialize object JSON from span.") from exc
+    elif LightningSpanAttributes.OBJECT_LITERAL.value in attributes:
+        literal = attributes[LightningSpanAttributes.OBJECT_LITERAL.value]
+        obj_type = attributes.get(LightningSpanAttributes.OBJECT_TYPE.value, "str")
+        if obj_type == "str":
+            return literal
+        elif obj_type == "int":
+            # Let it raise errors if there are any
+            return int(literal)  # type: ignore
+        elif obj_type == "float":
+            return float(literal)  # type: ignore
+        elif obj_type == "bool":
+            return literal.lower() == "true"  # type: ignore
+        elif obj_type == "bytes":
+            return base64.b64decode(literal.encode("utf-8"))  # type: ignore
+        else:
+            raise RuntimeError(f"Unsupported object type for literal deserialization: {obj_type}")
+    else:
+        return None