From 795ce2674bd35d9ea1c4440c0a5d50ec72e50b8b Mon Sep 17 00:00:00 2001 From: Mufaddal Rohawala Date: Tue, 3 Nov 2026 01:18:59 +0000 Subject: [PATCH 1/3] fix: create AI Registry boto3 clients lazily to avoid import-time region 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 --- .../src/sagemaker/ai_registry/air_hub.py | 41 ++++++++--- .../tests/unit/ai_registry/test_air_hub.py | 68 +++++++++++++++++++ 2 files changed, 98 insertions(+), 11 deletions(-) diff --git a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py index 701b397d93..b80200c599 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py +++ b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py @@ -32,9 +32,26 @@ class AIRHub: """AI Registry Hub class for managing hub content operations.""" - # Use production SageMaker endpoint (default) - _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 + # and does not require a configured region. See issue #6542. + _sagemaker_client = None + _s3_client = None + + @classmethod + 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) + cls._sagemaker_client = boto3.client("sagemaker") + return cls._sagemaker_client + + @classmethod + def _get_s3_client(cls): + """Return the shared S3 client, creating it on first use.""" + if cls._s3_client is None: + cls._s3_client = boto3.client("s3") + return cls._s3_client @classmethod def _generate_hub_names(cls, region: str, account_id: str) -> None: @@ -76,7 +93,7 @@ def _create_airegistry_hub_if_not_exists(cls, client=None) -> None: cls._ensure_hub_name_initialized() if client is None: - client = cls._sagemaker_client + client = cls._get_sagemaker_client() try: client.describe_hub(HubName=cls.hubName) @@ -118,7 +135,7 @@ def import_hub_content( Returns: Response from import_hub_content API call """ - client = session.sagemaker_client if session is not None else cls._sagemaker_client + client = session.sagemaker_client if session is not None else cls._get_sagemaker_client() cls._create_airegistry_hub_if_not_exists(client) @@ -156,7 +173,7 @@ def list_hub_content( """ cls._ensure_hub_name_initialized() - client = session.sagemaker_client if session is not None else cls._sagemaker_client + client = session.sagemaker_client if session is not None else cls._get_sagemaker_client() request = { "HubName": cls.hubName, @@ -204,7 +221,7 @@ def describe_hub_content( """ cls._ensure_hub_name_initialized() - client = session.sagemaker_client if session is not None else cls._sagemaker_client + client = session.sagemaker_client if session is not None else cls._get_sagemaker_client() request = { "HubName": cls.hubName, @@ -234,7 +251,7 @@ def list_hub_content_versions( """ cls._ensure_hub_name_initialized() - client = session.sagemaker_client if session is not None else cls._sagemaker_client + client = session.sagemaker_client if session is not None else cls._get_sagemaker_client() request = { "HubName": cls.hubName, @@ -265,7 +282,7 @@ def delete_hub_content( """ cls._ensure_hub_name_initialized() - client = session.sagemaker_client if session is not None else cls._sagemaker_client + client = session.sagemaker_client if session is not None else cls._get_sagemaker_client() request = { "HubName": cls.hubName, @@ -311,7 +328,9 @@ def upload_to_s3(bucket: str, prefix: str, local_file_path: str) -> str: S3 URI of uploaded file """ extra_args = AIRHub._default_bucket_expected_owner_args(bucket) - AIRHub._s3_client.upload_file(local_file_path, bucket, prefix, ExtraArgs=extra_args or None) + AIRHub._get_s3_client().upload_file( + local_file_path, bucket, prefix, ExtraArgs=extra_args or None + ) return f"s3://{bucket}/{prefix}" @staticmethod @@ -327,4 +346,4 @@ def download_from_s3(s3_uri: str, local_path: str) -> None: bucket = parsed.netloc key = parsed.path.lstrip("/") extra_args = AIRHub._default_bucket_expected_owner_args(bucket) - AIRHub._s3_client.download_file(bucket, key, local_path, ExtraArgs=extra_args or None) + AIRHub._get_s3_client().download_file(bucket, key, local_path, ExtraArgs=extra_args or None) diff --git a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py index 8cef29af6e..12be5d40ff 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py +++ b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py @@ -11,12 +11,80 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. +import os +import subprocess +import sys +import textwrap from unittest.mock import patch, MagicMock from sagemaker.ai_registry.air_constants import AIR_DEFAULT_PAGE_SIZE from sagemaker.ai_registry.air_hub import AIRHub +class TestAIRHubLazyClients: + """Importing the module must have no AWS side effects (issue #6542).""" + + def test_import_without_region_has_no_aws_side_effects(self): + """Importing sagemaker.mlops must not require a configured AWS region. + + Reproduces issue #6542: AIRHub used to create boto3 clients as class + attributes, which ran at import time and raised NoRegionError when no + region was configured. Clients are now created lazily, so the import + must succeed in a subprocess with every region source scrubbed. + """ + script = textwrap.dedent(""" + import sagemaker.mlops # noqa: F401 + import sagemaker.serve # noqa: F401 + from sagemaker.ai_registry.air_hub import AIRHub + + assert AIRHub._sagemaker_client is None, "sagemaker client created at import" + assert AIRHub._s3_client is None, "s3 client created at import" + print("OK") + """) + env = {k: v for k, v in os.environ.items() if not k.startswith("AWS_")} + # Point AWS config/credentials at a non-existent path so no region leaks in. + env["AWS_CONFIG_FILE"] = os.devnull + env["AWS_SHARED_CREDENTIALS_FILE"] = os.devnull + + result = subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + ) + + assert ( + result.returncode == 0 + ), f"import failed without a region (issue #6542):\n{result.stderr}" + assert "OK" in result.stdout + + @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() + mock_boto3.client.return_value = sentinel + + first = AIRHub._get_sagemaker_client() + second = AIRHub._get_sagemaker_client() + + assert first is sentinel + assert second is sentinel + mock_boto3.client.assert_called_once_with("sagemaker") + + @patch("sagemaker.ai_registry.air_hub.boto3") + def test_get_s3_client_created_lazily_and_cached(self, mock_boto3): + AIRHub._s3_client = None + sentinel = MagicMock() + mock_boto3.client.return_value = sentinel + + first = AIRHub._get_s3_client() + second = AIRHub._get_s3_client() + + assert first is sentinel + assert second is sentinel + mock_boto3.client.assert_called_once_with("s3") + + class TestAIRHub: @patch("sagemaker.ai_registry.air_hub.boto3") def test_import_hub_content(self, mock_boto3): From 78252a50469e5602544c764f29bf69f1bad6a17c Mon Sep 17 00:00:00 2001 From: Mufaddal Rohawala Date: Wed, 4 Nov 2026 01:29:44 +0000 Subject: [PATCH 2/3] test: make AIRHub import test tolerate absent mlops/serve packages 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 #6542 defect actually lived. The lazy-client assertions are unchanged. --- .../tests/unit/ai_registry/test_air_hub.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py index 12be5d40ff..6fd3d1a369 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py +++ b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py @@ -25,16 +25,25 @@ class TestAIRHubLazyClients: """Importing the module must have no AWS side effects (issue #6542).""" def test_import_without_region_has_no_aws_side_effects(self): - """Importing sagemaker.mlops must not require a configured AWS region. + """Importing AIRHub must not require a configured AWS region. Reproduces issue #6542: AIRHub used to create boto3 clients as class attributes, which ran at import time and raised NoRegionError when no region was configured. Clients are now created lazily, so the import must succeed in a subprocess with every region source scrubbed. + + Imports sagemaker.mlops / sagemaker.serve too when they are installed + (the transitive path from the original report), but tolerates their + absence: the sagemaker-train unit-test environment installs only + sagemaker-train and sagemaker-core, and the defect lives in air_hub + itself, which the direct import below exercises. """ script = textwrap.dedent(""" - import sagemaker.mlops # noqa: F401 - import sagemaker.serve # noqa: F401 + for _mod in ("sagemaker.mlops", "sagemaker.serve"): + try: + __import__(_mod) + except ImportError: + pass from sagemaker.ai_registry.air_hub import AIRHub assert AIRHub._sagemaker_client is None, "sagemaker client created at import" From b9442290ad2fc9e4b7d2c0ec4b6700ba86492899 Mon Sep 17 00:00:00 2001 From: Mufaddal Rohawala Date: Fri, 6 Nov 2026 03:18:49 +0000 Subject: [PATCH 3/3] ci: re-trigger checks (retry flaky Read the Docs / serve-deps install)