如何在AWS Batch Job启动前从Step Function获取Job ID并更新API

我正在搭建一套数据管道：从SQS获取Celery消息，解码后启动AWS Step Function，在Step Function内部提交运行Python脚本的AWS Batch Job。

我需要在AWS Batch Job的Docker容器启动前（即Job提交时）更新DRF API端点的job_id字段，但不知道怎么在Batch Job进入Running状态前从Step Function中获取其AWS Batch Job ID。

我不想直接在Lambda函数中提交AWS Batch Job，因为其他管道都基于Step Function运行，需要保持一致性。

现有代码 #

Lambda函数（lambda_function.py）

def start_step_function(step_function_input):
    state_machine_arn = os.environ['STEP_FUNCTION_ARN_MYAPP']

    response = client.start_execution(
        stateMachineArn=state_machine_arn,
        input=json.dumps(step_function_input)
    )

    return response
def lambda_handler(event, context):
    try:
        for record in event['Records']:
            # Step 1: Decode the body
            body_json = decode_body(record['body'])
            
            # Step 2: Extract parameters
            step_function_input = extract_parameters(body_json)

            logger.info(f"Step Function input: {step_function_input}")
            
            # Step 3: Start the Step Function execution
            start_step_function(step_function_input)

        
        return {
            'statusCode': 200,
            'body': json.dumps('Step Function executed successfully!')
        }
    except Exception as e:
        logger.error(f"Error processing record: {e}")
        return {
            'statusCode': 400,
            'body': json.dumps(f'Error processing record: {str(e)}')
        }

Step Function定义

{
  "Comment": "A description of my state machine",
  "StartAt": "submit_my_job",
  "States": {
    "submit_my_job": {
      "Type": "Task",
      "Resource": "arn:aws:states:::batch:submitJob.sync",
      "Parameters": {
        "JobName": "my_job",
        "JobQueue": "arn:aws:batch:eu-west-1:380921374981:job-queue/my_job-worker-queue",
        "Parameters": {
          "TASK_ID.$": "$.task_id",
          "SPIDER_ID.$": "$.spider_id"
        },
        "JobDefinition": "arn:aws:batch:eu-west-1:380921374981:job-definition/my_job"
      },
      "Next": "Success",
      "TimeoutSeconds": 5400,
      "Catch": [
        {
          "ErrorEquals": [
            "States.TaskFailed"
          ],
          "Next": "Fail"
        }
      ]
    },
    "Success": {
      "Type": "Succeed"
    },
    "Fail": {
      "Type": "Fail"
    }
  },
  "TimeoutSeconds": 5400
}

尝试过的无效方法 #

我曾在Lambda中尝试提取ID，但只能获取到Step Function的ARN，得到的是Step Function执行ID而非Batch Job ID：

...
# Step 3: Start the Step Function execution
response = start_step_function(step_function_input)

# Extract jobID from the response (assuming it's in the response ARN)
job_id = response['executionArn'].split(':')[-1] # 这是Step Function ID，不是Batch Job ID :(
task_id = step_function_input['task_id']

# Step 4: Update the task with the new job_id
update_task_in_my_api(task_id, {"job_id": job_id}, reason="Step Function execution completed")

期望的输出 #

Step Function提交Batch Job后，输出中包含Batch Job ID：

{
  "JobArn": "arn:aws:batch:eu-west-1:380921374981:job/53425042-22b2-4cad-a6a1-a7d881001702",
  "JobId": "53425042-22b2-4cad-a6a1-a7d881001702",
  "JobName": "my_job",
  "SdkHttpMetadata": {
    "AllHttpHeaders": {
      "Access-Control-Expose-Headers": [
        "X-amzn-errortype,X-amzn-requestid,X-amzn-errormessage,X-amzn-trace-id,X-amz-apigw-id,date"
      ],
      "x-amz-apigw-id": [
        "YesN4HTFDoEEXSw="
      ],
      "Access-Control-Allow-Origin": [
        "*"
      ],
      "Connection": [
        "keep-alive"
      ],
      "x-amzn-RequestId": [
        "91eab790-5e04-45bf-a94c-fa4891cda6d5"
      ],
      "Content-Length": [
        "169"
      ],
      "Date": [
        "Tue, 28 May 2024 11:32:15 GMT"
      ],
      "X-Amzn-Trace-Id": [
        "Root=1-6655c0bd-520650647e897a7fa51a62b7"
      ],
      "Content-Type": [
        "application/json"
      ]
    },
    "HttpHeaders": {
      "Access-Control-Allow-Origin": "*",
      "Access-Control-Expose-Headers": "X-amzn-errortype,X-amzn-requestid,X-amzn-errormessage,X-amzn-trace-id,X-amz-apigw-id,date",
      "Connection": "keep-alive",
      "Content-Length": "169",
      "Content-Type": "application/json",
      "Date": "Tue, 28 May 2024 11:32:15 GMT",
      "x-amz-apigw-id": "YesN4HTFDoEEXSw=",
      "x-amzn-RequestId": "91eab790-5e04-45bf-a94c-fa4891cda6d5",
      "X-Amzn-Trace-Id": "Root=1-6655c0bd-520650647e897a7fa51a62b7"
    },
    "HttpStatusCode": 200
  },
  "SdkResponseMetadata": {
    "RequestId": "91eab790-5e04-45bf-a94c-fa4891cda6d5"
  }
}

我的问题是：如何在AWS Batch Job启动前获取其Job ID并更新我的API端点？

方案1：修改Step Function，添加更新API的Lambda任务 #

将原有的batch:submitJob.sync（等待Job完成的同步任务）替换为batch:submitJob（异步任务），它会在提交Job成功后立即返回JobId，此时Job还处于Submitted状态，未进入Running，完全符合需求。之后添加一个Lambda任务调用DRF API更新job_id，最后可选择添加等待Job完成的步骤以保持原有流程一致性。

修改后的Step Function定义：

{
  "Comment": "A description of my state machine",
  "StartAt": "submit_my_job",
  "States": {
    "submit_my_job": {
      "Type": "Task",
      "Resource": "arn:aws:states:::batch:submitJob",
      "Parameters": {
        "JobName": "my_job",
        "JobQueue": "arn:aws:batch:eu-west-1:380921374981:job-queue/my_job-worker-queue",
        "Parameters": {
          "TASK_ID.$": "$.task_id",
          "SPIDER_ID.$": "$.spider_id"
        },
        "JobDefinition": "arn:aws:batch:eu-west-1:380921374981:job-definition/my_job"
      },
      "Next": "update_api_job_id",
      "Catch": [
        {
          "ErrorEquals": [
            "States.TaskFailed"
          ],
          "Next": "Fail"
        }
      ]
    },
    "update_api_job_id": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:eu-west-1:380921374981:function:update-drf-api-job-id",
      "Parameters": {
        "task_id.$": "$.task_id",
        "batch_job_id.$": "$.JobId"
      },
      "Next": "wait_for_job_completion"
    },
    "wait_for_job_completion": {
      "Type": "Task",
      "Resource": "arn:aws:states:::batch:waitForJob",
      "Parameters": {
        "JobId.$": "$.JobId",
        "JobQueue": "arn:aws:batch:eu-west-1:380921374981:job-queue/my_job-worker-queue"
      },
      "Next": "Success",
      "Catch": [
        {
          "ErrorEquals": [
            "States.TaskFailed"
          ],
          "Next": "Fail"
        }
      ]
    },
    "Success": {
      "Type": "Succeed"
    },
    "Fail": {
      "Type": "Fail"
    }
  },
  "TimeoutSeconds": 5400
}

对应的更新API Lambda函数示例：

import json
import requests

def lambda_handler(event, context):
    task_id = event['task_id']
    batch_job_id = event['batch_job_id']
    
    # 调用DRF API更新job_id字段
    api_url = f"https://your-drf-api-endpoint/tasks/{task_id}/"
    headers = {"Content-Type": "application/json"}
    payload = {"job_id": batch_job_id}
    
    try:
        response = requests.patch(api_url, json=payload, headers=headers)
        response.raise_for_status()
        return {
            'statusCode': 200,
            'body': json.dumps('API updated successfully')
        }
    except Exception as e:
        raise Exception(f"Failed to update API: {str(e)}")

关键说明 #

• batch:submitJob异步任务提交成功后立即返回JobId，此时Job未进入Running状态，满足更新API的时间要求。
• 后续的wait_for_job_completion任务可以保持Step Function等待Job完成的原有行为，保证流程一致性。

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