CognitoでSagemaker Studioにログインできるようにする

はじめに

D2Cのエンジニアの橋ヶ谷です。

D2Cでは社内のデータサイエンティスト向けにJupyterHubを提供しています。
JupyterHubにはCognitoを使用してログインできるようにしています。
JupyterHubからSageMaker Studioに移行することになったのですが、JupyterHubで使用していたCognitoをSageMaker Studioでも使えるようにしたので、その方法について書きたいと思います。

SageMaker Studioへの通常のアクセス方法

StudioにはAWSのマネジメントコンソールからアクセスできます。
ユーザプロファイル一覧から自分のユーザプロファイルを探し、Studioを起動します。

しかしこの方法では他のユーザのStudioにアクセスできてしまう問題があります。
またユーザ側のアカウントの管理の手間を減らすためデータサイエンティスト向けのサービスはできる限り同じ認証方法を利用できるようにしたいです。既にJupyterHubやAirflowで使用していたCognitoを使用したかったのですが、マネジメントコンソールからのアクセスだとそれができません。
これらの問題を解決するためにCognitoによる認証を用いた構成を考えて環境を構築しました。

Cognitoによる認証を利用した構成

Cognitoによる認証をCognito、ALB、Lambdaを使用して実現しました。
ユーザがALBにアクセスするとCognitoのログインを求められます。
ログインするとLambdaでCognitoのユーザ名を取得します。
そのユーザのユーザプロファイルの署名付きURLを発行しリダイレクトすることでユーザは自分の環境にアクセスすることができます。

ALBのリスナーにはCognitoによる認証後にLambdaのTargetGroupに接続するように設定をします。

以下Lambdaのコードです。

import base64
import json
import os

import boto3

DOMAIN_ID = os.environ["DOMAIN_ID"]


def lambda_handler(event, context):
    # ユーザ名の取得
    encoded_data = event["headers"]["x-amzn-oidc-data"]
    headers = encoded_data.split(".")[1]
    decoded_headers = base64.b64decode(headers)
    decoded_json = json.loads(decoded_headers)
    username = decoded_json["username"]

    # 署名付きURLを発行
    sagemaker = boto3.client("sagemaker")
    response = sagemaker.create_presigned_domain_url(
        DomainId=DOMAIN_ID, UserProfileName=username
    )

    # 署名付きURLにリダイレクト
    return {"statusCode": 303, "headers": {"Location": response["AuthorizedUrl"]}}

実際に使用してみます。
ALBにアクセスするとCognitoのログイン画面が表示されます。

ログインするとSageMaker Studioを起動することができました。

おわりに

Lambdaで署名付きURLを発行することでCognitoでSageMaker Studioにログインできるようにすることができました。
これによってユーザは今までJupyterHubやAirflowで使用していた認証をそのまま使用でき、SageMaker Studioを使用するためにAWSのマネジメントコンソールにログインする必要もなくなりました。

参考

https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreatePresignedDomainUrl.html

https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_presigned_domain_url.html