🐫

AWS ECSからVertex AIを呼び出す; Workload Identityは使用しない

に公開

概要

  • サービスアカウントキーファイルの中身をAWS Secret Managerに登録する
  • Secret Managerの値をConfigurationPropertiesに紐付ける
  • ByteArrayInputStreamに変換し、Credentialsを作成
  • Credentialsを使用してサービスアカウントとして認証、Vertex AIを呼び出す

1. Google Cloudの準備

1.1 サービスアカウントを作成、JSONキーファイルをダウンロード

公式ドキュメントに従い、作業を行う。
https://cloud.google.com/iam/docs/service-accounts-create?hl=ja
https://cloud.google.com/iam/docs/keys-create-delete?hl=ja

1.2 サービスアカウントにIAMロールを紐づける

Vertex AI Userを紐づける。
https://cloud.google.com/vertex-ai/generative-ai/docs/access-control?hl=ja

2. コードの作成

2.1 Secret Managerの値をConfigurationPropertiesに紐付ける

application.yml
application:
  gemini-api:
    project-id: ${GEMINI_PROJECT_ID}
    service-account-key: ${GEMINI_SERVICE_ACCOUNT_KEY}
GeminiApiProperties.java
import lombok.Data;
import org.springframework.boot.context.properties.ConfigurationProperties;

@Data
@ConfigurationProperties(prefix = "application.gemini-api")
public class GeminiApiProperties {
    private String projectId;
    private String serviceAccountKey;
}
GeminiApiConfig.java
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Configuration;

@Configuration
@EnableConfigurationProperties(GeminiApiProperties.class)
public class GeminiApiConfig {
}

2.2 Factoryを使用してClientを作成

GeminiApiClientFactory.java
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

@Component
@RequiredArgsConstructor
public class GeminiApiClientFactory {
    private final GeminiApiProperties properties;

    public GeminiApiClient create(final String systemInstruction) {
        return new GeminiApiClient(properties.getProjectId(), properties.getServiceAccountKey(), systemInstruction);
    }
}

2.3 GenerativeModelの作成、生成用メソッドの準備

GeminiApiClient.java
import com.google.auth.Credentials;
import com.google.auth.oauth2.ServiceAccountCredentials;
import com.google.cloud.vertexai.VertexAI;
import com.google.cloud.vertexai.api.GenerateContentResponse;
import com.google.cloud.vertexai.api.GenerationConfig;
import com.google.cloud.vertexai.generativeai.ContentMaker;
import com.google.cloud.vertexai.generativeai.GenerativeModel;
import com.google.cloud.vertexai.generativeai.ResponseHandler;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;

public class GeminiApiClient {
    private final GenerativeModel generativeModel;

    public GeminiApiClient(final String projectId,
                           final String serviceAccountKey,
                           final String systemInstruction) {
        final String location = "us-central1";
        final String modelName = "gemini-2.0-flash-001";
        final VertexAI.Builder vertexAI = new VertexAI.Builder();

        try {
            final ByteArrayInputStream serviceAccountKeyStream = new ByteArrayInputStream(serviceAccountKey.getBytes(
                StandardCharsets.UTF_8));
            final Credentials credentials = ServiceAccountCredentials.fromStream(serviceAccountKeyStream).createScoped(
                "https://www.googleapis.com/auth/cloud-platform");
            vertexAI.setProjectId(projectId)
                    .setCredentials(credentials)
                    .setLocation(location);
        } catch (final IOException e) {
            throw new RuntimeException(e);
        }

        final GenerationConfig config = GenerationConfig.newBuilder().setResponseMimeType("application/json").build();
        this.generativeModel = new GenerativeModel(modelName, vertexAI.build())
            .withSystemInstruction(ContentMaker.fromString(systemInstruction))
            .withGenerationConfig(config);
    }

    public final String generate(final String userPrompt) {
        String output = "";
        try {
            final GenerateContentResponse response = generativeModel.generateContent(userPrompt);
            output = ResponseHandler.getText(response);
        } catch (final IOException e) {
            throw new RuntimeException(e);
        }

        return output;
    }
}

3. AWSの作業

Secret Managerのシークレットの作成、ECSのタスク定義で参照できるように設定する。
https://docs.aws.amazon.com/ja_jp/AmazonECS/latest/developerguide/specifying-sensitive-data-tutorial.html

4. 他のやり方との比較

4.1 サービスアカウントキーファイルを使用するやり方

Secret Managerから値を読み取り、コンテナ起動時にファイルを作成、GOOGLE_APPLICATION_CREDENTIALS 環境変数にファイルパスを設定する。あとはサービスアカウント使用時に勝手に認証してくれる。記事で紹介したやり方とほぼ同じだが、少しまわりくどい&Dockerfileを修正する必要がある。

4.2 Workload Identityを使用するやり方

Google Cloudで推奨されているやり方だが、EC2を前提しており、ECSだとメタデータURLをカスタマイズしたり、AWS_ACCESS_KEY_ID、AWS_SECRET_ACCESS_KEY、AWS_SESSION_TOKENを設定する必要がある。また、アクセスキーをローテーションしている場合、一度の設定だけではまかなえない。さらに、プラットフォーム依存の設定になるため、Workload Identityに対応していないプライベートクラウドなどに移行する際には修正しなければいけない。

https://cloud.google.com/iam/docs/workload-identity-federation-with-other-clouds?hl=ja
https://zenn.dev/ohsawa0515/articles/gcp-workload-identity-federation
https://zenn.dev/1mono2/articles/1fd85d17e862e3

Discussion