refactor(feature-store): Remove client caching from FeatureGroupManager

Remove _lf_client_cache and _s3_client_cache instance caches from
_get_lake_formation_client and _get_s3_client. Each call now creates
a fresh boto3 client directly. Remove corresponding cache-specific
unit tests (cache reuse and different-region tests).
---
X-AI-Prompt: remove client caching for lf and s3 in feature_group_manager and update tests
X-AI-Tool: kiro-cli

Author: BassemHalim

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_group_manager.py

@@ -49,9 +49,6 @@ class LakeFormationConfig(Base):
 class FeatureGroupManager(FeatureGroup):
     """FeatureGroup with extended management capabilities."""
 
-    _lf_client_cache: dict
-    _s3_client_cache: dict
-
     @staticmethod
     def _s3_uri_to_arn(s3_uri: str, region: Optional[str] = None) -> str:
         """
@@ -190,11 +187,7 @@ def _get_lake_formation_client(
         region: Optional[str] = None,
     ):
         """
-        Get a Lake Formation client, reusing a cached client when possible.
-
-        The client is cached on the instance keyed by (session, region). Subsequent
-        calls with the same arguments return the existing client instead of creating
-        a new one.
+        Get a Lake Formation client.
 
         Args:
             session: Boto3 session. If not provided, a new session will be created.
@@ -203,29 +196,16 @@ def _get_lake_formation_client(
         Returns:
             A boto3 Lake Formation client.
         """
-        cache_key = (id(session), region)
-        if not hasattr(self, "_lf_client_cache") or self._lf_client_cache is None:
-            self._lf_client_cache = {}
-
-        if cache_key not in self._lf_client_cache:
-            boto_session = session or Session()
-            self._lf_client_cache[cache_key] = boto_session.client(
-                "lakeformation", region_name=region
-            )
-
-        return self._lf_client_cache[cache_key]
+        boto_session = session or Session()
+        return boto_session.client("lakeformation", region_name=region)
 
     def _get_s3_client(
         self,
         session: Optional[Session] = None,
         region: Optional[str] = None,
     ):
         """
-        Get an S3 client, reusing a cached client when possible.
-
-        The client is cached on the instance keyed by (session, region). Subsequent
-        calls with the same arguments return the existing client instead of creating
-        a new one.
+        Get an S3 client.
 
         Args:
             session: Boto3 session. If not provided, a new session will be created.
@@ -234,17 +214,8 @@ def _get_s3_client(
         Returns:
             A boto3 S3 client.
         """
-        cache_key = (id(session), region)
-        if not hasattr(self, "_s3_client_cache") or self._s3_client_cache is None:
-            self._s3_client_cache = {}
-
-        if cache_key not in self._s3_client_cache:
-            boto_session = session or Session()
-            self._s3_client_cache[cache_key] = boto_session.client(
-                "s3", region_name=region
-            )
-
-        return self._s3_client_cache[cache_key]
+        boto_session = session or Session()
+        return boto_session.client("s3", region_name=region)
 
     def _register_s3_with_lake_formation(
         self,

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_feature_group_manager.py

@@ -70,37 +70,6 @@ def test_creates_client_with_provided_session(self):
         mock_session.client.assert_called_with("lakeformation", region_name="us-west-2")
         assert client == mock_client
 
-    def test_caches_client_for_same_session_and_region(self):
-        """Test that repeated calls with the same session and region reuse the cached client."""
-        mock_session = MagicMock()
-        mock_client = MagicMock()
-        mock_session.client.return_value = mock_client
-
-        fg = MagicMock(spec=FeatureGroupManager)
-        fg._get_lake_formation_client = FeatureGroupManager._get_lake_formation_client.__get__(fg)
-
-        client1 = fg._get_lake_formation_client(session=mock_session, region="us-west-2")
-        client2 = fg._get_lake_formation_client(session=mock_session, region="us-west-2")
-
-        assert client1 is client2
-        mock_session.client.assert_called_once()
-
-    def test_creates_new_client_for_different_region(self):
-        """Test that a different region produces a new client."""
-        mock_session = MagicMock()
-        mock_client_west = MagicMock()
-        mock_client_east = MagicMock()
-        mock_session.client.side_effect = [mock_client_west, mock_client_east]
-
-        fg = MagicMock(spec=FeatureGroupManager)
-        fg._get_lake_formation_client = FeatureGroupManager._get_lake_formation_client.__get__(fg)
-
-        client1 = fg._get_lake_formation_client(session=mock_session, region="us-west-2")
-        client2 = fg._get_lake_formation_client(session=mock_session, region="us-east-1")
-
-        assert client1 is not client2
-        assert mock_session.client.call_count == 2
-
 
 class TestGetS3Client:
     """Tests for _get_s3_client method."""
@@ -135,37 +104,6 @@ def test_creates_client_with_provided_session(self):
         mock_session.client.assert_called_with("s3", region_name="us-west-2")
         assert client == mock_client
 
-    def test_caches_client_for_same_session_and_region(self):
-        """Test that repeated calls with the same session and region reuse the cached client."""
-        mock_session = MagicMock()
-        mock_client = MagicMock()
-        mock_session.client.return_value = mock_client
-
-        fg = MagicMock(spec=FeatureGroupManager)
-        fg._get_s3_client = FeatureGroupManager._get_s3_client.__get__(fg)
-
-        client1 = fg._get_s3_client(session=mock_session, region="us-west-2")
-        client2 = fg._get_s3_client(session=mock_session, region="us-west-2")
-
-        assert client1 is client2
-        mock_session.client.assert_called_once()
-
-    def test_creates_new_client_for_different_region(self):
-        """Test that a different region produces a new client."""
-        mock_session = MagicMock()
-        mock_client_west = MagicMock()
-        mock_client_east = MagicMock()
-        mock_session.client.side_effect = [mock_client_west, mock_client_east]
-
-        fg = MagicMock(spec=FeatureGroupManager)
-        fg._get_s3_client = FeatureGroupManager._get_s3_client.__get__(fg)
-
-        client1 = fg._get_s3_client(session=mock_session, region="us-west-2")
-        client2 = fg._get_s3_client(session=mock_session, region="us-east-1")
-
-        assert client1 is not client2
-        assert mock_session.client.call_count == 2
-
 
 class TestRegisterS3WithLakeFormation:
     """Tests for _register_s3_with_lake_formation method."""