Refactoring: Organize functions from core.py into classes

Author: IrvingMg

src/xpk/commands/batch.py

@@ -14,19 +14,24 @@
 limitations under the License.
 """
 
+import re
 from argparse import Namespace
 
-from ..core.cluster import setup_k8s_env, create_k8s_service_account
+from ..core.cluster import ClusterManager
 from ..core.commands import run_command_for_value
-from ..core.config import GCS_FUSE_ANNOTATION_KEY, GCS_FUSE_ANNOTATION_VALUE, XPK_SA, DEFAULT_NAMESPACE
-from ..core.gcloud_context import add_zone_and_project
+from ..core.config import (
+    DEFAULT_NAMESPACE,
+    GCS_FUSE_ANNOTATION_KEY,
+    GCS_FUSE_ANNOTATION_VALUE,
+    XPK_SA,
+)
+from ..core.gcloud_context import GCloudContextManager
+from ..core.kjob import AppProfileDefaults, Kueue_TAS_annotation, prepare_kjob
 from ..core.kueue import LOCAL_QUEUE_NAME
 from ..core.storage import get_auto_mount_gcsfuse_storages
 from ..utils.console import xpk_exit, xpk_print
 from .common import set_cluster_command
-from ..core.kjob import AppProfileDefaults, prepare_kjob, Kueue_TAS_annotation
 from .kind import set_local_cluster_command
-import re
 
 
 def batch(args: Namespace) -> None:
@@ -38,7 +43,7 @@ def batch(args: Namespace) -> None:
     None
   """
   if not args.kind_cluster:
-    add_zone_and_project(args)
+    GCloudContextManager.add_zone_and_project(args)
     set_cluster_command_code = set_cluster_command(args)
   else:
     set_cluster_command_code = set_local_cluster_command(args)
@@ -54,8 +59,9 @@ def batch(args: Namespace) -> None:
 
 
 def submit_job(args: Namespace) -> None:
-  k8s_api_client = setup_k8s_env(args)
-  create_k8s_service_account(XPK_SA, DEFAULT_NAMESPACE)
+  cluster_manager = ClusterManager(args, None)
+  k8s_api_client = cluster_manager.setup_k8s_env()
+  cluster_manager.create_k8s_service_account(XPK_SA, DEFAULT_NAMESPACE)
   gcs_fuse_storages = get_auto_mount_gcsfuse_storages(k8s_api_client)
 
   cmd = (

src/xpk/commands/cluster.py

@@ -16,25 +16,12 @@
 
 from tabulate import tabulate
 
-from ..core.capacity import H100_DEVICE_TYPE
-from ..core.cluster import (
-    get_all_clusters_programmatic,
-    get_cluster_credentials,
-    install_nccl_on_cluster,
-    set_jobset_on_cluster,
-    setup_k8s_env,
-    update_cluster_with_gcsfuse_driver_if_necessary,
-    update_cluster_with_workload_identity_if_necessary,
-)
+from ..core.capacity import CapacityManager, DeviceType
+from ..core.cluster import ClusterManager
 from ..core.cluster_private import authorize_private_cluster_access_if_necessary
 from ..core.commands import run_command_for_value, run_command_with_updates
 from ..core.config import VERTEX_TENSORBOARD_FEATURE_FLAG
-from ..core.gcloud_context import (
-    add_zone_and_project,
-    get_gke_control_plane_version,
-    get_gke_server_config,
-    zone_to_region,
-)
+from ..core.gcloud_context import GCloudContextManager, GKEVersionManager
 from ..core.kjob import apply_kjob_crds, prepare_kjob, verify_kjob_installed
 from ..core.kueue import (
     cluster_preheat_yml,
@@ -43,28 +30,23 @@
     wait_for_kueue_available,
 )
 from ..core.nap import enable_autoprovisioning_on_cluster
-from ..core.network import (
-    create_cluster_network_config,
-    delete_cluster_subnets,
-    set_up_cluster_network_for_gpu,
-)
-from ..core.nodepool import get_gke_node_pool_version, run_gke_node_pool_create_command
+from ..core.network import ClusterNetworkManager
+from ..core.nodepool import NodePoolManager
 from ..core.ray import install_ray_cluster
-from ..core.resources import create_cluster_configmaps
+from ..core.resources import ResourceManager
 from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
     AcceleratorType,
     AcceleratorTypeToAcceleratorCharacteristics,
     SystemCharacteristics,
     get_system_characteristics,
 )
-from ..core.vertex import create_vertex_tensorboard
+from ..core.vertex import VertexAI
 from ..core.workload import get_workload_list
 from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from . import cluster_gcluster
 from .common import set_cluster_command
-from ..core.cluster import update_cluster_with_gcpfilestore_driver_if_necessary
 
 
 def cluster_create(args) -> None:
@@ -77,13 +59,25 @@ def cluster_create(args) -> None:
     0 if successful and 1 otherwise.
   """
   system, return_code = get_system_characteristics(args)
-
   if return_code > 0:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
   xpk_print(f'Starting cluster create for cluster {args.cluster}:', flush=True)
-  add_zone_and_project(args)
+  GCloudContextManager.add_zone_and_project(args)
+  cluster_manager = ClusterManager(args, system)
+  capacity_manager = CapacityManager(args)
+  resource_manager = ResourceManager(args, capacity_manager, system)
+  nodepools_manager = NodePoolManager(
+      args, system, resource_manager, capacity_manager
+  )
+  network_manager = ClusterNetworkManager(args)
+  vertex_ai = VertexAI(args, resource_manager)
+  try:
+    version_manager = GKEVersionManager(args)
+  except RuntimeError as e:
+    xpk_print(e)
+    xpk_exit(1)
 
   if system.device_type in cluster_gcluster.supported_device_types:
     xpk_print(
@@ -93,12 +87,8 @@ def cluster_create(args) -> None:
     cluster_gcluster.cluster_create(args)
     xpk_exit(0)
 
-  return_code, gke_server_config = get_gke_server_config(args)
-  if return_code != 0:
-    xpk_exit(return_code)
-
-  return_code, gke_control_plane_version = get_gke_control_plane_version(
-      args, gke_server_config
+  return_code, gke_control_plane_version = (
+      version_manager.get_gke_control_plane_version()
   )
   if return_code != 0:
     xpk_exit(return_code)
@@ -123,60 +113,61 @@ def cluster_create(args) -> None:
       or args.enable_gcpfilestore_csi_driver
   ):
     update_cluster_command_code = (
-        update_cluster_with_workload_identity_if_necessary(args)
+        cluster_manager.update_cluster_with_workload_identity_if_necessary()
     )
     if update_cluster_command_code != 0:
       xpk_exit(update_cluster_command_code)
 
   # Enable GCSFuse CSI Driver if not enabled already.
   if args.enable_gcsfuse_csi_driver:
     update_cluster_command_code = (
-        update_cluster_with_gcsfuse_driver_if_necessary(args)
+        cluster_manager.update_cluster_with_gcsfuse_driver_if_necessary()
     )
     if update_cluster_command_code != 0:
       xpk_exit(update_cluster_command_code)
 
   if args.enable_gcpfilestore_csi_driver:
     update_cluster_command_code = (
-        update_cluster_with_gcpfilestore_driver_if_necessary(args)
+        cluster_manager.update_cluster_with_gcpfilestore_driver_if_necessary()
     )
     if update_cluster_command_code != 0:
       xpk_exit(update_cluster_command_code)
 
   # Update Pathways clusters with CloudDNS if not enabled already.
 
-  get_cluster_credentials(args)
+  cluster_manager.get_cluster_credentials()
 
   # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
   tensorboard_config = {}
   if VERTEX_TENSORBOARD_FEATURE_FLAG and args.create_vertex_tensorboard:
-    tensorboard_config = create_vertex_tensorboard(args)
+    tensorboard_config = vertex_ai.create_vertex_tensorboard()
     # exit if failed to create Tensorboard in Vertex AI
     if not tensorboard_config:
       xpk_exit(1)
 
   if system.accelerator_type == AcceleratorType['GPU']:
     xpk_print('Setting up Network for cluster')
-    set_up_cluster_network_code = set_up_cluster_network_for_gpu(args, system)
+    set_up_cluster_network_code = network_manager.set_up_network_for_gpu(system)
     if set_up_cluster_network_code != 0:
       xpk_exit(set_up_cluster_network_code)
 
-  if system.device_type == H100_DEVICE_TYPE:
+  if system.device_type == DeviceType.H100.value:
     xpk_print('Creating Network Config for cluster')
-    create_cluster_network_config_code = create_cluster_network_config(args)
+    create_cluster_network_config_code = network_manager.create_network_config()
     if create_cluster_network_config_code != 0:
       xpk_exit(create_cluster_network_config_code)
 
   # Check the control plane version of the cluster and determine the node pool
   # version to use.
-  return_code, gke_node_pool_version = get_gke_node_pool_version(
-      args, gke_server_config
+
+  return_code, gke_node_pool_version = (
+      nodepools_manager.get_gke_node_pool_version(version_manager)
   )
   if return_code != 0:
     xpk_exit(return_code)
 
-  run_gke_node_pool_create_command_code = run_gke_node_pool_create_command(
-      args, system, gke_node_pool_version
+  run_gke_node_pool_create_command_code = (
+      nodepools_manager.run_gke_node_pool_create_command(gke_node_pool_version)
   )
   if run_gke_node_pool_create_command_code != 0:
     xpk_exit(run_gke_node_pool_create_command_code)
@@ -193,8 +184,8 @@ def cluster_create(args) -> None:
       xpk_exit(return_code)
 
   xpk_print('Creating ConfigMap for cluster')
-  create_cluster_configmaps_code = create_cluster_configmaps(
-      args, system, tensorboard_config, autoprovisioning_config
+  create_cluster_configmaps_code = resource_manager.create_cluster_configmaps(
+      tensorboard_config, autoprovisioning_config
   )
   if create_cluster_configmaps_code != 0:
     xpk_exit(create_cluster_configmaps_code)
@@ -203,7 +194,7 @@ def cluster_create(args) -> None:
       'Enabling the jobset API on our cluster, to be deprecated when Jobset is'
       ' globally available'
   )
-  set_jobset_on_cluster_code = set_jobset_on_cluster(args)
+  set_jobset_on_cluster_code = cluster_manager.set_jobset_on_cluster()
   if set_jobset_on_cluster_code != 0:
     xpk_exit(set_jobset_on_cluster_code)
 
@@ -226,7 +217,7 @@ def cluster_create(args) -> None:
   if err_code > 0:
     xpk_exit(err_code)
 
-  k8s_client = setup_k8s_env(args)
+  k8s_client = cluster_manager.setup_k8s_env()
   install_storage_crd(k8s_client)
 
   xpk_print('Wait for Kueue to be fully available')
@@ -243,7 +234,7 @@ def cluster_create(args) -> None:
 
   if system.accelerator_type == AcceleratorType['GPU']:
     xpk_print('Installing NCCL Plugin for cluster')
-    install_nccl_code = install_nccl_on_cluster(args, system)
+    install_nccl_code = cluster_manager.install_nccl_on_cluster()
     if install_nccl_code != 0:
       xpk_exit(install_nccl_code)
 
@@ -257,7 +248,7 @@ def cluster_create(args) -> None:
   xpk_print(
       'See your GKE Cluster here:'
       # pylint: disable=line-too-long
-      f' https://console.cloud.google.com/kubernetes/clusters/details/{zone_to_region(args.zone)}/{args.cluster}/details?project={args.project}'
+      f' https://console.cloud.google.com/kubernetes/clusters/details/{GCloudContextManager.zone_to_region(args.zone)}/{args.cluster}/details?project={args.project}'
   )
   xpk_exit(0)
 
@@ -272,7 +263,7 @@ def cluster_delete(args) -> None:
     0 if successful and 1 otherwise.
   """
   xpk_print(f'Starting cluster delete for cluster: {args.cluster}', flush=True)
-  add_zone_and_project(args)
+  GCloudContextManager.add_zone_and_project(args)
 
   if cluster_gcluster.created_by_gcluster(args):
     xpk_print(f'Deleting {args.cluster} cluster using Cluster Toolkit...')
@@ -303,14 +294,13 @@ def cluster_cacheimage(args) -> None:
   xpk_print(
       f'Starting cluster cacheimage for cluster: {args.cluster}', flush=True
   )
-  add_zone_and_project(args)
-
-  get_cluster_credentials(args)
+  GCloudContextManager.add_zone_and_project(args)
   system, return_code = get_system_characteristics(args)
-
   if return_code > 0:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
+  cluster_manager = ClusterManager(args, system)
+  cluster_manager.get_cluster_credentials()
 
   node_selector_key = AcceleratorTypeToAcceleratorCharacteristics[
       system.accelerator_type
@@ -352,9 +342,13 @@ def cluster_describe(args) -> None:
     0 if successful and 1 otherwise.
   """
   xpk_print(f'Starting nodepool list for cluster: {args.cluster}', flush=True)
-  add_zone_and_project(args)
-
-  get_cluster_credentials(args)
+  GCloudContextManager.add_zone_and_project(args)
+  system, return_code = get_system_characteristics(args)
+  if return_code > 0:
+    xpk_print('Fetching system characteristics failed!')
+    xpk_exit(return_code)
+  cluster_manager = ClusterManager(args, system)
+  cluster_manager.get_cluster_credentials()
 
   return_code, data_table = nodepools_build_table(args)
   if return_code != 0:
@@ -583,7 +577,7 @@ def cluster_list(args) -> None:
   Returns:
     0 if successful and 1 otherwise.
   """
-  add_zone_and_project(args)
+  GCloudContextManager.add_zone_and_project(args)
   xpk_print(f'For project {args.project} and zone {args.zone}:', flush=True)
   if run_gke_clusters_list_command(args):
     xpk_exit(1)
@@ -631,7 +625,8 @@ def create_cluster_if_necessary(
   Returns:
     0 if successful and 1 otherwise.
   """
-  all_clusters, return_code = get_all_clusters_programmatic(args)
+  cluster_manager = ClusterManager(args, system)
+  all_clusters, return_code = cluster_manager.get_all_clusters_programmatic()
   if return_code > 0:
     xpk_print('Listing all clusters failed!')
     return 1
@@ -653,6 +648,8 @@ def run_gke_cluster_delete_command(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
+  network_manager = ClusterNetworkManager(args)
+
   if not args.force:
     xpk_print('Get the name of the workloads in the cluster.')
     args.filter_by_status = 'EVERYTHING'
@@ -675,15 +672,15 @@ def run_gke_cluster_delete_command(args) -> int:
   command = (
       'gcloud beta container clusters delete'
       f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)} --quiet'
+      f' --region={GCloudContextManager.zone_to_region(args.zone)} --quiet'
   )
 
   return_code = run_command_with_updates(command, 'Cluster Delete', args)
   if return_code != 0:
     xpk_print(f'Cluster delete request returned ERROR {return_code}')
     return 1
 
-  return_code = delete_cluster_subnets(args)
+  return_code = network_manager.delete_subnets()
   if return_code != 0:
     return return_code
 
@@ -701,7 +698,7 @@ def run_gke_clusters_list_command(args) -> int:
   """
   command = (
       'gcloud container clusters list'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} --region={GCloudContextManager.zone_to_region(args.zone)}'
   )
   return_code = run_command_with_updates(command, 'Cluster List', args)
   if return_code != 0:
@@ -748,7 +745,7 @@ def run_gke_cluster_create_command(
   command = (
       'gcloud beta container clusters create'
       f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
+      f' --region={GCloudContextManager.zone_to_region(args.zone)}'
       f' --node-locations={args.zone}'
       f' --cluster-version={gke_control_plane_version}'
       f' --machine-type={machine_type}'