如何在Amazon SageMaker端点中使用multiprocessing的Pool？

针对你遇到的SageMaker端点调用时多进程函数超时的问题，结合SageMaker容器环境的特性，给出以下实用解决建议：

1. 切换multiprocessing启动方式为spawn #

SageMaker sklearn容器的主进程是长期运行的WSGI服务（如gunicorn），默认的fork启动模式会继承主进程的资源（如网络连接、锁），导致子进程卡住或异常。切换为spawn模式可以避免这类问题：

import multiprocessing

def multi_json(data, cores=32):
    # 强制设置spawn启动方式，避免继承主进程资源
    multiprocessing.set_start_method('spawn', force=True)
    with multiprocessing.Pool(processes=cores) as pool:
        json_recom = pool.map(create_json, data)
    return json_recom

2. 限制主服务进程数，避免CPU资源竞争 #

Sklearn官方容器默认会根据CPU核心数设置gunicorn的工作进程数（workers），如果主服务有多个工作进程，再加上你的多进程任务，会导致CPU过载，任务超时。在模型部署时通过环境变量将主服务进程数设为1：

model = sagemaker.model.Model(
    # 其他参数不变
    env={
        'PROJECT': 'Test',
        'ENVIRONMENT': 'Test',
        # 设置gunicorn仅启动1个工作进程
        'SAGEMAKER_SKLEARN_SERVER_ARGS': '--workers=1'
    }, 
    sagemaker_session=sagemaker_session
)

3. 调整异步推理超时配置 #

异步推理任务有默认超时限制（通常15分钟），如果你的多进程任务耗时超过这个时间，需要显式设置更长的超时时间：

from sagemaker.async_inference import AsyncInferenceConfig

asyncConfig = AsyncInferenceConfig(
    output_path=f"s3://{BUCKET}/async-output",
    # 设置任务超时时间（单位：秒，示例为1小时）
    invocation_timeout=3600,
    max_concurrent_invocations_per_instance=1
)

4. 验证函数的可序列化性 #

spawn模式下，multiprocessing需要将函数和数据序列化为pickle格式。如果create_json函数是嵌套函数或包含不可序列化的对象，会导致子进程启动失败。你可以在本地用spawn模式测试函数是否正常：

import multiprocessing
multiprocessing.set_start_method('spawn', force=True)

# 本地测试multi_json函数
test_data = [...]  # 模拟你的输入数据
result = multi_json(test_data)

如果测试失败，需将create_json改为顶层函数，或使用dill替代默认pickler（需在requirements.txt中添加dill依赖）：

import multiprocessing
import dill

# 替换默认pickler
multiprocessing.reduction.ForkingPickler = dill.Pickler
multiprocessing.set_start_method('spawn', force=True)

def multi_json(data, cores=32):
    with multiprocessing.Pool(processes=cores) as pool:
        json_recom = pool.map(create_json, data)
    return json_recom

5. 查看CloudWatch日志定位具体错误 #

登录AWS CloudWatch，找到对应SageMaker端点的日志组（通常命名为/aws/sagemaker/Endpoints/test-endpoint），查看模型容器的日志，可获取子进程启动时的具体报错（如资源不足、序列化失败等），帮助精准定位问题。

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