基于 Amazon MWAA 和 Amazon Batch，实现载荷仿真任务弹性调度和计算

from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.amazon.aws.operators.batch import BatchOperator # 不同版本的apache-airflow-providers-amazon的import路径会不同
from airflow.operators.bash_operator import BashOperator
from airflow.operators.dummy_operator import DummyOperator

# 设置 DAG 的默认参数
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2024, 1, 19),
    'retries': 0,
    'retry_delay': timedelta(minutes=1),
}

# 创建 DAG
dag = DAG(
    'china_bladed_process_dag', # DAG名字，唯一
    default_args=default_args,
)

start_task = DummyOperator(task_id='start_task', dag=dag)
end_task = DummyOperator(task_id='end_task', dag=dag)

JOB_OVERRIDES = {} # 参考SDK说明--->https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/batch/client/，可以Override默认行为，比如containerOverrides中的command可以覆盖原有Job Definition中定义的命令

for i in range(30): # 30并发，可调整
    submit_batch_job = BatchOperator(
        task_id='submit_batch_job_'+str(i),
        job_name='batch_job_name_'+str(i),
        job_queue='airflow-bladed-getting-started-wizard-job-queue', # 第3步创建的batch queue
        job_definition='airflow-bladed-getting-started-wizard-job-definition:2', # 第3步创建的batch definition
        overrides=JOB_OVERRIDES,
        dag=dag,
    )

    start_task >> submit_batch_job >> end_task

import boto3
import os

client = boto3.client('fsx', region_name="cn-norhtwest-1")

def lambda_handler(event, context):
    # 得到参数值
    filesystem_id = os.environ.get("filesystem_id")
    
    # 删除Amazon FSx for Lustre
    response = client.delete_file_system(FileSystemId=filesystem_id)

import boto3
import time
import base64

# boto3初始化
client = boto3.client('fsx', region_name="cn-northwest-1")
client_lambda = boto3.client('lambda', region_name="cn-northwest-1")
client_ec2 = boto3.client('ec2', region_name="cn-northwest-1")
client_batch = boto3.client('batch', region_name="cn-northwest-1")

def lambda_handler(event, context):
    # 创建Amazon FSx for Lustre
    response = client.create_file_system(
        FileSystemType='LUSTRE',
        SecurityGroupIds=[
            'sg-08653e515b27628ac', # 需要修改，参考Amazon Batch服务compute environment中的Networking部分
        ],
        SubnetIds=[
            'subnet-099097bfc300e9550', # 需要修改，参考Amazon Batch服务compute environment中的Networking部分
        ],
        StorageCapacity=1200, # 按照需求修改
        LustreConfiguration={
            'DeploymentType': 'SCRATCH_2',
       # 1. 通过指定ImportPath实现Amazon Lustre和Amazon S3资源映射，通过Lazyloading方式进行访问；通过指定AutoImportPolicy实现Amazon Lustre和Amazon S3文件同步机制，如果未指定，则默认是NONE，即只在文件系统创建时进行同步--->请参考References 9
       # 2. 除了Lazyloading，Amazon Lustre还支持Preloading，提前加载Amazon S3中的资源到Lustre--->请参考 References 10
            'ImportPath': 's3://airflow-bladed/dags', # 模拟前处理任务结果存放目录
        }
    )
    
    filesystem_id = response['FileSystem']['FileSystemId']
    mount_name = response['FileSystem']['LustreConfiguration']['MountName']
    
    # 修改Amazon Lambda函数(delete_lustre)的filesystem_id环境变量
    response_lambda = client_lambda.update_function_configuration(
        FunctionName='delete_lustre', # 需要修改，删除Amazon Lustre文件系统的Amazon Lambda函数名字
        Environment={
            'Variables': {
                'filesystem_id': filesystem_id
            }
        }
    )
    
    # 更新后处理Launch template userData，先赋值若干参数
    # ----------从这里开始----------
    launch_template_name = 'postprocess-launch-template' # 步骤3.6创建的Launch template名字
    region='cn-northwest-1'
    fsx_directory='/scratch'
    file_system_id_01=filesystem_id
    fsx_mount_name=mount_name
    # ----------到这里结束----------
    # 编写UserData
    user_data_script = f"""MIME-Version: 1.0
Content-Type: multipart/mixed; boundary="==BOUNDARY==" 

--==BOUNDARY==
MIME-Version: 1.0
Content-Type: text/x-shellscript; charset="us-ascii"

#!/bin/bash
amazon-linux-extras install -y lustre2.10
mkdir -p {fsx_directory}
mount -t lustre -o noatime,flock {file_system_id_01}.fsx.{region}.amazonaws.com@tcp:/{fsx_mount_name} {fsx_directory}

--==BOUNDARY==--
    """
    # Base64-encoded UserData
    encoded_user_data = base64.b64encode(user_data_script.encode('utf-8')).decode('utf-8')
    
    # 创建新版本的Launch Template
    response_create_lt_version = client_ec2.create_launch_template_version(
        LaunchTemplateName=launch_template_name,
        SourceVersion="1", # 参考版本
        LaunchTemplateData={'UserData': encoded_user_data}
    )

    version_number = response_create_lt_version['LaunchTemplateVersion']['VersionNumber'] # 提取Launch template version number
    
    response_batch_update_compute_env = client_batch.update_compute_environment(
        computeEnvironment='postprocess-airflow-bladed-getting-started-wizard-compute-env', # 步骤3.5创建的Compute environment名字
        computeResources={
            'launchTemplate': {
                'launchTemplateName': launch_template_name,
                'version': str(version_number)
            }
        }
    )

# 定义变量
lambda_function_name_create_lustre = 'create_lustre' # 步骤3.8创建的Amazon Lambda函数名字
lambda_function_name_delete_lustre = 'delete_lustre' # 步骤3.7创建的Amazon Lambda函数名字
payload = '{"key1": "value1", "key2": "value2"}'

# 定义Amazon Lambda Operator
lambda_task_create_lustre = LambdaInvokeFunctionOperator(
    task_id='invoke_lambda_task_create_lustre',
    function_name=lambda_function_name_create_lustre,
    invocation_type='RequestResponse',
    payload=payload,
    dag=dag,
)

lambda_task_delete_lustre = LambdaInvokeFunctionOperator(
    task_id='invoke_lambda_task_delete_lustre',
    function_name=lambda_function_name_delete_lustre,
    invocation_type='RequestResponse',
    payload=payload,
    dag=dag,
)

# 定义Bash Operator，等待Amazon Lustre创建完成
bash_task = BashOperator(
    task_id='bash_task',
    bash_command='sleep 600',
    dag=dag,
)

for i in range(30):
    submit_batch_job = BatchOperator(
        task_id='submit_batch_job_'+str(i),
        job_name='batch_job_name_'+str(i),
        job_queue='airflow-bladed-getting-started-wizard-job-queue',
        job_definition='airflow-bladed-getting-started-wizard-job-definition:2',
        overrides=JOB_OVERRIDES,
        dag=dag,
    )

    start_task >> submit_batch_job >> lambda_task_create_lustre >> bash_task >> lambda_task_delete_lustre >> end_task

# 定义后处理BatchOperator
submit_batch_job_postprocess = BatchOperator(
    task_id='submit_batch_job_postprocess',
    job_name='batch_job_name_postprocess',
    job_queue='postprocess-airflow-bladed-getting-started-wizard-job-queue', # 步骤3.5创建的Job queue
    job_definition='postprocess-airflow-bladed-getting-started-wizard-job-definition:2', # 步骤3.5创建的Job definition
    overrides=JOB_OVERRIDES,
    dag=dag,
)

for i in range(30):
    # 定义前处理BatchOperator
    submit_batch_job = BatchOperator(
        task_id='submit_batch_job_'+str(i),
        job_name='batch_job_name_'+str(i),
        job_queue='airflow-bladed-getting-started-wizard-job-queue',
        job_definition='airflow-bladed-getting-started-wizard-job-definition:2',
        overrides=JOB_OVERRIDES,
        dag=dag,
    )

start_task >> submit_batch_job >> lambda_task_create_lustre >> bash_task >> submit_batch_job_postprocess >> lambda_task_delete_lustre >> end_task