AWS SageMaker Step Decorator模式下注册部署Scikit-learn线性回归模型时ModelBuilder参数冲突问题排查

我们团队正在基于AWS SageMaker的Step Decorator方案搭建MLOps管道，用来实现从数据获取到模型部署的全流程自动化。目前管道已经完成了Athena数据拉取、特征工程、Scikit-learn线性回归模型训练、评估报告生成等环节，所有步骤都能正常运行，中间产物也都正确存储在S3中。

不过在添加最后一步——注册并部署训练好的模型时，遇到了SageMaker ModelBuilder类抛出的矛盾错误，折腾了很久也没找到原因，特来求助：

我们的注册部署步骤代码 #

我们用Step Decorator定义了如下的注册部署步骤：

@step(
    name="register",
    instance_type=instance_type,
    keep_alive_period_in_seconds=300,
)
def register(pipeline_execution_base_path, model, pickled_model_path, eval_report_path, model_approval_status, test_s3_path):
    import json
    import numpy as np
    import pandas as pd
    from pathlib import Path
    from sagemaker import MetricsSource, ModelMetrics
    from sagemaker.serve.builder.model_builder import ModelBuilder
    from sagemaker.serve.builder.model_builder import ModelServer
    from sagemaker.serve.builder.schema_builder import SchemaBuilder
    from sagemaker.serve.spec.inference_spec import InferenceSpec
    from sagemaker.utils import unique_name_from_base
    from s3fs import S3FileSystem
    from sklearn.linear_model import LinearRegression
    import pickle

    s3 = S3FileSystem()

    class ModelInferenceSpec(InferenceSpec):
        def load(self, model_dir: str):
            print(model_dir)
            model = pickle.load(s3.open(model_dir + "/model.pkl", 'rb'))
            return model

        def invoke(self, input_object: object, model: object):
            predictions = model.predict(input_object)
            return predictions

    # 基于S3中的评估报告创建模型指标
    model_metrics = ModelMetrics(
        model_statistics=MetricsSource(
            s3_uri=eval_report_path,
            content_type="application/json",
        )
    )

    # 用测试集样本生成Schema
    x_cols = ['x_variable_1', 'x_variable_2', 'x_variable_3', 'x_variable_4']
    y_col = 'y_variable'
    sample_data = pd.read_csv(test_s3_path, nrows=50)
    sample_data.pop("y_variable")
    schema_builder = SchemaBuilder(
        sample_input=sample_data[x_cols].to_numpy(),
        sample_output=model.predict(sample_data[x_cols]),
    )

    # 临时保存模型到本地目录
    model_path = Path("/tmp/model/")
    model_path.mkdir(parents=True, exist_ok=True)
    with open(f"{model_path}/model.pkl", 'wb') as f:
        pickle.dump(model, f)

    # 构建、注册模型
    model_package_path = f"{pipeline_execution_base_path}/model_package/model-artifacts"
    model_builder = ModelBuilder(
        model_path=str(model_path),
        inference_spec=ModelInferenceSpec(),
        schema_builder=schema_builder,
        role_arn=role,
        s3_model_data_url=model_package_path,
        image_uri="141502667606.dkr.ecr.eu-west-1.amazonaws.com/sagemaker-scikit-learn:0.23-1-cpu-py3",
    )
    
    model_package = model_builder.build().register(
        model_package_group_name=model_package_group_name,
        approval_status=model_approval_status,
        model_metrics=model_metrics,
    )

    return model_package.model_package_arn

补充说明：

• 这里的model参数是上游训练步骤输出的sklearn.linear_model.LinearRegression实例
• 我们没有用SageMaker内置算法，所以通过image_uris.retrieve()获取了Scikit-learn的官方镜像作为image_uri

遇到的矛盾错误 #

1. 添加inference_spec时的错误
   运行管道时，SageMaker控制台抛出如下错误：

1740658570967 | [ 2025-02-27T12:16:10.967Z ] ValueError: Can only set one of the following: model, inference_spec.

1740658570967 | [ 2025-02-27T12:16:10.967Z ] Traceback (most recent call last): File "/opt/conda/lib/python3.11/site-packages/sagemaker/remote_function/invoke_function.py", line 144, in main _execute_remote_function( File "/opt/conda/lib/python3.11/site-packages/sagemaker/remote_function/invoke_function.py", line 119, in _execute_remote_function stored_function.load_and_invoke() File "/opt/conda/lib/python3.11/site-packages/sagemaker/remote_function/core/stored_function.py", line 183, in load_and_invoke result = func(*resolved_args, **resolved_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/tmp/ipykernel_21287/2654039261.py", line 86, in register File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/utils/telemetry_logger.py", line 116, in wrapper raise e File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/utils/telemetry_logger.py", line 104, in wrapper response = func(self, *args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/builder/model_builder.py", line 929, in build self._build_validations() File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/builder/model_builder.py", line 998, in _build_validations raise ValueError("Can only set one of the following: model, inference_spec.")

1740658570967 | [ 2025-02-27T12:16:10.967Z ] 2025-02-27 12:16:10,566 sagemaker.remote_function ERROR Error encountered while invoking the remote function.

这个错误很奇怪，因为我们明明没有给ModelBuilder传递model参数，只传了inference_spec，却被提示不能同时设置两者。

2. 移除inference_spec后的错误
   为了排查，我们注释掉了inference_spec=ModelInferenceSpec()这一行，结果又抛出了新的错误：

1740657870421 | [ 2025-02-27T12:04:30.421Z ] ValueError: Cannot detect required model or inference spec

1740657870421 | [ 2025-02-27T12:04:30.421Z ] Traceback (most recent call last): File "/opt/conda/lib/python3.11/site-packages/sagemaker/remote_function/invoke_function.py", line 144, in main _execute_remote_function( File "/opt/conda/lib/python3.11/site-packages/sagemaker/remote_function/invoke_function.py", line 119, in _execute_remote_function stored_function.load_and_invoke() File "/opt/conda/lib/python3.11/site-packages/sagemaker/remote_function/core/stored_function.py", line 183, in load_and_invoke result = func(*resolved_args, **resolved_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/tmp/ipykernel_20133/2156326203.py", line 83, in register File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/utils/telemetry_logger.py", line 116, in wrapper raise e File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/utils/telemetry_logger.py", line 104, in wrapper response = func(self, *args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/builder/model_builder.py", line 990, in build self.built_model = self._build_for_torchserve() ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/builder/model_builder.py", line 658, in _build_for_torchserve self._save_model_inference_spec() File "/opt/conda/lib/python3.11/site-packages/sagemaker/serve/builder/model_builder.py", line 351, in _save_model_inference_spec raise ValueError("Cannot detect required model or inference spec")

1740657870421 | [ 2025-02-27 12:04:29,842 sagemaker.remote_function ERROR Error encountered while invoking the remote function.

现在又提示找不到model或inference_spec，陷入了两难的矛盾中。

我们的环境是：Scikit-learn 1.5.2，Python 3.11.11，SageMaker SDK用的是最新稳定版。已经翻了很多AWS官方文档，但都没找到对应的解决方案，希望能得到指点。

备注：内容来源于stack exchange，提问作者Jimmy