feature: add enable_sagemaker_metrics flag (aws#254)

Add the enable_sagemaker_metrics flag for Estimators. This corresponds to the new EnableSageMakerMetricsTimeSeries parameter.

By default, if the flag is unset then we don't pass on the EnableSageMakerMetricsTimeSeries parameter unless you're using the following framework versions:

MXNet >= 1.6
TensorFlow >=1.15
Pytorch >=1.3

Author: ishaaq

src/sagemaker/estimator.py

@@ -98,6 +98,7 @@ def __init__(
         rules=None,
         debugger_hook_config=None,
         tensorboard_output_config=None,
+        enable_sagemaker_metrics=None,
     ):
         """Initialize an ``EstimatorBase`` instance.
 
@@ -195,6 +196,10 @@ def __init__(
                 started. If the path is unset then SageMaker assumes the
                 checkpoints will be provided under `/opt/ml/checkpoints/`.
                 (default: ``None``).
+            enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
+                Series. For more information see:
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
+                (default: ``None``).
         """
         self.role = role
         self.train_instance_count = train_instance_count
@@ -250,6 +255,8 @@ def __init__(
         self.debugger_rule_configs = None
         self.collection_configs = None
 
+        self.enable_sagemaker_metrics = enable_sagemaker_metrics
+
     @abstractmethod
     def train_image(self):
         """Return the Docker image to use for training.
@@ -958,6 +965,9 @@ def start_new(cls, estimator, inputs):
 
         cls._add_spot_checkpoint_args(local_mode, estimator, train_args)
 
+        if estimator.enable_sagemaker_metrics is not None:
+            train_args["enable_sagemaker_metrics"] = estimator.enable_sagemaker_metrics
+
         estimator.sagemaker_session.train(**train_args)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
@@ -1060,6 +1070,7 @@ def __init__(
         rules=None,
         debugger_hook_config=None,
         tensorboard_output_config=None,
+        enable_sagemaker_metrics=None,
     ):
         """Initialize an ``Estimator`` instance.
 
@@ -1171,6 +1182,10 @@ def __init__(
                 user entry script for training. The user entry script, files in
                 source_dir (if specified), and dependencies will be uploaded in
                 a tar to S3. Also known as internet-free mode (default: ``False``).
+            enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
+                Series. For more information see:
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
+                (default: ``None``).
         """
         self.image_name = image_name
         self.hyperparam_dict = hyperparameters.copy() if hyperparameters else {}
@@ -1201,6 +1216,7 @@ def __init__(
             rules=rules,
             debugger_hook_config=debugger_hook_config,
             tensorboard_output_config=tensorboard_output_config,
+            enable_sagemaker_metrics=enable_sagemaker_metrics,
         )
 
     def enable_network_isolation(self):
@@ -1354,6 +1370,7 @@ def __init__(
         git_config=None,
         checkpoint_s3_uri=None,
         checkpoint_local_path=None,
+        enable_sagemaker_metrics=None,
         **kwargs
     ):
         """Base class initializer. Subclasses which override ``__init__`` should
@@ -1500,6 +1517,10 @@ def __init__(
                 started. If the path is unset then SageMaker assumes the
                 checkpoints will be provided under `/opt/ml/checkpoints/`.
                 (default: ``None``).
+            enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
+                Series. For more information see:
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
+                (default: ``None``).
             **kwargs: Additional kwargs passed to the ``EstimatorBase``
                 constructor.
         """
@@ -1530,6 +1551,7 @@ def __init__(
         self._hyperparameters = hyperparameters or {}
         self.checkpoint_s3_uri = checkpoint_s3_uri
         self.checkpoint_local_path = checkpoint_local_path
+        self.enable_sagemaker_metrics = enable_sagemaker_metrics
 
     def enable_network_isolation(self):
         """Return True if this Estimator can use network isolation to run.

src/sagemaker/mxnet/estimator.py

@@ -21,6 +21,7 @@
     framework_version_from_tag,
     empty_framework_version_warning,
     python_deprecation_warning,
+    is_version_equal_or_higher,
 )
 from sagemaker.mxnet.defaults import MXNET_VERSION
 from sagemaker.mxnet.model import MXNetModel
@@ -103,6 +104,11 @@ def __init__(
             logger.warning(empty_framework_version_warning(MXNET_VERSION, self.LATEST_VERSION))
         self.framework_version = framework_version or MXNET_VERSION
 
+        if "enable_sagemaker_metrics" not in kwargs:
+            # enable sagemaker metrics for MXNet v1.6 or greater:
+            if is_version_equal_or_higher([1, 6], self.framework_version):
+                kwargs["enable_sagemaker_metrics"] = True
+
         super(MXNet, self).__init__(
             entry_point, source_dir, hyperparameters, image_name=image_name, **kwargs
         )

src/sagemaker/pytorch/estimator.py

@@ -21,6 +21,7 @@
     framework_version_from_tag,
     empty_framework_version_warning,
     python_deprecation_warning,
+    is_version_equal_or_higher,
 )
 from sagemaker.pytorch.defaults import PYTORCH_VERSION, PYTHON_VERSION
 from sagemaker.pytorch.model import PyTorchModel
@@ -98,6 +99,11 @@ def __init__(
             logger.warning(empty_framework_version_warning(PYTORCH_VERSION, PYTORCH_VERSION))
         self.framework_version = framework_version or PYTORCH_VERSION
 
+        if "enable_sagemaker_metrics" not in kwargs:
+            # enable sagemaker metrics for PT v1.3 or greater:
+            if is_version_equal_or_higher([1, 3], self.framework_version):
+                kwargs["enable_sagemaker_metrics"] = True
+
         super(PyTorch, self).__init__(
             entry_point, source_dir, hyperparameters, image_name=image_name, **kwargs
         )

src/sagemaker/session.py

@@ -370,6 +370,7 @@ def train(  # noqa: C901
         debugger_rule_configs=None,
         debugger_hook_config=None,
         tensorboard_output_config=None,
+        enable_sagemaker_metrics=None,
     ):
         """Create an Amazon SageMaker training job.
 
@@ -432,6 +433,10 @@ def train(  # noqa: C901
                 started. If the path is unset then SageMaker assumes the
                 checkpoints will be provided under `/opt/ml/checkpoints/`.
                 (default: ``None``).
+            enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
+                Series. For more information see:
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
+                (default: ``None``).
 
         Returns:
             str: ARN of the training job, if it is created.
@@ -467,6 +472,11 @@ def train(  # noqa: C901
         if metric_definitions is not None:
             train_request["AlgorithmSpecification"]["MetricDefinitions"] = metric_definitions
 
+        if enable_sagemaker_metrics is not None:
+            train_request["AlgorithmSpecification"][
+                "EnableSageMakerMetricsTimeSeries"
+            ] = enable_sagemaker_metrics
+
         if hyperparameters and len(hyperparameters) > 0:
             train_request["HyperParameters"] = hyperparameters
 

src/sagemaker/tensorflow/estimator.py

@@ -286,6 +286,11 @@ def __init__(
         if not py_version:
             py_version = "py3" if self._only_python_3_supported() else "py2"
 
+        if "enable_sagemaker_metrics" not in kwargs:
+            # enable sagemaker metrics for TF v1.15 or greater:
+            if fw.is_version_equal_or_higher([1, 15], self.framework_version):
+                kwargs["enable_sagemaker_metrics"] = True
+
         super(TensorFlow, self).__init__(image_name=image_name, **kwargs)
         self.checkpoint_path = checkpoint_path
 

tests/unit/test_estimator.py

@@ -205,6 +205,7 @@ def test_framework_all_init_args(sagemaker_session):
         encrypt_inter_container_traffic=True,
         checkpoint_s3_uri="s3://bucket/checkpoint",
         checkpoint_local_path="file://local/checkpoint",
+        enable_sagemaker_metrics=True,
     )
     _TrainingJob.start_new(f, "s3://mydata")
     sagemaker_session.train.assert_called_once()
@@ -241,6 +242,7 @@ def test_framework_all_init_args(sagemaker_session):
         "encrypt_inter_container_traffic": True,
         "checkpoint_s3_uri": "s3://bucket/checkpoint",
         "checkpoint_local_path": "file://local/checkpoint",
+        "enable_sagemaker_metrics": True,
     }
 
 
@@ -1835,6 +1837,59 @@ def test_generic_to_fit_with_network_isolation(sagemaker_session):
     assert args["enable_network_isolation"]
 
 
+def test_generic_to_fit_with_sagemaker_metrics_missing(sagemaker_session):
+    e = Estimator(
+        IMAGE_NAME,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+    )
+
+    e.fit()
+
+    sagemaker_session.train.assert_called_once()
+    args = sagemaker_session.train.call_args[1]
+    assert "enable_sagemaker_metrics" not in args
+
+
+def test_generic_to_fit_with_sagemaker_metrics_enabled(sagemaker_session):
+    e = Estimator(
+        IMAGE_NAME,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+        enable_sagemaker_metrics=True,
+    )
+
+    e.fit()
+
+    sagemaker_session.train.assert_called_once()
+    args = sagemaker_session.train.call_args[1]
+    assert args["enable_sagemaker_metrics"]
+
+
+def test_generic_to_fit_with_sagemaker_metrics_disabled(sagemaker_session):
+    e = Estimator(
+        IMAGE_NAME,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+        enable_sagemaker_metrics=False,
+    )
+
+    e.fit()
+
+    sagemaker_session.train.assert_called_once()
+    args = sagemaker_session.train.call_args[1]
+    assert not args["enable_sagemaker_metrics"]
+
+
 def test_generic_to_deploy(sagemaker_session):
     e = Estimator(
         IMAGE_NAME,

tests/unit/test_mxnet.py

@@ -736,3 +736,53 @@ def test_create_model_with_custom_hosting_image(sagemaker_session):
     model = mx.create_model(image_name=custom_hosting_image)
 
     assert model.image == custom_hosting_image
+
+
+def test_mx_enable_sm_metrics(sagemaker_session):
+    mx = MXNet(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        train_instance_count=INSTANCE_COUNT,
+        train_instance_type=INSTANCE_TYPE,
+        enable_sagemaker_metrics=True,
+    )
+    assert mx.enable_sagemaker_metrics
+
+
+def test_mx_disable_sm_metrics(sagemaker_session):
+    mx = MXNet(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        train_instance_count=INSTANCE_COUNT,
+        train_instance_type=INSTANCE_TYPE,
+        enable_sagemaker_metrics=False,
+    )
+    assert not mx.enable_sagemaker_metrics
+
+
+def test_mx_disable_sm_metrics_if_pt_ver_is_less_than_1_6(sagemaker_session):
+    for fw_version in ["1.1", "1.2", "1.3", "1.4", "1.5"]:
+        mx = MXNet(
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            train_instance_count=INSTANCE_COUNT,
+            train_instance_type=INSTANCE_TYPE,
+            framework_version=fw_version,
+        )
+        assert mx.enable_sagemaker_metrics is None
+
+
+def test_mx_enable_sm_metrics_if_fw_ver_is_at_least_1_6(sagemaker_session):
+    for fw_version in ["1.6", "1.7", "2.0", "2.1"]:
+        mx = MXNet(
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            train_instance_count=INSTANCE_COUNT,
+            train_instance_type=INSTANCE_TYPE,
+            framework_version=fw_version,
+        )
+        assert mx.enable_sagemaker_metrics

tests/unit/test_pytorch.py

@@ -523,3 +523,25 @@ def test_empty_framework_version(warning, sagemaker_session):
 
     assert estimator.framework_version == defaults.PYTORCH_VERSION
     warning.assert_called_with(defaults.PYTORCH_VERSION, defaults.PYTORCH_VERSION)
+
+
+def test_pt_enable_sm_metrics(sagemaker_session):
+    pytorch = _pytorch_estimator(sagemaker_session, enable_sagemaker_metrics=True)
+    assert pytorch.enable_sagemaker_metrics
+
+
+def test_pt_disable_sm_metrics(sagemaker_session):
+    pytorch = _pytorch_estimator(sagemaker_session, enable_sagemaker_metrics=False)
+    assert not pytorch.enable_sagemaker_metrics
+
+
+def test_pt_disable_sm_metrics_if_pt_ver_is_less_than_1_15(sagemaker_session):
+    for fw_version in ["1.1", "1.2"]:
+        pytorch = _pytorch_estimator(sagemaker_session, framework_version=fw_version)
+        assert pytorch.enable_sagemaker_metrics is None
+
+
+def test_pt_enable_sm_metrics_if_fw_ver_is_at_least_1_15(sagemaker_session):
+    for fw_version in ["1.3", "1.4", "2.0", "2.1"]:
+        pytorch = _pytorch_estimator(sagemaker_session, framework_version=fw_version)
+        assert pytorch.enable_sagemaker_metrics

tests/unit/test_session.py

@@ -665,6 +665,7 @@ def test_train_pack_to_request(sagemaker_session):
         tags=None,
         vpc_config=VPC_CONFIG,
         metric_definitions=None,
+        enable_sagemaker_metrics=None,
     )
 
     assert sagemaker_session.sagemaker_client.method_calls[0] == (
@@ -1172,6 +1173,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
         train_use_spot_instances=True,
         checkpoint_s3_uri="s3://mybucket/checkpoints/",
         checkpoint_local_path="/tmp/checkpoints",
+        enable_sagemaker_metrics=True,
     )
 
     _, _, actual_train_args = sagemaker_session.sagemaker_client.method_calls[0]
@@ -1180,6 +1182,7 @@ def test_train_pack_to_request_with_optional_params(sagemaker_session):
     assert actual_train_args["HyperParameters"] == hyperparameters
     assert actual_train_args["Tags"] == TAGS
     assert actual_train_args["AlgorithmSpecification"]["MetricDefinitions"] == METRIC_DEFINITONS
+    assert actual_train_args["AlgorithmSpecification"]["EnableSageMakerMetricsTimeSeries"] is True
     assert actual_train_args["EnableInterContainerTrafficEncryption"] is True
     assert actual_train_args["EnableManagedSpotTraining"] is True
     assert actual_train_args["CheckpointConfig"]["S3Uri"] == "s3://mybucket/checkpoints/"

tests/unit/test_tf_estimator.py

@@ -1200,3 +1200,29 @@ def test_tf_script_mode_attach(sagemaker_session, tf_version):
     assert estimator.hyperparameters() is not None
     assert estimator.source_dir == "s3://some/sourcedir.tar.gz"
     assert estimator.entry_point == "iris-dnn-classifier.py"
+
+
+@patch("sagemaker.utils.create_tar_file", MagicMock())
+def test_tf_enable_sm_metrics(sagemaker_session):
+    tf = _build_tf(sagemaker_session, enable_sagemaker_metrics=True)
+    assert tf.enable_sagemaker_metrics
+
+
+@patch("sagemaker.utils.create_tar_file", MagicMock())
+def test_tf_disable_sm_metrics(sagemaker_session):
+    tf = _build_tf(sagemaker_session, enable_sagemaker_metrics=False)
+    assert not tf.enable_sagemaker_metrics
+
+
+@patch("sagemaker.utils.create_tar_file", MagicMock())
+def test_tf_disable_sm_metrics_if_fw_ver_is_less_than_1_15(sagemaker_session):
+    for fw_version in ["1.11", "1.12", "1.13", "1.14"]:
+        tf = _build_tf(sagemaker_session, framework_version=fw_version)
+        assert tf.enable_sagemaker_metrics is None
+
+
+@patch("sagemaker.utils.create_tar_file", MagicMock())
+def test_tf_enable_sm_metrics_if_fw_ver_is_at_least_1_15(sagemaker_session):
+    for fw_version in ["1.15", "1.16", "2.0", "2.1"]:
+        tf = _build_tf(sagemaker_session, framework_version=fw_version)
+        assert tf.enable_sagemaker_metrics