From 12456b1f05ba843a5aac48240b71b9919d18dac0 Mon Sep 17 00:00:00 2001 From: pintaoz Date: Tue, 25 Feb 2025 23:36:47 -0800 Subject: [PATCH 1/2] Fix error when there is no session to call _create_model_request() --- src/sagemaker/pipeline.py | 16 ++++++++++++++++ src/sagemaker/test.py | 1 + 2 files changed, 17 insertions(+) create mode 100644 src/sagemaker/test.py diff --git a/src/sagemaker/pipeline.py b/src/sagemaker/pipeline.py index 1d1ece5965..eb88ff55fb 100644 --- a/src/sagemaker/pipeline.py +++ b/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,17 @@ 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) diff --git a/src/sagemaker/test.py b/src/sagemaker/test.py new file mode 100644 index 0000000000..f22a3bd2a7 --- /dev/null +++ b/src/sagemaker/test.py @@ -0,0 +1 @@ +from sagemaker.jumpstart import models \ No newline at end of file From 7a8cd343491e9de9568c22bb8a308cef4176d68a Mon Sep 17 00:00:00 2001 From: pintaoz Date: Wed, 26 Feb 2025 09:50:04 -0800 Subject: [PATCH 2/2] Fix codestyle --- src/sagemaker/pipeline.py | 1 - src/sagemaker/test.py | 1 - src/sagemaker/workflow/steps.py | 2 ++ 3 files changed, 2 insertions(+), 2 deletions(-) delete mode 100644 src/sagemaker/test.py diff --git a/src/sagemaker/pipeline.py b/src/sagemaker/pipeline.py index eb88ff55fb..b36cd4e917 100644 --- a/src/sagemaker/pipeline.py +++ b/src/sagemaker/pipeline.py @@ -563,7 +563,6 @@ def delete_model(self): 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. diff --git a/src/sagemaker/test.py b/src/sagemaker/test.py deleted file mode 100644 index f22a3bd2a7..0000000000 --- a/src/sagemaker/test.py +++ /dev/null @@ -1 +0,0 @@ -from sagemaker.jumpstart import models \ No newline at end of file diff --git a/src/sagemaker/workflow/steps.py b/src/sagemaker/workflow/steps.py index a80b5440c7..f49e457bc6 100644 --- a/src/sagemaker/workflow/steps.py +++ b/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,