Use list endpoints to check endpoint type

Author: pintaoz

src/sagemaker/hyperpod/inference/hp_endpoint.py

@@ -222,21 +222,24 @@ def list_pods(cls, namespace=None):
             namespace = get_default_namespace()
 
         v1 = client.CoreV1Api()
-        response = v1.list_namespaced_pod(namespace=namespace)
+        list_pods_response = v1.list_namespaced_pod(namespace=namespace)
+
+        list_response = cls.call_list_api(
+            kind=INFERENCE_ENDPOINT_CONFIG_KIND,
+            namespace=namespace,
+        )
+
+        endpoints = set()
+        if list_response and list_response["items"]:
+            for item in list_response["items"]:
+                endpoints.add(item["metadata"]["name"])
 
         pods = []
-        for item in response.items:
+        for item in list_pods_response.items:
             app_name = item.metadata.labels.get("app", None)
-            try:
+            if app_name in endpoints:
                 # list_namespaced_pod will return all pods in the namespace, so we need to filter
                 # out the pods that are created by custom endpoint
-                cls.call_get_api(
-                    name=app_name,
-                    kind=INFERENCE_ENDPOINT_CONFIG_KIND,
-                    namespace=namespace,
-                )
                 pods.append(item.metadata.name)
-            except Exception:
-                continue
 
         return pods

src/sagemaker/hyperpod/inference/hp_jumpstart_endpoint.py

@@ -251,21 +251,24 @@ def list_pods(cls, namespace=None):
             namespace = get_default_namespace()
 
         v1 = client.CoreV1Api()
-        response = v1.list_namespaced_pod(namespace=namespace)
+        list_pods_response = v1.list_namespaced_pod(namespace=namespace)
+
+        list_response = cls.call_list_api(
+            kind=JUMPSTART_MODEL_KIND,
+            namespace=namespace,
+        )
+
+        endpoints = set()
+        if list_response and list_response["items"]:
+            for item in list_response["items"]:
+                endpoints.add(item["metadata"]["name"])
 
         pods = []
-        for item in response.items:
+        for item in list_pods_response.items:
             app_name = item.metadata.labels.get("app", None)
-            try:
+            if app_name in endpoints:
                 # list_namespaced_pod will return all pods in the namespace, so we need to filter
                 # out the pods that are created by jumpstart endpoint
-                cls.call_get_api(
-                    name=app_name,
-                    kind=JUMPSTART_MODEL_KIND,
-                    namespace=namespace,
-                )
                 pods.append(item.metadata.name)
-            except Exception:
-                continue
 
         return pods

test/unit_tests/inference/test_hp_endpoint.py

@@ -195,10 +195,10 @@ def test_invoke(self, mock_endpoint_get, mock_get_cluster_context):
         )
         self.assertEqual(result, "response")
 
-    @patch.object(HPEndpoint, "call_get_api")
+    @patch.object(HPEndpoint, "call_list_api")
     @patch("kubernetes.client.CoreV1Api")
     @patch.object(HPEndpoint, "verify_kube_config")
-    def test_list_pods(self, mock_verify_config, mock_core_api, mock_get_api):
+    def test_list_pods(self, mock_verify_config, mock_core_api, mock_list_api):
         mock_pod1 = MagicMock()
         mock_pod1.metadata.name = "custom-endpoint-pod1"
         mock_pod1.metadata.labels = {"app": "custom-endpoint"}
@@ -214,12 +214,13 @@ def test_list_pods(self, mock_verify_config, mock_core_api, mock_get_api):
             mock_pod3,
         ]
 
-        def mock_behavior(name, kind, namespace):
-            if name.startswith("custom-endpoint"):
-                return
-            else:
-                raise Exception("Endpoint not found")
-        mock_get_api.side_effect = mock_behavior
+        mock_list_api.return_value = {
+            "items": [
+                {
+                    "metadata": {"name": "custom-endpoint"}
+                }
+            ]
+        }
 
         result = self.endpoint.list_pods(namespace="test-ns")
 

test/unit_tests/inference/test_hp_jumpstart_endpoint.py

@@ -141,10 +141,10 @@ def test_invoke(self, mock_endpoint_get, mock_get_cluster_context):
         )
         self.assertEqual(result, "response")
 
-    @patch.object(HPJumpStartEndpoint, "call_get_api")
+    @patch.object(HPJumpStartEndpoint, "call_list_api")
     @patch("kubernetes.client.CoreV1Api")
     @patch.object(HPJumpStartEndpoint, "verify_kube_config")
-    def test_list_pods(self, mock_verify_config, mock_core_api, mock_get_api):
+    def test_list_pods(self, mock_verify_config, mock_core_api, mock_list_api):
         mock_pod1 = MagicMock()
         mock_pod1.metadata.name = "js-endpoint-pod1"
         mock_pod1.metadata.labels = {"app": "js-endpoint"}
@@ -160,12 +160,13 @@ def test_list_pods(self, mock_verify_config, mock_core_api, mock_get_api):
             mock_pod3,
         ]
 
-        def mock_behavior(name, kind, namespace):
-            if name.startswith("js-endpoint"):
-                return
-            else:
-                raise Exception("Endpoint not found")
-        mock_get_api.side_effect = mock_behavior
+        mock_list_api.return_value = {
+            "items": [
+                {
+                    "metadata": {"name": "js-endpoint"}
+                }
+            ]
+        }
 
         result = self.endpoint.list_pods(namespace="test-ns")