Add telemetry annotation for estimator relied functions to track esti…

src/sagemaker/chainer/estimator.py

@@ -28,6 +28,9 @@
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -165,6 +168,7 @@ def hyperparameters(self):
         )
         return hyperparameters
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "chainer.estimator.create_model")
     def create_model(
         self,
         model_server_workers=None,

src/sagemaker/estimator.py

@@ -106,6 +106,9 @@
 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__)
 
 
@@ -1527,6 +1530,7 @@ def logs(self):
         """
         self.sagemaker_session.logs_for_job(self.latest_training_job.name, wait=True)
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "sagemaker.estimator.deploy")
     def deploy(
         self,
         initial_instance_count=None,

src/sagemaker/huggingface/estimator.py

@@ -28,6 +28,9 @@
 from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig
 from sagemaker.workflow.entities import PipelineVariable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -342,6 +345,7 @@ def hyperparameters(self):
 
         return hyperparameters
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "huggingface.estimator.create_model")
     def create_model(
         self,
         model_server_workers=None,

src/sagemaker/job.py

@@ -20,6 +20,9 @@
 from sagemaker.local import file_input
 from sagemaker.workflow import is_pipeline_variable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 
 class _Job(object):
     """Handle creating, starting and waiting for Amazon SageMaker jobs to finish.
@@ -36,6 +39,7 @@ def __init__(self, sagemaker_session, job_name):
         self.job_name = job_name
 
     @abstractmethod
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "job.start_new")
     def start_new(self, estimator, inputs):
         """Create a new Amazon SageMaker job from the estimator.
 
@@ -63,6 +67,7 @@ def stop(self):
         """Stop the job."""
 
     @staticmethod
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "job._load_config")
     def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
         """Placeholder docstring"""
         input_config = _Job._format_inputs_to_input_config(inputs, validate_uri)

src/sagemaker/jumpstart/estimator.py

@@ -47,6 +47,9 @@
 from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
 from sagemaker.workflow.entities import PipelineVariable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 
 class JumpStartEstimator(Estimator):
     """JumpStartEstimator class.
@@ -762,6 +765,7 @@ def attach(
             additional_kwargs=additional_kwargs,
         )
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "jumpstart.estimator.deploy")
     def deploy(
         self,
         initial_instance_count: Optional[int] = None,

src/sagemaker/mxnet/estimator.py

@@ -32,6 +32,9 @@
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -213,6 +216,7 @@ def _configure_distribution(self, distribution):
                     "custom_mpi_options", ""
                 )
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "mxnet.estimator.create_model")
     def create_model(
         self,
         model_server_workers=None,

src/sagemaker/pytorch/estimator.py

@@ -33,6 +33,9 @@
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -376,6 +379,7 @@ def hyperparameters(self):
 
         return hyperparameters
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "pytorch.estimator.create_model")
     def create_model(
         self,
         model_server_workers=None,

src/sagemaker/rl/estimator.py

@@ -26,6 +26,9 @@
 from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 SAGEMAKER_ESTIMATOR = "sagemaker_estimator"
@@ -171,6 +174,7 @@ def __init__(
             **kwargs
         )
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "rl.estimator.create_model")
     def create_model(
         self,
         role=None,

src/sagemaker/sklearn/estimator.py

@@ -30,6 +30,9 @@
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow import is_pipeline_variable
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -166,6 +169,7 @@ def __init__(
                 instance_type=instance_type,
             )
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "sklearn.estimator.create_model")
     def create_model(
         self,
         model_server_workers=None,

src/sagemaker/tensorflow/estimator.py

@@ -31,6 +31,9 @@
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.utils import format_tags
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -310,6 +313,7 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
 
         return init_params
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "tensorflow.estimator.create_model")
     def create_model(
         self,
         role=None,

src/sagemaker/workflow/airflow.py

@@ -24,7 +24,11 @@
 from sagemaker.processing import Processor
 from sagemaker.utils import format_tags
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
 
+
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.prepare_framework")
 def prepare_framework(estimator, s3_operations):
     """Prepare S3 operations and environment variables related to framework.
 
@@ -77,6 +81,7 @@ def prepare_framework(estimator, s3_operations):
     )
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.prepare_amazon_algorithm_estimator")
 def prepare_amazon_algorithm_estimator(estimator, inputs, mini_batch_size=None):
     """Sets up amazon algorithm estimator.
 
@@ -108,6 +113,7 @@ def prepare_amazon_algorithm_estimator(estimator, inputs, mini_batch_size=None):
     estimator.mini_batch_size = mini_batch_size
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.training_base_config")
 def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size=None):  # noqa: C901
     """Export Airflow base training config from an estimator
 
@@ -218,6 +224,7 @@ def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size=
     return train_config
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.training_config")
 def training_config(estimator, inputs=None, job_name=None, mini_batch_size=None):
     """Export Airflow training config from an estimator
 
@@ -470,6 +477,7 @@ def _merge_s3_operations(s3_operations_list):
     return s3_operations_merged
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.update_submit_s3_uri")
 def update_submit_s3_uri(estimator, job_name):
     """Updated the S3 URI of the framework source directory in given estimator.
 
@@ -494,6 +502,7 @@ def update_submit_s3_uri(estimator, job_name):
     estimator.uploaded_code = fw_utils.UploadedCode(submit_uri, script_name)
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.update_estimator_from_task")
 def update_estimator_from_task(estimator, task_id, task_type):
     """Update training job of the estimator from a task in the DAG
 
@@ -630,6 +639,7 @@ def model_config(model, instance_type=None, role=None, image_uri=None):
     return config
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.model_config_from_estimator")
 def model_config_from_estimator(
     estimator,
     task_id,
@@ -815,6 +825,7 @@ def transform_config(
     return config
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.transform_config_from_estimator")
 def transform_config_from_estimator(
     estimator,
     task_id,
@@ -1057,6 +1068,7 @@ def deploy_config(model, initial_instance_count, instance_type, endpoint_name=No
     return config
 
 
+@_telemetry_emitter(Feature.SDK_DEFAULTS, "airflow.deploy_config_from_estimator")
 def deploy_config_from_estimator(
     estimator,
     task_id,

src/sagemaker/xgboost/estimator.py

@@ -31,6 +31,9 @@
 from sagemaker.xgboost.model import XGBoostModel
 from sagemaker.xgboost.utils import validate_py_version, validate_framework_version
 
+from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
+from sagemaker.telemetry.constants import Feature
+
 logger = logging.getLogger("sagemaker")
 
 
@@ -129,6 +132,7 @@ def __init__(
                 image_scope="training",
             )
 
+    @_telemetry_emitter(Feature.SDK_DEFAULTS, "xgboost.estimator.create_model")
     def create_model(
         self,
         model_server_workers=None,

tests/unit/test_chainer.py

@@ -336,6 +336,7 @@ def test_create_model_with_custom_image(sagemaker_session):
 @patch("sagemaker.utils.create_tar_file", MagicMock())
 @patch("time.time", return_value=TIME)
 @patch("time.strftime", return_value=TIMESTAMP)
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_chainer(strftime, time, sagemaker_session, chainer_version, chainer_py_version):
     chainer = Chainer(
         entry_point=SCRIPT_PATH,

tests/unit/test_estimator.py

@@ -4228,6 +4228,7 @@ def test_generic_deploy_accelerator_type(sagemaker_session):
     assert args["production_variants"][0]["InstanceType"] == INSTANCE_TYPE
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_deploy_with_model_name(sagemaker_session):
     estimator = Estimator(
         IMAGE_URI,
@@ -4247,6 +4248,7 @@ def test_deploy_with_model_name(sagemaker_session):
     assert kwargs["name"] == model_name
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_deploy_with_no_model_name(sagemaker_session):
     estimator = Estimator(
         IMAGE_URI,
@@ -4266,6 +4268,7 @@ def test_deploy_with_no_model_name(sagemaker_session):
 
 
 @patch("sagemaker.estimator.Estimator.create_model")
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_deploy_with_customized_volume_size_timeout(create_model, sagemaker_session):
     estimator = Estimator(
         IMAGE_URI,

tests/unit/test_job.py

@@ -24,6 +24,7 @@
 from sagemaker.job import _Job
 from sagemaker.model import FrameworkModel
 from sagemaker.workflow.parameters import ParameterString
+from mock import patch
 
 BUCKET_NAME = "s3://mybucket/train"
 S3_OUTPUT_PATH = "s3://bucket/prefix"
@@ -136,6 +137,7 @@ def framework(sagemaker_session):
     )
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config(estimator):
     inputs = TrainingInput(BUCKET_NAME)
 
@@ -152,6 +154,7 @@ def test_load_config(estimator):
     assert config["stop_condition"]["MaxRuntimeInSeconds"] == MAX_RUNTIME
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config_with_output_compression_disabled(estimator):
     inputs = TrainingInput(BUCKET_NAME)
     estimator.disable_output_compression = True
@@ -168,6 +171,7 @@ def test_load_config_with_output_compression_disabled(estimator):
     assert config["stop_condition"]["MaxRuntimeInSeconds"] == MAX_RUNTIME
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config_with_model_channel(estimator):
     inputs = TrainingInput(BUCKET_NAME)
 
@@ -188,6 +192,7 @@ def test_load_config_with_model_channel(estimator):
     assert config["stop_condition"]["MaxRuntimeInSeconds"] == MAX_RUNTIME
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config_with_model_channel_no_inputs(estimator):
     estimator.model_uri = MODEL_URI
     estimator.model_channel_name = MODEL_CHANNEL_NAME
@@ -205,6 +210,7 @@ def test_load_config_with_model_channel_no_inputs(estimator):
     assert config["stop_condition"]["MaxRuntimeInSeconds"] == MAX_RUNTIME
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config_with_code_channel(framework):
     inputs = TrainingInput(BUCKET_NAME)
 
@@ -225,6 +231,7 @@ def test_load_config_with_code_channel(framework):
     assert config["resource_config"]["InstanceType"] == INSTANCE_TYPE
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config_with_code_channel_no_code_uri(framework):
     inputs = TrainingInput(BUCKET_NAME)
 
@@ -242,6 +249,7 @@ def test_load_config_with_code_channel_no_code_uri(framework):
     assert config["resource_config"]["InstanceType"] == INSTANCE_TYPE
 
 
+@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None)
 def test_load_config_with_role_as_pipeline_parameter(estimator):
     inputs = TrainingInput(BUCKET_NAME)
     estimator.role = ParameterString(name="Role")