Skip to content

Commit 0ea7780

Browse files
committed
feat: pass name from modelbuilder constructor to created model
1 parent 7aa39f9 commit 0ea7780

File tree

9 files changed

+28
-9
lines changed

9 files changed

+28
-9
lines changed

src/sagemaker/serve/builder/djl_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ def _create_djl_model(self) -> Type[Model]:
130130
huggingface_hub_token=self.env_vars.get("HF_TOKEN"),
131131
image_config=self.image_config,
132132
vpc_config=self.vpc_config,
133+
name=self.name
133134
)
134135

135136
if not self.image_uri:

src/sagemaker/serve/builder/jumpstart_builder.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,10 @@ def _is_jumpstart_model_id(self) -> bool:
147147
def _create_pre_trained_js_model(self) -> Type[Model]:
148148
"""Placeholder docstring"""
149149
pysdk_model = JumpStartModel(
150-
self.model, vpc_config=self.vpc_config, sagemaker_session=self.sagemaker_session
150+
self.model,
151+
name=self.name,
152+
vpc_config=self.vpc_config,
153+
sagemaker_session=self.sagemaker_session
151154
)
152155

153156
self._original_deploy = pysdk_model.deploy

src/sagemaker/serve/builder/model_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,7 @@ def _create_model(self):
492492
env=self.env_vars,
493493
sagemaker_session=self.sagemaker_session,
494494
predictor_cls=self._get_predictor,
495+
name=self.name,
495496
)
496497

497498
# store the modes in the model so that we may

src/sagemaker/serve/builder/tei_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ def _create_tei_model(self, **kwargs) -> Type[Model]:
105105
env=self.env_vars,
106106
role=self.role_arn,
107107
sagemaker_session=self.sagemaker_session,
108+
name=self.name
108109
)
109110

110111
logger.info("Detected %s. Proceeding with the the deployment.", self.image_uri)

src/sagemaker/serve/builder/tf_serving_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ def _create_tensorflow_model(self):
9797
env=self.env_vars,
9898
sagemaker_session=self.sagemaker_session,
9999
predictor_cls=self._get_tensorflow_predictor,
100+
name=self.name
100101
)
101102

102103
self.pysdk_model.mode = self.mode

src/sagemaker/serve/builder/tgi_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ def _create_tgi_model(self) -> Type[Model]:
142142
env=self.env_vars,
143143
role=self.role_arn,
144144
sagemaker_session=self.sagemaker_session,
145+
name=self.name
145146
)
146147

147148
self._original_deploy = pysdk_model.deploy

src/sagemaker/serve/builder/transformers_builder.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ def _create_transformers_model(self) -> Type[Model]:
105105
env=self.env_vars,
106106
role=self.role_arn,
107107
sagemaker_session=self.sagemaker_session,
108+
name=self.name
108109
)
109110

110111
logger.info("Detected %s. Proceeding with the the deployment.", self.image_uri)

tests/unit/sagemaker/serve/builder/test_model_builder.py

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -317,7 +317,7 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
317317
)
318318

319319
mock_model_obj = Mock()
320-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
320+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
321321
mock_model_obj
322322
if image_uri == mock_image_uri
323323
and image_config == MOCK_IMAGE_CONFIG
@@ -326,6 +326,7 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_byoc(
326326
and role == mock_role_arn
327327
and env == ENV_VARS
328328
and sagemaker_session == mock_session
329+
and "model-name-" in name
329330
else None
330331
)
331332

@@ -425,13 +426,14 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_1p_dlc_as_byoc(
425426
)
426427

427428
mock_model_obj = Mock()
428-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
429+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
429430
mock_model_obj
430431
if image_uri == mock_1p_dlc_image_uri
431432
and model_data == model_data
432433
and role == mock_role_arn
433434
and env == ENV_VARS
434435
and sagemaker_session == mock_session
436+
and "model-name-" in name
435437
else None
436438
)
437439

@@ -532,13 +534,14 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_and_inference_spec(
532534
)
533535

534536
mock_model_obj = Mock()
535-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
537+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
536538
mock_model_obj
537539
if image_uri == mock_image_uri
538540
and model_data == model_data
539541
and role == mock_role_arn
540542
and env == ENV_VARS_INF_SPEC
541543
and sagemaker_session == mock_session
544+
and "model-name-" in name
542545
else None
543546
)
544547

@@ -633,13 +636,14 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_model(
633636
)
634637

635638
mock_model_obj = Mock()
636-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
639+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
637640
mock_model_obj
638641
if image_uri == mock_image_uri
639642
and model_data == model_data
640643
and role == mock_role_arn
641644
and env == ENV_VARS
642645
and sagemaker_session == mock_session
646+
and "model-name-" in name
643647
else None
644648
)
645649

@@ -742,13 +746,14 @@ def test_build_happy_path_with_sagemakerEndpoint_mode_and_xgboost_model(
742746
)
743747

744748
mock_model_obj = Mock()
745-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
749+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
746750
mock_model_obj
747751
if image_uri == mock_image_uri
748752
and model_data == model_data
749753
and role == mock_role_arn
750754
and env == ENV_VARS
751755
and sagemaker_session == mock_session
756+
and "model-name-" in name
752757
else None
753758
)
754759

@@ -847,13 +852,14 @@ def test_build_happy_path_with_local_container_mode(
847852
mock_mode.prepare.side_effect = lambda: None
848853

849854
mock_model_obj = Mock()
850-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
855+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
851856
mock_model_obj
852857
if image_uri == mock_image_uri
853858
and model_data is None
854859
and role == mock_role_arn
855860
and env == {}
856861
and sagemaker_session == mock_session
862+
and "model-name-" in name
857863
else None
858864
)
859865

@@ -968,13 +974,14 @@ def test_build_happy_path_with_localContainer_mode_overwritten_with_sagemaker_mo
968974
)
969975

970976
mock_model_obj = Mock()
971-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
977+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
972978
mock_model_obj
973979
if image_uri == mock_image_uri
974980
and model_data is None
975981
and role == mock_role_arn
976982
and env == {}
977983
and sagemaker_session == mock_session
984+
and "model-name-" in name
978985
else None
979986
)
980987

@@ -1119,13 +1126,14 @@ def test_build_happy_path_with_sagemaker_endpoint_mode_overwritten_with_local_co
11191126
)
11201127

11211128
mock_model_obj = Mock()
1122-
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls: ( # noqa E501
1129+
mock_sdk_model.side_effect = lambda image_uri, image_config, vpc_config, model_data, role, env, sagemaker_session, predictor_cls, name: ( # noqa E501
11231130
mock_model_obj
11241131
if image_uri == mock_image_uri
11251132
and model_data == model_data
11261133
and role == mock_role_arn
11271134
and env == ENV_VARS
11281135
and sagemaker_session == mock_session
1136+
and "model-name-" in name
11291137
else None
11301138
)
11311139

tests/unit/sagemaker/serve/builder/test_tensorflow_serving_builder.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ def setUp(self):
3333
self.instance.image_config = {}
3434
self.instance.vpc_config = {}
3535
self.instance.modes = {}
36+
self.instance.name = "model-name-mock-uuid-hex"
3637

3738
@patch("os.makedirs")
3839
@patch("os.path.exists")
@@ -71,5 +72,6 @@ def test_create_tensorflow_model(self, mock_model):
7172
env=self.instance.env_vars,
7273
sagemaker_session=self.instance.sagemaker_session,
7374
predictor_cls=self.instance._get_tensorflow_predictor,
75+
name="model-name-mock-uuid-hex"
7476
)
7577
self.assertEqual(model, mock_model.return_value)

0 commit comments

Comments
 (0)