[Part-1] Refactor k8s related code into a separate module

varshaprasad96 · varshaprasad96 · commit f493678d4685 · 2024-10-01T09:34:20.000-07:00
This commit is first of the chain in refactoring the cluster/
package.
diff --git a/.gitignore b/.gitignore
@@ -12,3 +12,4 @@ node_modules
 .DS_Store
 ui-tests/playwright-report
 ui-tests/test-results
+/src/codeflare_sdk.egg-info/
diff --git a/poetry.lock b/poetry.lock
diff --git a/src/codeflare_sdk.egg-info/SOURCES.txt b/src/codeflare_sdk.egg-info/SOURCES.txt
@@ -7,17 +7,19 @@ src/codeflare_sdk.egg-info/SOURCES.txt
 src/codeflare_sdk.egg-info/dependency_links.txt
 src/codeflare_sdk.egg-info/top_level.txt
 src/codeflare_sdk/cluster/__init__.py
-src/codeflare_sdk/cluster/auth.py
 src/codeflare_sdk/cluster/awload.py
 src/codeflare_sdk/cluster/cluster.py
 src/codeflare_sdk/cluster/config.py
 src/codeflare_sdk/cluster/model.py
 src/codeflare_sdk/cluster/widgets.py
+src/codeflare_sdk/common/__init__.py
+src/codeflare_sdk/common/kubernetes_cluster/__init__.py
+src/codeflare_sdk/common/kubernetes_cluster/auth.py
+src/codeflare_sdk/common/kubernetes_cluster/kube_api_helpers.py
 src/codeflare_sdk/job/__init__.py
 src/codeflare_sdk/job/ray_jobs.py
 src/codeflare_sdk/utils/__init__.py
 src/codeflare_sdk/utils/demos.py
 src/codeflare_sdk/utils/generate_cert.py
 src/codeflare_sdk/utils/generate_yaml.py
-src/codeflare_sdk/utils/kube_api_helpers.py
 src/codeflare_sdk/utils/pretty_print.py
diff --git a/src/codeflare_sdk/__init__.py b/src/codeflare_sdk/__init__.py
@@ -1,8 +1,4 @@
 from .cluster import (
-    Authentication,
-    KubeConfiguration,
-    TokenAuthentication,
-    KubeConfigFileAuthentication,
     AWManager,
     Cluster,
     ClusterConfiguration,
@@ -17,6 +13,13 @@
     view_clusters,
 )
 
+from .common import (
+    Authentication,
+    KubeConfiguration,
+    TokenAuthentication,
+    KubeConfigFileAuthentication,
+)
+
 from .job import RayJobClient
 
 from .utils import generate_cert
diff --git a/src/codeflare_sdk/cluster/__init__.py b/src/codeflare_sdk/cluster/__init__.py
@@ -1,10 +1,3 @@
-from .auth import (
-    Authentication,
-    KubeConfiguration,
-    TokenAuthentication,
-    KubeConfigFileAuthentication,
-)
-
 from .model import (
     RayClusterStatus,
     AppWrapperStatus,
diff --git a/src/codeflare_sdk/cluster/awload.py b/src/codeflare_sdk/cluster/awload.py
@@ -22,9 +22,12 @@
 import os
 import yaml
 
-from kubernetes import client, config
-from ..utils.kube_api_helpers import _kube_api_error_handling
-from .auth import config_check, get_api_client
+from kubernetes import client
+from ..common import _kube_api_error_handling
+from ..common.kubernetes_cluster.auth import (
+    config_check,
+    get_api_client,
+)
 
 
 class AWManager:
diff --git a/src/codeflare_sdk/cluster/cluster.py b/src/codeflare_sdk/cluster/cluster.py
@@ -18,21 +18,21 @@
 cluster setup queue, a list of all existing clusters, and the user's working namespace.
 """
 
-import re
-import subprocess
 from time import sleep
 from typing import List, Optional, Tuple, Dict
 
-from kubernetes import config
 from ray.job_submission import JobSubmissionClient
 
-from .auth import config_check, get_api_client
+from ..common.kubernetes_cluster.auth import (
+    config_check,
+    get_api_client,
+)
 from ..utils import pretty_print
 from ..utils.generate_yaml import (
     generate_appwrapper,
     head_worker_gpu_count_from_cluster,
 )
-from ..utils.kube_api_helpers import _kube_api_error_handling
+from ..common import _kube_api_error_handling
 from ..utils.generate_yaml import is_openshift_cluster
 
 from .config import ClusterConfiguration
@@ -47,8 +47,7 @@
     cluster_up_down_buttons,
     is_notebook,
 )
-from kubernetes import client, config
-from kubernetes.utils import parse_quantity
+from kubernetes import client
 import yaml
 import os
 import requests
diff --git a/src/codeflare_sdk/cluster/widgets.py b/src/codeflare_sdk/cluster/widgets.py
@@ -28,8 +28,11 @@
 import pandas as pd
 from .config import ClusterConfiguration
 from .model import RayClusterStatus
-from ..utils.kube_api_helpers import _kube_api_error_handling
-from .auth import config_check, get_api_client
+from ..common import _kube_api_error_handling
+from ..common.kubernetes_cluster.auth import (
+    config_check,
+    get_api_client,
+)
 
 
 def cluster_up_down_buttons(cluster: "codeflare_sdk.cluster.Cluster") -> widgets.Button:
diff --git a/src/codeflare_sdk/common/__init__.py b/src/codeflare_sdk/common/__init__.py
@@ -0,0 +1,8 @@
+# Importing everything from the kubernetes_cluster module
+from .kubernetes_cluster import (
+    Authentication,
+    KubeConfiguration,
+    TokenAuthentication,
+    KubeConfigFileAuthentication,
+    _kube_api_error_handling,
+)
diff --git a/src/codeflare_sdk/common/kubernetes_cluster/__init__.py b/src/codeflare_sdk/common/kubernetes_cluster/__init__.py
@@ -0,0 +1,10 @@
+from .auth import (
+    Authentication,
+    KubeConfiguration,
+    TokenAuthentication,
+    KubeConfigFileAuthentication,
+    config_check,
+    get_api_client,
+)
+
+from .kube_api_helpers import _kube_api_error_handling
diff --git a/src/codeflare_sdk/common/kubernetes_cluster/auth.py b/src/codeflare_sdk/common/kubernetes_cluster/auth.py
@@ -23,7 +23,7 @@
 from kubernetes import client, config
 import os
 import urllib3
-from ..utils.kube_api_helpers import _kube_api_error_handling
+from .kube_api_helpers import _kube_api_error_handling
 
 from typing import Optional
 
diff --git a/src/codeflare_sdk/common/kubernetes_cluster/kube_api_helpers.py b/src/codeflare_sdk/common/kubernetes_cluster/kube_api_helpers.py
diff --git a/src/codeflare_sdk/utils/generate_cert.py b/src/codeflare_sdk/utils/generate_cert.py
@@ -19,9 +19,12 @@
 from cryptography import x509
 from cryptography.x509.oid import NameOID
 import datetime
-from ..cluster.auth import config_check, get_api_client
-from kubernetes import client, config
-from .kube_api_helpers import _kube_api_error_handling
+from ..common.kubernetes_cluster.auth import (
+    config_check,
+    get_api_client,
+)
+from kubernetes import client
+from ..common import _kube_api_error_handling
 
 
 def generate_ca_cert(days: int = 30):
diff --git a/src/codeflare_sdk/utils/generate_yaml.py b/src/codeflare_sdk/utils/generate_yaml.py
@@ -21,16 +21,14 @@
 from typing import Optional
 import typing
 import yaml
-import sys
 import os
-import argparse
 import uuid
-from kubernetes import client, config
-from .kube_api_helpers import _kube_api_error_handling
-from ..cluster.auth import get_api_client, config_check
-from os import urandom
-from base64 import b64encode
-from urllib3.util import parse_url
+from kubernetes import client
+from ..common import _kube_api_error_handling
+from ..common.kubernetes_cluster.auth import (
+    get_api_client,
+    config_check,
+)
 from kubernetes.client.exceptions import ApiException
 import codeflare_sdk
 
diff --git a/tests/e2e/support.py b/tests/e2e/support.py
@@ -4,7 +4,9 @@
 import subprocess
 from kubernetes import client, config
 import kubernetes.client
-from codeflare_sdk.utils.kube_api_helpers import _kube_api_error_handling
+from codeflare_sdk.common.kubernetes_cluster.kube_api_helpers import (
+    _kube_api_error_handling,
+)
 
 
 def get_ray_image():
diff --git a/tests/unit_test.py b/tests/unit_test.py
@@ -39,7 +39,7 @@
     _app_wrapper_status,
     _ray_cluster_status,
 )
-from codeflare_sdk.cluster.auth import (
+from codeflare_sdk.common.kubernetes_cluster import (
     TokenAuthentication,
     Authentication,
     KubeConfigFileAuthentication,
@@ -71,7 +71,7 @@
     get_package_and_version,
 )
 
-import codeflare_sdk.utils.kube_api_helpers
+import codeflare_sdk.common.kubernetes_cluster.kube_api_helpers
 from codeflare_sdk.utils.generate_yaml import (
     gen_names,
     is_openshift_cluster,
@@ -198,8 +198,8 @@ def test_token_auth_login_tls(mocker):
 def test_config_check_no_config_file(mocker):
     mocker.patch("os.path.expanduser", return_value="/mock/home/directory")
     mocker.patch("os.path.isfile", return_value=False)
-    mocker.patch("codeflare_sdk.cluster.auth.config_path", None)
-    mocker.patch("codeflare_sdk.cluster.auth.api_client", None)
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_path", None)
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.api_client", None)
 
     with pytest.raises(PermissionError) as e:
         config_check()
@@ -210,8 +210,8 @@ def test_config_check_with_incluster_config(mocker):
     mocker.patch("os.path.isfile", return_value=False)
     mocker.patch.dict(os.environ, {"KUBERNETES_PORT": "number"})
     mocker.patch("kubernetes.config.load_incluster_config", side_effect=None)
-    mocker.patch("codeflare_sdk.cluster.auth.config_path", None)
-    mocker.patch("codeflare_sdk.cluster.auth.api_client", None)
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_path", None)
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.api_client", None)
 
     result = config_check()
     assert result == None
@@ -221,16 +221,18 @@ def test_config_check_with_existing_config_file(mocker):
     mocker.patch("os.path.expanduser", return_value="/mock/home/directory")
     mocker.patch("os.path.isfile", return_value=True)
     mocker.patch("kubernetes.config.load_kube_config", side_effect=None)
-    mocker.patch("codeflare_sdk.cluster.auth.config_path", None)
-    mocker.patch("codeflare_sdk.cluster.auth.api_client", None)
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_path", None)
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.api_client", None)
 
     result = config_check()
     assert result == None
 
 
 def test_config_check_with_config_path_and_no_api_client(mocker):
-    mocker.patch("codeflare_sdk.cluster.auth.config_path", "/mock/config/path")
-    mocker.patch("codeflare_sdk.cluster.auth.api_client", None)
+    mocker.patch(
+        "codeflare_sdk.common.kubernetes_cluster.auth.config_path", "/mock/config/path"
+    )
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.api_client", None)
     result = config_check()
     assert result == "/mock/config/path"
 
@@ -2170,7 +2172,8 @@ def test_map_to_ray_cluster(mocker):
 
     mock_api_client = mocker.MagicMock(spec=client.ApiClient)
     mocker.patch(
-        "codeflare_sdk.cluster.auth.get_api_client", return_value=mock_api_client
+        "codeflare_sdk.common.kubernetes_cluster.auth.get_api_client",
+        return_value=mock_api_client,
     )
 
     mock_routes = {
diff --git a/tests/upgrade/raycluster_sdk_upgrade_sleep_test.py b/tests/upgrade/raycluster_sdk_upgrade_sleep_test.py
@@ -12,7 +12,7 @@
 from tests.e2e.support import *
 
 
-from codeflare_sdk.utils.kube_api_helpers import _kube_api_error_handling
+from codeflare_sdk.common import _kube_api_error_handling
 
 namespace = "test-ns-rayupgrade-sleep"
 # Global variables for kueue resources
diff --git a/tests/upgrade/raycluster_sdk_upgrade_test.py b/tests/upgrade/raycluster_sdk_upgrade_test.py
@@ -7,7 +7,7 @@
 from tests.e2e.support import *
 from codeflare_sdk.cluster.cluster import get_cluster
 
-from codeflare_sdk.utils.kube_api_helpers import _kube_api_error_handling
+from codeflare_sdk.common import _kube_api_error_handling
 
 namespace = "test-ns-rayupgrade"
 # Global variables for kueue resources
