如何在SageMaker Notebook中为训练任务的CloudWatch日志流添加标签？

问题背景 #

已通过SageMaker Notebook启动TensorFlow训练任务，并为任务添加了Project: my-project-1标签，现在希望直接在Notebook中为该任务对应的CloudWatch日志流添加相同标签，无需进入CloudWatch UI操作。

实现步骤 #

1. 获取训练任务的核心信息 #

训练任务启动后，通过Estimator对象获取任务的完整名称（该名称即为CloudWatch中对应的日志流名称），同时SageMaker训练任务的日志组是固定的：

# 获取训练任务的完整名称
job_name = tf_estimator.latest_training_job.job_name
# SageMaker训练任务默认日志组
log_group_name = '/aws/sagemaker/TrainingJobs'

2. 创建CloudWatch Logs客户端 #

通过boto3创建CloudWatch Logs客户端，用于执行标签添加操作：

import boto3
# 使用之前定义的region变量
logs_client = boto3.client('logs', region_name=region)

3. 为日志流添加标签 #

调用CloudWatch Logs的tag_log_stream接口添加标签。由于日志流可能在训练任务启动后需要几秒才会创建，建议添加简单的等待逻辑确保日志流存在：

import time

# 定义要添加的标签，与训练任务标签保持一致
tags = {'Project': 'my-project-1'}

# 等待日志流创建完成
max_retries = 5
retry_count = 0
log_stream_exists = False

while retry_count < max_retries and not log_stream_exists:
    try:
        # 检查日志流是否存在
        logs_client.describe_log_streams(
            logGroupName=log_group_name,
            logStreamNamePrefix=job_name
        )
        log_stream_exists = True
    except logs_client.exceptions.ResourceNotFoundException:
        retry_count += 1
        time.sleep(2)

# 日志流存在则添加标签
if log_stream_exists:
    logs_client.tag_log_stream(
        logGroupName=log_group_name,
        logStreamName=job_name,
        tags=tags
    )
    print(f"已成功为日志流 {job_name} 添加标签")
else:
    print(f"重试{max_retries}次后仍未找到日志流 {job_name}")

补充说明 #

• 若要给整个训练任务日志组添加统一标签，可调用tag_log_group接口，参数替换为日志组名称和目标标签即可。
• 需确保Notebook实例的IAM角色拥有logs:TagLogStream（或logs:TagLogGroup）权限，否则会触发权限错误。

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