calling register() on SKLearn Estimator fails because of SKLearnModel().register() is called with too many arguments

[thisislb]

Describe the bug When trying to register a model package by calling register() on the SKLearn Estimator which was already used for training and deploying a model to an endpoint, the registration of the model fails because SKLearn().register() calls SKLearnModel().register() with too many arguments since the signature of SKLearnModel.register() does not allow additional arguments.

To reproduce

1. fit and deploy some model using sagemaker.sklearn.estimator.SKLearn (works)
2. create model package group (works)
3. try to register a model package from estimator (FAILS)

estimator.register(
    content_types=[content_types.JSON],
    response_types=[content_types.JSON, content_types.CSV, content_types.NPY],
    inference_instances=["ml.t2.medium"],
    transform_instances=["ml.t2.medium"],
    model_package_group_name=model_pgr_arn,
    approval_status="Approved",
)

Expected behavior That a model package is created from the SKLearn estimator.

Screenshots or logs

[sagemaker/lib/python3.9/site-packages/sagemaker/estimator.py:1324, in EstimatorBase.register(self, content_types, 
response_types, inference_instances, transform_instances, image_uri, model_package_name, model_package_group_name, 
model_metrics, metadata_properties, marketplace_cert, approval_status, description, compile_model_family, model_name, 
drift_check_baselines, customer_metadata_properties, **kwargs)
   ]()[1322](sagemaker/lib/python3.9/site-packages/sagemaker/estimator.py?line=1321)[     model = self.create_model(image_uri=image_uri, **kwargs)
   ]()[1323](.../estimator.py?line=1322)[ model.name = model_name
-> ]()[1324](.../estimator.py?line=1323)[ return model.register(
   ]()[1325](.../estimator.py?line=1324)[     content_types,
   ]()[1326](.../estimator.py?line=1325)[     response_types,
   ]()[1327](.../estimator.py?line=1326)[     inference_instances,
   ]()[1328](.../estimator.py?line=1327)[     transform_instances,
   ]()[1329](.../estimator.py?line=1328)[     model_package_name,
   ]()[1330](.../estimator.py?line=1329)[     model_package_group_name,
   ]()[1331](.../estimator.py?line=1330)[     image_uri,
   ]()[1332](.../estimator.py?line=1331)[     model_metrics,
   ]()[1333](.../estimator.py?line=1332)[     metadata_properties,
   ]()[1334](.../estimator.py?line=1333)[     marketplace_cert,
   ]()[1335](.../estimator.py?line=1334)[     approval_status,
   ]()[1336](.../estimator.py?line=1335)[     description,
   ]()[1337](.../estimator.py?line=1336)[     drift_check_baselines=drift_check_baselines,
   ]()[1338](.../estimator.py?line=1337)[     customer_metadata_properties=customer_metadata_properties,
   ]()[1339](.../estimator.py?line=1338)[ )
TypeError: register() got an unexpected keyword argument 'drift_check_baselines']()

https://github.com/aws/sagemaker-python-sdk/blob/acadfd342abeeb9a3f0d4a32f595a03af6d1f6eb/src/sagemaker/estimator.py#L1324-L1339

imho calls

https://github.com/aws/sagemaker-python-sdk/blob/acadfd342abeeb9a3f0d4a32f595a03af6d1f6eb/src/sagemaker/sklearn/model.py#L140-L154

which doesn't accept additional arguments. maybe adding **kwargs could fix it?

System information A description of your system. Please provide:

• SageMaker Python SDK version:: 2.86.2
• Framework name: scikit-learn
• Framework version: 0.23-1
• Python version: py3 for estimator / locally 3.9.10
• CPU or GPU: CPU
• Custom Docker image (Y/N): No / used script mode

Additional context None

Maybe someone has time to confirm this error, or tell me if it is some user error by me? Thanks.

[davnn]

Maybe someone has time to confirm this error, or tell me if it is some user error by me?

I'm getting the same error.

[portisto]

I reproduced the error and in order for the cell to work it needs to be changed to:

inference_instance_type = "ml.m5.xlarge"
model_package_1 = sklearn_estimator_1.register(
    model_package_group_name=model_package_arn,
    inference_instances=[inference_instance_type],
    transform_instances=[inference_instance_type],
    content_types=["text/csv"],
    response_types=["text/csv"],
    approval_status="PendingManualApproval",
    
)


model_package_arn_1 = model_package_1.model_package_arn
print("Model Package ARN : ", model_package_arn_1)

This is because the model group that is created in the Sagemaker Registry gets created without a version, and the default status is "PendingManualApproval" . This change will make your cell work, however, the newly created model version after this step should be changed to "Validated" in order for the rest of the notebook to work. HTH.

[alfredolozano]

hi this error still persists, even with @portisto suggestion