部署SageMaker端点时如何传递环境变量？PyTorch模型读取失败问题

问题原因 #

1. 环境变量读取时机错误 #

SageMaker推理容器的启动流程是先加载代码包，再注入自定义环境变量。如果你在脚本的全局作用域（函数外的顶层代码）直接调用os.environ['MODEL_ARCHITECTURE']，读取逻辑会在环境变量注入前执行，自然读不到参数。

2. 旧版SDK参数传递逻辑问题 #

你使用的PyTorch 1.8.1对应的旧版本SageMaker Python SDK，PyTorchModel初始化时传入的env参数不会自动同步到端点部署配置，需要部署时显式二次传入。

解决方案 #

方案一：调整读取时机+显式传参（推荐，改动最小） #

1. 把所有环境变量读取逻辑移到model_fn函数内部，SageMaker会在环境变量全部注入完成后才调用model_fn加载模型，此时读取不会失败：

# inference.py 示例
import os
from model import load_model

# 不要在全局作用域读取环境变量，执行时机太早
# ARCH = os.environ['MODEL_ARCHITECTURE']

def model_fn(model_dir):
    # 放到model_fn内读取
    arch = os.environ['MODEL_ARCHITECTURE']
    model = load_model(arch, model_dir)
    return model

2. 调用deploy方法部署端点时，显式传入环境变量，兼容旧版SDK逻辑：

predictor = inference_model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.xlarge", # 替换为你的实例类型
    env={'MODEL_ARCHITECTURE': best_architecture_name}
)

方案二：动态写入配置文件（兼容性最强，适合交付场景） #

如果担心环境变量受容器、SDK版本影响，可直接把架构名称动态写入代码包的配置文件，完全绕过环境变量依赖：

1. 定义PyTorchModel前，动态生成配置文件到你的serve目录：

import json
import os

config = {"model_architecture": best_architecture_name}
with open(os.path.join("serve", "arch_config.json"), "w", encoding="utf-8") as f:
    json.dump(config, f)

# 再正常定义PyTorchModel，source_dir会自动打包生成的配置文件
inference_model = PyTorchModel(
    entry_point="inference.py",
    source_dir="serve",
    role=role,
    model_data=model_data_s3_path,
    framework_version="1.8.1",
    py_version="py36",
)

2. 推理脚本中直接读取配置文件即可：

import json
import os

def model_fn(model_dir):
    config_path = os.path.join(os.path.dirname(__file__), "arch_config.json")
    with open(config_path, "r", encoding="utf-8") as f:
        arch = json.load(f)["model_architecture"]
    # 后续模型加载逻辑

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