はじめに
ひきつづき EC2 インスタンスの情報収集を題材に Agents for Amazon Bedrock を触っています。今回はアクショングループに紐づいた Lambda 関数を Python から Go に書き直してパフォーマンスを改善してみました。
以下を試してどの程度パフォーマンスが改善するか確認します。コードはこちらのブランチに置いています。
- goroutine でリージョンごとに並行処理する
- 以下の処理を init 関数で実行し、結果を再利用する
- Config の初期化
- EC2 クライアントの初期化
- DescribeRegions API
これまでの経緯は過去記事をご覧ください。
- Agents for Amazon Bedrock にインスタンスの情報を収集させる
- Agents for Amazon Bedrock で単一のアクショングループから複数の機能を呼び分ける
- Agents for Amazon Bedrock にパラメーターを配列で渡したい
改善後の Lambda 関数
全体像はこちら。
. ├── Dockerfile ├── go.mod ├── go.sum └── main.go
まず各メソッドを生やす構造体を定義し、EC2 のクライアントとリージョンのデフォルト値を保持するようにします。そしてグローバル変数を用意し、init 関数によって値が代入されるようにします。
var d *Describer
type Describer struct {
ctx context.Context
client *ec2.Client
regions []string
}
init 関数
Go の init 関数はパッケージ内の定数と変数宣言が評価された後に実行されます。Lambda 関数の場合、ウォームスタート時には init は実行されず、コールドスタート時に実行された結果を再利用します。
これを利用して DescribeRegions を init 関数で行うことで API 実行のオーバーヘッドを軽減します。Config の初期化、EC2 クライアントの初期化も同様にコールドスタート時の実行結果を再利用するようにします。
init 関数は引数を取ることができず戻り値もないので、グローバル変数に結果を代入する必要があります。
func init() {
d = &Describer{
ctx: context.Background(),
}
cfg, err := config.LoadDefaultConfig(
d.ctx,
config.WithRetryMode(aws.RetryModeStandard),
config.WithRetryMaxAttempts(10),
)
if err != nil {
log.Fatal(err)
}
d.client = ec2.NewFromConfig(cfg)
out, err := d.client.DescribeRegions(d.ctx, &ec2.DescribeRegionsInput{})
if err != nil {
log.Fatal(err)
}
d.regions = make([]string, len(out.Regions))
for i, region := range out.Regions {
d.regions[i] = aws.ToString(region.RegionName)
}
}
全リージョンのインスタンス総数、実行中インスタンス数を調べる関数
API の機能を担う各関数はリージョンごとに並行で処理するように改善します。errgroup と sync.Mutex による排他制御を使ってオーソドックスに記述します。今回はページネーターも考慮しています。
type CountInfo struct {
Region string `json:"region"`
TotalInstances int `json:"totalInstances"`
RunningInstances int `json:"runningInstances"`
}
func GetInstancesCount(regions []string) ([]CountInfo, error) {
var mu sync.Mutex
var infos []CountInfo
eg, ctx := errgroup.WithContext(d.ctx)
for _, region := range regions {
region := region
eg.Go(func() error {
total := 0
running := 0
var token *string
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
out, err := d.client.DescribeInstances(
ctx,
&ec2.DescribeInstancesInput{
NextToken: token,
},
func(o *ec2.Options) {
o.Region = region
},
)
if err != nil {
return fmt.Errorf("%s: %w", region, err)
}
for _, r := range out.Reservations {
total += len(r.Instances)
for _, i := range r.Instances {
if i.State.Name == types.InstanceStateNameRunning {
running++
}
}
}
token = out.NextToken
if token == nil {
break
}
}
info := CountInfo{
Region: region,
TotalInstances: total,
RunningInstances: running,
}
mu.Lock()
infos = append(infos, info)
mu.Unlock()
return nil
})
}
if err := eg.Wait(); err != nil {
return nil, err
}
return infos, nil
}
全リージョンで Owner タグのついていないインスタンスを調べる関数
上記と同様の手法で改善します。
type InstanceInfo struct {
Region string `json:"region"`
InstanceID string `json:"instanceId"`
InstanceName string `json:"instanceName"`
State types.InstanceStateName `json:"state"`
}
func GetInstancesWithoutOwner(regions []string) ([]InstanceInfo, error) {
var mu sync.Mutex
var infos []InstanceInfo
eg, ctx := errgroup.WithContext(d.ctx)
for _, region := range regions {
region := region
eg.Go(func() error {
var token *string
var regionalInfos []InstanceInfo
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
out, err := d.client.DescribeInstances(
ctx,
&ec2.DescribeInstancesInput{
NextToken: token,
},
func(o *ec2.Options) {
o.Region = region
},
)
if err != nil {
return fmt.Errorf("%s: %w", region, err)
}
for _, r := range out.Reservations {
for _, i := range r.Instances {
if getInstanceTagValue("Owner", i.Tags) == "" {
regionalInfos = append(regionalInfos, InstanceInfo{
Region: region,
InstanceID: aws.ToString(i.InstanceId),
InstanceName: getInstanceTagValue("Name", i.Tags),
State: i.State.Name,
})
}
}
}
token = out.NextToken
if token == nil {
break
}
}
mu.Lock()
infos = append(infos, regionalInfos...)
mu.Unlock()
return nil
})
}
if err := eg.Wait(); err != nil {
return nil, err
}
return infos, nil
}
...
func getInstanceTagValue(key string, tags []types.Tag) string {
for _, t := range tags {
if t.Key != nil && strings.EqualFold(aws.ToString(t.Key), key) && t.Value != nil {
return *t.Value
}
}
return ""
}
全リージョンでインバウンドが解放された (0.0.0.0/0 が許可された) インスタンスを調べる関数
セキュリティグループの配列 (スライス) をループで回すオーバーヘッドを軽減するため、map を使います。
type PermissionInfo struct {
IpProtocol string `json:"ipProtocol"`
FromPort int32 `json:"fromPort"`
ToPort int32 `json:"toPort"`
AllowFrom string `json:"allowFrom"`
}
type InstanceSecurityGroupInfo struct {
Region string `json:"region"`
InstanceID string `json:"instanceId"`
InstanceName string `json:"instanceName"`
State types.InstanceStateName `json:"state"`
Permissions []PermissionInfo `json:"permissions"`
}
func GetInstancesWithOpenPermission(regions []string) ([]InstanceSecurityGroupInfo, error) {
var mu sync.Mutex
var infos []InstanceSecurityGroupInfo
eg, ctx := errgroup.WithContext(d.ctx)
for _, region := range regions {
region := region
eg.Go(func() error {
sgmap, err := getOpenSecurityGroups(ctx, region)
if err != nil {
return fmt.Errorf("%s: %w", region, err)
}
if len(sgmap) == 0 {
return nil
}
var sgids []string
for sgid := range sgmap {
sgids = append(sgids, sgid)
}
var token *string
var regionalInfos []InstanceSecurityGroupInfo
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
out, err := d.client.DescribeInstances(
ctx,
&ec2.DescribeInstancesInput{
NextToken: token,
Filters: []types.Filter{
{
Name: aws.String("instance.group-id"),
Values: sgids,
},
},
},
func(o *ec2.Options) {
o.Region = region
},
)
if err != nil {
return fmt.Errorf("%s: %w", region, err)
}
for _, r := range out.Reservations {
for _, i := range r.Instances {
var permissions []PermissionInfo
for _, sg := range i.SecurityGroups {
if perms, ok := sgmap[aws.ToString(sg.GroupId)]; ok {
permissions = append(permissions, perms...)
}
}
regionalInfos = append(regionalInfos, InstanceSecurityGroupInfo{
Region: region,
InstanceID: aws.ToString(i.InstanceId),
InstanceName: getInstanceTagValue("Name", i.Tags),
State: i.State.Name,
Permissions: permissions,
})
}
}
token = out.NextToken
if token == nil {
break
}
}
mu.Lock()
infos = append(infos, regionalInfos...)
mu.Unlock()
return nil
})
}
if err := eg.Wait(); err != nil {
return nil, err
}
return infos, nil
}
func getOpenSecurityGroups(ctx context.Context, region string) (map[string][]PermissionInfo, error) {
m := make(map[string][]PermissionInfo, defaultMapSize)
var token *string
for {
openSgs, err := d.client.DescribeSecurityGroups(
ctx,
&ec2.DescribeSecurityGroupsInput{
NextToken: token,
Filters: []types.Filter{
{
Name: aws.String("ip-permission.cidr"),
Values: []string{"0.0.0.0/0"},
},
},
},
func(o *ec2.Options) {
o.Region = region
},
)
if err != nil {
return nil, fmt.Errorf("%s: %w", region, err)
}
for _, sg := range openSgs.SecurityGroups {
var permissions []PermissionInfo
for _, p := range sg.IpPermissions {
for _, ip := range p.IpRanges {
if aws.ToString(ip.CidrIp) == "0.0.0.0/0" {
permissions = append(permissions, PermissionInfo{
IpProtocol: aws.ToString(p.IpProtocol),
FromPort: aws.ToInt32(p.FromPort),
ToPort: aws.ToInt32(p.ToPort),
AllowFrom: aws.ToString(sg.GroupName),
})
}
}
}
m[aws.ToString(sg.GroupId)] = permissions
}
token = openSgs.NextToken
if token == nil {
break
}
}
return m, nil
}
ハンドラー
ここからが Agents for Amazon Bedrock 固有の対応で、肝の部分です。github.com/aws/aws-lambda-go/events には Bedrock 関連イベントの型情報は現時点で存在しないため、自前で定義する必要があります。なお github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/types の型は Lambda の入出力には使えませんでした。
ResponseBody はキーが Content-Type で可変のためどのように型を定義するか迷いましたが、map をネストさせて対応しました。Body は実行する API によって型が決まるため any にしています。
type EventRequest struct {
ActionGroup string `json:"actionGroup"`
APIPath string `json:"apiPath"`
HTTPMethod string `json:"httpMethod"`
Parameters []struct {
Name string `json:"name"`
Type string `json:"type"`
Value string `json:"value"`
} `json:"parameters"`
ResponseBody map[string]map[string]any `json:"responseBody"`
}
type EventResponse struct {
MessageVersion string `json:"messageVersion"`
Response Response `json:"response"`
}
type Response struct {
ActionGroup string `json:"actionGroup"`
APIPath string `json:"apiPath"`
HTTPMethod string `json:"httpMethod"`
HTTPStatusCode int `json:"httpStatusCode"`
ResponseBody map[string]map[string]any `json:"responseBody"`
}
以下はパラメーターから regions の値を取り出してスライスに変換するヘルパー関数です。値が "all" であれば、デフォルト値として init 関数で取得した DescribeRegions API の結果をそのまま使います。
func getRegionNames(event EventRequest) []string {
for _, parameter := range event.Parameters {
if parameter.Name == "regions" && parameter.Value != "all" {
return strings.Split(parameter.Value, ",")
}
}
return d.regions
}
メインのハンドラーです。apiPath で機能を呼び分け、結果をシリアライズせずに lambda.Start に渡します。
func handle(event *EventRequest) (*EventResponse, error) {
fmt.Println("processing by golang")
regions := getRegionNames(*event)
apiPath := event.APIPath
var body any
var err error
switch apiPath {
case "/count/{regions}":
body, err = GetInstancesCount(regions)
case "/check-without-owner/{regions}":
body, err = GetInstancesWithoutOwner(regions)
case "/check-open-permission/{regions}":
body, err = GetInstancesWithOpenPermission(regions)
default:
return nil, fmt.Errorf("api path \"%s\" not suppported", apiPath)
}
if err != nil {
return nil, err
}
resp := &EventResponse{
MessageVersion: "1.0",
Response: Response{
ActionGroup: event.ActionGroup,
APIPath: event.APIPath,
HTTPMethod: event.HTTPMethod,
HTTPStatusCode: 200,
ResponseBody: map[string]map[string]any{
"application/json": {
"body": body,
},
},
},
}
if isDebug {
result, err := json.Marshal(resp)
if err != nil {
return nil, err
}
fmt.Println(string(result))
}
return resp, nil
}
func main() {
lambda.Start(handle)
}
Dockerfile
CDK の lambda.DockerImageFunction に Dockerfile を渡して cdk deploy 時にビルドさせる構成は変更しません。以下のような Dockerfile にしました。
FROM golang:1.22.4-alpine3.20 as build
ARG HTTP_PROXY
ARG HTTPS_PROXY
WORKDIR /agent-go
COPY go.mod go.sum ./
COPY main.go .
RUN HTTP_PROXY=${HTTP_PROXY} HTTPS_PROXY=${HTTPS_PROXY} go build -o main main.go
FROM alpine:3.20
COPY --from=build /agent-go/main /main
ENTRYPOINT [ "/main" ]
公式ドキュメントに倣って alpine を使ってみました。いつもどおり弊社環境に合わせてプロキシの URL を外から渡しますが、Go の場合は HTTP_PROXY に加えて HTTPS_PROXY も設定することでうまくいきました。
検証
ではどの程度パフォーマンスが改善したか見ていきます。
Python
まずは従来の Python 製スクリプトの実行ログです。
| Duration | Billed Duration | Memory Size | Max Memory Used | Init Duration | |
|---|---|---|---|---|---|
| Cold Start | 21296.26 ms | 23866 ms | 128 MB | 92 MB | 2569.61 ms |
| Warm Start | 21626.68 ms | 21627 ms | 128 MB | 95 MB | – |
雑に書いた直列のスクリプトでなんの最適化もしていないので、全リージョンを処理するのに 20 秒以上かかっています。またメモリも思いのほか消費しているようです。私が熟練の Pythonista だったらこんなことにはならなかったでしょう。
Go
次に Go 製コードの実行ログです。
| Duration | Billed Duration | Memory Size | Max Memory Used | Init Duration | |
|---|---|---|---|---|---|
| Cold Start | 1071.00 ms | 2238 ms | 128 MB | 25 MB | 1166.35 ms |
| Warm Start | 452.11 ms | 453 ms | 128 MB | 27 MB | – |
もともとパフォーマンスの改善を目的としてリファクタリングしたため、すべての数値において大幅な改善が見られました。特に init 関数の効果でウォームスタート時の Duration がコールドスタート時の半分以下になっています。課金面でもかなりのアドバンテージがありますね。
重要
私はこの記事において Python よりも Go が優れているというつもりは全然ありません。私の場合 Python のコードを的確にチューニングする実力がなかったため、Go を使って書き直したほうがパフォーマンスが出ました。が、もし Python のまま徹底的にリファクタリングした場合によりよい結果となる可能性は否定できません。
おわりに
Agents for Amazon Bedrock のアクショングループに Go で書いた Lambda 関数を紐づけるパターンはこれまであまりなかったと思います。Bedrock 成分薄めですが、ご参考になれば幸いです。
簡単30秒