fix: port forward for local interactive tests

Author: kryanbeane

.github/workflows/e2e_tests.yaml

@@ -86,6 +86,36 @@ jobs:
         with:
           user-name: sdk-user
 
+      - name: Grant sdk-user port-forwarding permissions
+        run: |
+          cat <<EOF | kubectl apply -f -
+          apiVersion: rbac.authorization.k8s.io/v1
+          kind: ClusterRole
+          metadata:
+            name: port-forward-permissions
+          rules:
+          - apiGroups: [""]
+            resources: ["services", "pods"]
+            verbs: ["get", "list", "watch"]
+          - apiGroups: [""]
+            resources: ["pods/portforward"]
+            verbs: ["create"]
+          ---
+          apiVersion: rbac.authorization.k8s.io/v1
+          kind: ClusterRoleBinding
+          metadata:
+            name: sdk-user-port-forward-binding
+          subjects:
+          - kind: User
+            name: sdk-user
+            apiGroup: rbac.authorization.k8s.io
+          roleRef:
+            kind: ClusterRole
+            name: port-forward-permissions
+            apiGroup: rbac.authorization.k8s.io
+          EOF
+        shell: bash
+
       - name: Configure RBAC for sdk user with limited permissions
         run: |
           kubectl create clusterrole list-ingresses --verb=get,list --resource=ingresses
@@ -117,7 +147,7 @@ jobs:
           pip install poetry
           poetry install --with test,docs
           echo "Running e2e tests..."
-          poetry run pytest -v -s ./tests/e2e -m 'kind and nvidia_gpu' > ${CODEFLARE_TEST_OUTPUT_DIR}/pytest_output.log 2>&1
+          poetry run pytest -v -s --log-cli-level=INFO ./tests/e2e/local_interactive_sdk_kind_test.py::TestRayLocalInteractiveOauth::test_local_interactives_nvidia_gpu > ${CODEFLARE_TEST_OUTPUT_DIR}/pytest_output.log 2>&1
         env:
           GRPC_DNS_RESOLVER: "native"
 

tests/e2e/local_interactive_sdk_kind_test.py

@@ -1,13 +1,14 @@
 from codeflare_sdk import (
     Cluster,
     ClusterConfiguration,
-    TokenAuthentication,
     generate_cert,
 )
 
 import pytest
 import ray
 import math
+import time
+import subprocess
 
 from support import *
 
@@ -16,8 +17,17 @@
 class TestRayLocalInteractiveOauth:
     def setup_method(self):
         initialize_kubernetes_client(self)
+        self.port_forward_process = None
 
     def teardown_method(self):
+        if self.port_forward_process:
+            self.port_forward_process.terminate()
+            try:
+                self.port_forward_process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                self.port_forward_process.kill()
+                self.port_forward_process.wait()
+            self.port_forward_process = None
         delete_namespace(self)
         delete_kueue_resources(self)
 
@@ -39,6 +49,8 @@ def run_local_interactives(
     ):
         cluster_name = "test-ray-cluster-li"
 
+        ray.shutdown()
+
         cluster = Cluster(
             ClusterConfiguration(
                 name=cluster_name,
@@ -49,45 +61,124 @@ def run_local_interactives(
                 head_memory_requests=2,
                 head_memory_limits=2,
                 worker_cpu_requests="500m",
-                worker_cpu_limits=1,
+                worker_cpu_limits="500m",
                 worker_memory_requests=1,
                 worker_memory_limits=4,
                 worker_extended_resource_requests={gpu_resource_name: number_of_gpus},
-                write_to_file=True,
                 verify_tls=False,
             )
         )
-        cluster.up()
-        cluster.wait_ready()
-
-        generate_cert.generate_tls_cert(cluster_name, self.namespace)
-        generate_cert.export_env(cluster_name, self.namespace)
-
-        print(cluster.local_client_url())
 
-        ray.shutdown()
-        ray.init(address=cluster.local_client_url(), logging_level="DEBUG")
-
-        @ray.remote(num_gpus=number_of_gpus / 2)
-        def heavy_calculation_part(num_iterations):
-            result = 0.0
-            for i in range(num_iterations):
-                for j in range(num_iterations):
-                    for k in range(num_iterations):
-                        result += math.sin(i) * math.cos(j) * math.tan(k)
-            return result
-
-        @ray.remote(num_gpus=number_of_gpus / 2)
-        def heavy_calculation(num_iterations):
-            results = ray.get(
-                [heavy_calculation_part.remote(num_iterations // 30) for _ in range(30)]
+        try:
+            cluster.up()
+
+            cluster.wait_ready()
+            cluster.status()
+
+            TIMEOUT = 300  # 5 minutes
+            END = time.time() + TIMEOUT
+
+            head_pod_name = None
+            worker_pod_name = None
+
+            while time.time() < END:
+                if not head_pod_name:
+                    head_pod_name = kubectl_get_pod_name_by_substring(
+                        self.namespace, cluster_name, "head"
+                    )
+                if not worker_pod_name:
+                    worker_pod_name = kubectl_get_pod_name_by_substring(
+                        self.namespace, cluster_name, "worker"
+                    )
+
+                head_status = (
+                    kubectl_get_pod_status(self.namespace, head_pod_name)
+                    if head_pod_name
+                    else "NotFound"
+                )
+                worker_status = (
+                    kubectl_get_pod_status(self.namespace, worker_pod_name)
+                    if worker_pod_name
+                    else "NotFound"
+                )
+
+                if (
+                    head_pod_name
+                    and worker_pod_name
+                    and "Running" in head_status
+                    and "Running" in worker_status
+                ):
+                    head_ready = kubectl_get_pod_ready(self.namespace, head_pod_name)
+                    worker_ready = kubectl_get_pod_ready(
+                        self.namespace, worker_pod_name
+                    )
+                    if head_ready and worker_ready:
+                        break
+                time.sleep(10)
+
+            generate_cert.generate_tls_cert(cluster_name, self.namespace)
+            generate_cert.export_env(cluster_name, self.namespace)
+
+            local_port = "20001"
+            ray_client_port = "10001"
+            head_service_name = f"{cluster_name}-head-svc"
+
+            port_forward_cmd = [
+                "kubectl",
+                "port-forward",
+                "-n",
+                self.namespace,
+                f"svc/{head_service_name}",
+                f"{local_port}:{ray_client_port}",
+            ]
+            self.port_forward_process = subprocess.Popen(
+                port_forward_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
             )
-            return sum(results)
-
-        ref = heavy_calculation.remote(3000)
-        result = ray.get(ref)
-        assert result == 1789.4644387076714
-        ray.cancel(ref)
-        ray.shutdown()
-
-        cluster.down()
+            time.sleep(5)
+
+            client_url = f"ray://localhost:{local_port}"
+            cluster.status()
+
+            ray.init(address=client_url, logging_level="INFO")
+
+            @ray.remote(num_gpus=number_of_gpus / 2)
+            def heavy_calculation_part(num_iterations):
+                result = 0.0
+                for i in range(num_iterations):
+                    for j in range(num_iterations):
+                        for k in range(num_iterations):
+                            result += math.sin(i) * math.cos(j) * math.tan(k)
+                return result
+
+            @ray.remote(num_gpus=number_of_gpus / 2)
+            def heavy_calculation(num_iterations):
+                results = ray.get(
+                    [
+                        heavy_calculation_part.remote(num_iterations // 30)
+                        for _ in range(30)
+                    ]
+                )
+                return sum(results)
+
+            ref = heavy_calculation.remote(3000)
+
+            try:
+                result = ray.get(ref)
+                assert result == 1789.4644387076714
+            except Exception as e:
+                raise
+            finally:
+                ray.cancel(ref)
+
+            ray.shutdown()
+
+        finally:
+            if self.port_forward_process:
+                self.port_forward_process.terminate()
+                try:
+                    self.port_forward_process.wait(timeout=10)
+                except subprocess.TimeoutExpired:
+                    self.port_forward_process.kill()
+                    self.port_forward_process.wait()
+                self.port_forward_process = None
+            cluster.down()