From 488b2bb3e3cdeb7649aa35edaf2d3fd0486debba Mon Sep 17 00:00:00 2001 From: ZealSV Date: Wed, 15 Jul 2026 11:18:42 -0700 Subject: [PATCH] fix(serve): dedicated INFERENCE_RECOMMENDER telemetry feature + type workload param Address review nits on the AI inference recommender: - Add Feature.INFERENCE_RECOMMENDER telemetry enum + FEATURE_TO_CODE mapping; tag generate_deployment_recommendations, start_benchmark, and show_result with it instead of MODEL_CUSTOMIZATION. - Type generate_deployment_recommendations(workload=...) as Optional[Union[Workload, str]] to match the other params. --- sagemaker-core/src/sagemaker/core/telemetry/constants.py | 1 + .../src/sagemaker/core/telemetry/telemetry_logging.py | 1 + .../serve/ai_inference_recommender/_model_builder_methods.py | 2 +- .../src/sagemaker/serve/ai_inference_recommender/jobs.py | 4 ++-- sagemaker-serve/src/sagemaker/serve/model_builder.py | 4 ++-- 5 files changed, 7 insertions(+), 5 deletions(-) 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,