Fix error when there is no session to call _create_model_request() (aws#5062)

* Fix error when there is no session to call _create_model_request()

* Fix codestyle

---------

Co-authored-by: pintaoz <[email protected]>

Author: 2 people

src/sagemaker/pipeline.py

@@ -17,6 +17,8 @@
 
 import sagemaker
 from sagemaker import ModelMetrics, Model
+from sagemaker import local
+from sagemaker import session
 from sagemaker.config import (
     ENDPOINT_CONFIG_KMS_KEY_ID_PATH,
     MODEL_VPC_CONFIG_PATH,
@@ -560,3 +562,16 @@ def delete_model(self):
             raise ValueError("The SageMaker model must be created before attempting to delete.")
 
         self.sagemaker_session.delete_model(self.name)
+
+    def _init_sagemaker_session_if_does_not_exist(self, instance_type=None):
+        """Set ``self.sagemaker_session`` to ``LocalSession`` or ``Session`` if it's not already.
+
+        The type of session object is determined by the instance type.
+        """
+        if self.sagemaker_session:
+            return
+
+        if instance_type in ("local", "local_gpu"):
+            self.sagemaker_session = local.LocalSession(sagemaker_config=self._sagemaker_config)
+        else:
+            self.sagemaker_session = session.Session(sagemaker_config=self._sagemaker_config)

src/sagemaker/workflow/steps.py

@@ -645,6 +645,7 @@ def arguments(self) -> RequestType:
             request_dict = self.step_args
         else:
             if isinstance(self.model, PipelineModel):
+                self.model._init_sagemaker_session_if_does_not_exist()
                 request_dict = self.model.sagemaker_session._create_model_request(
                     name="",
                     role=self.model.role,
@@ -653,6 +654,7 @@ def arguments(self) -> RequestType:
                     enable_network_isolation=self.model.enable_network_isolation,
                 )
             else:
+                self.model._init_sagemaker_session_if_does_not_exist()
                 request_dict = self.model.sagemaker_session._create_model_request(
                     name="",
                     role=self.model.role,