Implementing Tracing with AWS X-Ray service and AWS X-Ray SDK for REST API Gateway and Lambda functions

Introduction: This post details the implementation of tracing with AWS X-Ray and the AWS X-Ray SDK within a serverless AWS infrastructure. The focus is on integrating tracing across various AWS services, including API Gateway (REST API), Lambda functions, Systems Manager Parameter Store, and the Amazon Bedrock model. About the Project: Based on my earlier article that covers interacting with Amazon Bedrock models using API Gateway and Lambda functions, this post extends the existing infrastructure by adding configurations for tracing with AWS X-Ray. Initially, the setup utilized an HTTP API, which lacks native X-Ray support; therefore, a REST API was adopted to enable tracing. Additional information on API Gateway tracing with X-Ray is available in the official documentation. The AWS X-Ray SDK documentation can be found here. The Main Lambda function and IAM role configuration in infrastructure/tracing-rest-api.yaml CloudFormation template: Parameters: BedrockModelId: Type: String Default: 'amazon.titan-text-express-v1' LambdaLayerVersionArn: Type: String Default: 'arn:aws:lambda:::layer:aws-xray-sdk-layer:' Resources: MainLambdaFunction: Type: AWS::Lambda::Function Properties: FunctionName: MainLambdaFunction Description: Make requests to Bedrock models Runtime: python3.12 Handler: index.lambda_handler Role: !GetAtt MainLambdaExecutionRole.Arn Timeout: 30 MemorySize: 512 TracingConfig: Mode: Active Layers: - !Ref LambdaLayerVersionArn Environment: Variables: BEDROCK_MODEL_ID: !Ref BedrockModelId Code: ZipFile: | import json import boto3 import os import logging from botocore.exceptions import ClientError from aws_xray_sdk.core import patch_all, xray_recorder # Initialize logging logger = logging.getLogger() logger.setLevel(logging.INFO) # Initialize the X-Ray SDK patch_all() # Initialize the Bedrock Runtime client bedrock_runtime = boto3.client('bedrock-runtime') def lambda_handler(event, context): try: # Retrieve the model ID from environment variables model_id = os.environ['BEDROCK_MODEL_ID'] # Validate the input input_text = event.get("queryStringParameters", {}).get("inputText") if not input_text: logger.error('Input text is missing in the request') raise ValueError("Input text is required in the request query parameters.") # Prepare the payload for invoking the Bedrock model payload = json.dumps({ "inputText": input_text, "textGenerationConfig": { "maxTokenCount": 8192, "stopSequences": [], "temperature": 0, "topP": 1 } }) logger.info('Payload for Bedrock model: %s', payload) # Create a subsegment for the Bedrock model invocation with xray_recorder.in_subsegment('Bedrock InvokeModel') as subsegment: # Invoke the Bedrock model response = bedrock_runtime.invoke_model( modelId=model_id, contentType="application/json", accept="application/json", body=payload ) logger.info('Response from Bedrock model: %s', response) # Check if the 'body' exists in the response and handle it correctly if 'body' not in response or not response['body']: logger.error('Response body is empty') raise ValueError("Response body is empty.") # Read and process the response response_body = json.loads(response['body'].read().decode('utf-8')) logger.info('Processed response body: %s', response_body) return { 'statusCode': 200, 'body': json.dumps(response_body) } except ClientError as e: logger.error('ClientError: %s', e) return { 'statusCode': 500, '

Mar 13, 2025 - 21:49

Implementing Tracing with AWS X-Ray service and AWS X-Ray SDK for REST API Gateway and Lambda functions

Introduction:

This post details the implementation of tracing with AWS X-Ray and the AWS X-Ray SDK within a serverless AWS infrastructure. The focus is on integrating tracing across various AWS services, including API Gateway (REST API), Lambda functions, Systems Manager Parameter Store, and the Amazon Bedrock model.

About the Project:

Based on my earlier article that covers interacting with Amazon Bedrock models using API Gateway and Lambda functions, this post extends the existing infrastructure by adding configurations for tracing with AWS X-Ray. Initially, the setup utilized an HTTP API, which lacks native X-Ray support; therefore, a REST API was adopted to enable tracing. Additional information on API Gateway tracing with X-Ray is available in the official documentation. The AWS X-Ray SDK documentation can be found here.

The Main Lambda function and IAM role configuration in infrastructure/tracing-rest-api.yaml CloudFormation template:

    Parameters:
      BedrockModelId:
        Type: String
        Default: 'amazon.titan-text-express-v1'
      LambdaLayerVersionArn:
        Type: String
        Default: 'arn:aws:lambda:::layer:aws-xray-sdk-layer:'

    Resources:   
      MainLambdaFunction:
        Type: AWS::Lambda::Function
        Properties:
          FunctionName: MainLambdaFunction
          Description: Make requests to Bedrock models
          Runtime: python3.12
          Handler: index.lambda_handler
          Role: !GetAtt MainLambdaExecutionRole.Arn
          Timeout: 30
          MemorySize: 512
          TracingConfig:
            Mode: Active
          Layers:
            - !Ref LambdaLayerVersionArn
          Environment:
            Variables:
              BEDROCK_MODEL_ID: !Ref BedrockModelId
          Code:
            ZipFile: |
              import json
              import boto3
              import os
              import logging
              from botocore.exceptions import ClientError
              from aws_xray_sdk.core import patch_all, xray_recorder

              # Initialize logging
              logger = logging.getLogger()
              logger.setLevel(logging.INFO)

              # Initialize the X-Ray SDK
              patch_all()

              # Initialize the Bedrock Runtime client
              bedrock_runtime = boto3.client('bedrock-runtime')

              def lambda_handler(event, context):
                  try:
                      # Retrieve the model ID from environment variables
                      model_id = os.environ['BEDROCK_MODEL_ID']

                      # Validate the input
                      input_text = event.get("queryStringParameters", {}).get("inputText")
                      if not input_text:
                          logger.error('Input text is missing in the request')
                          raise ValueError("Input text is required in the request query parameters.")

                      # Prepare the payload for invoking the Bedrock model
                      payload = json.dumps({
                          "inputText": input_text,
                          "textGenerationConfig": {
                              "maxTokenCount": 8192,
                              "stopSequences": [],
                              "temperature": 0,
                              "topP": 1
                          }
                      })

                      logger.info('Payload for Bedrock model: %s', payload)

                      # Create a subsegment for the Bedrock model invocation
                      with xray_recorder.in_subsegment('Bedrock InvokeModel') as subsegment:
                          # Invoke the Bedrock model
                          response = bedrock_runtime.invoke_model(
                              modelId=model_id,
                              contentType="application/json",
                              accept="application/json",
                              body=payload
                          )

                          logger.info('Response from Bedrock model: %s', response)

                          # Check if the 'body' exists in the response and handle it correctly
                          if 'body' not in response or not response['body']:
                              logger.error('Response body is empty')
                              raise ValueError("Response body is empty.")

                          # Read and process the response
                          response_body = json.loads(response['body'].read().decode('utf-8'))
                          logger.info('Processed response body: %s', response_body)

                          return {
                              'statusCode': 200,
                              'body': json.dumps(response_body)
                          }

                  except ClientError as e:
                      logger.error('ClientError: %s', e)
                      return {
                          'statusCode': 500,
                          'body': json.dumps({"error": "Error interacting with the Bedrock API"})
                      }
                  except ValueError as e:
                      logger.error('ValueError: %s', e)
                      return {
                          'statusCode': 400,
                          'body': json.dumps({"error": str(e)})
                      }
                  except Exception as e:
                      logger.error('Exception: %s', e)
                      return {
                          'statusCode': 500,
                          'body': json.dumps({"error": "Internal Server Error"})
                      }

      MainLambdaExecutionRole:
        Type: AWS::IAM::Role
        Properties:
          RoleName: MainLambdaExecutionRole
          AssumeRolePolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Principal:
                  Service:
                    - lambda.amazonaws.com
                Action:
                  - sts:AssumeRole
          ManagedPolicyArns:
            - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
          Policies:
            - PolicyName: BedrockAccessPolicy
              PolicyDocument:
                Version: '2012-10-17'
                Statement:
                  - Effect: Allow
                    Action:
                      - bedrock:InvokeModel
                      - bedrock:ListFoundationModels
                    Resource: '*'
            - PolicyName: XRayAccessPolicy
              PolicyDocument:
                Version: '2012-10-17'
                Statement:
                  - Effect: Allow
                    Action:
                      - xray:PutTelemetryRecords
                      - xray:PutAnnotation
                      - xray:PutTraceSegments
                    Resource: '*'

Prerequisites:

Ensure the following prerequisites are in place:

An AWS account with sufficient permissions to create and manage resources.
The AWS CLI installed on the local machine.

Deployment:

Before Starting: If the infrastructure from my earlier post has not been set up, follow steps 1 – 4 from this article before proceeding with the deployment steps below.

1.Create a Lambda layer that includes the AWS X-Ray SDK. Use the layer’s version ARN in the tracing-rest-api.yaml CloudFormation template.

    aws lambda publish-layer-version --layer-name aws-xray-sdk-layer --zip-file fileb://infrastructure/aws-xray-sdk-layer-layer/aws-xray-sdk-layer-layer.zip --compatible-runtimes python3.12

2.Update the existing CloudFormation stack to enable X-Ray tracing and transition from an HTTP API to a REST API.

    aws cloudformation update-stack \
        --stack-name apigw-lambda-bedrock \
        --template-body file://infrastructure/tracing-rest-api.yaml \
        --capabilities CAPABILITY_NAMED_IAM \
        --disable-rollback

3.Retrieve the Invoke URL for the API Gateway Stage using the retrieve_invoke_url_rest_api.sh script. Test the API using CURL.

    ./scripts/retrieve_invoke_url_rest_api.sh

    export APIGW_TOKEN='token_value'

    curl -s -X GET -H "Authorization: Bearer $APIGW_TOKEN" "https://api_id.execute-api.eu-central-1.amazonaws.com/dev/invoke?inputText=your_question" | jq -r '.results[0].outputText'

4.To view the X-Ray trace map and segments timeline, navigate to the AWS Console: CloudWatch -> X-Ray traces -> Traces. In the Traces section, a list of trace IDs will be displayed. Clicking on a trace ID will reveal the Trace map, Segments timeline, and associated Logs.


![X-Ray Trace Map](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/voy3jootl0gy0ufj080m.png)

![X-Ray Segments Timeline](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/ev6b0d86fjdt7kv12ptg.png)

5.Clean Up Resources. After testing, clean up resources by deleting the Lambda layer version, the token from the Parameter Store, and the CloudFormation stack.

aws lambda delete-layer-version --layer-name aws-xray-sdk-layer --version-number 

aws ssm delete-parameter --name "AuthorizationLambdaToken"

aws cloudformation delete-stack --stack-name apigw-lambda-bedrock

Conclusion:

The integration of AWS X-Ray into a serverless infrastructure provides deep insights into the performance of APIs and Lambda functions, as well as their interactions with other AWS services such as Systems Manager and Bedrock. This enhanced visibility facilitates effective debugging, performance optimization, and comprehensive system monitoring, contributing to the smooth and efficient operation of applications.

If you found this post helpful and interesting, please click the button to show your support. Feel free to use and share this post. You can also support me with a virtual coffee