Skip to content

Commit fe8357f

Browse files
committed
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
1 parent 8d674b7 commit fe8357f

File tree

1 file changed

+17
-1
lines changed

1 file changed

+17
-1
lines changed

src/sagemaker_core/main/utils.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,11 +320,19 @@ def __call__(cls, *args, **kwargs):
320320
return cls._instances[cls]
321321

322322

323-
class SageMakerClient(metaclass=SingletonMeta):
323+
class SageMakerClient:
324324
"""
325325
A singleton class for creating a SageMaker client.
326326
"""
327327

328+
_instance = None
329+
_initialized = False
330+
331+
def __new__(cls, *args, **kwargs):
332+
if cls._instance is None:
333+
cls._instance = super().__new__(cls)
334+
return cls._instance
335+
328336
def __init__(
329337
self,
330338
session: Session = None,
@@ -335,6 +343,12 @@ def __init__(
335343
Initializes the SageMakerClient with a boto3 session, region name, and service name.
336344
Creates a boto3 client using the provided session, region, and service.
337345
"""
346+
347+
# Only skip reinitialization if region is the same
348+
if self._initialized and session is None and config is None:
349+
if region_name is None or region_name == self._instance.region_name:
350+
return
351+
338352
if session is None:
339353
logger.warning("No boto3 session provided. Creating a new session.")
340354
session = Session(region_name=region_name)
@@ -361,6 +375,8 @@ def __init__(
361375
"sagemaker-metrics", region_name, config=self.config
362376
)
363377

378+
self._initialized = True
379+
364380
def get_client(self, service_name: str) -> Any:
365381
"""
366382
Get the client of corresponding service

0 commit comments

Comments
 (0)