From fe8357f68238b2ff756d89516c4084b61db32530 Mon Sep 17 00:00:00 2001 From: jzhaoqwa Date: Wed, 25 Jun 2025 12:27:43 -0700 Subject: [PATCH] Remove singleton metaclass dependency for SageMakerClient Fix the bug that we cannot update the region once SageMakerClient is initialized Unit tests pass and tested locally. Region and configs can be updated --- src/sagemaker_core/main/utils.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/sagemaker_core/main/utils.py b/src/sagemaker_core/main/utils.py index 126aa556..700df524 100644 --- a/src/sagemaker_core/main/utils.py +++ b/src/sagemaker_core/main/utils.py @@ -320,11 +320,19 @@ def __call__(cls, *args, **kwargs): return cls._instances[cls] -class SageMakerClient(metaclass=SingletonMeta): +class SageMakerClient: """ A singleton class for creating a SageMaker client. """ + _instance = None + _initialized = False + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + def __init__( self, session: Session = None, @@ -335,6 +343,12 @@ def __init__( Initializes the SageMakerClient with a boto3 session, region name, and service name. Creates a boto3 client using the provided session, region, and service. """ + + # Only skip reinitialization if region is the same + if self._initialized and session is None and config is None: + if region_name is None or region_name == self._instance.region_name: + return + if session is None: logger.warning("No boto3 session provided. Creating a new session.") session = Session(region_name=region_name) @@ -361,6 +375,8 @@ def __init__( "sagemaker-metrics", region_name, config=self.config ) + self._initialized = True + def get_client(self, service_name: str) -> Any: """ Get the client of corresponding service