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..6fd3d1a369 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,89 @@ # 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 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(""" + 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" + 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):