diff --git a/src/sagemaker/chainer/estimator.py b/src/sagemaker/chainer/estimator.py index 09addf9910..8f5e8a574c 100644 --- a/src/sagemaker/chainer/estimator.py +++ b/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, diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index b6af6cf5de..e2bb872938 100644 --- a/src/sagemaker/estimator.py +++ b/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, diff --git a/src/sagemaker/huggingface/estimator.py b/src/sagemaker/huggingface/estimator.py index 86df43d4e9..d9ea36b5b0 100644 --- a/src/sagemaker/huggingface/estimator.py +++ b/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, diff --git a/src/sagemaker/job.py b/src/sagemaker/job.py index 7040c376ab..5a3857a199 100644 --- a/src/sagemaker/job.py +++ b/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) diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py index f53d109dc8..8b48062c9e 100644 --- a/src/sagemaker/jumpstart/estimator.py +++ b/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, diff --git a/src/sagemaker/mxnet/estimator.py b/src/sagemaker/mxnet/estimator.py index 7d0ce2d494..8c5310f358 100644 --- a/src/sagemaker/mxnet/estimator.py +++ b/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, diff --git a/src/sagemaker/pytorch/estimator.py b/src/sagemaker/pytorch/estimator.py index 412926279c..3e5f01eb53 100644 --- a/src/sagemaker/pytorch/estimator.py +++ b/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, diff --git a/src/sagemaker/rl/estimator.py b/src/sagemaker/rl/estimator.py index 3ed539fa2e..f923900b6c 100644 --- a/src/sagemaker/rl/estimator.py +++ b/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, diff --git a/src/sagemaker/sklearn/estimator.py b/src/sagemaker/sklearn/estimator.py index 9f4b25f214..ff179718b3 100644 --- a/src/sagemaker/sklearn/estimator.py +++ b/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, diff --git a/src/sagemaker/tensorflow/estimator.py b/src/sagemaker/tensorflow/estimator.py index df2bc74935..f4559c07aa 100644 --- a/src/sagemaker/tensorflow/estimator.py +++ b/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, diff --git a/src/sagemaker/workflow/airflow.py b/src/sagemaker/workflow/airflow.py index 3678c3d97e..e2e4dddd4a 100644 --- a/src/sagemaker/workflow/airflow.py +++ b/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, diff --git a/src/sagemaker/xgboost/estimator.py b/src/sagemaker/xgboost/estimator.py index dfd7145e93..c7577ce76c 100644 --- a/src/sagemaker/xgboost/estimator.py +++ b/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, diff --git a/tests/unit/test_chainer.py b/tests/unit/test_chainer.py index 8ad2ae0bab..9a41ccb9be 100644 --- a/tests/unit/test_chainer.py +++ b/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, diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index 295f1a8d24..2344b40137 100644 --- a/tests/unit/test_estimator.py +++ b/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, diff --git a/tests/unit/test_job.py b/tests/unit/test_job.py index 603b494e5a..cab35c60fc 100644 --- a/tests/unit/test_job.py +++ b/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") diff --git a/tests/unit/test_mxnet.py b/tests/unit/test_mxnet.py index 4a584dfae4..e739563ec9 100644 --- a/tests/unit/test_mxnet.py +++ b/tests/unit/test_mxnet.py @@ -209,6 +209,7 @@ def _create_compilation_job(input_shape, output_location): @patch("sagemaker.estimator.name_from_base") @patch("sagemaker.utils.create_tar_file", MagicMock()) +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model( name_from_base, sagemaker_session, mxnet_inference_version, mxnet_inference_py_version ): @@ -249,6 +250,7 @@ def test_create_model( name_from_base.assert_called_with(base_job_name) +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_optional_params( sagemaker_session, mxnet_inference_version, mxnet_inference_py_version ): @@ -291,6 +293,7 @@ def test_create_model_with_optional_params( @patch("sagemaker.estimator.name_from_base") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_custom_image(name_from_base, sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py index 618d0d7ea8..14fc2d4abe 100644 --- a/tests/unit/test_pytorch.py +++ b/tests/unit/test_pytorch.py @@ -191,6 +191,7 @@ def _get_environment(submit_directory, model_url, image_uri): @patch("sagemaker.estimator.name_from_base") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model( name_from_base, sagemaker_session, pytorch_inference_version, pytorch_inference_py_version ): @@ -230,6 +231,7 @@ def test_create_model( name_from_base.assert_called_with(base_job_name) +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_optional_params( sagemaker_session, pytorch_inference_version, pytorch_inference_py_version ): @@ -272,6 +274,7 @@ def test_create_model_with_optional_params( @patch("sagemaker.estimator.name_from_base") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_custom_image(name_from_base, sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" diff --git a/tests/unit/test_rl.py b/tests/unit/test_rl.py index d9c4129cf6..5a0947369a 100644 --- a/tests/unit/test_rl.py +++ b/tests/unit/test_rl.py @@ -246,6 +246,7 @@ def test_create_mxnet_model(name_from_base, sagemaker_session, coach_mxnet_versi name_from_base.assert_called_with("sagemaker-rl-mxnet") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_optional_params(sagemaker_session, coach_mxnet_version): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" @@ -279,6 +280,7 @@ def test_create_model_with_optional_params(sagemaker_session, coach_mxnet_versio @patch("sagemaker.estimator.name_from_base") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_custom_image(name_from_base, sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" @@ -316,6 +318,7 @@ def test_create_model_with_custom_image(name_from_base, sagemaker_session): @patch("sagemaker.utils.create_tar_file", MagicMock()) @patch("time.strftime", return_value=TIMESTAMP) @patch("time.time", return_value=TIME) +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_rl(time, strftime, sagemaker_session, coach_mxnet_version): rl = RLEstimator( entry_point=SCRIPT_PATH, diff --git a/tests/unit/test_sklearn.py b/tests/unit/test_sklearn.py index b0df31fee1..b1eb45e64b 100644 --- a/tests/unit/test_sklearn.py +++ b/tests/unit/test_sklearn.py @@ -172,6 +172,7 @@ def test_training_image_uri(sagemaker_session, sklearn_version): assert _get_full_cpu_image_uri(sklearn_version) == sklearn.training_image_uri() +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model(sagemaker_session, sklearn_version): source_dir = "s3://mybucket/source" @@ -188,6 +189,7 @@ def test_create_model(sagemaker_session, sklearn_version): @patch("sagemaker.model.FrameworkModel._upload_code") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_network_isolation(upload, sagemaker_session, sklearn_version): source_dir = "s3://mybucket/source" repacked_model_data = "s3://mybucket/prefix/model.tar.gz" @@ -208,6 +210,7 @@ def test_create_model_with_network_isolation(upload, sagemaker_session, sklearn_ @patch("sagemaker.estimator.name_from_base") +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_from_estimator(name_from_base, sagemaker_session, sklearn_version): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" @@ -246,6 +249,7 @@ def test_create_model_from_estimator(name_from_base, sagemaker_session, sklearn_ name_from_base.assert_called_with(base_job_name) +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_optional_params(sagemaker_session, sklearn_version): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source" @@ -291,6 +295,7 @@ def test_create_model_with_optional_params(sagemaker_session, sklearn_version): assert model.name == model_name +@patch("sagemaker.telemetry.telemetry_logging.resolve_value_from_config", side_effect=None) def test_create_model_with_custom_image(sagemaker_session): container_log_level = '"logging.INFO"' source_dir = "s3://mybucket/source"