[Python] Cannot register and deploy scikit-learn linear regression model on AWS Sagemaker AI

We are evaluating AWS Sagemaker AI pipeline for MLOps, using the step decorator approach as introduced here, and mostly following the example here.

Since we are just starting out we'd prefer to have full flexibility in how we script our models, so we are not using any of the built-in algos as shown in the example above, but instead we import algos directly from scikit-learn, specifically we are just using a basic linear regression model (i.e. sklearn.linear_model)

Using the step decorator approach, we have built out a working pipeline that fetches data via Athena, performs feature engineering, trains a basic linear regression model, generates an model evaluation report, storing any outputs in S3 as we go, all of which are working fine.

However, when we try to add a final step to the pipeline that registers and deploys the trained model on AWS Sagemaker using the Sagemaker ModelBuilder class, we are getting errors that don't seem to make sense.

Our final step that registers and deploys the trained model looks like this:

@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

# Create model_metrics as per evaluation report in s3
model_metrics = ModelMetrics(
model_statistics=MetricsSource(
s3_uri=eval_report_path,
content_type="application/json",
)
)

# take a small sample from test, generate predictions off it, and pass those into SchemaBuilder which infers the I/O 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)

# Build the trained model and register it
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


As mentioned, model here is an instance of sklearn.linear_model, which we trained in a model training step that happens before this final step. Also, since we are not using one of the built-in algos, we understood that we have to supply an image_uri parameter to ModelBuilder, which we fetched via the Sagemaker SDK by:

from sagemaker import image_uris
image_uris.retrieve(framework='sklearn',region='eu-west-1',version='0.23-1',image_scope='inference')


However, when we add this final step to the pipeline and execute the whole thing, we are getting the following error off the Sagemaker pipeline UI:


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.


It seems to be complaining that we cannot supply both the model and inference_spec parameters to ModelBuilder, which is strange since we didn't use the model parameter. If we comment out the inference_spec parameter and execute the pipeline, we get this error instead:


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-27T12:04:30.421Z ] 2025-02-27 12:04:29,842 sagemaker.remote_function ERROR Error encountered while invoking the remote function.


So now it is complaining that we are supplying neither of model and inference_spec even though we only took out the inference_spec parameter.

We have looked at tons of AWS documentation but nothing seems to help.

We are using scikit-learn 1.5.2 and python 3.11.11.

Continue reading...