xpk 0.9.0__py3-none-any.whl → 0.10.1__py3-none-any.whl

xpk/core/blueprint/blueprint_generator.py

@@ -22,7 +22,9 @@ from ruamel import yaml
 
 from ...utils.console import xpk_exit, xpk_print
 from ...utils.file import ensure_directory_exists
+
 from ..capacity import (
+    H100_DEVICE_TYPE,
     B200_DEVICE_TYPE,
     H100_MEGA_DEVICE_TYPE,
     H200_DEVICE_TYPE,
@@ -30,10 +32,11 @@ from ..capacity import (
 )
 from ..system_characteristics import get_system_characteristics_by_device_type
 from .blueprint_definitions import Blueprint, DeploymentGroup, DeploymentModule
-
+from ..kueue import KUEUE_VERSION
 
 yaml = yaml.YAML()
 
+a3high_device_type = H100_DEVICE_TYPE
 a3mega_device_type = H100_MEGA_DEVICE_TYPE
 a3ultra_device_type = H200_DEVICE_TYPE
 a4_device_type = B200_DEVICE_TYPE
@@ -49,7 +52,7 @@ blueprint_dependencies_dir = {
 }
 
 cluster_toolkit_url = "github.com/GoogleCloudPlatform/cluster-toolkit"
-cluster_toolkit_version = "v1.48.0"
+cluster_toolkit_version = "v1.57.1"
 
 
 class BlueprintGeneratorOutput:
@@ -92,6 +95,8 @@ class BlueprintGenerator:
       group_placement_max_distance: int = 2,
       subnetwork_cidr_suffix: int = 24,
       reservation: str | None = None,
+      reservation_placement_policy: dict[str, str] | None = None,
+      reservation_maintenance_interval: str = "PERIODIC",
       gcs_bucket: Optional[str | None] = None,
       capacity_type: CapacityType = CapacityType.ON_DEMAND,
       system_node_pool_min_node_count: int = 2,
@@ -142,7 +147,6 @@ class BlueprintGenerator:
         source="modules/scheduler/gke-cluster",
         use=[primary_vpc_name, gpu_subnets_name],
         settings={
-            "release_channel": "RAPID",
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "enable_private_endpoint": False,
@@ -176,34 +180,42 @@ class BlueprintGenerator:
             "group_placement_max_distance": group_placement_max_distance,
         },
     )
+    nodepool_used_deps = ["gke_cluster", gpu_subnets_name]
 
     a3_megagpu_pool_0 = DeploymentModule(
         id="a3_megagpu_pool_0",
         source="modules/compute/gke-node-pool",
-        use=["gke_cluster", gpu_subnets_name],
+        use=nodepool_used_deps,
         settings={
             "name": f"{cluster_name}-a3-megagpu-pool-0",
             "machine_type": system.gce_machine_type,
-            "static_node_count": num_nodes,
             "zones": [zone],
-            "host_maintenance_interval": (
-                None
-                if capacity_type == CapacityType.RESERVATION
-                else "PERIODIC"
-            ),
+            "host_maintenance_interval": reservation_maintenance_interval,
             "reservation_affinity": self._getblock_reservation_affinity(
                 reservation
             ),
             "run_workload_script": False,
             "spot": capacity_type == CapacityType.SPOT,
             "max_pods_per_node": 32,
-            "auto_upgrade": True,
+            "guest_accelerator": [{
+                "type": "nvidia-h100-mega-80gb",
+                "count": 8,
+                "gpu_driver_installation_config": {
+                    "gpu_driver_version": "LATEST"
+                },
+            }],
+            "auto_upgrade": (
+                True if capacity_type != CapacityType.FLEX_START else False
+            ),
         },
         outputs=["instructions"],
     )
+    if capacity_type == CapacityType.FLEX_START:
+      a3_megagpu_pool_0.settings.update(self.get_dws_flex_start())
+    else:
+      a3_megagpu_pool_0.settings.update({"static_node_count": num_nodes})
 
     set_placement_policy = capacity_type != CapacityType.SPOT
-    tas_name = "topologyName: 'gke-default'" if set_placement_policy else ""
     num_chips = num_nodes * system.chips_per_vm
     workload = DeploymentModule(
         id="workload_component_install",
@@ -212,11 +224,16 @@ class BlueprintGenerator:
         settings={
             "kueue": {
                 "install": True,
-                "version": "v0.10.0",  # TAS feature-gates is enabled in CT
+                "version": KUEUE_VERSION,  # TAS feature-gates is enabled in CT
                 "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
                 "config_template_vars": {
                     "num_chips": num_chips,
-                    "tas_name": tas_name,
+                    "reservation": (
+                        1 if capacity_type == CapacityType.RESERVATION else 0
+                    ),
+                    "flex_start": (
+                        1 if capacity_type == CapacityType.FLEX_START else 0
+                    ),
                 },
             },
             "jobset": {"install": True, "version": "v0.7.2"},
@@ -247,6 +264,13 @@ class BlueprintGenerator:
             }]
         },
     )
+
+    print(reservation_placement_policy)
+    if reservation_placement_policy is not None:
+      a3_megagpu_pool_0.settings["placement_policy"] = (
+          reservation_placement_policy
+      )
+
     primary_group = DeploymentGroup(
         group="primary",
         modules=[
@@ -258,11 +282,9 @@ class BlueprintGenerator:
             workload_configmap,
         ],
     )
-
-    if set_placement_policy:
+    if set_placement_policy and reservation_placement_policy is None:
       a3_megagpu_pool_0.use.append(group_placement_0.id)
       primary_group.modules.append(group_placement_0)
-
     a3_mega_blueprint = Blueprint(
         terraform_backend_defaults=self._getblock_terraform_backend(
             gcs_bucket, cluster_name, prefix
@@ -478,14 +500,22 @@ class BlueprintGenerator:
         source="modules/scheduler/gke-cluster",
         use=[net_0_id],
         settings={
-            "release_channel": "RAPID",
-            "version_prefix": "1.31.",
-            "maintenance_exclusions": [{
-                "name": "no-minor-or-node-upgrades-indefinite",
-                "start_time": "2024-12-01T00:00:00Z",
-                "end_time": "2025-12-22T00:00:00Z",
-                "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
-            }],
+            "release_channel": (
+                "UNSPECIFIED"
+                if capacity_type == CapacityType.FLEX_START
+                else "RAPID"
+            ),
+            "version_prefix": "1.32.",
+            "maintenance_exclusions": (
+                []
+                if capacity_type == CapacityType.FLEX_START
+                else [{
+                    "name": "no-minor-or-node-upgrades-indefinite",
+                    "start_time": "2024-12-01T00:00:00Z",
+                    "end_time": "2025-12-22T00:00:00Z",
+                    "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+                }]
+            ),
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "system_node_pool_machine_type": system_node_pool_machine_type,
@@ -534,9 +564,10 @@ class BlueprintGenerator:
         use=[cluster_id],
         settings={
             "machine_type": system.gce_machine_type,
-            "auto_upgrade": True,
+            "auto_upgrade": (
+                True if capacity_type != CapacityType.FLEX_START else False
+            ),
             "zones": [zone],
-            "static_node_count": num_nodes,
             "spot": capacity_type == CapacityType.SPOT,
             "reservation_affinity": self._getblock_reservation_affinity(
                 reservation
@@ -562,6 +593,10 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+    if capacity_type == CapacityType.FLEX_START:
+      gpu_pool.settings.update(self.get_dws_flex_start())
+    else:
+      gpu_pool.settings.update({"static_node_count": num_nodes})
 
     num_chips = num_nodes * system.chips_per_vm
     workload_manager_install_id = "workload-manager-install"
@@ -572,9 +607,14 @@ class BlueprintGenerator:
         settings={
             "kueue": {
                 "install": True,
-                "version": "v0.10.0",  # TAS feature-gates is enabled in CT
+                "version": KUEUE_VERSION,  # TAS feature-gates is enabled in CT
                 "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {"num_chips": num_chips},
+                "config_template_vars": {
+                    "num_chips": num_chips,
+                    "flex_start": (
+                        1 if capacity_type == CapacityType.FLEX_START else 0
+                    ),
+                },
             },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [
@@ -777,13 +817,21 @@ class BlueprintGenerator:
                 f" {cluster_name}-rdma-net.subnetwork_interfaces_gke))"
             ),
             "version_prefix": "1.32.",
-            "release_channel": "RAPID",
-            "maintenance_exclusions": [{
-                "name": "no-minor-or-node-upgrades-indefinite",
-                "start_time": "2024-12-01T00:00:00Z",
-                "end_time": "2025-12-22T00:00:00Z",
-                "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
-            }],
+            "release_channel": (
+                "UNSPECIFIED"
+                if capacity_type == CapacityType.FLEX_START
+                else "RAPID"
+            ),
+            "maintenance_exclusions": (
+                []
+                if capacity_type == CapacityType.FLEX_START
+                else [{
+                    "name": "no-minor-or-node-upgrades-indefinite",
+                    "start_time": "2024-12-01T00:00:00Z",
+                    "end_time": "2025-12-22T00:00:00Z",
+                    "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+                }]
+            ),
         },
         outputs=["instructions"],
     )
@@ -800,10 +848,11 @@ class BlueprintGenerator:
         use=[cluster_id],
         settings={
             "machine_type": system.gce_machine_type,
-            "auto_upgrade": True,
+            "auto_upgrade": (
+                True if capacity_type != CapacityType.FLEX_START else False
+            ),
             "zones": [zone],
             "disk_type": "hyperdisk-balanced",
-            "static_node_count": num_nodes,
             "local_ssd_count_ephemeral_storage": 32,
             "spot": capacity_type == CapacityType.SPOT,
             "reservation_affinity": self._getblock_reservation_affinity(
@@ -830,6 +879,10 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+    if capacity_type == CapacityType.FLEX_START:
+      gpu_pool.settings.update(self.get_dws_flex_start())
+    else:
+      gpu_pool.settings.update({"static_node_count": num_nodes})
 
     num_chips = num_nodes * system.chips_per_vm
     workload_manager_install_id = "workload-manager-install"
@@ -840,9 +893,14 @@ class BlueprintGenerator:
         settings={
             "kueue": {
                 "install": True,
-                "version": "v0.10.0",  # TAS feature-gates is enabled in CT
+                "version": KUEUE_VERSION,  # TAS feature-gates is enabled in CT
                 "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {"num_chips": num_chips},
+                "config_template_vars": {
+                    "num_chips": num_chips,
+                    "flex_start": (
+                        1 if capacity_type == CapacityType.FLEX_START else 0
+                    ),
+                },
             },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [
@@ -992,6 +1050,16 @@ class BlueprintGenerator:
     )
     return deployment_files_path
 
+  def get_dws_flex_start(self) -> dict:
+    return {
+        "enable_flex_start": True,
+        "enable_queued_provisioning": True,
+        "autoscaling_total_min_nodes": 0,
+        "release_channel": "UNSPECIFIED",
+        "auto_repair": False,
+        "auto_upgrade": False,
+    }
+
 
 yaml.register_class(Blueprint)
 yaml.register_class(DeploymentGroup)

xpk/core/capacity.py

@@ -16,8 +16,8 @@ limitations under the License.
 
 import enum
 
-from ..utils.console import xpk_print
-from .commands import run_command_with_updates
+from ..utils.console import xpk_print, xpk_exit
+from .commands import run_command_with_updates, run_command_for_value
 
 AUTOPROVISIONING_CONFIG_VALUE = 'AUTOPROVISION'
 AUTOPROVISIONING_CONFIG_MINIMUM_KEY = 'minimum_chips'
@@ -36,6 +36,7 @@ class CapacityType(enum.Enum):
   RESERVATION = 'reservation'
   SPOT = 'spot'
   UNKNOWN = 'unknown'
+  FLEX_START = 'flex_start'
 
 
 def print_reservations(args) -> int:
@@ -84,6 +85,9 @@ def get_capacity_type(args) -> tuple[CapacityType, int]:
   if args.spot:
     capacity_type = CapacityType.SPOT
     num_types += 1
+  if args.flex:
+    capacity_type = CapacityType.FLEX_START
+    num_types += 1
 
   # Check that the number of user arguments provided is valid.
   if num_types == 0:
@@ -91,14 +95,62 @@ def get_capacity_type(args) -> tuple[CapacityType, int]:
   elif num_types != 1:
     xpk_print(
         'ERROR: User specified more than one of the following arguments. Please'
-        ' specify only one of `--reservation=$RESERVATION_NAME`, `--on-demand`'
-        ' or `--spot`.'
+        ' specify only one of `--reservation=$RESERVATION_NAME`, `--on-demand`,'
+        ' `--flex` or `--spot`.'
     )
     return_code = 1
 
   return capacity_type, return_code
 
 
+def get_reservation_maintenance_interval(
+    reservation: str, zone: str, project: str
+) -> str:
+  """Get reservation maintenance interval.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      f'gcloud beta compute reservations describe {reservation}'
+      f' --project={project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
+  )
+  return_code, output = run_command_for_value(
+      command, 'Get reservation maintenance interval', None
+  )
+  if return_code != 0:
+    xpk_print(f'Get reservation maintenance interval ERROR {return_code}')
+    xpk_exit(1)
+  return output.strip()
+
+
+def get_reservation_placement_policy(
+    reservation: str, zone: str, project: str
+) -> str:
+  """Get reservation placement policy.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      f'gcloud beta compute reservations describe {reservation}'
+      f' --project={project} --zone={zone} --format="value(resourcePolicies.policy)"'
+  )
+  return_code, output = run_command_for_value(
+      command, 'Get reservation placement policy', None
+  )
+  if return_code != 0:
+    xpk_print(f'Get reservation placement policy ERROR {return_code}')
+    xpk_exit(1)
+  return output.strip()
+
+
 def verify_reservation_exists(args) -> int:
   """Verify the reservation exists.
 
@@ -121,9 +173,9 @@ def verify_reservation_exists(args) -> int:
 
 
 def get_capacity_arguments_from_capacity_type(
-    args, capacity_type: CapacityType
+    args, capacity_type: CapacityType, max_nodes: int
 ) -> tuple[str, int]:
-  """Determine the TPU Nodepool creation capacity arguments needed.
+  """Determine the Nodepool creation capacity arguments needed.
 
   Args:
     args: user provided arguments for running the command.
@@ -141,6 +193,12 @@ def get_capacity_arguments_from_capacity_type(
       capacity_args = ''
     case CapacityType.SPOT:
       capacity_args = '--spot'
+    case CapacityType.FLEX_START:
+      capacity_args = (
+          ' --flex-start --enable-queued-provisioning --enable-autoscaling'
+          ' --location-policy=ANY --reservation-affinity=none'
+          f' --no-enable-autorepair --max-nodes={max_nodes}'
+      )
     case CapacityType.RESERVATION:
       capacity_args = (
           f'--reservation-affinity=specific --reservation={args.reservation}'
@@ -173,6 +231,8 @@ def get_capacity_node_selectors_from_capacity_type(
   match capacity_type:
     case CapacityType.ON_DEMAND.name:
       node_selector = ''
+    case CapacityType.FLEX_START.name:
+      node_selector = 'cloud.google.com/gke-queued="true"'
     case CapacityType.SPOT.name:
       node_selector = 'cloud.google.com/gke-spot="true"'
     case CapacityType.RESERVATION.name:

xpk/core/cluster.py

@@ -38,7 +38,7 @@ from .resources import get_cluster_system_characteristics
 from .system_characteristics import SystemCharacteristics
 
 JOBSET_VERSION = 'v0.8.0'
-PATHWAYS_JOB_VERSION = 'v0.1.1'
+PATHWAYS_JOB_VERSION = 'v0.1.2'
 INSTALLER_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
 INSTALLER_NCCL_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
 INSTALLER_NCCL_RDMA = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer.yaml'
@@ -315,28 +315,60 @@ def update_cluster_with_pd_driver_if_necessary(args) -> int:
   return 0
 
 
-def is_driver_enabled_on_cluster(args, driver: str) -> bool:
+def update_cluster_with_lustre_driver_if_necessary(args) -> int:
+  """Updates a GKE cluster to enable Lustre CSI driver, if not enabled already.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and error code otherwise.
+  """
+  if is_driver_enabled_on_cluster(
+      args, driver='lustreCsiDriver'
+  ) and is_driver_enabled_on_cluster(
+      args, driver='lustreCsiDriver', config_key='enableLegacyLustrePort'
+  ):
+    return 0
+  cluster_update_return_code = update_gke_cluster_with_lustre_driver_enabled(
+      args
+  )
+  if cluster_update_return_code > 0:
+    xpk_print(
+        'Updating GKE cluster to enable PersistentDisk CSI driver failed!'
+    )
+    return cluster_update_return_code
+
+  return 0
+
+
+def is_driver_enabled_on_cluster(
+    args, driver: str, config_key: str = 'enabled', config_val: str = 'true'
+) -> bool:
   """Checks if the CSI driver is enabled on the cluster.
   Args:
     args: user provided arguments for running the command.
     driver (str) : name of the driver
+    config (str): the config to look for; by default looks for "enabled" parameter
+    config_val (str): the value indicating the enabled; default vale is "true"
   Returns:
     True if driver is enabled on the cluster and False otherwise.
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
       f' --project={args.project} --region={zone_to_region(args.zone)}'
-      f' --format="value(addonsConfig.{driver}Config.enabled)"'
+      f' --format="value(addonsConfig.{driver}Config.{config_key})"'
   )
   return_code, driver_enabled = run_command_for_value(
       command,
-      f'Checks if {driver} driver is enabled in cluster describe.',
+      f"Checks if {driver} driver's {config_key} is enabled in cluster"
+      ' describe.',
       args,
   )
   if return_code != 0:
     xpk_exit(return_code)
-  if driver_enabled.strip().lower() == 'true':
-    xpk_print(f'{driver} driver is enabled on the cluster, no update needed.')
+  if driver_enabled.strip().lower() == config_val.lower():
+    xpk_print(
+        f"{driver} driver's {config_key} config is {config_val} on the cluster."
+    )
     return True
   return False
 
@@ -423,6 +455,19 @@ def get_gpu_type_from_cluster(args) -> str:
   return ''
 
 
+def setup_k8s_service_accounts() -> None:
+  """
+  Creates/sets up SAs and the roles for them
+  """
+  default_sa = 'default'
+
+  create_xpk_k8s_service_account()
+
+  role_name = create_pod_reader_role()
+  create_role_binding(default_sa, role_name)
+  create_role_binding(XPK_SA, role_name)
+
+
 def create_xpk_k8s_service_account() -> None:
   k8s_core_client = k8s_client.CoreV1Api()
   sa = k8s_client.V1ServiceAccount(
@@ -441,6 +486,94 @@ def create_xpk_k8s_service_account() -> None:
     )
 
 
+def create_pod_reader_role() -> str:
+  """
+  Creates the 'pod-reader' Role in the default namespace.
+  """
+  k8s_rbac_client = k8s_client.RbacAuthorizationV1Api()
+  role_name = 'pod-reader'
+
+  role = k8s_client.V1Role(
+      metadata=k8s_client.V1ObjectMeta(
+          name=role_name, namespace=DEFAULT_NAMESPACE
+      ),
+      rules=[
+          k8s_client.V1PolicyRule(
+              api_groups=[''],
+              resources=['pods', 'services'],
+              verbs=['get', 'list', 'watch'],
+          ),
+          k8s_client.V1PolicyRule(
+              api_groups=['batch'],
+              resources=['jobs'],
+              verbs=['get', 'list', 'watch'],
+          ),
+      ],
+  )
+
+  xpk_print(
+      f'Attempting to create Role: {role_name} in namespace:'
+      f' {DEFAULT_NAMESPACE}'
+  )
+  try:
+    k8s_rbac_client.create_namespaced_role(DEFAULT_NAMESPACE, role, pretty=True)
+    xpk_print(f'Successfully created Role: {role_name}')
+    return role_name
+  except ApiException as e:
+    if e.status == 409:  # Conflict, meaning it already exists
+      xpk_print(f'Role: {role_name} already exists. Skipping its creation.')
+      return role_name
+    else:
+      xpk_print(f'Error creating Role {role_name}: {e}')
+      xpk_exit(1)
+
+
+def create_role_binding(sa: str, role_name: str) -> None:
+  """
+  Creates a RoleBinding to associate the Service Account
+  with the Role in the default namespace.
+  Assumes the Service Account and the Role already exist.
+  """
+  k8s_rbac_client = k8s_client.RbacAuthorizationV1Api()
+  role_binding_name = f'{sa}-{role_name}-binding'
+
+  role_binding = k8s_client.V1RoleBinding(
+      metadata=k8s_client.V1ObjectMeta(
+          name=role_binding_name, namespace=DEFAULT_NAMESPACE
+      ),
+      subjects=[
+          k8s_client.RbacV1Subject(
+              kind='ServiceAccount', name=sa, namespace=DEFAULT_NAMESPACE
+          )
+      ],
+      role_ref=k8s_client.V1RoleRef(
+          kind='Role', name=role_name, api_group='rbac.authorization.k8s.io'
+      ),
+  )
+
+  xpk_print(
+      f'Attempting to create RoleBinding: {role_binding_name} for Service'
+      f' Account: {XPK_SA} to Role: {role_name} in namespace:'
+      f' {DEFAULT_NAMESPACE}'
+  )
+  try:
+    k8s_rbac_client.create_namespaced_role_binding(
+        DEFAULT_NAMESPACE, role_binding, pretty=True
+    )
+    xpk_print(
+        f'Successfully created RoleBinding: {role_binding_name} for {XPK_SA}'
+    )
+  except ApiException as e:
+    if e.status == 409:  # Conflict, meaning it already exists
+      xpk_print(
+          f'RoleBinding: {role_binding_name} already exists. Skipping its'
+          ' creation.'
+      )
+    else:
+      xpk_print(f'Error creating RoleBinding {role_binding_name}: {e}')
+      xpk_exit(1)
+
+
 def update_gke_cluster_with_clouddns(args) -> int:
   """Run the GKE cluster update command for existing clusters and enable CloudDNS.
 
@@ -522,6 +655,32 @@ def update_gke_cluster_with_gcsfuse_driver_enabled(args) -> int:
   return 0
 
 
+def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
+  """Run the GKE cluster update command for existing cluster and enable Lustre CSI driver.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'gcloud container clusters update'
+      f' {args.cluster} --project={args.project}'
+      f' --region={zone_to_region(args.zone)}'
+      ' --enable-legacy-lustre-port'
+      ' --quiet'
+  )
+  xpk_print(
+      'Updating GKE cluster to enable Lustre CSI driver, may take a while!'
+  )
+  return_code = run_command_with_updates(
+      command, 'GKE Cluster Update to enable Lustre CSI driver', args
+  )
+  if return_code != 0:
+    xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
+    return 1
+  return 0
+
+
 def upgrade_gke_control_plane_version(args, default_rapid_gke_version) -> int:
   """Upgrade GKE cluster's control plane version before updating nodepools to use CloudDNS.
 
@@ -731,7 +890,6 @@ def get_cluster_credentials(args) -> None:
   command = (
       'gcloud container clusters get-credentials'
       f' {args.cluster} --region={zone_to_region(args.zone)}'
-      ' --dns-endpoint'
       f' --project={args.project} &&'
       ' kubectl config view && kubectl config set-context --current'
       ' --namespace=default'