DEV Community

Andrii Shykhov
Andrii Shykhov

Posted on • Originally published at Medium

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

#awsxray #awsserverless

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:<region>:<account_id>:layer:aws-xray-sdk-layer:<version_number>' 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: '*'
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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'
Enter fullscreen mode Exit fullscreen mode

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)
Enter fullscreen mode Exit fullscreen mode

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 <number> aws ssm delete-parameter --name "AuthorizationLambdaToken" aws cloudformation delete-stack --stack-name apigw-lambda-bedrock
Enter fullscreen mode Exit fullscreen mode

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 🙂

Top comments (0)