[Updated] Add telemetry to ModelTrainer, Estimator and ModelBuilder (#1608)

* [Updated] Add telemetry to ModelTrainer, Estimator and ModelBuilder

* fix unit test name

* Fix unit test

* Add handshake telemetry

* restyle code format

* remove debug logs

* remove debug code

Author: zhaoqizqwang

src/sagemaker/estimator.py

@@ -106,7 +106,8 @@
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
-
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 
 logger = logging.getLogger(__name__)
 
@@ -1297,6 +1298,7 @@ def latest_job_profiler_artifacts_path(self):
             )
         return None
 
+    @_telemetry_emitter(feature=Feature.ESTIMATOR, func_name="estimator.fit")
     @runnable_by_pipeline
     def fit(
         self,

src/sagemaker/modules/train/model_trainer.py

@@ -86,6 +86,8 @@
     EXEUCTE_TORCHRUN_DRIVER,
     EXECUTE_BASIC_SCRIPT_DRIVER,
 )
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 from sagemaker.modules import logger
 from sagemaker.modules.train.sm_recipes.utils import _get_args_from_recipe, _determine_device_type
 
@@ -117,7 +119,7 @@ class ModelTrainer(BaseModel):
     ```
 
     Attributes:
-        session (Optiona(Session)):
+        sagemaker_session (Optiona(Session)):
             The SageMaker session.
             If not specified, a new session will be created.
         role (Optional(str)):
@@ -181,7 +183,7 @@ class ModelTrainer(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
 
     training_mode: Mode = Mode.SAGEMAKER_TRAINING_JOB
-    session: Optional[Session] = None
+    sagemaker_session: Optional[Session] = None
     role: Optional[str] = None
     base_job_name: Optional[str] = None
     source_code: Optional[SourceCode] = None
@@ -299,12 +301,12 @@ def model_post_init(self, __context: Any):
         self._validate_distributed_runner(self.source_code, self.distributed_runner)
 
         if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB:
-            if self.session is None:
-                self.session = Session()
-                logger.warning("Session not provided. Using default Session.")
+            if self.sagemaker_session is None:
+                self.sagemaker_session = Session()
+                logger.warning("SageMaker session not provided. Using default Session.")
 
             if self.role is None:
-                self.role = get_execution_role(sagemaker_session=self.session)
+                self.role = get_execution_role(sagemaker_session=self.sagemaker_session)
                 logger.warning(f"Role not provided. Using default role:\n{self.role}")
 
         if self.base_job_name is None:
@@ -333,7 +335,7 @@ def model_post_init(self, __context: Any):
             )
 
         if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB and self.output_data_config is None:
-            session = self.session
+            session = self.sagemaker_session
             base_job_name = self.base_job_name
             self.output_data_config = OutputDataConfig(
                 s3_output_path=f"s3://{session.default_bucket()}/{base_job_name}",
@@ -348,6 +350,7 @@ def model_post_init(self, __context: Any):
         if self.training_image:
             logger.info(f"Training image URI: {self.training_image}")
 
+    @_telemetry_emitter(feature=Feature.MODEL_TRAINER, func_name="model_trainer.train")
     @validate_call
     def train(
         self,
@@ -451,7 +454,7 @@ def train(
                 resource_config=resource_config,
                 vpc_config=vpc_config,
                 # Public Instance Attributes
-                session=self.session.boto_session,
+                session=self.sagemaker_session.boto_session,
                 role_arn=self.role,
                 tags=self.tags,
                 stopping_condition=self.stopping_condition,
@@ -494,7 +497,7 @@ def train(
                 instance_count=resource_config.instance_count,
                 image=algorithm_specification.training_image,
                 container_root=self.local_container_root,
-                sagemaker_session=self.session,
+                sagemaker_session=self.sagemaker_session,
                 container_entrypoint=algorithm_specification.container_entrypoint,
                 container_arguments=algorithm_specification.container_arguments,
                 input_data_config=input_data_config,
@@ -539,9 +542,9 @@ def create_input_data_channel(self, channel_name: str, data_source: DataSourceTy
                         input_mode="File",
                     )
                 else:
-                    s3_uri = self.session.upload_data(
+                    s3_uri = self.sagemaker_session.upload_data(
                         path=data_source,
-                        bucket=self.session.default_bucket(),
+                        bucket=self.sagemaker_session.default_bucket(),
                         key_prefix=f"{self.base_job_name}/input/{channel_name}",
                     )
                     channel = Channel(
@@ -821,7 +824,7 @@ def from_recipe(
         training_input_mode: Optional[str] = "File",
         environment: Optional[Dict[str, str]] = None,
         tags: Optional[List[Tag]] = None,
-        session: Optional[Session] = None,
+        sagemaker_session: Optional[Session] = None,
         role: Optional[str] = None,
         base_job_name: Optional[str] = None,
     ) -> "ModelTrainer":
@@ -863,7 +866,7 @@ def from_recipe(
             tags (Optional[List[Tag]]):
                 An array of key-value pairs. You can use tags to categorize your AWS resources
                 in different ways, for example, by purpose, owner, or environment.
-            session (Optional[Session]):
+            sagemaker_session (Optional[Session]):
                 The SageMaker session.
                 If not specified, a new session will be created.
             role (Optional[str]):
@@ -885,9 +888,9 @@ def from_recipe(
                 + "Please provide a GPU or Tranium instance type."
             )
 
-        if session is None:
-            session = Session()
-            logger.warning("Session not provided. Using default Session.")
+        if sagemaker_session is None:
+            sagemaker_session = Session()
+            logger.warning("SageMaker session not provided. Using default Session.")
         if role is None:
             role = get_execution_role(sagemaker_session=session)
             logger.warning(f"Role not provided. Using default role:\n{role}")
@@ -903,13 +906,13 @@ def from_recipe(
             recipe_overrides=recipe_overrides,
             requirements=requirements,
             compute=compute,
-            region_name=session.boto_region_name,
+            region_name=sagemaker_session.boto_region_name,
         )
         if training_image is not None:
             model_trainer_args["training_image"] = training_image
 
         model_trainer = cls(
-            session=session,
+            sagemaker_session=sagemaker_session,
             role=role,
             base_job_name=base_job_name,
             training_image_config=training_image_config,

src/sagemaker/serve/builder/model_builder.py

@@ -836,11 +836,24 @@ def _initialize_for_mlflow(self, artifact_path: str) -> None:
         self.env_vars.update({"MLFLOW_MODEL_FLAVOR": f"{deployment_flavor}"})
         self.dependencies.update({"requirements": mlflow_model_dependency_path})
 
+    @_capture_telemetry("ModelBuilder.build_training_job")
+    def _collect_training_job_model_telemetry(self):
+        return
+
+    @_capture_telemetry("ModelBuilder.build_model_trainer")
+    def _collect_model_trainer_model_telemetry(self):
+        return
+
+    @_capture_telemetry("ModelBuilder.build_estimator")
+    def _collect_estimator_model_telemetry(self):
+        return
+
     # Model Builder is a class to build the model for deployment.
     # It supports three modes of deployment
     # 1/ SageMaker Endpoint
     # 2/ Local launch with container
     # 3/ In process mode with Transformers server in beta release
+    @_capture_telemetry("ModelBuilder.build")
     def build(  # pylint: disable=R0911
         self,
         mode: Type[Mode] = None,
@@ -868,15 +881,20 @@ def build(  # pylint: disable=R0911
         if role_arn:
             self.role_arn = role_arn
 
+        self.serve_settings = self._get_serve_setting()
+
         if isinstance(self.model, TrainingJob):
             self.model_path = self.model.model_artifacts.s3_model_artifacts
             self.model = None
+            self._collect_training_job_model_telemetry()
         elif isinstance(self.model, ModelTrainer):
             self.model_path = self.model._latest_training_job.model_artifacts.s3_model_artifacts
             self.model = None
+            self._collect_model_trainer_model_telemetry()
         elif isinstance(self.model, Estimator):
             self.model_path = self.model.output_path
             self.model = None
+            self._collect_estimator_model_telemetry()
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
 
@@ -899,7 +917,6 @@ def build(  # pylint: disable=R0911
             self.sagemaker_session.sagemaker_client._user_agent_creator.to_string
         )
 
-        self.serve_settings = self._get_serve_setting()
         self._is_custom_image_uri = self.image_uri is not None
 
         self._handle_mlflow_input()
@@ -1017,6 +1034,7 @@ def _build_for_model_server(self):  # pylint: disable=R0911, R1710
         if self.model_server == ModelServer.MMS:
             return self._build_for_transformers()
 
+    @_capture_telemetry("ModelBuilder.save")
     def save(
         self,
         save_path: Optional[str] = None,
@@ -1566,6 +1584,7 @@ def _optimize_prepare_for_hf(self):
         )
         self.pysdk_model.env.update(env)
 
+    @_capture_telemetry("ModelBuilder.deploy")
     def deploy(
         self,
         endpoint_name: str = None,

src/sagemaker/serve/utils/telemetry_logger.py

@@ -171,7 +171,7 @@ def wrapper(self, *args, **kwargs):
 
             extra += f"&x-latency={round(elapsed, 2)}"
 
-            if not self.serve_settings.telemetry_opt_out:
+            if hasattr(self, "serve_settings") and not self.serve_settings.telemetry_opt_out:
                 _send_telemetry(
                     status,
                     MODE_TO_CODE[str(self.mode)],

src/sagemaker/telemetry/constants.py

@@ -25,6 +25,8 @@ class Feature(Enum):
     SDK_DEFAULTS = 1
     LOCAL_MODE = 2
     REMOTE_FUNCTION = 3
+    MODEL_TRAINER = 4
+    ESTIMATOR = 5
 
     def __str__(self):  # pylint: disable=E0307
         """Return the feature name."""

src/sagemaker/telemetry/telemetry_logging.py

@@ -52,6 +52,8 @@
     str(Feature.SDK_DEFAULTS): 1,
     str(Feature.LOCAL_MODE): 2,
     str(Feature.REMOTE_FUNCTION): 3,
+    str(Feature.MODEL_TRAINER): 4,
+    str(Feature.ESTIMATOR): 5,
 }
 
 STATUS_TO_CODE = {
@@ -61,7 +63,13 @@
 
 
 def _telemetry_emitter(feature: str, func_name: str):
-    """Decorator to emit telemetry logs for SageMaker Python SDK functions"""
+    """
+    Decorator to emit telemetry logs for SageMaker Python SDK functions. This class needs
+    sagemaker_session object as a member. Default session object is a pysdk v2 Session object
+    in this repo. When collecting telemetry for classes using sagemaker-core Session object,
+    we should be aware of its differences, such as sagemaker_session.sagemaker_config does not
+    exist in new Session class.
+    """
 
     def decorator(func):
         @functools.wraps(func)
@@ -95,10 +103,18 @@ def wrapper(*args, **kwargs):
                 # Construct the feature list to track feature combinations
                 feature_list: List[int] = [FEATURE_TO_CODE[str(feature)]]
 
-                if sagemaker_session.sagemaker_config and feature != Feature.SDK_DEFAULTS:
+                if (
+                    hasattr(sagemaker_session, "sagemaker_config")
+                    and sagemaker_session.sagemaker_config
+                    and feature != Feature.SDK_DEFAULTS
+                ):
                     feature_list.append(FEATURE_TO_CODE[str(Feature.SDK_DEFAULTS)])
 
-                if sagemaker_session.local_mode and feature != Feature.LOCAL_MODE:
+                if (
+                    hasattr(sagemaker_session, "local_mode")
+                    and sagemaker_session.local_mode
+                    and feature != Feature.LOCAL_MODE
+                ):
                     feature_list.append(FEATURE_TO_CODE[str(Feature.LOCAL_MODE)])
 
                 # Construct the extra info to track platform and environment usage metadata
@@ -111,7 +127,7 @@ def wrapper(*args, **kwargs):
                 )
 
                 # Add endpoint ARN to the extra info if available
-                if sagemaker_session.endpoint_arn:
+                if hasattr(sagemaker_session, "endpoint_arn") and sagemaker_session.endpoint_arn:
                     extra += f"&x-endpointArn={sagemaker_session.endpoint_arn}"
 
                 start_timer = perf_counter()
@@ -171,8 +187,9 @@ def _send_telemetry_request(
 ) -> None:
     """Make GET request to an empty object in S3 bucket"""
     try:
-        accountId = _get_accountId(session)
-        region = _get_region_or_default(session)
+        accountId = _get_accountId(session) if session else "NotAvailable"
+        # telemetry will be sent to us-west-2 if no session availale
+        region = _get_region_or_default(session) if session else DEFAULT_AWS_REGION
         url = _construct_url(
             accountId,
             region,

src/sagemaker/utils.py

@@ -1160,7 +1160,7 @@ def get_sagemaker_config_value(sagemaker_session, key, sagemaker_config: dict =
     Returns:
         object: The corresponding default value in the configuration file.
     """
-    if sagemaker_session:
+    if sagemaker_session and hasattr(sagemaker_session, "sagemaker_config"):
         config_to_check = sagemaker_session.sagemaker_config
     else:
         config_to_check = sagemaker_config

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -146,9 +146,9 @@ def model_trainer():
 def test_model_trainer_param_validation(test_case, modules_session):
     if test_case["should_throw"]:
         with pytest.raises(ValueError):
-            ModelTrainer(**test_case["init_params"], session=modules_session)
+            ModelTrainer(**test_case["init_params"], sagemaker_session=modules_session)
     else:
-        trainer = ModelTrainer(**test_case["init_params"], session=modules_session)
+        trainer = ModelTrainer(**test_case["init_params"], sagemaker_session=modules_session)
         assert trainer is not None
         assert trainer.training_image == DEFAULT_IMAGE
         assert trainer.compute == DEFAULT_COMPUTE_CONFIG
@@ -261,7 +261,7 @@ def test_metric_settings(mock_training_job, modules_session):
 
     model_trainer = ModelTrainer(
         training_image=image_uri,
-        session=modules_session,
+        sagemaker_session=modules_session,
         role=role,
     ).with_metric_settings(
         enable_sage_maker_metrics_time_series=True, metric_definitions=[metric_definition]
@@ -306,7 +306,7 @@ def test_debugger_settings(mock_training_job, modules_session):
 
     model_trainer = ModelTrainer(
         training_image=image_uri,
-        session=modules_session,
+        sagemaker_session=modules_session,
         role=role,
     ).with_debugger_settings(
         debug_hook_config=debug_hook_config,
@@ -367,7 +367,7 @@ def test_additional_settings(mock_training_job, modules_session):
     )
     model_trainer = ModelTrainer(
         training_image=image_uri,
-        session=modules_session,
+        sagemaker_session=modules_session,
         role=role,
     ).with_additional_settings(
         retry_strategy=retry_strategy,
@@ -467,7 +467,7 @@ def test_train_with_distributed_runner(
 
     try:
         model_trainer = ModelTrainer(
-            session=modules_session,
+            sagemaker_session=modules_session,
             training_image=DEFAULT_IMAGE,
             source_code=test_case["source_code"],
             distributed_runner=test_case["distributed_runner"],