diff --git a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py index 701b397d93..0350386db1 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py +++ b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py @@ -32,9 +32,24 @@ 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") + # Default clients are created on first use rather than at import time, so that + # importing the SDK does not require AWS region/credentials to be configured. + _sagemaker_client = None + _s3_client = None + + @classmethod + def _get_sagemaker_client(cls): + """Return the default SageMaker client, creating it on first use.""" + if cls._sagemaker_client is None: + cls._sagemaker_client = boto3.client("sagemaker") + return cls._sagemaker_client + + @classmethod + def _get_s3_client(cls): + """Return the default 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 +91,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 +133,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 +171,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 +219,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 +249,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 +280,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 +326,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 +344,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..a02d0f2475 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py +++ b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py @@ -11,6 +11,9 @@ # 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 from unittest.mock import patch, MagicMock from sagemaker.ai_registry.air_constants import AIR_DEFAULT_PAGE_SIZE @@ -225,3 +228,39 @@ def test_generate_hub_names_no_padding(self): assert ( AIRHub.hubName.isalnum() ), f"Hub name should only contain alphanumeric characters: {AIRHub.hubName}" + + def test_import_does_not_require_aws_region(self, tmp_path): + """Importing AIRHub must not create boto3 clients (fails without a region).""" + env = { + k: v + for k, v in os.environ.items() + if k not in ("AWS_REGION", "AWS_DEFAULT_REGION", "AWS_PROFILE") + } + env["AWS_CONFIG_FILE"] = str(tmp_path / "config") + env["AWS_SHARED_CREDENTIALS_FILE"] = str(tmp_path / "credentials") + + result = subprocess.run( + [sys.executable, "-c", "import sagemaker.ai_registry.air_hub"], + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + @patch("sagemaker.ai_registry.air_hub.boto3") + def test_clients_created_lazily_and_cached(self, mock_boto3): + AIRHub._sagemaker_client = None + AIRHub._s3_client = None + try: + sm_client = AIRHub._get_sagemaker_client() + s3_client = AIRHub._get_s3_client() + + assert AIRHub._get_sagemaker_client() is sm_client + assert AIRHub._get_s3_client() is s3_client + mock_boto3.client.assert_any_call("sagemaker") + mock_boto3.client.assert_any_call("s3") + assert mock_boto3.client.call_count == 2 + finally: + AIRHub._sagemaker_client = None + AIRHub._s3_client = None