feat: script mode support for estimator class

Author: evakravi

src/sagemaker/algorithm.py

@@ -174,7 +174,7 @@ def __init__(
         self.validate_train_spec()
         self.hyperparameter_definitions = self._parse_hyperparameters()
 
-        self.hyperparam_dict = {}
+        self._hyperparameters = {}
         if hyperparameters:
             self.set_hyperparameters(**hyperparameters)
 
@@ -215,7 +215,7 @@ def set_hyperparameters(self, **kwargs):
         """Placeholder docstring"""
         for k, v in kwargs.items():
             value = self._validate_and_cast_hyperparameter(k, v)
-            self.hyperparam_dict[k] = value
+            self._hyperparameters[k] = value
 
         self._validate_and_set_default_hyperparameters()
 
@@ -225,7 +225,7 @@ def hyperparameters(self):
         The fit() method, that does the model training, calls this method to
         find the hyperparameters you specified.
         """
-        return self.hyperparam_dict
+        return self._hyperparameters
 
     def training_image_uri(self):
         """Returns the docker image to use for training.
@@ -464,10 +464,10 @@ def _validate_and_set_default_hyperparameters(self):
         # Check if all the required hyperparameters are set. If there is a default value
         # for one, set it.
         for name, definition in self.hyperparameter_definitions.items():
-            if name not in self.hyperparam_dict:
+            if name not in self._hyperparameters:
                 spec = definition["spec"]
                 if "DefaultValue" in spec:
-                    self.hyperparam_dict[name] = spec["DefaultValue"]
+                    self._hyperparameters[name] = spec["DefaultValue"]
                 elif "IsRequired" in spec and spec["IsRequired"]:
                     raise ValueError("Required hyperparameter: %s is not set" % name)
 

src/sagemaker/chainer/estimator.py

@@ -15,7 +15,7 @@
 
 import logging
 
-from sagemaker.estimator import Framework
+from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.fw_utils import (
     framework_name_from_image,
     framework_version_from_tag,
@@ -158,7 +158,9 @@ def hyperparameters(self):
 
         # remove unset keys.
         additional_hyperparameters = {k: v for k, v in additional_hyperparameters.items() if v}
-        hyperparameters.update(Framework._json_encode_hyperparameters(additional_hyperparameters))
+        hyperparameters.update(
+            EstimatorBase._json_encode_hyperparameters(additional_hyperparameters)
+        )
         return hyperparameters
 
     def create_model(

src/sagemaker/estimator.py

@@ -17,7 +17,7 @@
 import re
 
 from sagemaker.deprecations import renamed_kwargs
-from sagemaker.estimator import Framework
+from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.fw_utils import (
     framework_name_from_image,
     warn_if_parameter_server_with_multi_gpu,
@@ -246,13 +246,13 @@ def hyperparameters(self):
             distribution=self.distribution
         )
         hyperparameters.update(
-            Framework._json_encode_hyperparameters(distributed_training_hyperparameters)
+            EstimatorBase._json_encode_hyperparameters(distributed_training_hyperparameters)
         )
 
         if self.compiler_config:
             training_compiler_hyperparameters = self.compiler_config._to_hyperparameter_dict()
             hyperparameters.update(
-                Framework._json_encode_hyperparameters(training_compiler_hyperparameters)
+                EstimatorBase._json_encode_hyperparameters(training_compiler_hyperparameters)
             )
 
         return hyperparameters

src/sagemaker/huggingface/estimator.py

@@ -18,7 +18,7 @@
 from packaging.version import Version
 
 from sagemaker.deprecations import renamed_kwargs
-from sagemaker.estimator import Framework
+from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.fw_utils import (
     framework_name_from_image,
     framework_version_from_tag,
@@ -192,7 +192,9 @@ def hyperparameters(self):
         additional_hyperparameters = self._distribution_configuration(
             distribution=self.distribution
         )
-        hyperparameters.update(Framework._json_encode_hyperparameters(additional_hyperparameters))
+        hyperparameters.update(
+            EstimatorBase._json_encode_hyperparameters(additional_hyperparameters)
+        )
         return hyperparameters
 
     def create_model(

src/sagemaker/pytorch/estimator.py

@@ -18,7 +18,7 @@
 import re
 
 from sagemaker import image_uris, fw_utils
-from sagemaker.estimator import Framework
+from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.model import FrameworkModel, SAGEMAKER_OUTPUT_LOCATION
 from sagemaker.mxnet.model import MXNetModel
 from sagemaker.tensorflow.model import TensorFlowModel
@@ -340,7 +340,9 @@ def hyperparameters(self):
             SAGEMAKER_ESTIMATOR: SAGEMAKER_ESTIMATOR_VALUE,
         }
 
-        hyperparameters.update(Framework._json_encode_hyperparameters(additional_hyperparameters))
+        hyperparameters.update(
+            EstimatorBase._json_encode_hyperparameters(additional_hyperparameters)
+        )
         return hyperparameters
 
     @classmethod

src/sagemaker/rl/estimator.py

@@ -19,7 +19,7 @@
 
 from sagemaker import image_uris, s3, utils
 from sagemaker.deprecations import renamed_kwargs
-from sagemaker.estimator import Framework
+from sagemaker.estimator import Framework, EstimatorBase
 import sagemaker.fw_utils as fw
 from sagemaker.tensorflow import defaults
 from sagemaker.tensorflow.model import TensorFlowModel
@@ -327,7 +327,9 @@ def hyperparameters(self):
             )
             additional_hyperparameters["model_dir"] = self.model_dir
 
-        hyperparameters.update(Framework._json_encode_hyperparameters(additional_hyperparameters))
+        hyperparameters.update(
+            EstimatorBase._json_encode_hyperparameters(additional_hyperparameters)
+        )
         return hyperparameters
 
     def _default_s3_path(self, directory, mpi=False):

src/sagemaker/tensorflow/estimator.py

@@ -17,6 +17,8 @@
 import os
 import subprocess
 from time import sleep
+from sagemaker.fw_utils import UploadedCode
+
 
 import pytest
 from botocore.exceptions import ClientError
@@ -3350,3 +3352,87 @@ def test_image_name_map(sagemaker_session):
     )
 
     assert e.image_uri == IMAGE_URI
+
+
+@patch("sagemaker.estimator.Estimator._stage_user_code_in_s3")
+def test_script_mode_estimator(patched_stage_user_code, sagemaker_session):
+    patched_stage_user_code.return_value = UploadedCode(
+        s3_prefix="s3://%s/%s" % ("bucket", "key"), script_name="script_name"
+    )
+    script_uri = "s3://codebucket/someprefix/sourcedir.tar.gz"
+    image_uri = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.9.0-gpu-py38"
+    model_uri = "s3://someprefix2/models/model.tar.gz"
+    t = Estimator(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        source_dir=script_uri,
+        image_uri=image_uri,
+        model_uri=model_uri,
+    )
+    t.fit("s3://bucket/mydata")
+
+    patched_stage_user_code.assert_called_once()
+    sagemaker_session.train.assert_called_once()
+
+
+@patch("time.time", return_value=TIME)
+@patch("sagemaker.estimator.tar_and_upload_dir")
+def test_script_mode_estimator_same_calls_as_framework(
+    patched_tar_and_upload_dir, sagemaker_session
+):
+
+    patched_tar_and_upload_dir.return_value = UploadedCode(
+        s3_prefix="s3://%s/%s" % ("bucket", "key"), script_name="script_name"
+    )
+    sagemaker_session.boto_region_name = REGION
+
+    script_uri = "s3://codebucket/someprefix/sourcedir.tar.gz"
+
+    instance_type = "ml.p2.xlarge"
+    instance_count = 1
+
+    model_uri = "s3://someprefix2/models/model.tar.gz"
+    training_data_uri = "s3://bucket/mydata"
+
+    generic_estimator = Estimator(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        region=REGION,
+        sagemaker_session=sagemaker_session,
+        instance_count=instance_count,
+        instance_type=instance_type,
+        source_dir=script_uri,
+        image_uri=IMAGE_URI,
+        model_uri=model_uri,
+        environment={"USE_SMDEBUG": "0"},
+        dependencies=[],
+        debugger_hook_config={},
+    )
+    generic_estimator.fit(training_data_uri)
+
+    generic_estimator_tar_and_upload_dir_args = patched_tar_and_upload_dir.call_args_list
+    generic_estimator_train_args = sagemaker_session.train.call_args_list
+
+    patched_tar_and_upload_dir.reset_mock()
+    sagemaker_session.train.reset_mock()
+
+    framework_estimator = DummyFramework(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        region=REGION,
+        source_dir=script_uri,
+        instance_count=instance_count,
+        instance_type=instance_type,
+        sagemaker_session=sagemaker_session,
+        model_uri=model_uri,
+        dependencies=[],
+        debugger_hook_config={},
+    )
+    framework_estimator.fit(training_data_uri)
+
+    assert len(generic_estimator_tar_and_upload_dir_args) == 1
+    assert generic_estimator_tar_and_upload_dir_args == patched_tar_and_upload_dir.call_args_list
+    assert generic_estimator_train_args == sagemaker_session.train.call_args_list