fix: Stabilize integration cluster tests

src/sagemaker/hyperpod/cli/clients/kubernetes_client.py

@@ -69,10 +69,20 @@ class KubernetesClient:
     def __new__(cls, config_file: Optional[str] = None) -> "KubernetesClient":
         if cls._instance is None:
             cls._instance = super(KubernetesClient, cls).__new__(cls)
+            cls._instance._config_file = config_file or KUBE_CONFIG_PATH
             config.load_kube_config(config_file=config_file or KUBE_CONFIG_PATH)
             cls._instance._kube_client = client.ApiClient()
         return cls._instance
 
+    @classmethod
+    def _from_config_file(cls, config_file: str) -> "KubernetesClient":
+        instance = super(KubernetesClient, cls).__new__(cls)
+        instance._config_file = config_file
+        instance._kube_client = config.new_client_from_config(
+            config_file=config_file
+        )
+        return instance
+
     def set_context(
         self,
         context_name: str,
@@ -110,8 +120,10 @@ def set_context(
         with open(KUBE_CONFIG_PATH, "w") as file:
             yaml.safe_dump(kubeconfig, file)
 
-        # Load the updated kubeconfig
+        # Load the updated kubeconfig and refresh this client's configuration.
         config.load_kube_config(config_file=KUBE_CONFIG_PATH)
+        self._config_file = KUBE_CONFIG_PATH
+        self._kube_client = client.ApiClient()
         logger.debug(f"Current context set to '{context_name}'")
 
     def get_core_v1_api(self) -> client.CoreV1Api:
@@ -182,20 +194,23 @@ def list_node_with_temp_config(
             file (str): The path to the kubeconfig file
             selector_str (str): Selector to filter
         """
-        config.load_kube_config(config_file=file)
-        v1Client = client.CoreV1Api()
+        api_client = self._kube_client
+        if getattr(self, "_config_file", None) != file:
+            api_client = config.new_client_from_config(config_file=file)
+
+        v1_client = client.CoreV1Api(api_client)
         _continue = None
 
         nodes = []
         while True:
-            response = v1Client.list_node(
+            response = v1_client.list_node(
                 label_selector=selector_str,
                 limit=200,  # Set a reasonable limit
                 _continue=_continue,
             )
             nodes.extend(response.items)
 
-            _continue = response._metadata._continue  # Get the continue token
+            _continue = response.metadata._continue  # Get the continue token
 
             if not _continue:  # If there's no more data, break the loop
                 break
@@ -236,7 +251,9 @@ def get_sagemaker_managed_namespace(self, namespace: Optional[str]):
         if namespace is None:
             return None
         try:
-            response = client.CoreV1Api().read_namespace(name=namespace)
+            response = client.CoreV1Api(self._kube_client).read_namespace(
+                name=namespace
+            )
             labels = response.metadata.labels
             if labels and SAGEMAKER_MANAGED_QUEUE_LABEL in labels and labels[SAGEMAKER_MANAGED_QUEUE_LABEL] == "true":
                 return response
@@ -254,19 +271,19 @@ def list_pods_with_labels(self, namespace: str, label_selector: str):
         )
 
     def list_pods_in_all_namespaces_with_labels(self, label_selector: str):
-        v1Client = client.CoreV1Api()
+        v1_client = client.CoreV1Api(self._kube_client)
         pods = []
         _continue = None
 
         while True:
-            response = v1Client.list_pod_for_all_namespaces(
+            response = v1_client.list_pod_for_all_namespaces(
                 label_selector=label_selector,
                 limit=200,  # Set a reasonable limit
                 _continue=_continue,
             )
             pods.extend(response.items)
 
-            _continue = response._metadata._continue  # Get the continue token
+            _continue = response.metadata._continue  # Get the continue token
 
             if not _continue:  # If there's no more data, break the loop
                 break
@@ -367,7 +384,9 @@ def list_workload_priority_classes(self):
         )
 
     def get_cluster_queue(self, cluster_queue_name: str):
-        return client.CustomObjectsApi().get_cluster_custom_object(
+        return client.CustomObjectsApi(
+            self._kube_client
+        ).get_cluster_custom_object(
             group=KUEUE_CUSTOM_OBJECT_GROUP,
             version=KUEUE_CUSTOM_OBJECT_VERSION,
             plural=CLUSTER_QUEUE_PRIORITY_CLASS_CUSTOM_OBJECT_PLURAL,

src/sagemaker/hyperpod/cli/commands/cluster.py

@@ -48,7 +48,7 @@
     get_user_agent_extra_suffix,
 )
 from sagemaker.hyperpod.cli.service.list_pods import (
-    ListPods,
+    _list_pods_and_get_requested_resources_group_by_node_name,
 )
 from sagemaker.hyperpod.cli.utils import (
     get_name_from_arn,
@@ -303,11 +303,11 @@ def rate_limited_operation(
             return None
         eks_cluster_name = get_name_from_arn(eks_cluster_arn)
         _update_kube_config(eks_cluster_name, region, temp_config_file)
-        k8s_client = KubernetesClient(config_file=temp_config_file)
+        k8s_client = KubernetesClient._from_config_file(temp_config_file)
         nodes = k8s_client.list_node_with_temp_config(
             temp_config_file, SAGEMAKER_HYPERPOD_NAME_LABEL
         )
-        nodes_info = _aggregate_nodes_info(nodes)
+        nodes_info = _aggregate_nodes_info(nodes, k8s_client)
 
         ns_nominal_quota = {}
         ns_quota_usage = {}
@@ -466,10 +466,10 @@ def _restructure_output(summary_list, namespaces):
 
 def _aggregate_nodes_info(
     nodes: List[client.V1Node],
+    k8s_client: KubernetesClient,
 ) -> Dict[str, Dict[str, Any]]:
-    list_pods_service = ListPods()
     nodes_resource_allocated_dict = (
-        list_pods_service.list_pods_and_get_requested_resources_group_by_node_name()
+        _list_pods_and_get_requested_resources_group_by_node_name(k8s_client)
     )
     nodes_summary: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int))
 

src/sagemaker/hyperpod/cli/service/list_pods.py

@@ -27,6 +27,36 @@
 from sagemaker.hyperpod.cli.constants.command_constants import NVIDIA_GPU_RESOURCE_LIMIT_KEY
 from sagemaker.hyperpod.cli.service.discover_namespaces import DiscoverNamespaces
 
+def _list_pods_and_get_requested_resources_group_by_node_name(
+    k8s_client: KubernetesClient,
+):
+    label_filter = "training.kubeflow.org/job-name"
+    pods = k8s_client.list_pods_in_all_namespaces_with_labels(label_filter)
+
+    accelerator_devices_requests_by_node = defaultdict(int)
+    for pod in pods:
+        node_name = pod.spec.node_name
+        if node_name:
+            for container in pod.spec.containers:
+                if container.resources and container.resources.requests:
+                    gpu_request = container.resources.requests.get(
+                        NVIDIA_GPU_RESOURCE_LIMIT_KEY
+                    )
+                    neuron_request = container.resources.requests.get(
+                        "aws.amazon.com/neurondevice"
+                    )
+                    if gpu_request:
+                        accelerator_devices_requests_by_node[node_name] += int(
+                            gpu_request
+                        )
+                    if neuron_request:
+                        accelerator_devices_requests_by_node[node_name] += int(
+                            neuron_request
+                        )
+
+    return accelerator_devices_requests_by_node
+
+
 class ListPods:
     def __init__(self):
         return
@@ -72,35 +102,9 @@ def list_pods_and_get_requested_resources_group_by_node_name(
         Group by the node_name of the pod and the value is the
         accelerator devices resources requested
         """
-        k8s_client = KubernetesClient()
-
-        label_filter = f"training.kubeflow.org/job-name"
-        pods = k8s_client.list_pods_in_all_namespaces_with_labels(label_filter)
-
-        # Dictionary to hold total GPU/Neuron requests per node
-        accelerator_devices_requests_by_node = defaultdict(int)
-
-        # Loop through each pod and aggregate the GPU resource requests per node
-        for pod in pods:
-            node_name = pod.spec.node_name
-            # Check if the pod has a node_name assigned, if pod is not scheduled, it may not have node_name
-            if node_name:
-                for container in pod.spec.containers:
-                    if container.resources and container.resources.requests:
-                        gpu_request = container.resources.requests.get(NVIDIA_GPU_RESOURCE_LIMIT_KEY)
-                        neuron_request = container.resources.requests.get(
-                            "aws.amazon.com/neurondevice"
-                        )
-                        if gpu_request:
-                            accelerator_devices_requests_by_node[node_name] += int(
-                                gpu_request
-                            )
-                        if neuron_request:
-                            accelerator_devices_requests_by_node[node_name] += int(
-                                neuron_request
-                            )
-
-        return accelerator_devices_requests_by_node
+        return _list_pods_and_get_requested_resources_group_by_node_name(
+            KubernetesClient()
+        )
 
     def _generate_list_pods_output(self, pods: V1PodList) -> Optional[str]:
         output_pods = {"pods": []}

src/sagemaker/hyperpod/cli/service/self_subject_access_review.py

@@ -38,7 +38,6 @@ def self_subject_access_review(
             spec=V1SelfSubjectAccessReviewSpec(
                 resource_attributes=resource_attributes,
                 non_resource_attributes=non_resource_attributes,
-                local_vars_configuration=local_vars_configuration,
             )
         )