feat: OpenTelemetry + X-Ray tracing (#103)

Author: hrodmn

.gitignore

@@ -105,3 +105,5 @@ cdk.out/
 node_modules
 cdk.context.json
 *.nc
+
+.test-deploy-env

infrastructure/aws/cdk/app.py

@@ -43,7 +43,6 @@ def __init__(
         id: str,
         memory: int = 1024,
         timeout: int = 30,
-        runtime: aws_lambda.Runtime = aws_lambda.Runtime.PYTHON_3_12,
         concurrent: Optional[int] = None,
         permissions: Optional[List[iam.PolicyStatement]] = None,
         environment: Optional[Dict] = None,
@@ -138,12 +137,19 @@ def __init__(
                 **environment,
                 "TITILER_MULTIDIM_ROOT_PATH": app_settings.root_path,
                 "TITILER_MULTIDIM_CACHE_HOST": redis_cluster.attr_redis_endpoint_address,
+                "OTEL_METRICS_EXPORTER": "none",  # Disable metrics - only using traces
+                "OTEL_PYTHON_DISABLED_INSTRUMENTATIONS": "aws-lambda,requests,urllib3,aiohttp-client",  # Disable aws-lambda auto-instrumentation (handled by otel_wrapper.py)
+                "OTEL_PROPAGATORS": "tracecontext,baggage,xray",
+                "OPENTELEMETRY_COLLECTOR_CONFIG_URI": "/opt/collector-config/config.yaml",
+                # AWS_LAMBDA_LOG_FORMAT not set - using custom JSON formatter in handler.py
+                "AWS_LAMBDA_EXEC_WRAPPER": "/opt/otel-instrument",  # Enable OTEL wrapper to avoid circular import
             },
             log_retention=logs.RetentionDays.ONE_WEEK,
             vpc=vpc,
             vpc_subnets=ec2.SubnetSelection(subnet_type=ec2.SubnetType.PUBLIC),
             allow_public_subnet=True,
             role=veda_reader_role,
+            tracing=aws_lambda.Tracing.ACTIVE,
         )
 
         for perm in permissions:
@@ -207,6 +213,17 @@ def __init__(
         )
     )
 
+# Add X-Ray permissions for tracing
+perms.append(
+    iam.PolicyStatement(
+        actions=[
+            "xray:PutTraceSegments",
+            "xray:PutTelemetryRecords",
+        ],
+        resources=["*"],
+    )
+)
+
 
 lambda_stack = LambdaStack(
     app,

infrastructure/aws/lambda/Dockerfile

@@ -1,32 +1,40 @@
 ARG PYTHON_VERSION=3.12
 
-# Build stage - includes all build tools and dependencies
-FROM public.ecr.aws/lambda/python:${PYTHON_VERSION} AS builder
+# Stage 1: OTEL
+# Download the OpenTelemetry layer
+# ref: https://github.com/athewsey/opentelemetry-lambda-container/blob/98069d5eb6d812ccd28d5c80e2f9d6c8a8c76fb9/python-example/lambda-function/Dockerfile
+FROM public.ecr.aws/lambda/python:${PYTHON_VERSION} as otel-builder
+RUN <<EOF
+dnf install -y unzip wget
+wget https://github.com/aws-observability/aws-otel-python-instrumentation/releases/download/v0.12.1/layer.zip -O /tmp/layer.zip
+mkdir -p /opt-builder
+unzip /tmp/layer.zip -d /opt-builder/
+EOF
 
-# Copy uv for faster dependency management
+# Stage 2: titiler-multidim application and dependencies
+FROM public.ecr.aws/lambda/python:${PYTHON_VERSION} AS builder
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
 
-# Install system dependencies needed for compilation
 RUN dnf install -y gcc-c++ && dnf clean all
 
-# Set working directory for build
 WORKDIR /build
 
-# Copy dependency files first for better caching
 COPY README.md uv.lock .python-version pyproject.toml ./
 COPY src/titiler/ ./src/titiler/
 
-# Install dependencies to temporary directory with Lambda-specific optimizations
-RUN uv export --locked --no-editable --no-dev --extra lambda --format requirements.txt -o requirements.txt && \
-  uv pip install \
+RUN <<EOF
+uv export --locked --no-editable --no-dev --extra lambda --format requirements.txt -o requirements.txt
+uv pip install \
   --compile-bytecode \
   --no-binary pydantic \
   --target /deps \
   --no-cache-dir \
   --disable-pip-version-check \
   -r requirements.txt
+EOF
 
 # Aggressive cleanup to minimize size and optimize for Lambda container
+# Clean up app dependencies in /deps
 WORKDIR /deps
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 RUN <<EOF
@@ -46,34 +54,27 @@ find . -name 'README*' -delete
 find . -name '*.md' -delete
 # Strip debug symbols from shared libraries (preserve numpy.libs)
 find . -type f -name '*.so*' -not -path "*/numpy.libs/*" -exec strip --strip-unneeded {} \; 2>/dev/null || true
-# Create a manifest file for debugging
-du -sh . > /tmp/package_size.txt
 EOF
 
-# Final runtime stage - minimal Lambda image optimized for container runtime
+# Stage 3: Final runtime stage - minimal Lambda image optimized for container runtime
 FROM public.ecr.aws/lambda/python:${PYTHON_VERSION}
 
-# Set Lambda-specific environment variables for optimal performance
-ENV PYTHONPATH=${LAMBDA_RUNTIME_DIR} \
-  PYTHONUNBUFFERED=1 \
-  PYTHONDONTWRITEBYTECODE=1 \
-  AWS_LWA_ENABLE_COMPRESSION=true
-
-# Copy only the cleaned dependencies from builder stage
-# Copy required system library
-COPY --from=builder /deps /usr/lib64/libexpat.so.1 ${LAMBDA_RUNTIME_DIR}/
+ENV PYTHONUNBUFFERED=1 \
+  PYTHONDONTWRITEBYTECODE=1
 
-# Copy application handler
+COPY --from=otel-builder /opt-builder/ /opt/
+COPY infrastructure/aws/lambda/collector-config.yaml /opt/collector-config/config.yaml
+COPY --from=builder /deps ${LAMBDA_RUNTIME_DIR}/
+COPY --from=builder /usr/lib64/libexpat.so.1 ${LAMBDA_RUNTIME_DIR}/
 COPY infrastructure/aws/lambda/handler.py ${LAMBDA_RUNTIME_DIR}/
 
-# Ensure handler is executable and optimize permissions
 RUN <<EOF
 chmod 644 "${LAMBDA_RUNTIME_DIR}"/handler.py
+chmod -R 755 /opt/
 # Pre-compile the handler for faster cold starts
 python -c "import py_compile; py_compile.compile('${LAMBDA_RUNTIME_DIR}/handler.py', doraise=True)"
 # Create cache directories with proper permissions
 mkdir -p /tmp/.cache && chmod 777 /tmp/.cache
 EOF
 
-# Set the Lambda handler
 CMD ["handler.lambda_handler"]

infrastructure/aws/lambda/collector-config.yaml

@@ -0,0 +1,38 @@
+extensions:
+  # AWS Proxy extension - forwards X-Ray segments to local X-Ray daemon via UDP
+  # This avoids the awsxray exporter making HTTPS API calls
+  awsproxy:
+    endpoint: 127.0.0.1:2000
+
+receivers:
+  otlp:
+    protocols:
+      grpc:
+        endpoint: localhost:4317
+      http:
+        endpoint: localhost:4318
+
+processors:
+  batch:
+    timeout: 1s
+    send_batch_size: 50
+
+exporters:
+  # Export to AWS X-Ray via local X-Ray daemon (UDP, no internet required)
+  # The awsproxy extension bridges the collector to the daemon at 127.0.0.1:2000
+  awsxray:
+    endpoint: http://127.0.0.1:2000
+    local_mode: true  # Use local X-Ray daemon instead of direct API calls
+    index_all_attributes: true
+
+  # Debug exporter to see traces in CloudWatch logs
+  debug:
+    verbosity: detailed
+
+service:
+  extensions: [awsproxy]
+  pipelines:
+    traces:
+      receivers: [otlp]
+      processors: [batch]
+      exporters: [debug, awsxray]

infrastructure/aws/lambda/handler.py

@@ -1,13 +1,141 @@
-"""AWS Lambda handler optimized for container runtime."""
+"""AWS Lambda handler optimized for container runtime with OTEL instrumentation."""
 
+import json
 import logging
+import os
 import warnings
+from datetime import datetime, timezone
 from typing import Any, Dict
 
 from mangum import Mangum
 
 from titiler.multidim.main import app
 
+
+def otel_trace_id_to_xray_format(otel_trace_id: str) -> str:
+    """
+    Convert OpenTelemetry trace ID to X-Ray format.
+
+    OTEL format: 32 hex chars (e.g., "68eeb2ec45b07caf760899f308d34ab6")
+    X-Ray format: "1-{first 8 chars}-{remaining 24 chars}" (e.g., "1-68eeb2ec-45b07caf760899f308d34ab6")
+
+    The first 8 hex chars represent the Unix timestamp, which is how X-Ray generates compatible IDs.
+    """
+    if len(otel_trace_id) == 32:
+        return f"1-{otel_trace_id[:8]}-{otel_trace_id[8:]}"
+    return otel_trace_id
+
+
+class XRayJsonFormatter(logging.Formatter):
+    """
+    Custom JSON formatter that includes X-Ray trace ID for log correlation.
+
+    This formatter outputs logs as JSON and includes:
+    - Standard log fields (timestamp, level, message, logger)
+    - X-Ray trace ID (converted from OTEL format)
+    - OTEL trace context fields (if present)
+    - Any extra fields passed via logger.info("msg", extra={...})
+    """
+
+    # Standard fields that shouldn't be duplicated in the output
+    RESERVED_ATTRS = {
+        "name",
+        "msg",
+        "args",
+        "created",
+        "filename",
+        "funcName",
+        "levelname",
+        "levelno",
+        "lineno",
+        "module",
+        "msecs",
+        "message",
+        "pathname",
+        "process",
+        "processName",
+        "relativeCreated",
+        "thread",
+        "threadName",
+        "exc_info",
+        "exc_text",
+        "stack_info",
+        "taskName",
+    }
+
+    def format(self, record: logging.LogRecord) -> str:  # noqa: C901
+        """Format log record as JSON with X-Ray trace ID."""
+        # Build base log object with standard fields
+        log_object = {
+            "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc)
+            .isoformat()
+            .replace("+00:00", "Z"),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "logger": record.name,
+        }
+
+        # Add X-Ray trace ID
+        xray_trace_id = None
+
+        # Method 1: Extract from Lambda's X-Ray environment variable (preferred)
+        trace_header = os.environ.get("_X_AMZN_TRACE_ID", "")
+        if trace_header:
+            for part in trace_header.split(";"):
+                if part.startswith("Root="):
+                    xray_trace_id = part.split("=", 1)[1]
+                    break
+
+        # Method 2: Convert OTEL trace ID if available (fallback)
+        if not xray_trace_id and hasattr(record, "otelTraceID"):
+            xray_trace_id = otel_trace_id_to_xray_format(record.otelTraceID)
+
+        if xray_trace_id:
+            log_object["xray_trace_id"] = xray_trace_id
+
+        # Add exception info if present
+        if record.exc_info:
+            log_object["exception"] = self.formatException(record.exc_info)
+
+        # Add OTEL fields if present
+        for attr in [
+            "otelSpanID",
+            "otelTraceID",
+            "otelTraceSampled",
+            "otelServiceName",
+        ]:
+            if hasattr(record, attr):
+                log_object[attr] = getattr(record, attr)
+
+        # Add AWS request ID if available
+        if hasattr(record, "aws_request_id"):
+            log_object["requestId"] = record.aws_request_id
+
+        # Add any extra fields from record.__dict__ that aren't standard
+        for key, value in record.__dict__.items():
+            if key not in self.RESERVED_ATTRS and key not in log_object:
+                log_object[key] = value
+
+        return json.dumps(log_object)
+
+
+# Configure root logger with custom JSON formatter that includes X-Ray trace ID
+root_logger = logging.getLogger()
+root_logger.setLevel(logging.WARN)
+
+# Remove any existing handlers
+for log_handler in root_logger.handlers[:]:
+    root_logger.removeHandler(log_handler)
+
+# Add StreamHandler with our custom JSON formatter
+json_handler = logging.StreamHandler()
+json_handler.setFormatter(XRayJsonFormatter())
+root_logger.addHandler(json_handler)
+
+# Set titiler loggers to INFO level
+logging.getLogger("titiler").setLevel(logging.INFO)
+
+# Keep specific loggers at ERROR/WARNING levels
 logging.getLogger("mangum.lifespan").setLevel(logging.ERROR)
 logging.getLogger("mangum.http").setLevel(logging.ERROR)
 logging.getLogger("botocore").setLevel(logging.WARNING)
@@ -16,15 +144,8 @@
 warnings.filterwarnings("ignore", category=UserWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
 
-
-# Pre-import commonly used modules for faster cold starts
-try:
-    import numpy  # noqa: F401
-    import pandas  # noqa: F401
-    import rioxarray  # noqa: F401
-    import xarray  # noqa: F401
-except ImportError:
-    pass
+# LoggingInstrumentor().instrument(set_logging_format=False)
+# FastAPIInstrumentor.instrument_app(app)
 
 handler = Mangum(
     app,
@@ -40,10 +161,5 @@
 
 
 def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
-    """Lambda handler with container-specific optimizations."""
-    response = handler(event, context)
-
-    return response
-
-
-handler.lambda_handler = lambda_handler
+    """Lambda handler with container-specific optimizations and OTEL tracing."""
+    return handler(event, context)

infrastructure/aws/lambda/requirements-lambda.txt

@@ -51,7 +51,7 @@ server = [
 ]
 lambda = [
     "mangum==0.19.0",
-    "aiobotocore>=2.24.0,<2.24.2"
+    "aiobotocore>=2.24.0,<2.24.2",
 ]
 
 [dependency-groups]