Change: Accept model_data as dictionary in the model deploy

Author: martinRenou

src/sagemaker/model.py

@@ -831,16 +831,10 @@ def _create_sagemaker_model(
                 # _base_name, model_name are not needed under PipelineSession.
                 # the model_data may be Pipeline variable
                 # which may break the _base_name generation
-                model_uri = None
-                if isinstance(self.model_data, (str, PipelineVariable)):
-                    model_uri = self.model_data
-                elif isinstance(self.model_data, dict):
-                    model_uri = self.model_data.get("S3DataSource", {}).get("S3Uri", None)
-
                 self._ensure_base_name_if_needed(
                     image_uri=container_def["Image"],
                     script_uri=self.source_dir,
-                    model_uri=model_uri,
+                    model_uri=self._get_model_uri(),
                 )
                 self._set_model_name_if_needed()
 
@@ -877,6 +871,14 @@ def _create_sagemaker_model(
             )
             self.sagemaker_session.create_model(**create_model_args)
 
+    def _get_model_uri(self):
+        model_uri = None
+        if isinstance(self.model_data, (str, PipelineVariable)):
+            model_uri = self.model_data
+        elif isinstance(self.model_data, dict):
+            model_uri = self.model_data.get("S3DataSource", {}).get("S3Uri", None)
+        return model_uri
+
     def _ensure_base_name_if_needed(self, image_uri, script_uri, model_uri):
         """Create a base name from the image URI if there is no model name provided.
 
@@ -1434,7 +1436,7 @@ def deploy(
             self._ensure_base_name_if_needed(
                 image_uri=self.image_uri,
                 script_uri=self.source_dir,
-                model_uri=self.model_data,
+                model_uri=self._get_model_uri(),
             )
             if self._base_name is not None:
                 self._base_name = "-".join((self._base_name, compiled_model_suffix))