Skip to content

Commit 844b558

Browse files
Aditi2424adishaa
andauthored
Fix: Object of type ModelLifeCycle is not JSON serializable (#5197)
* Fix: Object of type ModelLifeCycle is not JSON serializable * Fix unit test * Fix integ tests * Revert "Fix integ tests" This reverts commit f6513fe. * Fix integration tests --------- Co-authored-by: adishaa <[email protected]>
1 parent 829030a commit 844b558

File tree

4 files changed

+9
-9
lines changed

4 files changed

+9
-9
lines changed

src/sagemaker/session.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7509,7 +7509,7 @@ def get_model_package_args(
75097509
if source_uri is not None:
75107510
model_package_args["source_uri"] = source_uri
75117511
if model_life_cycle is not None:
7512-
model_package_args["model_life_cycle"] = model_life_cycle
7512+
model_package_args["model_life_cycle"] = model_life_cycle._to_request_dict()
75137513
if model_card is not None:
75147514
original_req = model_card._create_request_args()
75157515
if original_req.get("ModelCardName") is not None:

tests/integ/sagemaker/workflow/test_model_create_and_registration.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
from sagemaker.s3 import S3Uploader
4949
from sagemaker.sklearn import SKLearnModel, SKLearnProcessor
5050
from sagemaker.mxnet.model import MXNetModel
51+
from sagemaker.model_life_cycle import ModelLifeCycle
5152
from sagemaker.workflow.condition_step import ConditionStep
5253
from sagemaker.workflow.parameters import ParameterInteger, ParameterString
5354
from sagemaker.workflow.pipeline import Pipeline
@@ -1005,11 +1006,11 @@ def test_model_registration_with_model_life_cycle_object(
10051006
py_version="py3",
10061007
role=role,
10071008
)
1008-
create_model_life_cycle = {
1009-
"Stage": "Development",
1010-
"StageStatus": "In-Progress",
1011-
"StageDescription": "Development In Progress",
1012-
}
1009+
create_model_life_cycle = ModelLifeCycle(
1010+
stage="Development",
1011+
stage_status="In-Progress",
1012+
stage_description="Development In Progress",
1013+
)
10131014

10141015
step_register = RegisterModel(
10151016
name="MyRegisterModelStep",

tests/integ/test_model_package.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ def test_update_model_life_cycle_model_package(sagemaker_session):
103103
inference_instances=["ml.m5.large"],
104104
transform_instances=["ml.m5.large"],
105105
model_package_group_name=model_group_name,
106-
model_life_cycle=create_model_life_cycle._to_request_dict(),
106+
model_life_cycle=create_model_life_cycle,
107107
)
108108

109109
desc_model_package = sagemaker_session.sagemaker_client.describe_model_package(

tests/unit/test_estimator.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4369,7 +4369,6 @@ def test_register_default_image(sagemaker_session):
43694369
stage_status="In-Progress",
43704370
stage_description="Sending for Staging Verification",
43714371
)
4372-
update_model_life_cycle_req = update_model_life_cycle._to_request_dict()
43734372

43744373
estimator.register(
43754374
content_types=content_types,
@@ -4384,7 +4383,7 @@ def test_register_default_image(sagemaker_session):
43844383
nearest_model_name=nearest_model_name,
43854384
data_input_configuration=data_input_config,
43864385
model_card=model_card,
4387-
model_life_cycle=update_model_life_cycle_req,
4386+
model_life_cycle=update_model_life_cycle,
43884387
)
43894388
sagemaker_session.create_model.assert_not_called()
43904389
exp_model_card = {

0 commit comments

Comments
 (0)