fix: use local updated args; use train_max_wait (#1945)

Co-authored-by: Ajay Karpur <[email protected]>

Author: metrizable

doc/v2.rst

@@ -231,7 +231,7 @@ The following estimator parameters have been renamed:
 +------------------------------+------------------------+
 | ``train_use_spot_instances`` | ``use_spot_instances`` |
 +------------------------------+------------------------+
-| ``train_max_run_wait``       | ``max_wait``           |
+| ``train_max_wait``           | ``max_wait``           |
 +------------------------------+------------------------+
 | ``train_volume_size``        | ``volume_size``        |
 +------------------------------+------------------------+

src/sagemaker/cli/compatibility/v2/modifiers/training_params.py

@@ -45,7 +45,7 @@
     "train_instance_count",
     "train_instance_type",
     "train_max_run",
-    "train_max_run_wait",
+    "train_max_wait",
     "train_use_spot_instances",
     "train_volume_size",
     "train_volume_kms_key",
@@ -63,7 +63,7 @@ def node_should_be_modified(self, node):
         - ``train_instance_count``
         - ``train_instance_type``
         - ``train_max_run``
-        - ``train_max_run_wait``
+        - ``train_max_wait``
         - ``train_use_spot_instances``
         - ``train_volume_kms_key``
         - ``train_volume_size``

src/sagemaker/estimator.py

@@ -242,7 +242,7 @@ def __init__(
         use_spot_instances = renamed_kwargs(
             "train_use_spot_instances", "use_spot_instances", use_spot_instances, kwargs
         )
-        max_wait = renamed_kwargs("train_max_run_wait", "max_wait", max_wait, kwargs)
+        max_wait = renamed_kwargs("train_max_wait", "max_wait", max_wait, kwargs)
         volume_size = renamed_kwargs("train_volume_size", "volume_size", volume_size, kwargs)
         volume_kms_key = renamed_kwargs(
             "train_volume_kms_key", "volume_kms_key", volume_kms_key, kwargs

src/sagemaker/mxnet/estimator.py

@@ -150,6 +150,9 @@ def __init__(
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
         distribution = renamed_kwargs("distributions", "distribution", distribution, kwargs)
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
         validate_version_or_image_args(framework_version, py_version, image_uri)
         if py_version == "py2":
             logger.warning(
@@ -168,7 +171,6 @@ def __init__(
         )
 
         if distribution is not None:
-            instance_type = kwargs.get("instance_type")
             warn_if_parameter_server_with_multi_gpu(
                 training_instance_type=instance_type, distribution=distribution
             )

src/sagemaker/sklearn/estimator.py

@@ -16,6 +16,7 @@
 import logging
 
 from sagemaker import image_uris
+from sagemaker.deprecations import renamed_kwargs
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import (
     framework_name_from_image,
@@ -107,6 +108,12 @@ def __init__(
             :class:`~sagemaker.estimator.Framework` and
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
+        instance_count = renamed_kwargs(
+            "train_instance_count", "instance_count", kwargs.get("instance_count"), kwargs
+        )
         validate_version_or_image_args(framework_version, py_version, image_uri)
         if py_version and py_version != "py3":
             raise AttributeError(
@@ -117,10 +124,8 @@ def __init__(
 
         # SciKit-Learn does not support distributed training or training on GPU instance types.
         # Fail fast.
-        instance_type = kwargs.get("instance_type")
         _validate_not_gpu_instance_type(instance_type)
 
-        instance_count = kwargs.get("instance_count")
         if instance_count:
             if instance_count != 1:
                 raise AttributeError(

src/sagemaker/tensorflow/estimator.py

@@ -19,6 +19,7 @@
 
 from sagemaker import image_uris, s3, utils
 from sagemaker.debugger import DebuggerHookConfig
+from sagemaker.deprecations import renamed_kwargs
 from sagemaker.estimator import Framework
 import sagemaker.fw_utils as fw
 from sagemaker.tensorflow import defaults
@@ -112,6 +113,9 @@ def __init__(
             :class:`~sagemaker.estimator.Framework` and
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
         fw.validate_version_or_image_args(framework_version, py_version, image_uri)
         if py_version == "py2":
             logger.warning(
@@ -121,7 +125,6 @@ def __init__(
         self.py_version = py_version
 
         if distribution is not None:
-            instance_type = kwargs.get("instance_type")
             fw.warn_if_parameter_server_with_multi_gpu(
                 training_instance_type=instance_type, distribution=distribution
             )

src/sagemaker/xgboost/estimator.py

@@ -16,6 +16,7 @@
 import logging
 
 from sagemaker import image_uris
+from sagemaker.deprecations import renamed_kwargs
 from sagemaker.estimator import Framework, _TrainingJob
 from sagemaker.fw_utils import (
     framework_name_from_image,
@@ -95,6 +96,9 @@ def __init__(
             :class:`~sagemaker.estimator.Framework` and
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
+        instance_type = renamed_kwargs(
+            "train_instance_type", "instance_type", kwargs.get("instance_type"), kwargs
+        )
         super(XGBoost, self).__init__(
             entry_point, source_dir, hyperparameters, image_uri=image_uri, **kwargs
         )
@@ -111,7 +115,7 @@ def __init__(
                 self.sagemaker_session.boto_region_name,
                 version=framework_version,
                 py_version=self.py_version,
-                instance_type=kwargs.get("instance_type"),
+                instance_type=instance_type,
                 image_scope="training",
             )
 

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_training_params.py

@@ -48,7 +48,7 @@
     "train_instance_count=1",
     "train_instance_type='ml.c4.xlarge'",
     "train_max_run=8 * 60 * 60",
-    "train_max_run_wait=1 * 60 * 60",
+    "train_max_wait=1 * 60 * 60",
     "train_use_spot_instances=True",
     "train_volume_size=30",
     "train_volume_kms_key='key'",