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

integration/gcluster_a3mega_test.py

@@ -18,6 +18,7 @@ from xpk.commands.cluster_gcluster import get_unique_name
 from xpk.core.docker_manager import DockerManager
 from xpk.core.gcluster_manager import GclusterManager
 from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
+from xpk.utils.versions import ReleaseChannel
 import pytest
 import os
 import shutil
@@ -28,6 +29,8 @@ region = os.getenv("REGION")
 zone = os.getenv("ZONE")
 auth_cidr = os.getenv("AUTH_CIDR")
 cluster_name = os.getenv("A3_MEGA_TEST_CLUSTER_NAME")
+release_channel = os.getenv("RELEASE_CHANNEL")
+cluster_version = os.getenv("CLUSTER_VERSION")
 
 uploads_dir = "uploads"
 
@@ -87,6 +90,8 @@ def test_create_a3_mega_deployment_files(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
 
   blueprint_name = f"{cluster_name}-a3-mega-xpk"
@@ -107,6 +112,8 @@ def test_create_a3_mega_deployment_files(setup_tests):
       auth_cidr=auth_cidr,
       zone=zone,
       system_node_pool_min_node_count=3,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, prefix, blueprint_name)
@@ -164,6 +171,8 @@ def create_test_a3_mega_deployment(docker_path: str, bp_path: str):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
 
   blueprint_name = f"{cluster_name}-a3-mega-xpk"
   prefix = "prefix"
@@ -183,6 +192,8 @@ def create_test_a3_mega_deployment(docker_path: str, bp_path: str):
       auth_cidr=auth_cidr,
       zone=zone,
       system_node_pool_min_node_count=3,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
 
   gcluster_manager = GclusterManager(

integration/gcluster_a3ultra_test.py

@@ -24,6 +24,7 @@ from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
 from xpk.core.capacity import CapacityType
 from xpk.core.docker_manager import DockerManager
 from xpk.core.gcluster_manager import GclusterManager
+from xpk.utils.versions import ReleaseChannel
 
 ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
 project_id = os.getenv("PROJECT_ID")
@@ -31,6 +32,8 @@ region = os.getenv("REGION")
 zone = os.getenv("ZONE")
 auth_cidr = os.getenv("AUTH_CIDR")
 cluster_name = os.getenv("A3_ULTRA_TEST_CLUSTER_NAME")
+release_channel = os.getenv("RELEASE_CHANNEL")
+cluster_version = os.getenv("CLUSTER_VERSION")
 
 
 @pytest.fixture(name="setup_tests")
@@ -60,6 +63,8 @@ def test_create_a3_ultra_deployment_files(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a3-ultra-xpk"
 
@@ -80,6 +85,8 @@ def test_create_a3_ultra_deployment_files(setup_tests):
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
       prefix=prefix,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)
@@ -125,6 +132,8 @@ def test_create_a3_ultra_deployment(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a3-ultra-xpk"
 
@@ -144,6 +153,8 @@ def test_create_a3_ultra_deployment(setup_tests):
       capacity_type=CapacityType.SPOT,
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)

integration/gcluster_a4_test.py

@@ -24,6 +24,7 @@ from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
 from xpk.core.capacity import CapacityType
 from xpk.core.docker_manager import DockerManager
 from xpk.core.gcluster_manager import GclusterManager
+from xpk.utils.versions import ReleaseChannel
 
 ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
 project_id = os.getenv("PROJECT_ID")
@@ -31,6 +32,8 @@ region = os.getenv("REGION")
 zone = os.getenv("ZONE")
 auth_cidr = os.getenv("AUTH_CIDR")
 cluster_name = os.getenv("A4_TEST_CLUSTER_NAME")
+release_channel = os.getenv("RELEASE_CHANNEL")
+cluster_version = os.getenv("CLUSTER_VERSION")
 
 
 @pytest.fixture(name="setup_tests")
@@ -60,6 +63,8 @@ def test_create_a4_deployment_files(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a4-xpk"
 
@@ -80,6 +85,8 @@ def test_create_a4_deployment_files(setup_tests):
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
       prefix=prefix,
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, prefix, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)
@@ -125,6 +132,8 @@ def test_create_a4_deployment(setup_tests):
   assert auth_cidr is not None
   assert ctk_gcloud_cfg is not None
   assert cluster_name is not None
+  assert release_channel is not None
+  assert cluster_version is not None
   docker_path, bp_path = setup_tests[0], setup_tests[1]
   blueprint_name = f"{cluster_name}-a4-xpk"
 
@@ -144,6 +153,8 @@ def test_create_a4_deployment(setup_tests):
       capacity_type=CapacityType.SPOT,
       num_nodes=1,
       system_node_pool_machine_type="e2-standard-16",
+      release_channel=ReleaseChannel(release_channel),
+      cluster_version=cluster_version,
   )
   blueprint_test_path = os.path.join(bp_path, f"{blueprint_name}.yaml")
   blueprint_deps_test_path = os.path.join(bp_path, blueprint_name)

xpk/commands/cluster.py

@@ -17,6 +17,7 @@ limitations under the License.
 from tabulate import tabulate
 
 from ..utils.feature_flags import FeatureFlags
+from ..utils.versions import ReleaseChannel
 from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE, get_reservation_deployment_type
 from ..core.cluster import (
     get_all_clusters_programmatic,
@@ -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
@@ -263,15 +266,15 @@ def cluster_create(args) -> None:
   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)
+  _log_cluster_create_telemetry(args)
 
-  return_code, gke_server_config = get_gke_server_config(args)
+  release_channel = (
+      ReleaseChannel.REGULAR if args.gke_version else ReleaseChannel.RAPID
+  )
+
+  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 +284,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)
@@ -1022,7 +1037,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 +1061,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 +1074,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 +1085,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 +1132,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 +1175,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
@@ -1285,7 +1306,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 +1320,7 @@ def _install_kueue(
           configure_sub_slicing=(
               FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing
           ),
-      ),
+      )
   )
 
 
@@ -1320,3 +1341,18 @@ def prepare_gpus(system: SystemCharacteristics):
     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

@@ -21,6 +21,7 @@ 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.utils.versions import ReleaseChannel
 
 
 @pytest.fixture
@@ -103,7 +104,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"]
@@ -153,7 +158,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)