change: add git test for estimator class, update typing

Author: evakravi

src/sagemaker/estimator.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import os
+from typing import Dict
 import uuid
 from abc import ABCMeta, abstractmethod
 
@@ -582,7 +583,7 @@ def _get_or_create_name(self, name=None):
         return name_from_base(self.base_job_name)
 
     @staticmethod
-    def _json_encode_hyperparameters(hyperparameters):
+    def _json_encode_hyperparameters(hyperparameters: dict) -> dict:
         """Applies Json encoding for certain Hyperparameter types, returns hyperparameters.
 
         Args:
@@ -661,14 +662,14 @@ def _prepare_for_training(self, job_name=None):
         self._prepare_debugger_for_training()
         self._prepare_profiler_for_training()
 
-    def _script_mode_hyperparam_update(self, code_dir, script):
+    def _script_mode_hyperparam_update(self, code_dir: str, script: str) -> None:
         """Applies in-place update to hyperparameters required for script mode with training.
 
         Args:
             code_dir (str): The directory hosting the training scripts.
             script (str): The relative filepath of the training entry-point script.
         """
-        hyperparams = {}
+        hyperparams: Dict[str, str] = {}
         hyperparams[DIR_PARAM_NAME] = code_dir
         hyperparams[SCRIPT_PARAM_NAME] = script
         hyperparams[CONTAINER_LOG_LEVEL_PARAM_NAME] = self.container_log_level
@@ -677,7 +678,7 @@ def _script_mode_hyperparam_update(self, code_dir, script):
 
         self._hyperparameters.update(EstimatorBase._json_encode_hyperparameters(hyperparams))
 
-    def _stage_user_code_in_s3(self):
+    def _stage_user_code_in_s3(self) -> str:
         """Upload the user training script to s3 and return the location.
 
         Returns: s3 uri
@@ -2615,14 +2616,14 @@ def _prepare_for_training(self, job_name=None):
 
         self._validate_and_set_debugger_configs()
 
-    def _script_mode_hyperparam_update(self, code_dir, script):
+    def _script_mode_hyperparam_update(self, code_dir: str, script: str) -> None:
         """Applies in-place update to hyperparameters required for script mode with training.
 
         Args:
             code_dir (str): The directory hosting the training scripts.
             script (str): The relative filepath of the training entry-point script.
         """
-        hyperparams = {}
+        hyperparams: Dict[str, str] = {}
         hyperparams[DIR_PARAM_NAME] = code_dir
         hyperparams[SCRIPT_PARAM_NAME] = script
         hyperparams[CONTAINER_LOG_LEVEL_PARAM_NAME] = self.container_log_level

tests/unit/test_estimator.py

@@ -3354,10 +3354,34 @@ def test_image_name_map(sagemaker_session):
     assert e.image_uri == IMAGE_URI
 
 
+@patch("sagemaker.git_utils.git_clone_repo")
+def test_git_support_with_branch_and_commit_succeed_estimator_class(
+    git_clone_repo, sagemaker_session
+):
+    git_clone_repo.side_effect = lambda gitconfig, entrypoint, source_dir=None, dependencies=None: {
+        "entry_point": "/tmp/repo_dir/entry_point",
+        "source_dir": None,
+        "dependencies": None,
+    }
+    git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
+    entry_point = "entry_point"
+    fw = Estimator(
+        entry_point=entry_point,
+        git_config=git_config,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        image_uri=IMAGE_URI,
+    )
+    fw.fit()
+    git_clone_repo.assert_called_once_with(git_config, entry_point, None, None)
+
+
 @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"
+        s3_prefix="s3://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"