LambdaのX-Rayを試す

---

LambdaでOpenAI APIを呼び出しつつ、AWS X-RayによるトレーシングをTerraformで設定する方法を解説しています。レイヤーとしてOpenAIとX-Ray SDKを追加し、初期化からリクエストまで詳細なプロファイリングを行う手順をまとめました。

はじめに

• ECS FargateとX-Rayの連携を試すで、ECSとX-Rayの統合は試したが、Lambdaでも同様のことはできるとわかったので試してみた
• ECSと同様にX-Ray SDK for Pythonを追加すればいいだけだが、Lambdaの場合はレイヤーとして定義する必要があるのが違い。

作るもの

• OpenAIのAPIに対してリクエストを送るLambda
• APIキーはパラメーターストアのSecureStringとして事前登録済み
• LambdaのレイヤーとしてOpenAIとX-Ray SDKを追加する
• ディレクトリ構成は以下の通り

./
├── lambda.tf                # Lambda関連のTerraform設定ファイル
├── main.tf                  # メインのTerraform設定ファイル（割愛）
├── openai.zip               # OpenAIとX-Ray SDKが入ったレイヤー用のZip
└── lambda_function.py       # Lambda関数のPythonスクリプト

レイヤーの作成

以下の方法で、OpenAIとX-Ray SDKが入ったLambdaのレイヤー（openai.zip）を作成する。requirements.txtは以下のようにする。

• DockerでAWS LambdaのレイヤーのZipファイルを一括生成する

openai==1.58.1
aws-xray-sdk==2.14.0

ソースファイル

lambda.tf

ポイントは以下の2点

• tracing_configのmode="Active"として、X-Rayトレーシングを有効化する
• AWSXRayDaemonWriteAccessのマネージドポリシーをアタッチする

# Lambda 用の IAM ロールを作成
resource "aws_iam_role" "lambda_role" {
  name = "lambda_xray_python_role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "lambda.amazonaws.com"
        }
      }
    ]
  })
}

# Lambda の基本実行権限をアタッチ
resource "aws_iam_role_policy_attachment" "lambda_basic_execution" {
  role       = aws_iam_role.lambda_role.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
}

# X-Ray の書き込み権限をアタッチ
resource "aws_iam_role_policy_attachment" "lambda_xray_access" {
  role       = aws_iam_role.lambda_role.name
  policy_arn = "arn:aws:iam::aws:policy/AWSXRayDaemonWriteAccess"
}

# Systems Manager Parameter Store から openai-api-key を取得
data "aws_ssm_parameter" "openai_api_key" {
  name            = "openai-api-key"
  with_decryption = true
}

# Lambdaの作成
data "archive_file" "lambda" {
  type        = "zip"
  source_file = "lambda_function.py"
  output_path = ".cache/lambda_function.zip"
}

# Lambda Layer を作成
resource "aws_lambda_layer_version" "openai_layer" {
  filename         = "openai.zip"
  layer_name       = "openai_layer"
  compatible_runtimes = ["python3.12"]
  source_code_hash = filebase64sha256("openai.zip")
}

# Lambdaの作成
resource "aws_lambda_function" "exmaple" {
  filename         = data.archive_file.lambda.output_path
  function_name    = "openai_lambda_with_xray"
  role             = aws_iam_role.lambda_role.arn
  handler          = "lambda_function.lambda_handler"
  runtime          = "python3.12"
  source_code_hash = data.archive_file.lambda.output_base64sha256
  memory_size      = 512
  timeout          = 20

  layers = [
    aws_lambda_layer_version.openai_layer.arn
  ]

  environment {
    variables = {
      OPENAI_API_KEY = data.aws_ssm_parameter.openai_api_key.value
    }
  }  

  # X-Ray トレーシングの有効化
  tracing_config {
    mode = "Active"
  }
}

lambda_function.py

import json
import openai
from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core import patch_all

# X-Ray SDK をパッチ
patch_all()

@xray_recorder.capture('Request openai')
def run_openai_message(user_input):
    # ChatGPT APIへのリクエスト
    response = openai.chat.completions.create(
        model="gpt-4o",  # または最新のモデルを指定
        messages=[
            {"role": "system", "content": "あなたはかわいい猫耳おじさんアシスタントです。おじさん特有の気持ち悪さを残しつつ、可愛く返信してください"},
            {"role": "user", "content": user_input},
        ],
        temperature=0.5,
    )

    # 応答の取得
    chat_response = response.choices[0].message.content.strip()

    return chat_response

def lambda_handler(event, context):
    user_input = event['user_input']
    chat_response = run_openai_message(user_input)

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

結果

かなり詳細にプロファイリングしてくれた。単なるprint文ではわからない、初期化のオーバーヘッド（init）も見てくれるのはとても良い。

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---