change: update _upload_code typing, add git test for model class

Author: evakravi

src/sagemaker/model.py

@@ -405,7 +405,7 @@ def prepare_container_def(
             self.image_uri, self.model_data, deploy_env, image_config=self.image_config
         )
 
-    def _upload_code(self, key_prefix, repack=False):
+    def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
         """Uploads code to S3 to be used with script mode with SageMaker inference.
 
         Args:

tests/unit/test_model.py

@@ -28,6 +28,9 @@
 SCRIPT_URI = "s3://codebucket/someprefix/sourcedir.tar.gz"
 IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
 MODEL_DATA = "s3://someprefix2/models/model.tar.gz"
+GIT_REPO = "https://github.com/aws/sagemaker-python-sdk.git"
+BRANCH = "test-branch-git-config"
+COMMIT = "ae15c9d7d5b97ea95ea451e4662ee43da3401d73"
 
 
 class DummyFrameworkModel(FrameworkModel):
@@ -97,3 +100,30 @@ def test_script_mode_model_same_calls_as_framework(repack_model, sagemaker_sessi
         == sagemaker_session.endpoint_from_production_variants.call_args_list
     )
     assert generic_model_repack_model_args == repack_model.call_args_list
+
+
+@patch("sagemaker.git_utils.git_clone_repo")
+@patch("sagemaker.model.fw_utils.tar_and_upload_dir")
+def test_git_support_succeed_model_class(tar_and_upload_dir, git_clone_repo, sagemaker_session):
+    git_clone_repo.side_effect = lambda gitconfig, entrypoint, sourcedir, dependency: {
+        "entry_point": "entry_point",
+        "source_dir": "/tmp/repo_dir/source_dir",
+        "dependencies": ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"],
+    }
+    entry_point = "entry_point"
+    source_dir = "source_dir"
+    dependencies = ["foo", "bar"]
+    git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
+    model = Model(
+        sagemaker_session=sagemaker_session,
+        entry_point=entry_point,
+        source_dir=source_dir,
+        dependencies=dependencies,
+        git_config=git_config,
+        image_uri=IMAGE_URI,
+    )
+    model.prepare_container_def(instance_type=INSTANCE_TYPE)
+    git_clone_repo.assert_called_with(git_config, entry_point, source_dir, dependencies)
+    assert model.entry_point == "entry_point"
+    assert model.source_dir == "/tmp/repo_dir/source_dir"
+    assert model.dependencies == ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"]