diff --git a/sagemaker-core/src/sagemaker/core/telemetry/constants.py b/sagemaker-core/src/sagemaker/core/telemetry/constants.py index 04231ef8a0..4b9b9d3e66 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/constants.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/constants.py @@ -32,6 +32,7 @@ class Feature(Enum): PROCESSING = 18 MODEL_CUSTOMIZATION_NOVA = 19 MODEL_CUSTOMIZATION_OSS = 20 + INFERENCE_RECOMMENDER = 21 def __str__(self): # pylint: disable=E0307 """Return the feature name.""" diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index 96196e1182..aa44bac0af 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -76,6 +76,7 @@ str(Feature.PROCESSING): 18, str(Feature.MODEL_CUSTOMIZATION_NOVA): 19, str(Feature.MODEL_CUSTOMIZATION_OSS): 20, + str(Feature.INFERENCE_RECOMMENDER): 21, } STATUS_TO_CODE = { diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/_model_builder_methods.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/_model_builder_methods.py index c4ce3ba29f..5bbe4907e6 100644 --- a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/_model_builder_methods.py +++ b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/_model_builder_methods.py @@ -81,7 +81,7 @@ def _map_service_error(error: Exception) -> Exception: @_telemetry_emitter( - feature=Feature.MODEL_CUSTOMIZATION, func_name="ai_inference_recommender.start_benchmark" + feature=Feature.INFERENCE_RECOMMENDER, func_name="ai_inference_recommender.start_benchmark" ) def start_benchmark( endpoint: Union[Endpoint, str], diff --git a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/jobs.py b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/jobs.py index 5b5c459128..ec06111740 100644 --- a/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/jobs.py +++ b/sagemaker-serve/src/sagemaker/serve/ai_inference_recommender/jobs.py @@ -34,7 +34,7 @@ class BenchmarkJob(AIBenchmarkJob): """ @_telemetry_emitter( - feature=Feature.MODEL_CUSTOMIZATION, func_name="BenchmarkJob.show_result" + feature=Feature.INFERENCE_RECOMMENDER, func_name="BenchmarkJob.show_result" ) def show_result(self): """Download the benchmark output from S3 and return a parsed result. @@ -58,7 +58,7 @@ class RecommendationJob(AIRecommendationJob): """ @_telemetry_emitter( - feature=Feature.MODEL_CUSTOMIZATION, func_name="RecommendationJob.show_result" + feature=Feature.INFERENCE_RECOMMENDER, func_name="RecommendationJob.show_result" ) def show_result(self) -> "_RecommendationsView": """Return the ranked recommendations produced by the job. diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 275b09707e..453b097b73 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -4427,12 +4427,12 @@ def _model_builder_optimize_wrapper( return self.built_model @_telemetry_emitter( - feature=Feature.MODEL_CUSTOMIZATION, + feature=Feature.INFERENCE_RECOMMENDER, func_name="ModelBuilder.generate_deployment_recommendations", ) def generate_deployment_recommendations( self, - workload=None, + workload: Optional[Union["Workload", str]] = None, performance_target: Optional[Union[str, "PerformanceTarget"]] = None, *, output_path: Optional[str] = None,