AWS SSM Automation Runbook 从 S3 运行 Python 脚本

2025-03-01 19:34:03

AWS SSM Automation Runbook 中使用 executeScript 从 S3 运行 Python 脚本

在使用 AWS Systems Manager (SSM) Automation Runbook 时，把 Python 脚本放在 Runbook 文档里会使文档变得臃肿。我想把脚本放到 S3，然后从 Runbook 引用。我已经尝试了各种方法，但是总是遇到类似下面的错误：

Failure message Step fails when it is Poll action status for completion. Traceback (most recent call last): AttributeError: module 'test' has no attribute 'main' Handler main is not found in provided script. Please refer to Automation Service Troubleshooting Guide for more diagnosis details.

接下来，我会分析这个问题，并提供几种解决方法。

一、问题原因分析

错误信息表明 SSM 在尝试执行脚本时找不到 main 函数。这通常有几个原因：

Handler 名称错误: aws:executeScript 的 Handler 参数指定了脚本中要执行的函数。如果 Handler 设置不正确，SSM 就找不到入口点。
脚本路径问题: 如果脚本没有正确上传到 S3 或者 Runbook 没有正确配置以访问 S3 中的脚本，也会导致找不到脚本或函数。
附件名称问题 : Runbook 附件名称设置有误。
依赖问题: 如果你的 Python 脚本依赖于其他库，而这些库没有在 SSM 执行环境中预安装，也会导致错误。
S3 权限 : SSM 执行角色没有权限访问s3。

二、解决方案

以下是解决这个问题的几种方法。

1. 使用 `aws:executeScript` 的 `Attachments` (不推荐从 S3 加载)

虽然问题中提到了使用 Attachments，但这通常指的是附加到 Runbook 文档本身的文件，而不是 S3 中的文件。这有很大的局限性，下面演示如何正确的设置：

原理: Runbook 会尝试在附件中找对应的脚本，并且通过Handler去执行指定函数。

Runbook 示例 (YAML):

description: A simple SSM runbook that calls a templated script.
schemaVersion: '0.3'
parameters:
  Message:
    type: String
    description: The message to print and output.
    default: "Hello from the runbook!"
mainSteps:
  - name: ExecuteTemplateScript
    action: aws:executeScript
    isEnd: true
    inputs:
      Runtime: python3.10
      Handler: test.main  # [file].[function] format
      InputPayload:
        Message: '{{ Message }}'
      Script: ''
      Attachments:
        - Key: "SourceUrl"
          Values:
            - "file://test.py" # 注意格式
    outputs:
      - Name: OutputMessage
        Selector: $.Payload.OutputMessage
        Type: String
files:
  test.py:
    checksums:
      sha256: 590708757b79b9438bf299ee496a121c98cf865899db8fea5d788d0cb616d1f5 #这个值必须匹配你文件的checksum

Python 脚本 (test.py):

#!/usr/bin/env python3
"""Simple templated script for SSM that prints and outputs a message."""

import json

def process_message(payload: dict) -> dict:
    """Process the input message and return it."""
    message = payload.get('Message', 'No message provided')
    print(f"Message received: {message}")  # Printed to SSM logs
    return {'OutputMessage': message}

def main(events, context):
    """Main function for SSM execution."""
    # SSM passes InputPayload as 'events'
    payload = events
    result = process_message(payload)
    return result  # SSM captures this as output

if __name__ == "__main__":
    # For local testing, simulate SSM input
    import sys
    if not sys.stdin.isatty():
        payload = json.load(sys.stdin)
    else:
        payload = {'Message': 'Hello, world!'}
    result = process_message(payload)
    print(json.dumps(result))

注意:

Handler 字段应该设置为 test.main，其中 test 是文件名（不带 .py 扩展名），main 是脚本中的入口函数。
Attachments设置必须是file://+ 文件名。
上传Runbook是，需要一同上传脚本。
需要保证Checksum 一致。

2. 通过 `aws:downloadContent` 和 `aws:executeScript` 从 S3 加载脚本 (推荐)

这种方法更灵活，允许你将脚本存储在 S3 中。

原理:

使用 aws:downloadContent 步骤从 S3 下载 Python 脚本到 SSM 执行实例的本地文件系统。
使用 aws:executeScript 步骤，通过指定本地文件路径来执行脚本。

Runbook 示例 (YAML):

description: SSM Runbook to download Python script from S3 and execute it.
schemaVersion: '0.3'
parameters:
  S3BucketName:
    type: String
    description: The name of the S3 bucket where the script is located.
  S3Key:
    type: String
    description: The S3 key (path) to the Python script.
  Message:
    type: String
    description: The message to print and output.
    default: "Hello from S3!"
mainSteps:
  - name: DownloadScript
    action: aws:downloadContent
    inputs:
      SourceType: S3
      SourceInfo:
        path: 'https://s3.amazonaws.com/{{ S3BucketName }}/{{ S3Key }}'  # 完整的 S3 URL
    nextStep: ExecuteScript

  - name: ExecuteScript
    action: aws:executeScript
    inputs:
      Runtime: python3.10
      Handler: main # 文件在当前路径,不需要加文件名称
      InputPayload:
        Message: '{{ Message }}'
      Script: |
        {{ DownloadScript.Output }}
    outputs:
      - Name: OutputMessage
        Selector: $.Payload.OutputMessage
        Type: String

Python 脚本 (test.py): 与之前相同。

步骤:

将 Python 脚本上传到 S3 存储桶的指定位置。
在创建 Runbook 时，填写 S3BucketName 和 S3Key 参数。

注意:

S3 URL 格式: SourceInfo 中的 path 参数需要使用完整的 S3 URL 格式，如 https://s3.amazonaws.com/your-bucket-name/your-script-path.py。
Handler : 由于 aws:downloadContent会将脚本下载到执行环境的当前工作目录下，因此 aws:executeScript 步骤的 Handler只需填写函数名即可。
权限: 确保 SSM 自动化角色具有从 S3 存储桶读取对象的权限 ( s3:GetObject )。

3. 使用 Lambda 作为中转

这种方式稍微复杂一些，但如果你的脚本需要更多依赖，或者你想利用 Lambda 的特性（比如自动扩展、版本控制等），这将是最佳选择。

原理:

将你的 Python 脚本和依赖项打包成一个 Lambda 函数。
在 Runbook 中，使用 aws:invokeLambdaFunction 步骤来调用 Lambda 函数。

Runbook 示例 (YAML):

description: SSM Runbook to invoke a Python Lambda function.
schemaVersion: '0.3'
parameters:
  LambdaFunctionName:
    type: String
    description: The name of the Lambda function.
  Message:
    type: String
    description: Message to pass to the Lambda function.
    default: "Hello from Lambda!"
mainSteps:
  - name: InvokeLambda
    action: aws:invokeLambdaFunction
    inputs:
      FunctionName: '{{ LambdaFunctionName }}'
      Payload: '{"Message": "{{ Message }}"}'
    outputs:
      - Name: OutputMessage
        Selector: $.Payload.OutputMessage
        Type: String

Lambda 函数 (Python): 基本结构和之前的 test.py 相同, 但不需要 if __name__ == "__main__": 部分。

import json

def process_message(payload: dict) -> dict:
    """Process the input message and return it."""
    message = payload.get('Message', 'No message provided')
    print(f"Message received: {message}")
    return {'OutputMessage': message}

def lambda_handler(event, context): #Lambda 函数名必须为 lambda_handler
    """Main function for Lambda execution."""
    payload = event
    result = process_message(payload)
    return result