fix: lazy-init boto3 clients in AIRHub to avoid NoRegionError at import time#5935
fix: lazy-init boto3 clients in AIRHub to avoid NoRegionError at import time#5935aviruthen wants to merge 1 commit into
Conversation
| def _get_sagemaker_client(cls): | ||
| """Get or create the default SageMaker client (lazy initialization).""" | ||
| if cls._sagemaker_client is None: | ||
| cls._sagemaker_client = boto3.client("sagemaker") |
There was a problem hiding this comment.
Can we use Session() instead of boto3?
sagemaker-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review
This PR correctly fixes a NoRegionError at import time by converting class-level boto3 client construction to lazy initialization. The approach is sound and the implementation is clean. However, there are a few concerns: thread safety of the lazy initialization, direct boto3 usage (which conflicts with V3 architecture tenets), missing type annotations, and the lack of unit tests for this change.
| _sagemaker_client = None | ||
| _s3_client = None | ||
|
|
||
| @classmethod |
There was a problem hiding this comment.
Thread safety concern: The lazy initialization pattern here is not thread-safe. If two threads call _get_sagemaker_client() concurrently when _sagemaker_client is None, both could create a client. While creating two boto3 clients isn't catastrophic (one will just be GC'd), it's worth protecting with a lock for correctness:
import threading
class AIRHub:
_sagemaker_client = None
_s3_client = None
_lock = threading.Lock()
@classmethod
def _get_sagemaker_client(cls):
if cls._sagemaker_client is None:
with cls._lock:
if cls._sagemaker_client is None:
cls._sagemaker_client = boto3.client("sagemaker")
return cls._sagemaker_client| _sagemaker_client = None | ||
| _s3_client = None | ||
|
|
||
| @classmethod |
There was a problem hiding this comment.
Missing type annotations: Per SDK conventions, new public/internal methods should have type annotations. These classmethods should declare return types:
@classmethod
def _get_sagemaker_client(cls) -> "botocore.client.SageMaker":Or more practically, since boto3 clients don't have great type stubs, at minimum use -> Any or the appropriate type from mypy_boto3_sagemaker.
| _s3_client = None | ||
|
|
||
| @classmethod | ||
| def _get_sagemaker_client(cls): |
There was a problem hiding this comment.
V3 Architecture concern (pre-existing): Per V3 tenets, all subpackages should depend on sagemaker-core for AWS API interactions rather than calling boto3 directly. This is a pre-existing issue not introduced by this PR, but worth noting — ideally these clients should come from a sagemaker_session or sagemaker-core resource layer rather than raw boto3.client(). Consider filing a follow-up to migrate this class to use sagemaker-core.
| _sagemaker_client = None | ||
| _s3_client = None | ||
|
|
||
| @classmethod |
There was a problem hiding this comment.
Missing unit tests: This PR should include a unit test that verifies importing AIRHub does not trigger boto3 client creation (i.e., no NoRegionError when region is unconfigured). A simple test could mock boto3.client and assert it's not called on import, only on first method invocation. This would prevent regressions of the exact issue being fixed.
Issue #, if available:
Description of changes:
The AIRHub class in sagemaker-train constructs boto3.client("sagemaker") and boto3.client("s3") as class-level attributes, which execute at import time. Since the MTRL Launch PR (v3.13.0) added from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer to model_builder.py, import sagemaker.serve now transitively imports AIRHub, causing a NoRegionError in any environment without AWS region configured (e.g., conda-forge build CI).
This change converts those class-level client constructions to lazy initialization via _get_sagemaker_client() and _get_s3_client() classmethods. The clients are still created once and cached — the only difference is they're created on first use rather than at import time.
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.