如何在SageMaker Notebook中获取当前SageMaker Studio用户配置文件名？

问题背景 #

已完成SageMaker域搭建并添加用户配置文件，登录Studio进入对应配置文件的Notebook后，已完成以下依赖安装：

!pip install --upgrade sagemaker
!pip install torch==1.9.1
!pip install torchvision==0.10.1

以及库导入：

import boto3
import numpy as np
import os
import sagemaker
import torch
import torchvision
from sagemaker.pytorch.model import PyTorchModel

用于创建域和用户配置文件的CloudFormation代码如下：

StudioDomain:
    Type: Custom::StudioDomain
    Properties:
      ServiceToken: !GetAtt StudioDomainFunction.Arn
      VPC: !Ref VPCId
      SubnetIds: !Ref SubnetIds
      DomainName: "MyDomainName"
      DefaultUserSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn

UserProfile:
    Type: Custom::UserProfile
    Properties:
      ServiceToken: !GetAtt UserProfileFunction.Arn
      DomainId: !GetAtt StudioDomain.DomainId
      UserProfileName: 'chuck'
      UserSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn

现需获取并打印当前使用的用户配置文件名称（示例为chuck）。

实现方案 #

方法一：直接读取环境变量（最快捷） #

SageMaker Studio会自动为Notebook会话设置包含用户信息的环境变量，直接读取即可：

import os

user_profile_name = os.environ.get('SAGEMAKER_USER_PROFILE_NAME')
print(f"当前用户配置文件名称：{user_profile_name}")

方法二：通过SageMaker API验证（确保准确性） #

如果需要通过API确认当前用户配置文件信息，可以结合Domain ID和执行角色ARN匹配：

import boto3
import sagemaker
import os

# 初始化SageMaker客户端
sagemaker_client = boto3.client('sagemaker')
# 获取当前会话的Domain ID
domain_id = os.environ.get('SAGEMAKER_DOMAIN_ID')
# 获取当前执行角色ARN
current_role_arn = sagemaker.get_execution_role()

# 列出当前Domain下所有用户配置文件
response = sagemaker_client.list_user_profiles(DomainIdEquals=domain_id)

# 匹配当前执行角色对应的用户配置文件
for profile in response['UserProfiles']:
    if profile['UserSettings']['ExecutionRole'] == current_role_arn:
        print(f"当前用户配置文件名称：{profile['UserProfileName']}")
        break

说明 #

• 方法一无需额外API调用，适合快速获取；
• 方法二通过API验证，适用于需要确认角色与用户配置文件关联关系的场景。

内容的提问来源于stack exchange，提问作者haju