Trusted Space の Codelab

1. 概要

GPU アクセラレーション ワークロードのセキュリティとプライバシーを強化する準備はできていますか？この Codelab では、機密性の高い AI/ML ワークロードに強力なオペレータ分離とアクセラレータ サポートを提供するサービスである Trusted Space の機能について説明します。

貴重なデータ、モデル、鍵を保護することは、これまで以上に重要です。Trusted Space は、ワークロード オペレーターでもアクセスできない安全で信頼できる環境内でワークロードが動作するようにすることで、ソリューションを提供します。

Trusted Space には次のような機能があります。

• プライバシーとセキュリティの強化: Trusted Space は、機密性の高いアセット（モデル、貴重なデータ、鍵など）を暗号証明によって保護された高信頼実行環境を提供します。
• オペレーターの分離: オペレーターの干渉に関する懸念を解消します。Trusted Space では、ワークロードのオペレーターもアクセスできないため、SSH 接続、データへのアクセス、ソフトウェアのインストール、コードの改ざんを防ぐことができます。
• アクセラレータのサポート: Trusted Space は、H100、A100、T4、L4 などの GPU など、さまざまなハードウェア アクセラレータとシームレスに連携するように設計されています。これにより、パフォーマンスが重要な AI/ML アプリケーションをスムーズに実行できます。

学習内容

• Trusted Space の主なサービスについて理解する。
• AI/ML ワークロードの貴重なアセットを保護するために Trusted Space 環境をデプロイして構成する方法について学びます。

必要なもの

• Google Cloud Platform プロジェクト
• Google Compute Engine とアクセラレータに関する基本的な知識。
• サービス アカウント、鍵管理、Workload Identity 連携、属性条件の基礎知識。
• コンテナと Artifact Registry に関する基本的な知識

Primus Company による機密性の高いコード生成プロンプトの保護

この Codelab では、従業員データのプライバシーとセキュリティを重視する Primus という企業の立場になって、Primus は、コード生成モデルをデプロイして、開発者のコーディング作業を支援したいと考えています。ただし、社員から送信されたプロンプトの機密性保護について懸念しています。これらのプロンプトには、機密性の高いコード スニペット、内部プロジェクトの詳細、独自のアルゴリズムが含まれていることがよくあります。

Primus 社が事業者を信頼していないのはなぜですか？

Primus 社は競争の激しい市場で事業を展開しています。コードベースには、競争上の優位性を提供する独自のアルゴリズムや機密性の高いコード スニペットなど、貴重な知的財産が含まれています。ワークロード オペレーターによる企業スパイ行為の可能性を懸念しています。また、従業員のプロンプトには、Primus Corp が保護したい機密性の高い「Need To Know」のコード部分が含まれている場合があります。

この懸念に対処するため、Primus Corp は Trusted Space を利用して、コード生成用にモデルを実行している推論サーバーを分離します。仕組みは次のとおりです。

• プロンプトの暗号化: 各社員は、推論サーバーにプロンプトを送信する前に、Google Cloud で Primus Corp が管理する KMS 鍵を使用してプロンプトを暗号化します。これにより、対応する復号鍵が利用可能な Trusted Space 環境のみが復号し、平文プロンプトにアクセスできるようになります。実際のシナリオでは、クライアントサイド暗号化は、利用可能なライブラリ（tink など）で処理できます。この Codelab では、エンベロープ暗号化を使用してこのサンプル クライアント アプリケーションを使用します。
• オペレータの分離: 暗号化に使用される鍵にアクセスできるのは、Trusted Space 環境内で実行されている推論サーバーのみであり、信頼できる環境でプロンプトを復号できます。暗号鍵へのアクセスは Workload Identity プールによって保護されます。Trusted Space の分離保証により、ワークロード オペレーターでさえ、暗号化に使用された鍵と復号されたコンテンツにアクセスできません。
• アクセラレータを使用したセキュア推論: 推論サーバーは、（信頼できる空間の設定の一部として）Shielded VM で起動されます。これにより、ワークロード インスタンスがブートレベルまたはカーネルレベルのマルウェアやルートキットによって侵害されていないことが保証されます。このサーバーは、Trusted Space 環境内でプロンプトを復号し、コード生成モデルを使用して推論を行い、生成されたコードを社員に返します。

2. Cloud リソースを設定する

始める前に

• 次のコマンドを使用して このリポジトリのクローンを作成し、この Codelab で使用する必要なスクリプトを取得します。

git clone https://github.com/GoogleCloudPlatform/confidential-space.git

• この Codelab のディレクトリを変更します。

cd confidential-space/codelabs/trusted_space_codelab/scripts

• 以下に示すように、必要なプロジェクト環境変数が設定されていることを確認します。GCP プロジェクトの設定の詳細については、 こちらの Codelab をご覧ください。プロジェクト ID の取得方法と、プロジェクト ID とプロジェクト名、プロジェクト番号の違いについては、こちらをご覧ください。

export PRIMUS_PROJECT_ID=<GCP project id of Primus>

• プロジェクトの課金を有効にします。
• 両方のプロジェクトで Confidential Computing API と次の API を有効にします。

gcloud services enable \
    cloudapis.googleapis.com \
    cloudresourcemanager.googleapis.com \
    cloudkms.googleapis.com \
    cloudshell.googleapis.com \
    container.googleapis.com \
    containerregistry.googleapis.com \
    iam.googleapis.com \
    confidentialcomputing.googleapis.com

• 次のコマンドを使用して、上記で指定したリソース名の変数に値を割り当てます。これらの変数を使用すると、必要に応じてリソース名をカスタマイズできます。また、既存のリソースがすでに作成されている場合は、そのリソースを使用することもできます。（例: export PRIMUS_SERVICE_ACCOUNT='my-service-account'）

1. Primus プロジェクトの既存のクラウド リソース名を使用して、次の変数を設定できます。変数が設定されている場合、Primus プロジェクトの対応する既存のクラウド リソースが使用されます。変数が設定されていない場合、クラウド リソース名はプロジェクト名から生成され、その名前で新しいクラウド リソースが作成されます。リソース名でサポートされている変数は次のとおりです。

$PRIMUS_PROJECT_REGION	Primus 社用にリージョン リソースを作成するリージョン。
$PRIMUS_SERVICE_LOCATION	Primus 社用にリソースが作成されるロケーション。
$PRIMUS_PROJECT_ZONE	Primus 社用にゾーンリソースが作成されるゾーン。
$PRIMUS_WORKLOAD_IDENTITY_POOL	クラウド リソースを保護するための Primus 社の Workload Identity プール。
$PRIMUS_WIP_PROVIDER	Primus 社の Workload Identity プール プロバイダ。構成には、構成証明検証サービスによって署名されたトークンに使用する承認条件が含まれています。
$PRIMUS_SERVICEACCOUNT	$PRIMUS_WORKLOAD_IDENTITY_POOL が保護されたリソースにアクセスするために使用する Primus 社のサービス アカウント。このステップでは、$PRIMUS_INPUT_STORAGE_BUCKET バケットに保存されている顧客データを表示する権限が付与されています。
$PRIMUS_ENC_KEY	KMS 鍵は、Primus 社の社員が提供するプロンプトを暗号化するために使用されます。
$PRIMUS_ENC_KEYRING	Primus 社の暗号鍵 $PRIMUS_ENC_KEY の作成に使用される KMS キーリング。
$PRIMUS_ENC_KEYVERSION	暗号鍵 $PRIMUS_ENC_KEY の KMS 鍵バージョン。デフォルト値は 1 です。過去にローテーションされ、バージョンが更新された既存の鍵を使用している場合は、この値を更新します。
$PRIMUS_ARTIFACT_REPOSITORY	ワークロード Docker イメージが push されるアーティファクト リポジトリ。
$PRIMUS_PROJECT_REPOSITORY_REGION	公開されたワークロード Docker イメージを含むアーティファクト リポジトリのリージョン。
$WORKLOAD_VM	ワークロード VM の名前。
$WORKLOAD_IMAGE_NAME	ワークロード Docker イメージの名前。
$WORKLOAD_IMAGE_TAG	ワークロード コンテナ イメージのタグ。
$WORKLOAD_SERVICEACCOUNT	ワークロードを実行する Confidential VM にアクセスする権限を持つサービス アカウント。
$CLIENT_VM	推論サーバーのクライアント アプリケーションを実行するクライアント VM の名前。
$CLIENT_SERVICEACCOUNT	$CLIENT_VM が使用するサービス アカウント

• プロジェクト $PRIMUS_PROJECT_ID に対して、ストレージ管理者、Artifact Registry 管理者、Cloud KMS 管理者、サービス アカウント管理者、IAM Workload Identity プール管理者のロールが必要です。GCP Console を使用して IAM ロールを付与する方法については、こちらのガイドをご覧ください。
• $PRIMUS_PROJECT_ID の場合、次のスクリプトを実行して、残りの変数名をリソース名のプロジェクト ID に基づく値に設定します。

source config_env.sh

注: シームレスな操作を行うには、この Codelab で Cloud Shell の使用を検討してください。通常のターミナルを使用している場合、ターミナルを閉じて再度ログインする場合は、この手順を再度実行する必要があります。Cloud Shell は再接続を自動的に処理します。

Primus の会社リソースを設定する

このステップでは、Primus に必要なクラウド リソースを設定します。次のスクリプトを実行して、Primus のリソースを設定します。スクリプトの実行の一環として、次のリソースが作成されます。

• Primus 社の顧客データ ファイルを暗号化するための KMS の暗号鍵（$PRIMUS_ENC_KEY）とキーリング（$PRIMUS_ENC_KEYRING）。
• Workload Identity プール（$PRIMUS_WORKLOAD_IDENTITY_POOL）: プロバイダで構成された属性条件に基づいてクレームを検証します。
• 上記のワークロード ID プール（$PRIMUS_WORKLOAD_IDENTITY_POOL）に接続されているサービス アカウント（$PRIMUS_SERVICE_ACCOUNT）には、KMS 鍵を使用してデータを復号する（roles/cloudkms.cryptoKeyDecrypter ロールを使用）、KMS 鍵を使用してデータを暗号化する（roles/cloudkms.cryptoKeyEncrypter ロールを使用）、Cloud Storage バケットからデータを読み取る（objectViewer ロールを使用）、サービス アカウントをワークロード ID プールに接続する（roles/iam.workloadIdentityUser を使用）権限があります。

./setup_primus_resources.sh

注: この Codelab を再実行中に、KMS 鍵または Workload Identity プールの作成中に ALREADY_EXISTS などのエラーが発生した場合は、以前の実行からリソースがソフト削除されていることが原因である可能性があります。これらのリソースは 30 日間再利用できないため、別の名前で作成する必要があります。

これを行うには、次のコマンドを使用して、Workload Identity プールと KMS 鍵に新しい名前を設定します。

export PRIMUS_WORKLOAD_IDENTITY_POOL=<new name of Workload Identity Pool>

export PRIMUS_ENC_KEY=<new name of KMS Key>

次に、前述の手順に沿って ./setup_primus_resources.sh スクリプトを再実行します。

3. ワークロードを作成する

ワークロード サービス アカウントを作成する

次に、必要なロールと権限を持つワークロードのサービス アカウントを作成します。次のスクリプトを実行して、Primus プロジェクトにワークロード サービス アカウントを作成します。このサービス アカウントは、推論サーバーを実行する VM によって使用されます。

このワークロード サービス アカウント（$WORKLOAD_SERVICEACCOUNT）には、次のロールが付与されます。

• confidentialcomputing.workloadUser: 構成証明トークンを取得する
• logging.logWriter: Cloud Logging にログを書き込みます。

./create_workload_service_account.sh

ワークロードを作成する

このステップでは、ワークロード Docker イメージを作成します。ワークロードは Primus 社によって作成されます。この Codelab で使用するワークロードは、（Vertex Model Garden の）一般公開されている GCS バケットにある codegemma モデルを使用する Python コードです。ワークロードは、codegemma モデルを読み込み、Primus のデベロッパーからのコード生成リクエストを処理する推論サーバーを起動します。

コード生成リクエストで、Workload は暗号化されたプロンプトとともにラップされた DEK を取得します。ワークロードは、DEK の復号のために KMS API 呼び出しを行い、この DEK を使用してプロンプトを復号します。暗号鍵（DEK 用）は Workload Identity プールによって保護され、属性条件を満たすワークロードにアクセス権が付与されます。これらの属性条件については、次のワークロードの承認に関するセクションで詳しく説明します。推論サーバーが復号されたプロンプトを取得すると、読み込まれたモデルを使用してコードを生成し、レスポンスを返します。

次のスクリプトを実行して、次の手順が実行されるワークロードを作成します。

• Primus が所有する Artifact Registry（$PRIMUS_ARTIFACT_REGISTRY）を作成します。
• 必要なリソース名でワークロード コードを更新します。
• 推論サーバー ワークロードをビルドし、ワークロード コードの Docker イメージをビルドするための Dockerfile を作成します。この Codelab で使用する Dockerfile はこちらです。
• Docker イメージをビルドして、Primus が所有する Artifact Registry（$PRIMUS_ARTIFACT_REGISTRY）に公開します。
• $WORKLOAD_SERVICEACCOUNT に $PRIMUS_ARTIFACT_REGISTRY の読み取り権限を付与します。これは、ワークロード コンテナが Artifact Registry からワークロード Docker イメージを pull するために必要です。

./create_workload.sh

なお、この Codelab で作成して使用しているワークロードの generate() メソッドは次のとおりです（ワークロードのコード全体はこちらで確認できます）。

def generate():
  try:
    data = request.get_json()
    ciphertext = base64.b64decode(data["ciphertext"])
    wrapped_dek = base64.b64decode(data["wrapped_dek"])
    unwrapped_dek_response = kms_client.decrypt(
        request={"name": key_name, "ciphertext": wrapped_dek}
    )
    unwrapped_dek = unwrapped_dek_response.plaintext
    f = Fernet(unwrapped_dek)
    plaintext = f.decrypt(ciphertext)
    prompt = plaintext.decode("utf-8")
    tokens = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**tokens, max_new_tokens=128)
    generated_code = tokenizer.decode(outputs[0])
    generated_code_bytes = generated_code.encode("utf-8")

    response = f.encrypt(generated_code_bytes)
    ciphertext_base64 = base64.b64encode(response).decode("utf-8")
    response = {"generated_code_ciphertext": ciphertext_base64}
    return jsonify(response)

  except (ValueError, TypeError, KeyError) as e:
    return jsonify({"error": str(e)}), 500

4. ワークロードを承認して実行する

ワークロードを承認する

Primus は、次のリソースの属性に基づいて、プロンプトの暗号化に使用される KMS 鍵にアクセスするワークロードを承認したいと考えています。

• 内容: 確認済みのコード
• Where: 安全な環境
• Who: 信頼できるオペレーター

Primus は Workload Identity 連携を使用して、これらの要件に基づいてアクセス ポリシーを適用します。Workload Identity 連携では、属性条件を指定できます。これらの条件により、Workload Identity プール（WIP）で認証できる ID が制限されます。証明書検証サービスは、Workload Identity プール プロバイダとして WIP に追加して、測定結果を提示し、ポリシーを適用できます。

Workload Identity プールは、クラウド リソースの設定手順の一環としてすでに作成されています。Primus が新しい OIDC Workload Identity プール プロバイダを作成します。指定された --attribute-condition は、ワークロード コンテナへのアクセスを承認します。この構成では、次のものが必要になります。

• 内容: $PRIMUS_ARTIFACT_REPOSITORY リポジトリにアップロードされた最新の $WORKLOAD_IMAGE_NAME。
• 使用条件: Confidential Space 高信頼実行環境が、完全にサポートされている Confidential Space VM イメージで実行されている。
• ユーザー: Primus $WORKLOAD_SERVICE_ACCOUNT サービス アカウント。

export WORKLOAD_IMAGE_DIGEST=$(gcloud artifacts docker images describe ${PRIMUS_PROJECT_REPOSITORY_REGION}-docker.pkg.dev/$PRIMUS_PROJECT_ID/$PRIMUS_ARTIFACT_REPOSITORY/$WORKLOAD_IMAGE_NAME:$WORKLOAD_IMAGE_TAG  --format="value(image_summary.digest)" --project ${PRIMUS_PROJECT_ID})

gcloud iam workload-identity-pools providers create-oidc $PRIMUS_WIP_PROVIDER \
  --location="global" \
  --project="$PRIMUS_PROJECT_ID" \
  --workload-identity-pool="$PRIMUS_WORKLOAD_IDENTITY_POOL" \
  --issuer-uri="https://confidentialcomputing.googleapis.com/" \
  --allowed-audiences="https://sts.googleapis.com" \
  --attribute-mapping="google.subject='assertion.sub'" \
  --attribute-condition="assertion.swname == 'HARDENED_SHIELDED' && assertion.hwmodel == 'GCP_SHIELDED_VM' && 
assertion.submods.container.image_digest == '${WORKLOAD_IMAGE_DIGEST}' &&
 assertion.submods.container.image_reference == '${PRIMUS_PROJECT_REPOSITORY_REGION}-docker.pkg.dev/$PRIMUS_PROJECT_ID/$PRIMUS_ARTIFACT_REPOSITORY/$WORKLOAD_IMAGE_NAME:$WORKLOAD_IMAGE_TAG' && 
'$WORKLOAD_SERVICEACCOUNT@$PRIMUS_PROJECT_ID.iam.gserviceaccount.com' in assertion.google_service_accounts"

上記のコマンドは、hwmodel が「GCP_SHIELDED_VM」に設定され、swname が「HARDENED_SHIELDED」に設定されていることを確認することで、ワークロードが信頼できる空間環境で実行されていることを確認します。さらに、image_digest や image_reference などのワークロード固有の断言が含まれているため、セキュリティを強化し、実行中のワークロードの整合性を確保できます。

ワークロードを実行する

このステップでは、アクセラレータが接続された Trusted Space VM でワークロードを実行します。必要な TEE 引数は、メタデータ フラグを使用して渡されます。ワークロード コンテナの引数は、フラグの「tee-cmd」部分を使用して渡されます。ワークロード VM に Nvidia Tesla T4 GPU を装備するには、--accelerator=type=nvidia-tesla-t4,count=1 フラグを使用します。これにより、1 つの GPU が VM にアタッチされます。また、適切な GPU ドライバのインストールをトリガーするために、メタデータ フラグに tee-install-gpu-driver=true を含める必要があります。

gcloud compute instances create ${WORKLOAD_VM} \
  --accelerator=type=nvidia-tesla-t4,count=1 \
  --machine-type=n1-standard-16 \
  --shielded-secure-boot \
  --image-project=conf-space-images-preview \
  --image=confidential-space-0-gpupreview-796705b \
  --zone=${PRIMUS_PROJECT_ZONE} \
  --maintenance-policy=TERMINATE \
  --boot-disk-size=40 \
  --scopes=cloud-platform \
  --service-account=${WORKLOAD_SERVICEACCOUNT}@${PRIMUS_PROJECT_ID}.iam.gserviceaccount.com \
  --metadata="^~^tee-image-reference=${PRIMUS_PROJECT_REPOSITORY_REGION}-docker.pkg.dev/${PRIMUS_PROJECT_ID}/${PRIMUS_ARTIFACT_REPOSITORY}/${WORKLOAD_IMAGE_NAME}:${WORKLOAD_IMAGE_TAG}~tee-install-gpu-driver=true~tee-restart-policy=Never"

注: インスタンスに使用できるイメージは、confidential-space-0-gpupreview-796705b と confidential-space-debug-0-gpupreview-796705b の 2 つです。どちらを使用してもかまいません。主な違いは、confidential-space-0-gpupreview-796705b ではサービスが無効になっているため、オペレーターがワークロードを妨害できないのに対し、confidential-space-debug-0-gpupreview-796705b では SSH が有効になり、オペレーターは VM に SSH 接続できることです。これにより、オペレーターはデバッグできます（ただし、構成証明サービスは、デバッグ状態が有効になっていることを構成証明トークンに記録します）。ただし、その後 VM を手動でシャットダウンする必要があります。

推論クエリを実行する

ワークロード推論サーバーが正常に起動すると、Primus 社の社員は推論サーバーにコード生成リクエストを送信できるようになります。

注: 推論サーバーの起動には約 8 ～ 10 分かかります。サーバーが完全に準備完了するまで待ってから、次のコマンドを実行してください。

この Codelab では、次のスクリプトを使用して、推論サーバーとやり取りするクライアント アプリケーションを設定します。このスクリプトを実行してクライアント VM を設定します。

./setup_client.sh

次の手順では、クライアント VM に SSH 接続し、Python 仮想環境内でサンプル クライアント アプリケーションを実行する方法を示します。このサンプル アプリケーションでは、Fernet ライブラリでエンベロープ暗号化を使用していますが、特定の暗号化ライブラリはさまざまなユースケースに合わせて調整できます。

gcloud compute ssh ${CLIENT_VM} --zone=${PRIMUS_PROJECT_ZONE}

次のコマンドを実行して、クライアント VM で Python 仮想環境を有効にし、クライアント アプリケーションを実行します。

source venv/bin/activate

python3 inference_client.py

このサンプル クライアント アプリケーションの出力には、暗号化プロンプトと平文プロンプトのリクエストと、それに対応する暗号化されたレスポンスと復号されたレスポンスが表示されます。

5. クリーンアップ

この Codelab で作成したリソースをクリーンアップするために使用できるスクリプトは、こちらにあります。このクリーンアップの一環として、次のリソースが削除されます。

• Primus サービス アカウント（$PRIMUS_SERVICEACCOUNT）。
• Primus 暗号鍵（$PRIMUS_ENC_KEY）。
• Primus のアーティファクト リポジトリ（$PRIMUS_ARTIFACT_REPOSITORY）。
• Primus Workload Identity プール（$PRIMUS_WORKLOAD_IDENTITY_POOL）とそのプロバイダ。
• Primus のワークロード サービス アカウント（$WORKLOAD_SERVICEACCOUNT）。
• ワークロード VM（$WORKLOAD_VM）とクライアント VM（$CLIENT_VM）。

./cleanup.sh

確認が完了したら、プロジェクトの削除を検討してください。

• Cloud Platform Console に移動します。
• シャットダウンするプロジェクトを選択し、上部の [削除] をクリックします。これにより、プロジェクトの削除がスケジュールされます。