Skip to content

Commit 679c58c

Browse files
gwang111gwang111
committed
add tags for inf and dep recommendations
1 parent 5e56f01 commit 679c58c

File tree

6 files changed

+280
-25
lines changed

6 files changed

+280
-25
lines changed

‎src/sagemaker/inference_recommender/inference_recommender_mixin.py

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@
3030

3131
LOGGER = logging.getLogger("sagemaker")
3232

33+
DEPLOYMENT_RECOMMENDATION_TAG = "PythonSDK-DeploymentRecommendation"
34+
35+
RIGHT_SIZE_TAG = "PythonSDK-RightSize"
3336

3437
class Phase:
3538
"""Used to store phases of a traffic pattern to perform endpoint load testing.
@@ -218,6 +221,7 @@ def _update_params(
218221
explainer_config = kwargs["explainer_config"]
219222
inference_recommendation_id = kwargs["inference_recommendation_id"]
220223
inference_recommender_job_results = kwargs["inference_recommender_job_results"]
224+
tags = kwargs["tags"]
221225
if inference_recommendation_id is not None:
222226
inference_recommendation = self._update_params_for_recommendation_id(
223227
instance_type=instance_type,
@@ -237,7 +241,11 @@ def _update_params(
237241
async_inference_config,
238242
explainer_config,
239243
)
240-
return inference_recommendation or (instance_type, initial_instance_count)
244+
245+
if inference_recommendation:
246+
tags = self._add_client_type_tag(tags, inference_recommendation[2])
247+
return (inference_recommendation[0], inference_recommendation[1], tags)
248+
return (instance_type, initial_instance_count, tags)
241249

242250
def _update_params_for_right_size(
243251
self,
@@ -301,7 +309,7 @@ def _update_params_for_right_size(
301309
initial_instance_count = self.inference_recommendations[0]["EndpointConfiguration"][
302310
"InitialInstanceCount"
303311
]
304-
return (instance_type, initial_instance_count)
312+
return (instance_type, initial_instance_count, "PythonSDK-RightSize")
305313

306314
def _update_params_for_recommendation_id(
307315
self,
@@ -401,7 +409,7 @@ def _update_params_for_recommendation_id(
401409
raise ValueError("Must specify model recommendation id and instance count.")
402410
self.env.update(model_recommendation["Environment"])
403411
instance_type = model_recommendation["InstanceType"]
404-
return (instance_type, initial_instance_count)
412+
return (instance_type, initial_instance_count, DEPLOYMENT_RECOMMENDATION_TAG)
405413

406414
# Update params based on default inference recommendation
407415
if bool(instance_type) != bool(initial_instance_count):
@@ -465,7 +473,7 @@ def _update_params_for_recommendation_id(
465473
"InitialInstanceCount"
466474
]
467475

468-
return (instance_type, initial_instance_count)
476+
return (instance_type, initial_instance_count, RIGHT_SIZE_TAG)
469477

470478
def _convert_to_endpoint_configurations_json(
471479
self, hyperparameter_ranges: List[Dict[str, CategoricalParameter]]
@@ -605,3 +613,11 @@ def _search_recommendation(self, recommendation_list, inference_recommendation_i
605613
),
606614
None,
607615
)
616+
617+
def _add_client_type_tag(self, tags, client_type):
618+
client_type_tag = {
619+
"Key": "ClientType",
620+
"Value": client_type
621+
}
622+
tags = tags.append(client_type_tag) if tags else [client_type_tag]
623+
return tags

‎src/sagemaker/model.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,7 +1232,7 @@ def deploy(
12321232
inference_recommendation_id is not None
12331233
or self.inference_recommender_job_results is not None
12341234
):
1235-
instance_type, initial_instance_count = self._update_params(
1235+
instance_type, initial_instance_count, tags = self._update_params(
12361236
instance_type=instance_type,
12371237
initial_instance_count=initial_instance_count,
12381238
accelerator_type=accelerator_type,
@@ -1241,6 +1241,7 @@ def deploy(
12411241
explainer_config=explainer_config,
12421242
inference_recommendation_id=inference_recommendation_id,
12431243
inference_recommender_job_results=self.inference_recommender_job_results,
1244+
tags=tags,
12441245
)
12451246

12461247
is_async = async_inference_config is not None
@@ -1735,7 +1736,7 @@ def _create_sagemaker_model(self, *args, **kwargs): # pylint: disable=unused-ar
17351736
17361737
Args:
17371738
args: Positional arguments coming from the caller. This class does not require
1738-
any so they are ignored.
1739+
any but will specifically look for Tags (3rd arg positionally) if specified
17391740
17401741
kwargs: Keyword arguments coming from the caller. This class does not require
17411742
any so they are ignored.
@@ -1766,6 +1767,7 @@ def _create_sagemaker_model(self, *args, **kwargs): # pylint: disable=unused-ar
17661767
container_def,
17671768
vpc_config=self.vpc_config,
17681769
enable_network_isolation=self.enable_network_isolation(),
1770+
tags=args[2],
17691771
)
17701772

17711773
def _ensure_base_name_if_needed(self, base_name):

‎tests/integ/test_inference_recommender.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,6 +303,30 @@ def default_right_sized_unregistered_base_model(sagemaker_session, cpu_instance_
303303
sagemaker_session.delete_model(ModelName=model.name)
304304

305305

306+
@pytest.fixture(scope="module")
307+
def created_base_model(sagemaker_session, cpu_instance_type):
308+
model_data = sagemaker_session.upload_data(path=IR_SKLEARN_MODEL)
309+
region = sagemaker_session._region_name
310+
image_uri = image_uris.retrieve(
311+
framework="sklearn", region=region, version="1.0-1", image_scope="inference"
312+
)
313+
314+
iam_client = sagemaker_session.boto_session.client("iam")
315+
role_arn = iam_client.get_role(RoleName="SageMakerRole")["Role"]["Arn"]
316+
317+
model = Model(
318+
model_data=model_data,
319+
role=role_arn,
320+
entry_point=IR_SKLEARN_ENTRY_POINT,
321+
image_uri=image_uri,
322+
sagemaker_session=sagemaker_session,
323+
)
324+
325+
model.create(instance_type=cpu_instance_type)
326+
327+
return model
328+
329+
306330
@pytest.mark.slow_test
307331
def test_default_right_size_and_deploy_registered_model_sklearn(
308332
default_right_sized_model, sagemaker_session
@@ -453,3 +477,48 @@ def test_deploy_inference_recommendation_id_with_registered_model_sklearn(
453477
)
454478
predictor.delete_model()
455479
predictor.delete_endpoint()
480+
481+
482+
@pytest.mark.slow_test
483+
def test_deploy_deployment_recommendation_id_with_model(created_base_model, sagemaker_session):
484+
with timeout(minutes=20):
485+
try:
486+
deployment_recommendation = poll_for_deployment_recommendation(created_base_model, sagemaker_session)
487+
488+
assert deployment_recommendation != None
489+
490+
real_time_recommendations = deployment_recommendation.get("RealTimeInferenceRecommendations")
491+
recommendation_id = real_time_recommendations[0].get('RecommendationId')
492+
493+
endpoint_name = unique_name_from_base("test-rec-id-deployment-default-sklearn")
494+
created_base_model.predictor_cls = SKLearnPredictor
495+
predictor = created_base_model.deploy(
496+
inference_recommendation_id=recommendation_id, initial_instance_count=1, endpoint_name=endpoint_name
497+
)
498+
499+
payload = pd.read_csv(IR_SKLEARN_DATA, header=None)
500+
501+
inference = predictor.predict(payload)
502+
assert inference is not None
503+
assert 26 == len(inference)
504+
finally:
505+
predictor.delete_model()
506+
predictor.delete_endpoint()
507+
508+
509+
def poll_for_deployment_recommendation(created_base_model, sagemaker_session):
510+
with timeout(minutes=1):
511+
try:
512+
completed = False
513+
while not completed:
514+
describe_model_response = sagemaker_session.sagemaker_client.describe_model(ModelName=created_base_model.name)
515+
deployment_recommendation = describe_model_response.get("DeploymentRecommendation")
516+
517+
completed = (
518+
deployment_recommendation is not None
519+
and "COMPLETED" == deployment_recommendation.get("RecommendationStatus")
520+
)
521+
return deployment_recommendation
522+
except Exception as e:
523+
created_base_model.delete_model()
524+
raise e

‎tests/unit/sagemaker/inference_recommender/constants.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,3 +152,44 @@
152152
"ModelArtifacts": {"S3ModelArtifacts": IR_COMPILATION_MODEL_DATA},
153153
"InferenceImage": IR_COMPILATION_IMAGE,
154154
}
155+
156+
IR_CONTAINER_DEF = {
157+
"Image": IR_IMAGE,
158+
"Environment": IR_ENV,
159+
"ModelDataUrl": IR_MODEL_DATA,
160+
}
161+
162+
DEPLOYMENT_RECOMMENDATION_CONTAINER_DEF = {
163+
"Image": IR_IMAGE,
164+
"Environment": MODEL_RECOMMENDATION_ENV,
165+
"ModelDataUrl": IR_MODEL_DATA,
166+
}
167+
168+
IR_COMPILATION_CONTAINER_DEF = {
169+
"Image": IR_COMPILATION_IMAGE,
170+
"Environment": {},
171+
"ModelDataUrl": IR_COMPILATION_MODEL_DATA,
172+
}
173+
174+
IR_MODEL_PACKAGE_CONTAINER_DEF = {
175+
"ModelPackageName": IR_MODEL_PACKAGE_VERSION_ARN,
176+
"Environment": IR_ENV,
177+
}
178+
179+
IR_COMPILATION_MODEL_PACKAGE_CONTAINER_DEF = {
180+
"ModelPackageName": IR_MODEL_PACKAGE_VERSION_ARN,
181+
}
182+
183+
IR_TAGS = [
184+
{
185+
"Key": "ClientType",
186+
"Value": "PythonSDK-RightSize",
187+
}
188+
]
189+
190+
DEPLOYMENT_RECOMMENDATION_TAGS = [
191+
{
192+
"Key": "ClientType",
193+
"Value": "PythonSDK-DeploymentRecommendation",
194+
}
195+
]

‎tests/unit/sagemaker/inference_recommender/test_inference_recommender_mixin.py

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,17 @@
177177
}
178178
]
179179

180+
IR_MODEL_PACKAGE_CONTAINER_DEF = {
181+
"ModelPackageName": MODEL_PACKAGE_ARN,
182+
}
183+
184+
IR_TAGS = [
185+
{
186+
"Key": "ClientType",
187+
"Value": "PythonSDK-RightSize",
188+
}
189+
]
190+
180191

181192
@pytest.fixture()
182193
def sagemaker_session():
@@ -371,6 +382,8 @@ def test_right_size_default_with_model_package_successful(sagemaker_session, mod
371382
framework=IR_SAMPLE_FRAMEWORK,
372383
)
373384

385+
sagemaker_session.create_model.assert_not_called()
386+
374387
# assert that the create api has been called with default parameters
375388
sagemaker_session.create_inference_recommendations_job.assert_called_with(
376389
role=IR_ROLE_ARN,
@@ -426,6 +439,8 @@ def test_right_size_advanced_list_instances_model_package_successful(
426439
max_parallel_tests=5,
427440
)
428441

442+
sagemaker_session.create_model.assert_not_called()
443+
429444
# assert that the create api has been called with advanced parameters
430445
sagemaker_session.create_inference_recommendations_job.assert_called_with(
431446
role=IR_ROLE_ARN,
@@ -481,6 +496,8 @@ def test_right_size_advanced_single_instances_model_package_successful(
481496
max_parallel_tests=5,
482497
)
483498

499+
sagemaker_session.create_model.assert_not_called()
500+
484501
# assert that the create api has been called with advanced parameters
485502
sagemaker_session.create_inference_recommendations_job.assert_called_with(
486503
role=IR_ROLE_ARN,
@@ -517,6 +534,8 @@ def test_right_size_advanced_model_package_partial_params_successful(
517534
model_latency_thresholds=IR_SAMPLE_MODEL_LATENCY_THRESHOLDS,
518535
)
519536

537+
sagemaker_session.create_model.assert_not_called()
538+
520539
# assert that the create api has been called with advanced parameters
521540
sagemaker_session.create_inference_recommendations_job.assert_called_with(
522541
role=IR_ROLE_ARN,
@@ -567,14 +586,23 @@ def test_deploy_right_size_with_model_package_succeeds(
567586
default_right_sized_model.name = MODEL_NAME
568587
default_right_sized_model.deploy(endpoint_name=IR_DEPLOY_ENDPOINT_NAME)
569588

589+
sagemaker_session.create_model.assert_called_with(
590+
MODEL_NAME,
591+
IR_ROLE_ARN,
592+
IR_MODEL_PACKAGE_CONTAINER_DEF,
593+
vpc_config=None,
594+
enable_network_isolation=False,
595+
tags=IR_TAGS,
596+
)
597+
570598
sagemaker_session.endpoint_from_production_variants.assert_called_with(
571599
async_inference_config_dict=None,
572600
data_capture_config_dict=None,
573601
explainer_config_dict=None,
574602
kms_key=None,
575603
name="ir-endpoint-test",
576604
production_variants=IR_PRODUCTION_VARIANTS,
577-
tags=None,
605+
tags=IR_TAGS,
578606
wait=True,
579607
)
580608

@@ -589,6 +617,15 @@ def test_deploy_right_size_with_both_overrides_succeeds(
589617
endpoint_name=IR_DEPLOY_ENDPOINT_NAME,
590618
)
591619

620+
sagemaker_session.create_model.assert_called_with(
621+
MODEL_NAME,
622+
IR_ROLE_ARN,
623+
IR_MODEL_PACKAGE_CONTAINER_DEF,
624+
vpc_config=None,
625+
enable_network_isolation=False,
626+
tags=None,
627+
)
628+
592629
sagemaker_session.endpoint_from_production_variants.assert_called_with(
593630
async_inference_config_dict=None,
594631
data_capture_config_dict=None,
@@ -637,6 +674,15 @@ def test_deploy_right_size_serverless_override(sagemaker_session, default_right_
637674
serverless_inference_config = ServerlessInferenceConfig()
638675
default_right_sized_model.deploy(serverless_inference_config=serverless_inference_config)
639676

677+
sagemaker_session.create_model.assert_called_with(
678+
MODEL_NAME,
679+
IR_ROLE_ARN,
680+
IR_MODEL_PACKAGE_CONTAINER_DEF,
681+
vpc_config=None,
682+
enable_network_isolation=False,
683+
tags=None,
684+
)
685+
640686
sagemaker_session.endpoint_from_production_variants.assert_called_with(
641687
name=MODEL_NAME,
642688
production_variants=IR_SERVERLESS_PRODUCTION_VARIANTS,
@@ -661,6 +707,15 @@ def test_deploy_right_size_async_override(sagemaker_session, default_right_sized
661707
async_inference_config=async_inference_config,
662708
)
663709

710+
sagemaker_session.create_model.assert_called_with(
711+
MODEL_NAME,
712+
IR_ROLE_ARN,
713+
IR_MODEL_PACKAGE_CONTAINER_DEF,
714+
vpc_config=None,
715+
enable_network_isolation=False,
716+
tags=None,
717+
)
718+
664719
sagemaker_session.endpoint_from_production_variants.assert_called_with(
665720
name=MODEL_NAME,
666721
production_variants=[ANY],
@@ -695,6 +750,15 @@ def test_deploy_right_size_explainer_config_override(sagemaker_session, default_
695750
explainer_config=explainer_config,
696751
)
697752

753+
sagemaker_session.create_model.assert_called_with(
754+
MODEL_NAME,
755+
IR_ROLE_ARN,
756+
IR_MODEL_PACKAGE_CONTAINER_DEF,
757+
vpc_config=None,
758+
enable_network_isolation=False,
759+
tags=None,
760+
)
761+
698762
sagemaker_session.endpoint_from_production_variants.assert_called_with(
699763
name=MODEL_NAME,
700764
production_variants=[ANY],

0 commit comments

Comments
 (0)