fix(train): lazily create AIRHub boto3 clients to avoid NoRegionError on import

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)

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