refactor: ray and appwrapper modules

Signed-off-by: Bobbins228 <[email protected]>

Author: Bobbins228

src/codeflare_sdk/__init__.py

@@ -1,29 +1,28 @@
-from .cluster import (
-    AWManager,
+from .ray import (
     Cluster,
     ClusterConfiguration,
     RayClusterStatus,
-    AppWrapperStatus,
     CodeFlareClusterStatus,
     RayCluster,
-    AppWrapper,
     get_cluster,
     list_all_queued,
     list_all_clusters,
-    view_clusters,
+    AWManager,
+    AppWrapperStatus,
+    RayJobClient,
 )
 
+from .cluster import view_clusters
+
 from .common import (
     Authentication,
     KubeConfiguration,
     TokenAuthentication,
     KubeConfigFileAuthentication,
 )
 
-from .job import RayJobClient
-
-from .utils import generate_cert
-from .utils.demos import copy_demo_nbs
+from .common.utils import generate_cert
+from .common.utils.demos import copy_demo_nbs
 
 from importlib.metadata import version, PackageNotFoundError
 

src/codeflare_sdk/cluster/__init__.py

@@ -1,21 +1,3 @@
-from .model import (
-    RayClusterStatus,
-    AppWrapperStatus,
-    CodeFlareClusterStatus,
-    RayCluster,
-    AppWrapper,
-)
-
-from .cluster import (
-    Cluster,
-    ClusterConfiguration,
-    get_cluster,
-    list_all_queued,
-    list_all_clusters,
-)
-
 from .widgets import (
     view_clusters,
 )
-
-from .awload import AWManager

src/codeflare_sdk/cluster/widgets.py

@@ -26,16 +26,18 @@
 import ipywidgets as widgets
 from IPython.display import display, HTML, Javascript
 import pandas as pd
-from .config import ClusterConfiguration
-from .model import RayClusterStatus
+from ..ray.cluster.config import ClusterConfiguration
+from ..ray.cluster.status import RayClusterStatus
 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:
+def cluster_up_down_buttons(
+    cluster: "codeflare_sdk.ray.cluster.cluster.Cluster",
+) -> widgets.Button:
     """
     The cluster_up_down_buttons function returns two button widgets for a create and delete button.
     The function uses the appwrapper bool to distinguish between resource type for the tool tip.
@@ -115,7 +117,7 @@ def view_clusters(namespace: str = None):
         )
         return  # Exit function if not in Jupyter Notebook
 
-    from .cluster import get_current_namespace
+    from ..ray.cluster.cluster import get_current_namespace
 
     if not namespace:
         namespace = get_current_namespace()
@@ -278,7 +280,7 @@ def _on_ray_dashboard_button_click(
     """
     _on_ray_dashboard_button_click handles the event when the Open Ray Dashboard button is clicked, opening the Ray Dashboard in a new tab
     """
-    from codeflare_sdk.cluster import Cluster
+    from codeflare_sdk.ray.cluster import Cluster
 
     cluster_name = classification_widget.value
     namespace = ray_clusters_df[ray_clusters_df["Name"] == classification_widget.value][
@@ -309,7 +311,7 @@ def _on_list_jobs_button_click(
     """
     _on_list_jobs_button_click handles the event when the View Jobs button is clicked, opening the Ray Jobs Dashboard in a new tab
     """
-    from codeflare_sdk.cluster import Cluster
+    from codeflare_sdk.ray.cluster import Cluster
 
     cluster_name = classification_widget.value
     namespace = ray_clusters_df[ray_clusters_df["Name"] == classification_widget.value][
@@ -342,7 +344,7 @@ def _delete_cluster(
     _delete_cluster function deletes the cluster with the given name and namespace.
     It optionally waits for the cluster to be deleted.
     """
-    from .cluster import _check_aw_exists
+    from ..ray.cluster.cluster import _check_aw_exists
 
     try:
         config_check()
@@ -400,7 +402,7 @@ def _fetch_cluster_data(namespace):
     """
     _fetch_cluster_data function fetches all clusters and their spec in a given namespace and returns a DataFrame.
     """
-    from .cluster import list_all_clusters
+    from ..ray.cluster.cluster import list_all_clusters
 
     rayclusters = list_all_clusters(namespace, False)
     if not rayclusters:

src/codeflare_sdk/utils/__init__.py renamed to src/codeflare_sdk/common/utils/__init__.py

@@ -19,12 +19,12 @@
 from cryptography import x509
 from cryptography.x509.oid import NameOID
 import datetime
-from ..common.kubernetes_cluster.auth import (
+from ..kubernetes_cluster.auth import (
     config_check,
     get_api_client,
 )
 from kubernetes import client
-from ..common import _kube_api_error_handling
+from .. import _kube_api_error_handling
 
 
 def generate_ca_cert(days: int = 30):

src/codeflare_sdk/utils/demos.py renamed to src/codeflare_sdk/common/utils/demos.py

@@ -0,0 +1,16 @@
+from .appwrapper import AppWrapper, AppWrapperStatus, AWManager
+
+from .client import (
+    RayJobClient,
+)
+
+from .cluster import (
+    Cluster,
+    ClusterConfiguration,
+    get_cluster,
+    list_all_queued,
+    list_all_clusters,
+    RayClusterStatus,
+    CodeFlareClusterStatus,
+    RayCluster,
+)

src/codeflare_sdk/utils/generate_cert.py renamed to src/codeflare_sdk/common/utils/generate_cert.py

@@ -0,0 +1,6 @@
+from .awload import AWManager
+
+from .status import (
+    AppWrapperStatus,
+    AppWrapper,
+)

src/codeflare_sdk/ray/__init__.py

@@ -23,8 +23,8 @@
 import yaml
 
 from kubernetes import client
-from ..common import _kube_api_error_handling
-from ..common.kubernetes_cluster.auth import (
+from ...common import _kube_api_error_handling
+from ...common.kubernetes_cluster.auth import (
     config_check,
     get_api_client,
 )

src/codeflare_sdk/ray/appwrapper/__init__.py

@@ -0,0 +1,46 @@
+# Copyright 2024 IBM, Red Hat
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+The status sub-module defines Enums containing information for
+AppWrapper states, as well as dataclasses to store information for AppWrappers.
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+
+
+class AppWrapperStatus(Enum):
+    """
+    Defines the possible reportable phases of an AppWrapper.
+    """
+
+    SUSPENDED = "suspended"
+    RESUMING = "resuming"
+    RUNNING = "running"
+    RESETTING = "resetting"
+    SUSPENDING = "suspending"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    TERMINATING = "terminating"
+
+
+@dataclass
+class AppWrapper:
+    """
+    For storing information about an AppWrapper.
+    """
+
+    name: str
+    status: AppWrapperStatus