xpk 0.14.4__py3-none-any.whl → 0.16.0__py3-none-any.whl

xpk/commands/cluster.py

@@ -17,7 +17,9 @@ limitations under the License.
 from tabulate import tabulate
 
 from ..utils.feature_flags import FeatureFlags
-from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE, get_reservation_deployment_type
+from ..utils.versions import ReleaseChannel
+from ..core.pathways import get_pathways_machine_types
+from ..core.capacity import H100_DEVICE_TYPE, get_reservation_deployment_type
 from ..core.cluster import (
     get_all_clusters_programmatic,
     get_cluster_credentials,
@@ -26,7 +28,6 @@ from ..core.cluster import (
     set_jobset_on_cluster,
     set_pathways_job_on_cluster,
     setup_k8s_env,
-    disable_mglru_on_cluster,
     count_nodes_on_cluster,
     update_cluster_with_gcpfilestore_driver_if_necessary,
     update_cluster_with_gcsfuse_driver_if_necessary,
@@ -38,6 +39,8 @@ from ..core.cluster import (
 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.telemetry import MetricsCollector, MetricsEventMetadataKey
+from ..core.capacity import get_capacity_type
 from ..core.gcloud_context import (
     add_zone_and_project,
     get_gke_control_plane_version,
@@ -71,9 +74,9 @@ from ..core.system_characteristics import (
 )
 from ..core.vertex import create_vertex_tensorboard
 from ..core.workload import get_workload_list
-from ..utils.console import get_user_input, xpk_exit, xpk_print
+from ..utils.console import ask_for_user_consent, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
-from ..utils.execution_context import is_dry_run
+from ..utils.execution_context import is_dry_run, is_quiet
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from . import cluster_gcluster
 from .common import set_cluster_command, validate_sub_slicing_system
@@ -81,6 +84,7 @@ from jinja2 import Environment, FileSystemLoader
 from ..utils.templates import get_templates_absolute_path
 import shutil
 import os
+from .managed_ml_diagnostics import install_mldiagnostics_prerequisites
 
 CLUSTER_PREHEAT_JINJA_FILE = 'cluster_preheat.yaml.j2'
 
@@ -207,6 +211,25 @@ def _validate_cluster_create_args(args, system: SystemCharacteristics):
   if FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing:
     validate_sub_slicing_system(system)
     _validate_sub_slicing_reservation(args)
+  if args.enable_pathways:
+    _validate_pathways_machine(args)
+
+
+def _validate_pathways_machine(args):
+  return_code, result = get_pathways_machine_types(
+      project=args.project, zone=args.zone
+  )
+  if return_code != 0:
+    xpk_print('Error: Unable to retrieve available pathways machine types')
+    xpk_exit(1)
+
+  if args.pathways_gce_machine_type not in result:
+    xpk_print(
+        'Error: Invalid --pathways-gce-machine-type. Specify machine type that'
+        ' has at least 100GB of memory and at least 49 CPUs.'
+    )
+    xpk_print(f'Available machine types: {", ".join(result)}')
+    xpk_exit(1)
 
 
 def _validate_sub_slicing_reservation(args):
@@ -258,20 +281,19 @@ def cluster_create(args) -> None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
-  _validate_cluster_create_args(args, system)
-
   xpk_print(f'Starting cluster create for cluster {args.cluster}:', flush=True)
   add_zone_and_project(args)
 
-  if system.device_type in cluster_gcluster.supported_device_types:
-    xpk_print(
-        'Creating the cluster using Cluster Toolkit. Machine Type:'
-        f' {system.gce_machine_type} ...'
-    )
-    cluster_gcluster.cluster_create(args)
-    xpk_exit(0)
+  _validate_cluster_create_args(args, system)
+  _log_cluster_create_telemetry(args)
+
+  release_channel = (
+      ReleaseChannel.REGULAR if args.gke_version else ReleaseChannel.RAPID
+  )
 
-  return_code, gke_server_config = get_gke_server_config(args)
+  return_code, gke_server_config = get_gke_server_config(
+      args, release_channel=release_channel
+  )
   if return_code != 0 or gke_server_config is None:
     xpk_exit(return_code)
 
@@ -281,8 +303,20 @@ def cluster_create(args) -> None:
   if return_code != 0 or gke_control_plane_version is None:
     xpk_exit(return_code)
 
+  if system.device_type in cluster_gcluster.supported_device_types:
+    xpk_print(
+        'Creating the cluster using Cluster Toolkit. Machine Type:'
+        f' {system.gce_machine_type} ...'
+    )
+    cluster_gcluster.cluster_create(
+        args,
+        gke_control_plane_version=gke_control_plane_version,
+        release_channel=release_channel,
+    )
+    xpk_exit(0)
+
   create_cluster_command_code = create_cluster_if_necessary(
-      args, gke_control_plane_version, system
+      args, gke_control_plane_version, system, release_channel=release_channel
   )
   if create_cluster_command_code != 0:
     xpk_exit(create_cluster_command_code)
@@ -407,6 +441,13 @@ def cluster_create(args) -> None:
       # pylint: disable=line-too-long
       f' https://console.cloud.google.com/kubernetes/clusters/details/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/details?project={args.project}'
   )
+
+  if args.managed_mldiagnostics:
+    return_code = install_mldiagnostics_prerequisites()
+    if return_code != 0:
+      xpk_print('Installation of MLDiagnostics failed.')
+      xpk_exit(return_code)
+
   xpk_exit(0)
 
 
@@ -964,7 +1005,7 @@ def update_coredns() -> int:
 
   # 6. Scale up coredns and verify readiness
   scale_up_coredns(replicas=15)
-  verify_coredns_readiness(timeout=120)
+  verify_coredns_readiness()
 
   xpk_print('The CoreDNS setup process has been completed.')
 
@@ -1022,7 +1063,10 @@ def update_coredns_if_necessary() -> int:
 
 
 def create_cluster_if_necessary(
-    args, gke_control_plane_version: str, system: SystemCharacteristics
+    args,
+    gke_control_plane_version: str,
+    system: SystemCharacteristics,
+    release_channel: ReleaseChannel,
 ) -> int:
   """Creates cluster if not present in the project.
 
@@ -1043,7 +1087,7 @@ def create_cluster_if_necessary(
     return 0
   else:
     return run_gke_cluster_create_command(
-        args, gke_control_plane_version, system
+        args, gke_control_plane_version, system, release_channel=release_channel
     )
 
 
@@ -1056,7 +1100,7 @@ def run_gke_cluster_delete_command(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
-  if not args.force:
+  if not is_quiet():
     xpk_print('Get the name of the workloads in the cluster.')
     args.filter_by_status = 'EVERYTHING'
     return_code, return_value = get_workload_list(args)
@@ -1067,10 +1111,9 @@ def run_gke_cluster_delete_command(args) -> int:
     # Ignore Column Names line.
     if len(return_value) > 1:
       workloads = [x.split(' ')[0] for x in return_value.splitlines()][1:]
-      if workloads and not get_user_input(
+      if workloads and not ask_for_user_consent(
           f'Planning to delete {len(workloads)} workloads in the cluster'
-          f' {args.cluster} including {workloads}. \nDo you wish to delete: y'
-          ' (yes) / n (no):\n'
+          f' {args.cluster} including {workloads}. \nDo you wish to delete?'
       ):
         xpk_print('Skipping delete command.')
         return 0
@@ -1115,7 +1158,10 @@ def run_gke_clusters_list_command(args) -> int:
 
 
 def run_gke_cluster_create_command(
-    args, gke_control_plane_version: str, system: SystemCharacteristics
+    args,
+    gke_control_plane_version: str,
+    system: SystemCharacteristics,
+    release_channel: ReleaseChannel,
 ) -> int:
   """Run the Create GKE Cluster request.
 
@@ -1155,9 +1201,10 @@ def run_gke_cluster_create_command(
       ' --enable-dns-access'
       ' --autoscaling-profile=optimize-utilization'
       ' --labels=gke_product_type=xpk'
+      f' --release-channel={release_channel.value.lower()}'
   )
 
-  if args.gke_version or system.accelerator_type == AcceleratorType.GPU:
+  if args.gke_version:
     command += ' --no-enable-autoupgrade'
 
   enable_ip_alias = False
@@ -1199,7 +1246,8 @@ def run_gke_cluster_create_command(
 
   if args.enable_lustre_csi_driver:
     addons.append('LustreCsiDriver')
-    command += ' --enable-legacy-lustre-port'
+    if args.enable_legacy_lustre_port:
+      command += ' --enable-legacy-lustre-port'
 
   if hasattr(args, 'enable_mtc') and args.enable_mtc:
     addons.append('HighScaleCheckpointing')
@@ -1285,7 +1333,7 @@ def _install_kueue(
   else:
     # Determine total chips based on user specified topology.
     total_chips = get_total_chips_requested_from_args(args, system)
-  kueue_manager = KueueManager()
+  kueue_manager = KueueManager(args.project, args.zone)
   return kueue_manager.install_or_upgrade(
       KueueConfig(
           system,
@@ -1299,7 +1347,7 @@ def _install_kueue(
           configure_sub_slicing=(
               FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing
           ),
-      ),
+      )
   )
 
 
@@ -1315,8 +1363,17 @@ def prepare_gpus(system: SystemCharacteristics):
     if install_nri_code != 0:
       xpk_exit(install_nri_code)
 
-  if system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
-    xpk_print('Disabling MGLRU')
-    err_code = disable_mglru_on_cluster()
-    if err_code > 0:
-      xpk_exit(err_code)
+
+def _log_cluster_create_telemetry(args) -> None:
+  if FeatureFlags.TELEMETRY_ENABLED:
+    capacity_type, _ = get_capacity_type(args)
+    MetricsCollector.log_custom(
+        name='cluster_create',
+        metadata={
+            MetricsEventMetadataKey.ZONE: args.zone,
+            MetricsEventMetadataKey.SYSTEM_CHARACTERISTICS: (
+                args.tpu_type if args.tpu_type else args.device_type
+            ),
+            MetricsEventMetadataKey.PROVISIONING_MODE: capacity_type.value,
+        },
+    )

xpk/commands/cluster_gcluster.py

@@ -17,6 +17,7 @@ limitations under the License.
 import os
 
 from ..utils.feature_flags import FeatureFlags
+from ..utils.versions import ReleaseChannel
 from ..utils.execution_context import is_dry_run
 from ..core.kueue_manager import KueueConfig, KueueManager
 from ..core.nap import enable_autoprovisioning_on_cluster
@@ -51,11 +52,15 @@ gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')
 gcloud_cfg_path = os.path.expanduser('~/.config/gcloud')
 
 
-def cluster_create(args) -> None:
+def cluster_create(
+    args, gke_control_plane_version: str, release_channel: ReleaseChannel
+) -> None:
   """Function around cluster creation using Cluster toolkit.
 
   Args:
     args: user provided arguments for running the command.
+    gke_control_plane_version: the GKE version used for the new cluster.
+    release_channel:t the release channel used for the new cluster.
 
   Returns:
     0 if successful and 1 otherwise.
@@ -79,7 +84,13 @@ def cluster_create(args) -> None:
     )
   gcm = prepare_gcluster_manager(remote_state_client)
 
-  bp = generate_blueprint(blueprint_name=unique_name, args=args, prefix=prefix)
+  bp = generate_blueprint(
+      blueprint_name=unique_name,
+      args=args,
+      prefix=prefix,
+      gke_control_plane_version=gke_control_plane_version,
+      release_channel=release_channel,
+  )
 
   # staging: sending the blueprint file(s) to gcluster's working directory
   if is_dry_run():
@@ -141,7 +152,7 @@ def __install_kueue(args) -> int:
   else:
     # Determine total chips based on user specified topology.
     total_chips = get_total_chips_requested_from_args(args, system)
-  kueue_manager = KueueManager()
+  kueue_manager = KueueManager(args.project, args.zone)
 
   tolerations = [{
       'key': 'components.gke.io/gke-managed-components',
@@ -149,7 +160,6 @@ def __install_kueue(args) -> int:
       'value': 'true',
       'effect': 'NoSchedule',
   }]
-
   kueue_manager.install_or_upgrade(
       KueueConfig(
           system,
@@ -287,7 +297,11 @@ def validate_state_gcs_bucket(args):
 
 
 def generate_blueprint(
-    blueprint_name, args, prefix=None
+    blueprint_name,
+    args,
+    gke_control_plane_version: str,
+    release_channel: ReleaseChannel,
+    prefix=None,
 ) -> BlueprintGeneratorOutput:
   capacity_type, return_code = get_capacity_type(args)
   if return_code != 0:
@@ -342,6 +356,8 @@ def generate_blueprint(
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
           gcs_bucket=args.cluster_state_gcs_bucket,
+          cluster_version=gke_control_plane_version,
+          release_channel=release_channel,
       )
     if args.device_type == a3ultra_device_type:
       num_nodes = args.num_nodes if not args.num_nodes is None else 2
@@ -360,6 +376,8 @@ def generate_blueprint(
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
           gcs_bucket=args.cluster_state_gcs_bucket,
+          cluster_version=gke_control_plane_version,
+          release_channel=release_channel,
       )
     if args.device_type == a4_device_type:
       num_nodes = args.num_nodes if not args.num_nodes is None else 2
@@ -376,6 +394,8 @@ def generate_blueprint(
           capacity_type=capacity_type,
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
+          cluster_version=gke_control_plane_version,
+          release_channel=release_channel,
       )
   xpk_print('Device type is not supported.')
   xpk_exit(1)

xpk/commands/cluster_gcluster_test.py

@@ -20,7 +20,8 @@ import pytest
 
 from xpk.commands.cluster_gcluster import cluster_create
 from xpk.core.kueue_manager import KueueConfig
-from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics
+from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics, DockerPlatform, GpuConfig
+from xpk.utils.versions import ReleaseChannel
 
 
 @pytest.fixture
@@ -96,6 +97,8 @@ def test_install_kueue_standard(
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-mega-80gb-8",
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
+      gpu_config=GpuConfig(requires_topology=True),
   )
   mock_cluster_create_deps["get_system_characteristics"].return_value = (
       mock_system,
@@ -103,7 +106,11 @@ def test_install_kueue_standard(
   )
   mock_get_total_chips.return_value = 16
 
-  cluster_create(mock_args)
+  cluster_create(
+      mock_args,
+      release_channel=ReleaseChannel.RAPID,
+      gke_control_plane_version="1.2.3",
+  )
 
   mock_cluster_create_deps["xpk_exit"].assert_called_with(0)
   mock_kueue_manager = mock_cluster_create_deps["KueueManager"]
@@ -143,6 +150,8 @@ def test_install_kueue_with_autoprovisioning(
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-mega-80gb-8",
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
+      gpu_config=GpuConfig(requires_topology=True),
   )
   mock_cluster_create_deps["get_system_characteristics"].return_value = (
       mock_system,
@@ -153,7 +162,11 @@ def test_install_kueue_with_autoprovisioning(
   mock_autoprovisioning_config.maximum_chips = 128
   mock_enable_autoprovisioning.return_value = (mock_autoprovisioning_config, 0)
 
-  cluster_create(mock_args)
+  cluster_create(
+      mock_args,
+      release_channel=ReleaseChannel.RAPID,
+      gke_control_plane_version="1.2.3",
+  )
 
   mock_cluster_create_deps["xpk_exit"].assert_called_with(0)
   mock_enable_autoprovisioning.assert_called_once_with(mock_args, mock_system)