fix: create AI Registry boto3 clients lazily to avoid import-time reg… by mufaddal-rohawala · Pull Request #6558 · aws/sagemaker-python-sdk

mufaddal-rohawala · 2026-11-03T20:04:28Z

…ion errors

AIRHub created its sagemaker and s3 boto3 clients as class attributes, so they were instantiated at import time. Importing sagemaker.mlops or sagemaker.serve (which transitively import sagemaker.ai_registry.air_hub) therefore raised botocore.exceptions.NoRegionError whenever no AWS region was configured, breaking CI runners, docs builds, and first-run setup.

Create both clients lazily on first use via _get_sagemaker_client() / _get_s3_client() so importing the SDK has no AWS side effects and does not require a region. Add a subprocess regression test that imports the SDK with every region source scrubbed, plus unit tests covering lazy creation and caching.

Fixes #6542

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

github-actions

🤖 AI code review

This PR converts AIRHub's class-level boto3 sagemaker and s3 clients to lazy initialization via _get_sagemaker_client() / _get_s3_client(), so importing sagemaker.mlops / sagemaker.serve no longer raises NoRegionError at import time (issue #6542). It updates all call sites and adds a subprocess regression test plus unit tests for lazy creation and caching. The change is correct, minimal, and preserves the None-initialized attributes so existing mock-assignment tests keep working. Only a minor thread-safety point worth considering.

_{Reviewed commit b12a514. Automated review; verify before acting.}

github-actions · 2026-11-03T20:35:41Z

+    def _get_sagemaker_client(cls):
+        """Return the shared SageMaker client, creating it on first use."""
+        if cls._sagemaker_client is None:
+            # Use production SageMaker endpoint (default)


🟡 Low: The lazy _get_sagemaker_client() / _get_s3_client() use an unguarded check-then-create pattern. If two threads call these concurrently before the client is cached, boto3.client(...) can run twice and one result is discarded. Prior lazy-init reviews on this repo recommended guarding with threading.Lock() (double-checked locking). Low severity since the discarded client is harmless, but worth matching the established pattern. Source: #5935

github-actions

🤖 AI code review

This PR converts AIRHub's class-level boto3 sagemaker/s3 clients from eager instantiation to lazy creation via _get_sagemaker_client()/_get_s3_client(), eliminating the import-time NoRegionError when importing sagemaker.mlops/sagemaker.serve. It updates all call sites and adds a subprocess regression test plus unit tests for lazy creation/caching. The change is consistent with prior precedent (#6401, #5935) and preserves the None-initialized class attributes for mock compatibility. The PR looks clean; no blocking issues found.
[]

_{Reviewed commit eb45821. Automated review; verify before acting.}

…ion errors AIRHub created its `sagemaker` and `s3` boto3 clients as class attributes, so they were instantiated at import time. Importing `sagemaker.mlops` or `sagemaker.serve` (which transitively import `sagemaker.ai_registry.air_hub`) therefore raised `botocore.exceptions.NoRegionError` whenever no AWS region was configured, breaking CI runners, docs builds, and first-run setup. Create both clients lazily on first use via `_get_sagemaker_client()` / `_get_s3_client()` so importing the SDK has no AWS side effects and does not require a region. Add a subprocess regression test that imports the SDK with every region source scrubbed, plus unit tests covering lazy creation and caching. Fixes aws#6542

The region-regression test imported sagemaker.mlops and sagemaker.serve in a subprocess, but the sagemaker-train unit-test job installs only sagemaker-train and sagemaker-core. The subprocess failed with ModuleNotFoundError before ever reaching the region check, so the test failed in CI. Import those packages only when present (the transitive path from the original report) and always import sagemaker.ai_registry.air_hub directly, where the aws#6542 defect actually lived. The lazy-client assertions are unchanged.

github-actions

🤖 AI code review

This PR converts AIRHub's class-level boto3.client() attributes to lazy initialization via _get_sagemaker_client() / _get_s3_client(), eliminating the import-time NoRegionError described in issue #6542. It adds a subprocess regression test (import with AWS env scrubbed) plus unit tests for lazy creation and caching. The change faithfully follows the established pattern from PR #6401 and preserves the None-initialized attributes so existing mock-assignment tests still work. The change is clean and correct; findings below are minor.

Findings not attached to a changed line:

sagemaker-train/tests/unit/ai_registry/test_air_hub.py:155 🟡 Low: The lazy-client tests (test_get_sagemaker_client_created_lazily_and_cached, test_get_s3_client_created_lazily_and_cached) set AIRHub._sagemaker_client/_s3_client to a MagicMock sentinel but never reset them after. Because these are class attributes shared across the whole test session, the cached sentinel leaks into any later test that relies on lazy creation. The existing TestAIRHub tests happen to reassign the attribute explicitly so they're unaffected today, but consider adding teardown (or a fixture) that restores the attributes to None to keep this isolated as tests are added.

_{Reviewed commit 78252a5. Automated review; verify before acting.}

github-actions · 2026-11-04T18:26:53Z

+    def _get_sagemaker_client(cls):
+        """Return the shared SageMaker client, creating it on first use."""
+        if cls._sagemaker_client is None:
+            # Use production SageMaker endpoint (default)


🟡 Low: Retrieved precedent for this exact fix repeatedly recommends guarding the lazy init with a threading.Lock (double-checked locking) so concurrent first use doesn't construct duplicate clients. The direct prior fix (PR #6401) also skipped the lock, so this is not blocking, but worth considering if AIRHub may be used from multiple threads. Source: #5935

github-actions

🤖 AI code review

The PR converts AIRHub's class-level boto3.client() attributes to lazy, cached _get_sagemaker_client()/_get_s3_client() classmethods so importing sagemaker.mlops/sagemaker.serve no longer raises NoRegionError at import time (issue #6542). The change is correct, preserves the None-initialized attributes for existing mock-assignment patterns, and adds a subprocess regression test plus unit tests for lazy creation/caching. Overall it's clean and matches established precedent; my notes are minor (optional thread-safety lock, missing return annotations, and test state leakage)."

_{Reviewed commit 2f0d244. Automated review; verify before acting.}

github-actions · 2026-11-05T02:25:55Z

+    _sagemaker_client = None
+    _s3_client = None
+
+    @classmethod


🟡 Low: _get_sagemaker_client/_get_s3_client use the simple if cls._x is None lazy pattern without a lock. Prior lazy-init work on this repo recommends guarding lazy boto3-client initialization with a threading.Lock() (double-checked locking) so concurrent first use doesn't create duplicate clients. Low priority since the clients are cached and duplicate creation is harmless, but worth matching the established pattern.

Source: #5935

github-actions · 2026-11-05T02:25:56Z

+    def _get_sagemaker_client(cls):
+        """Return the shared SageMaker client, creating it on first use."""
+        if cls._sagemaker_client is None:
+            # Use production SageMaker endpoint (default)


🟡 Low: The new lazy-init classmethods lack return type annotations. Precedent on this change recommends annotating lazy-init classmethods (e.g. mypy_boto3_sagemaker types or Any if stubs are unavailable).

Source: #5935

github-actions · 2026-11-05T02:25:56Z

+    @patch("sagemaker.ai_registry.air_hub.boto3")
+    def test_get_sagemaker_client_created_lazily_and_cached(self, mock_boto3):
+        AIRHub._sagemaker_client = None
+        sentinel = MagicMock()


🟡 Low: test_get_sagemaker_client_created_lazily_and_cached and test_get_s3_client_created_lazily_and_cached set AIRHub._sagemaker_client/_s3_client to None but never restore them. Because these are class-level attributes, the None state (or the sentinel MagicMock from the lazy call) leaks into any later test that doesn't explicitly reassign the client. Existing tests in this file happen to reassign the attributes, so they're unaffected today, but consider resetting state in a fixture/teardown to avoid order-dependent fragility.

github-actions

🤖 AI code review

This PR converts AIRHub's class-level boto3 sagemaker and s3 clients to lazy initialization via _get_sagemaker_client() / _get_s3_client(), eliminating NoRegionError at import time (issue #6542). It adds a subprocess regression test that imports with AWS env scrubbed, plus unit tests for lazy creation and caching. The change closely follows the established precedent (PRs #5935, #6401) and looks correct. A couple of minor points worth considering below.

_{Reviewed commit b944229. Automated review; verify before acting.}

github-actions · 2026-11-06T03:21:10Z

+    def _get_sagemaker_client(cls):
+        """Return the shared SageMaker client, creating it on first use."""
+        if cls._sagemaker_client is None:
+            # Use production SageMaker endpoint (default)


🟡 Low: The lazy-init classmethods use an unguarded check-then-set, so two threads reaching _get_sagemaker_client() concurrently while the client is still None could each call boto3.client(...) (the later one overwrites the cache). Prior precedent for this exact pattern recommends protecting the double-checked initialization with threading.Lock() when concurrent first use is possible. Source: #5935

github-actions · 2026-11-06T03:21:10Z

+    @patch("sagemaker.ai_registry.air_hub.boto3")
+    def test_get_sagemaker_client_created_lazily_and_cached(self, mock_boto3):
+        AIRHub._sagemaker_client = None
+        sentinel = MagicMock()


🟡 Low: test_get_sagemaker_client_created_lazily_and_cached sets AIRHub._sagemaker_client = None at the start but doesn't restore it afterward, leaving the class attribute set to the MagicMock sentinel (with boto3 patched). Since AIRHub state is class-level and shared across tests, consider restoring/resetting it in teardown to avoid order-dependent leakage into other tests that rely on _sagemaker_client/_s3_client state. The same applies to the s3 variant.

github-actions · 2026-11-06T03:21:10Z

-    _sagemaker_client = boto3.client("sagemaker")
-    _s3_client = boto3.client("s3")
+    # Boto3 clients are created lazily on first use so that importing this module
+    # (and, transitively, sagemaker.mlops / sagemaker.serve) has no AWS side effects


🟡 Low: Issue #6542 also notes that clients are bound to the default credential chain rather than a caller-provided session. The non-session code paths here still fall back to boto3.client(...) directly (default credential chain). This PR addresses the import-time NoRegionError, but the session-binding concern from the issue remains unaddressed; worth confirming it's intentionally out of scope. Source: #6542

mufaddal-rohawala temporarily deployed to auto-approve November 3, 2026 20:26 — with GitHub Actions Inactive

mufaddal-rohawala temporarily deployed to auto-approve November 3, 2026 20:27 — with GitHub Actions Inactive

github-actions Bot reviewed Nov 3, 2026

View reviewed changes

mufaddal-rohawala temporarily deployed to auto-approve November 3, 2026 22:52 — with GitHub Actions Inactive

mufaddal-rohawala temporarily deployed to auto-approve November 3, 2026 22:53 — with GitHub Actions Inactive

github-actions Bot reviewed Nov 3, 2026

View reviewed changes

mufiAmazon added 2 commits November 4, 2026 18:24

mufaddal-rohawala force-pushed the fix/air-hub-lazy-clients-6542 branch from eb45821 to 78252a5 Compare November 4, 2026 18:25

mufaddal-rohawala temporarily deployed to auto-approve November 4, 2026 18:25 — with GitHub Actions Inactive

github-actions Bot reviewed Nov 4, 2026

View reviewed changes

Merge branch 'master' into fix/air-hub-lazy-clients-6542

2f0d244

mufaddal-rohawala temporarily deployed to auto-approve November 5, 2026 02:24 — with GitHub Actions Inactive

github-actions Bot reviewed Nov 5, 2026

View reviewed changes

zhaoqizqwang approved these changes Nov 5, 2026

View reviewed changes

mohamedzeidan2021 approved these changes Nov 5, 2026

View reviewed changes

ci: re-trigger checks (retry flaky Read the Docs / serve-deps install)

b944229

mufaddal-rohawala temporarily deployed to auto-approve November 6, 2026 03:20 — with GitHub Actions Inactive

mufaddal-rohawala deployed to auto-approve November 6, 2026 03:20 — with GitHub Actions Active

mufaddal-rohawala temporarily deployed to auto-approve November 6, 2026 03:20 — with GitHub Actions Inactive

github-actions Bot reviewed Nov 6, 2026

View reviewed changes

Conversation

mufaddal-rohawala commented Nov 3, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🤖 AI code review

Uh oh!

github-actions Bot Nov 3, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🤖 AI code review

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🤖 AI code review

Uh oh!

github-actions Bot Nov 4, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🤖 AI code review

Uh oh!

github-actions Bot Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🤖 AI code review

Uh oh!

github-actions Bot Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants