Zenn
🙆

SageMaker TrainingからWorkload Identity連携で認証切れを気にせずGoogle Cloudにアクセスする

に公開
AWS
Google Cloud
SageMaker
Workload Identity
tech

はじめに

Amazon SageMaker Training で AWS のリソースと Google Cloud のリソースにアクセスしたい場面がありました。
Google Cloud リソースへのアクセスに　Service Account Key を使う場合は、特に問題にぶち当たることはないですが、今回は Workload Identity 連携を用いて、Service Account Key より安全に Google Cloud へアクセスすることを目指します。

問題の背景

Google Cloud のリソースにアクセスするライブラリには、google-auth-library-python を使っています。

EC2 インスタンスで、このライブラリを使いつつ Workload Identity 連携を実施する場合、認証構成ファイルを読み込ませることによって接続が可能です。

次に示す認証構成ファイルの credential_source.url あたりにアクセスしながらアクセスキーを取得して認証を行っています。

認証構成ファイル
{
  "universe_domain": "googleapis.com",
  "type": "external_account",
  "audience": "//iam.googleapis.com/projects/<project_number>/locations/global/workloadIdentityPools/<pool_name>/providers/<provider_name>",
  "subject_token_type": "urn:ietf:params:aws:token-type:aws4_request",
  "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/<name>@<project>.iam.gserviceaccount.com:generateAccessToken",
  "token_url": "https://sts.googleapis.com/v1/token",
  "credential_source": {
    "environment_id": "aws1",
    "region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone",
    "url": "http://169.254.169.254/latest/meta-data/iam/security-credentials",
    "regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15"
  }
}

これが ECS となると話はちょっと変わってきます。 credential_source.url を見ると http://169.254.169.254/latest/meta-data/iam/security-credentials となっています。

EC2 の場合はこれが正しいアクセス先ですが、 ECS の場合は http://169.254.170.2$AWS_CONTAINER_CREDENTIALS_RELATIVE_URI がメタデータの取得先です。 AWS_CONTAINER_CREDENTIALS_RELATIVE_URI は環境変数にセットされており、固定ではなく動的です。

実際にこのURLへアクセスしてみると EC2 のときと同じような感じで認証情報を取ってくることはできます。つまりは認証情報を取得して環境変数に設定するところを自前実装するため、上に記載した認証構成ファイルは不要になります。

しかし、 ここで取得してきた認証情報は有効期限があるということがハマりポイント です。

初回はここから取得できるキーを使えば、 Google 側の認証が通るので問題なく Google Cloud のサービスにアクセスすることができます。しかし、有効期限があるということは Google 側でも再認証を走らせないと行けないタイミングが来るので定期的にメタデータURLへアクセスする必要があります。

また、同時に AWS のクライアントライブラリを使用している場合はもっと注意です。 Google Cloud の認証のために、環境変数へ認証情報を登録しましたが、これは AWS のクライアントライブラリも同じ仕様で環境変数にセットされている場合はそれらを先に見に行きます。インスタンスに用意したロールよりも先にです。

つまりは、環境変数に登録している認証情報の有効期限が切れた瞬間に Google Cloud のサービスにも AWS のサービスにもアクセスできなくなるという罠が存在しています。

この課題は、ECS だけでなく SageMaker Training も同様でした。

長くなりましたが、これらをどうにか楽に回避したいというモチベーションで調査し、解決しましたというのがこの記事の結論です。

問題の詳細

話はズレましたが、今回は SageMaker Training で AWS と Google Cloud 両方にアクセスしたい。AWS は IAM Role で。Google Cloud は Workload Identity 連携で。です。

SageMaker Training の後ろでは EC2 が起動しています。では EC2 と同じように、認証構成ファイルを使えば問題なくアクセスできるのはないかと思い、Training Job を実行してみますがエラーを吐いてアクセスできませんでした。

もしやと思い、Training Job のなかに入り、 env コマンドを叩いてみると ECS と同じメタデータエンドポイントがありました。

env
...
AWS_CONTAINER_CREDENTIALS_RELATIVE_URI=....
...

どうやら、Training Job は ECS/EC2 で動いているようです。
ECS と同じようにやればよいことがわかったのですが、問題の背景で述べたように愚直にやると認証情報の定期的な更新をする実装が必要になってきますが、やりたくないです。

EC2 では認証構成ファイルがあれば特別な設定なく Google Cloud へのアクセスができることから、このファイルでどうにかして突破したいところから解決の糸口を探し始めました。

解決策

google-auth-library-python がどのように認証構成ファイルを利用しているのか実装を追っていきました。

AWS環境で認証を行う流れは以下のようになっています：

  1. メタデータサービスURL（http://169.254.169.254/latest/meta-data/iam/security-credentials）にアクセスしてロール名を取得
  2. 取得したロール名をパスに追加した新しいURL（http://169.254.169.254/latest/meta-data/iam/security-credentials/{ロール名}）にアクセスして認証情報を取得

この処理は以下で実装されています：

https://github.com/googleapis/google-auth-library-python/blob/77ad53eb00c74b3badc486c8207a16dbc49f37e5/google/auth/aws.py#L421-L449

実装を読み解いていくと、 SageMaker Training Job 環境での問題の核心は、メタデータサービスのURLに対してロール名を追加するステップが入ることで正しく認証情報を取得できなくなることだとわかりました。

この問題を解決するために クエリパラメータを使って　URL　のパス部分への影響を回避 しました。

解決策の仕組み

URL にクエリパラメータ区切り文字（?）を追加すると、それ以降の文字列は URL のパスではなくクエリパラメータとして扱われます。これを利用することで、クライアントライブラリがロール名をパスに追加しても、エンドポイントのベース URL を変更せずに済みます。

実装方法

SageMaker Training Job 環境で Google Cloud の認証情報を正しく取得するための実装例を示します。

def _google_application_credentials_for_sagemaker():
    credentials_path = os.getenv("GOOGLE_APPLICATION_CREDENTIALS", None)
    if credentials_path is None:
        return

    with open(credentials_path) as f:
        credentials = json.loads(f.read())

    # SageMaker 用の認証情報に差し替え
    ip_address = "169.254.170.2"
    uri_path = os.environ["AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"]
    credentials["credential_source"]["url"] = f"http://{ip_address}{uri_path}?"

    with open(credentials_path, "w") as f:
        json.dump(credentials, f, indent=2)

重要なポイントは、credential_source.urlの末尾に?を追加していることです。これにより、GCPクライアントライブラリがロール名をパスに追加しても、クエリパラメータとして扱われるため、正しいエンドポイントにアクセスできます。

この関数をアプリケーションの起動時に呼んであげるだけで OK です。

技術的な詳細解説

google-auth-library-pythonの処理フローと問題点

通常の AWS 環境では：

  1. _get_metadata_role_name がメタデータURLにアクセスしてロール名を取得
  2. _get_metadata_security_credentials が {base_url}/{role_name} にアクセスして認証情報を取得

しかし SageMaker Training Job 環境では：

  • 認証情報を取得するエンドポイントは http://{ip_address}{uri_path} の形式
  • このエンドポイント自体が認証情報を直接返す
  • ロール名取得と認証情報取得のステップが分かれていることが問題

クエリパラメータによる解決策の詳細

URLにクエリパラメータ区切り文字（?）を追加すると、それ以降に追加される文字列はすべてクエリパラメータとして扱われます。

例えば、以下のように設定すると：

http://169.254.169.254/latest/meta-data/iam/security-credentials?

GCPクライアントライブラリが内部でロール名をパスに追加しようとしても：

http://169.254.169.254/latest/meta-data/iam/security-credentials?sample-role

sample-role はクエリパラメータとして扱われ、実際のリクエスト先のパスは変わりません。これにより、SageMaker環境でもGCPの認証情報を正しく取得できるようになります。

注意点

  • このアプローチは SageMaker Training 特有の問題を無理やり解決したものです。
  • クライアントライブラリのバージョンによって挙動が変わる可能性があります。

まとめ

  • SageMaker Training Job の環境変数を確認することで、認証情報の取得には ECS と同様の仕組みが使われていることがわかった。
  • クライアントライブラリの実装を追って、認証の流れを確認し、途中の処理をクエリパラメータ化によってキャンセルし、認証構成ファイルを少しいじるだけで、AWSにもGoogle Cloudにもアクセスできる方法を示した。

Discussion