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

xpk/commands/cluster_test.py

@@ -14,16 +14,20 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 from argparse import Namespace
 from dataclasses import dataclass
 from typing import Any
 from unittest.mock import MagicMock, patch
 import pytest
 
-from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command
+from xpk.core.telemetry import MetricsCollector
+from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command, cluster_create, _log_cluster_create_telemetry
+from xpk.core.capacity import CapacityType
 from xpk.core.system_characteristics import SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
 from xpk.utils.feature_flags import FeatureFlags
+from xpk.utils.versions import ReleaseChannel
 
 
 @dataclass
@@ -34,6 +38,30 @@ class _Mocks:
   commands_tester: CommandsTester
 
 
+@dataclass
+class _ClusterCreateMocks:
+  """Holds all the mocked dependencies for the cluster_create function."""
+
+  get_all_clusters_programmatic: MagicMock
+  get_gke_server_config: MagicMock
+  get_gke_control_plane_version: MagicMock
+  get_system_characteristics: MagicMock
+  authorize_private_cluster_access_if_necessary: MagicMock
+  update_coredns_if_necessary: MagicMock
+  get_cluster_credentials: MagicMock
+  setup_k8s_env: MagicMock
+  get_gke_node_pool_version: MagicMock
+  run_gke_node_pool_create_command: MagicMock
+  create_cluster_configmaps: MagicMock
+  set_jobset_on_cluster: MagicMock
+  get_cluster_location: MagicMock
+  install_kjob: MagicMock
+  xpk_exit: MagicMock
+  update_jobset_resources_if_necessary: MagicMock
+  _install_kueue: MagicMock
+  set_pathways_job_on_cluster: MagicMock
+
+
 @pytest.fixture
 def mocks(mocker) -> _Mocks:
   common_print_mock = mocker.patch(
@@ -65,6 +93,10 @@ def construct_args(**kwargs: Any) -> Namespace:
       project='project',
       zone='us-central1-a',
       reservation='',
+      on_demand=False,
+      tpu_type=None,
+      device_type=None,
+      spot=False,
       default_pool_cpu_machine_type='test-machine-type',
       cluster='test-cluster',
       default_pool_cpu_num_nodes='100',
@@ -87,11 +119,82 @@ def construct_args(**kwargs: Any) -> Namespace:
       memory_limit='100Gi',
       cpu_limit=100,
       cluster_cpu_machine_type='',
+      create_vertex_tensorboard=False,
+      enable_autoprovisioning=False,
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)
 
 
+@pytest.fixture
+def cluster_create_mocks(mocker) -> _ClusterCreateMocks:
+  """Mocks all dependencies for the cluster_create function."""
+  # This fixture patches all the functions called by cluster_create, allowing
+  # tests to focus on specific logic paths without executing external commands
+  # or complex sub-functions. Each mock can be configured within the test
+  # itself if a specific return value or behavior is needed.
+  return _ClusterCreateMocks(
+      get_all_clusters_programmatic=mocker.patch(
+          'xpk.commands.cluster.get_all_clusters_programmatic',
+          return_value=([], 0),
+      ),
+      get_gke_server_config=mocker.patch(
+          'xpk.commands.cluster.get_gke_server_config',
+          return_value=(0, MagicMock()),
+      ),
+      get_gke_control_plane_version=mocker.patch(
+          'xpk.commands.cluster.get_gke_control_plane_version'
+      ),
+      get_system_characteristics=mocker.patch(
+          'xpk.commands.cluster.get_system_characteristics',
+          return_value=(TPU_TEST_SYSTEM, 0),
+      ),
+      authorize_private_cluster_access_if_necessary=mocker.patch(
+          'xpk.commands.cluster.authorize_private_cluster_access_if_necessary',
+          return_value=0,
+      ),
+      update_coredns_if_necessary=mocker.patch(
+          'xpk.commands.cluster.update_coredns_if_necessary', return_value=0
+      ),
+      get_cluster_credentials=mocker.patch(
+          'xpk.commands.cluster.get_cluster_credentials', return_value=0
+      ),
+      setup_k8s_env=mocker.patch('xpk.commands.cluster.setup_k8s_env'),
+      get_gke_node_pool_version=mocker.patch(
+          'xpk.commands.cluster.get_gke_node_pool_version',
+          return_value=(0, '1.2.3'),
+      ),
+      run_gke_node_pool_create_command=mocker.patch(
+          'xpk.commands.cluster.run_gke_node_pool_create_command',
+          return_value=0,
+      ),
+      create_cluster_configmaps=mocker.patch(
+          'xpk.commands.cluster.create_cluster_configmaps', return_value=0
+      ),
+      set_jobset_on_cluster=mocker.patch(
+          'xpk.commands.cluster.set_jobset_on_cluster', return_value=0
+      ),
+      get_cluster_location=mocker.patch(
+          'xpk.commands.cluster.get_cluster_location',
+          return_value='us-central1',
+      ),
+      install_kjob=mocker.patch(
+          'xpk.commands.cluster.install_kjob', return_value=0
+      ),
+      xpk_exit=mocker.patch('xpk.commands.cluster.xpk_exit'),
+      update_jobset_resources_if_necessary=mocker.patch(
+          'xpk.commands.cluster.update_jobset_resources_if_necessary',
+          return_value=0,
+      ),
+      _install_kueue=mocker.patch(
+          'xpk.commands.cluster._install_kueue', return_value=0
+      ),
+      set_pathways_job_on_cluster=mocker.patch(
+          'xpk.commands.cluster.set_pathways_job_on_cluster', return_value=0
+      ),
+  )
+
+
 GPU_TEST_SYSTEM: SystemCharacteristics = UserFacingNameToSystemCharacteristics[
     'l4-1'
 ]
@@ -106,7 +209,7 @@ TPU_TEST_SYSTEM: SystemCharacteristics = UserFacingNameToSystemCharacteristics[
 def test_validate_cluster_create_args_for_correct_args_pass(
     mocks: _Mocks,
 ):
-  args = Namespace()
+  args = construct_args()
 
   _validate_cluster_create_args(args, GPU_TEST_SYSTEM)
 
@@ -209,6 +312,7 @@ def test_run_gke_cluster_create_command_specifies_custom_cluster_arguments_last(
       ),
       gke_control_plane_version='1.2.3',
       system=TPU_TEST_SYSTEM,
+      release_channel=ReleaseChannel.STABLE,
   )
 
   assert result == 0
@@ -226,12 +330,16 @@ def test_run_gke_cluster_create_command_without_gke_version_does_not_have_no_aut
       args=construct_args(gke_version=''),
       gke_control_plane_version='1.2.3',
       system=TPU_TEST_SYSTEM,
+      release_channel=ReleaseChannel.RAPID,
   )
 
   assert result == 0
   mocks.commands_tester.assert_command_not_run(
       'clusters create', ' --no-enable-autoupgrade'
   )
+  mocks.commands_tester.assert_command_run(
+      'clusters create', ' --release-channel=rapid'
+  )
 
 
 def test_run_gke_cluster_create_command_with_gke_version_has_no_autoupgrade_flag(
@@ -241,24 +349,137 @@ def test_run_gke_cluster_create_command_with_gke_version_has_no_autoupgrade_flag
       args=construct_args(gke_version='1.2.3'),
       gke_control_plane_version='1.2.3',
       system=TPU_TEST_SYSTEM,
+      release_channel=ReleaseChannel.REGULAR,
   )
 
   assert result == 0
   mocks.commands_tester.assert_command_run(
-      'clusters create', ' --no-enable-autoupgrade'
+      'clusters create', '--release-channel=regular', ' --no-enable-autoupgrade'
   )
 
 
-def test_run_gke_cluster_create_command_with_gpu_system_has_no_enable_autoupgrade(
-    mocks: _Mocks,
+def test_log_cluster_create_telemetry_does_not_log_when_feature_flag_is_disabled():
+  FeatureFlags.TELEMETRY_ENABLED = False
+  _log_cluster_create_telemetry(construct_args())
+  events = json.loads(MetricsCollector.flush())['log_event']
+  assert len(events) == 0
+
+
+def test_log_cluster_create_telemetry_logs_correct_event_when_tpu_type_is_provided(
+    mocker: MagicMock,
 ):
-  result = run_gke_cluster_create_command(
-      args=construct_args(gke_version=''),
-      gke_control_plane_version='1.2.3',
-      system=GPU_TEST_SYSTEM,
+  FeatureFlags.TELEMETRY_ENABLED = True
+  mocker.patch(
+      'xpk.commands.cluster.get_capacity_type',
+      return_value=(CapacityType.SPOT, 0),
+  )
+  _log_cluster_create_telemetry(construct_args(device_type='test-device-type'))
+  event = json.loads(MetricsCollector.flush())['log_event'][0]
+  payload = json.loads(event['source_extension_json'])
+  event_metadata = payload['event_metadata']
+  assert payload['event_name'] == 'cluster_create'
+  assert (
+      _get_event_metadata_value_by_key(
+          event_metadata,
+          'XPK_ZONE',
+      )
+      == 'us-central1-a'
+  )
+  assert (
+      _get_event_metadata_value_by_key(
+          event_metadata,
+          'XPK_SYSTEM_CHARACTERISTICS',
+      )
+      == 'test-device-type'
+  )
+  assert (
+      _get_event_metadata_value_by_key(
+          event_metadata,
+          'XPK_PROVISIONING_MODE',
+      )
+      == 'spot'
   )
 
-  assert result == 0
-  mocks.commands_tester.assert_command_run(
-      'clusters create', ' --no-enable-autoupgrade'
+
+def test_log_cluster_create_telemetry_logs_correct_event_when_device_type_is_provided(
+    mocker: MagicMock,
+):
+  FeatureFlags.TELEMETRY_ENABLED = True
+  mocker.patch(
+      'xpk.commands.cluster.get_capacity_type',
+      return_value=(CapacityType.SPOT, 0),
+  )
+  _log_cluster_create_telemetry(construct_args(tpu_type='test-tpu-type'))
+  event = json.loads(MetricsCollector.flush())['log_event'][0]
+  payload = json.loads(event['source_extension_json'])
+  event_metadata = payload['event_metadata']
+  assert payload['event_name'] == 'cluster_create'
+  assert (
+      _get_event_metadata_value_by_key(
+          event_metadata,
+          'XPK_ZONE',
+      )
+      == 'us-central1-a'
+  )
+  assert (
+      _get_event_metadata_value_by_key(
+          event_metadata,
+          'XPK_SYSTEM_CHARACTERISTICS',
+      )
+      == 'test-tpu-type'
   )
+  assert (
+      _get_event_metadata_value_by_key(
+          event_metadata,
+          'XPK_PROVISIONING_MODE',
+      )
+      == 'spot'
+  )
+
+
+def _get_event_metadata_value_by_key(
+    event_metadata: list[dict[str, str]], key: str
+) -> str | None:
+  return next(
+      (meta['value'] for meta in event_metadata if meta['key'] == key),
+      None,
+  )
+
+
+@pytest.mark.parametrize(
+    'gke_version_arg, expected_channel, expected_version',
+    [
+        (None, ReleaseChannel.RAPID, '1.2.4'),  # No version, should use RAPID
+        (
+            '1.2.3',
+            ReleaseChannel.REGULAR,
+            '1.2.3',
+        ),  # Version provided, should use REGULAR
+    ],
+)
+def test_cluster_create_calls_run_command_with_correct_channel_and_version(
+    gke_version_arg,
+    expected_channel,
+    expected_version,
+    mocks: _Mocks,
+    cluster_create_mocks: _ClusterCreateMocks,
+):
+  """
+  Verifies that cluster_create calls run_gke_cluster_create_command with the correct
+  release channel and GKE version based on whether a version is provided.
+  """
+  cluster_create_mocks.get_gke_control_plane_version.return_value = (
+      0,
+      expected_version,
+  )
+
+  args = construct_args(gke_version=gke_version_arg)
+  cluster_create(args)
+
+  expected_command_parts = [
+      'clusters create',
+      f'--cluster-version={expected_version}',
+      f'--release-channel={expected_channel.value.lower()}',
+  ]
+
+  mocks.commands_tester.assert_command_run(*expected_command_parts)

xpk/commands/config.py

@@ -14,16 +14,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.config import XpkConfig
+from ..core.config import xpk_config
 from ..utils.console import xpk_print
 
-xpk_cfg = XpkConfig()
-
 
 def set_config(args):
-  xpk_cfg.set(args.set_config_args[0], args.set_config_args[1])
+  xpk_config.set(args.set_config_args[0], args.set_config_args[1])
 
 
 def get_config(args):
-  value = xpk_cfg.get(args.get_config_key[0])
+  value = xpk_config.get(args.get_config_key[0])
   xpk_print(value)

xpk/commands/kind.py

@@ -99,7 +99,7 @@ def cluster_create(args) -> None:
       supports_sub_slicing=False,
   )
 
-  kueue_manager = KueueManager()
+  kueue_manager = KueueManager(project='', zone='')
   kueue_manager.install_or_upgrade(
       KueueConfig(
           system,

xpk/commands/storage.py

@@ -56,7 +56,7 @@ from ..core.storage import (
     list_storages,
     print_storages_for_cluster,
 )
-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.kubectl import apply_kubectl_manifest
 from ..utils.execution_context import is_dry_run
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -133,15 +133,13 @@ def storage_delete(args: Namespace) -> None:
       if storage.bucket.startswith(filestore_instance_name)
   ]
 
-  if children and not args.force:
-    detach = get_user_input(
-        "Deleting a filestore storage will destroy your filestore instance and"
-        " all its data in all volumes will be lost. Do you wish to delete the"
-        f" filestore instance {filestore_instance_name}?\n y (yes) / n (no):\n'"
-    )
-    if not detach:
-      xpk_print("Deleting storage canceled.")
-      xpk_exit(0)
+  if children and not ask_for_user_consent(
+      "Deleting a filestore storage will destroy your filestore instance and"
+      " all its data in all volumes will be lost. Do you wish to delete the"
+      f" filestore instance {filestore_instance_name}?"
+  ):
+    xpk_print("Deleting storage canceled.")
+    xpk_exit(0)
 
   for child in children:
     delete_storage_resources(k8s_api_client, child)

xpk/commands/workload.py

@@ -27,15 +27,14 @@ from ..core.cluster import (
     setup_k8s_env,
 )
 from ..core.commands import run_command_with_updates, run_commands
-from ..core.kueue_manager import KueueManager, has_sub_slicing_enabled
 from ..core.config import (VERTEX_TENSORBOARD_FEATURE_FLAG, XPK_CURRENT_VERSION)
 from ..core.docker_container import (
     get_main_container_docker_image,
     get_user_workload_container,
 )
+from ..core.kueue_manager import has_sub_slicing_enabled, get_installed_kueue_version, LOCAL_QUEUE_NAME
 from ..core.docker_resources import get_volumes, parse_env_config
 from ..core.gcloud_context import add_zone_and_project
-from ..core.kueue_manager import LOCAL_QUEUE_NAME
 from ..core.monitoring import get_gke_outlier_dashboard
 from ..core.nap import (
     get_autoprovisioning_node_selector_args,
@@ -64,6 +63,8 @@ from ..core.scheduling import (
     get_cpu_affinity,
     get_gpu_scheduler,
     create_sub_slicing_annotations,
+    create_placement_policy_label,
+    is_placement_policy_supported,
 )
 from ..core.storage import (
     GCE_PD_TYPE,
@@ -77,6 +78,7 @@ from ..core.storage import (
     get_storages_to_mount,
 )
 from ..core.system_characteristics import (
+    SUB_SLICING_TOPOLOGIES,
     AcceleratorType,
     get_system_characteristics,
     compute_vms_per_slice,
@@ -95,7 +97,7 @@ from ..core.workload_decorators import (
     tcpx_decorator,
     tcpxo_decorator,
 )
-from ..utils.console import get_user_input, xpk_exit, xpk_print
+from ..utils.console import ask_for_user_consent, xpk_exit, xpk_print
 from packaging.version import Version
 from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
@@ -144,6 +146,7 @@ spec:
               nodeSelector:
                 {accelerator_label}
                 {machine_label}
+                {placement_policy_label}
                 {autoprovisioning_args}
               priorityClassName: {args.priority}
               hostNetwork: true
@@ -193,6 +196,8 @@ spec:
               {gpu_scheduler}
               priorityClassName: {args.priority}
               restartPolicy: Never
+              nodeSelector:
+                {placement_policy_label}
               imagePullSecrets:
               - name: {args.docker_image_pull_secret}
               hostNetwork: true
@@ -238,6 +243,8 @@ spec:
             spec:
               priorityClassName: {args.priority}
               restartPolicy: Never
+              nodeSelector:
+                {placement_policy_label}
               imagePullSecrets:
               - name: {args.docker_image_pull_secret}
               dnsPolicy: ClusterFirstWithHostNet
@@ -273,6 +280,7 @@ PW_WORKLOAD_CREATE_YAML = """
         terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
         priorityClassName: {args.priority}
         nodeSelector:
+          {placement_policy_label}
           {autoprovisioning_args}
       pathwaysDir: {args.pathways_gcs_location} #This bucket needs to be created in advance.
       controller:
@@ -284,7 +292,6 @@ PW_WORKLOAD_CREATE_YAML = """
       {user_workload}
 """
 
-SUB_SLICING_TOPOLOGIES = ['2x2', '2x4', '4x4', '4x8', '8x8', '8x16', '16x16']
 SUB_SLICING_MINIMUM_KUEUE_VERSION = Version('0.13.0')
 
 
@@ -486,6 +493,12 @@ def workload_create(args) -> None:
                 operator: NotIn
                 values: [{restart_on_exit_codes}]"""
 
+  placement_policy_label = (
+      create_placement_policy_label(system)
+      if is_placement_policy_supported(system)
+      else ''
+  )
+
   # Create the workload file based on accelerator type or workload type.
   if system.accelerator_type == AcceleratorType.GPU:
     container, debugging_dashboard_id = get_user_workload_container(
@@ -519,6 +532,7 @@ def workload_create(args) -> None:
           failure_policy_rules=failure_policy_rules,
           pod_failure_policy=pod_failure_policy,
           annotations=annotations,
+          placement_policy_label=placement_policy_label,
       )
 
       sub_networks = get_cluster_subnetworks()
@@ -543,6 +557,7 @@ def workload_create(args) -> None:
           service_account=service_account,
           failure_policy_rules=failure_policy_rules,
           pod_failure_policy=pod_failure_policy,
+          placement_policy_label=placement_policy_label,
       )
 
   elif args.use_pathways and ensure_pathways_workload_prerequisites(
@@ -560,6 +575,7 @@ def workload_create(args) -> None:
         user_workload=get_user_workload_for_pathways(args, system),
         local_queue_name=LOCAL_QUEUE_NAME,
         autoprovisioning_args=autoprovisioning_args,
+        placement_policy_label=placement_policy_label,
     )
   else:
     container, debugging_dashboard_id = get_user_workload_container(
@@ -587,6 +603,7 @@ def workload_create(args) -> None:
                 create_sub_slicing_annotations(args.sub_slicing_topology)
             )
         ),
+        placement_policy_label=placement_policy_label,
         machine_label=create_machine_label(system.accelerator_type, system),
         local_queue_name=LOCAL_QUEUE_NAME,
         autoprovisioning_args=autoprovisioning_args,
@@ -697,9 +714,10 @@ def _validate_sub_slicing_availability():
     )
     xpk_exit(1)
 
-  kueue_manager = KueueManager()
-  return_code, current_version = kueue_manager.get_installed_kueue_version()
-  if return_code != 0:
+  return_code, current_version = get_installed_kueue_version(
+      dry_run_version=Version('0.13')
+  )
+  if return_code != 0 or not current_version:
     xpk_print(
         'Error: Unable to validate sub-slicing support on a given cluster.'
     )
@@ -784,11 +802,10 @@ def workload_delete(args) -> None:
       xpk_exit(return_code)
     # Skip the header
     workloads = [x.split(' ')[0] for x in return_value.splitlines()][1:]
-    if workloads and not args.force:
-      will_delete = get_user_input(
+    if workloads:
+      will_delete = 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?'
       )
   else:
     workloads = [args.workload]

xpk/commands/workload_test.py

@@ -107,7 +107,7 @@ def test_validate_sub_slicing_availability_exits_when_kueue_version_cannot_be_de
       return_value=(0, True),
   )
   mocker.patch(
-      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      'xpk.commands.workload.get_installed_kueue_version',
       return_value=(1, None),
   )
   with pytest.raises(SystemExit):
@@ -124,7 +124,7 @@ def test_validate_sub_slicing_availability_exits_when_kueue_version_does_not_mee
       return_value=(0, True),
   )
   mocker.patch(
-      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      'xpk.commands.workload.get_installed_kueue_version',
       return_value=(0, Version('0.0.0')),
   )
   with pytest.raises(SystemExit):
@@ -141,7 +141,7 @@ def test_validate_sub_slicing_availability_does_nothing_when_cluster_is_correctl
       return_value=(0, True),
   )
   mocker.patch(
-      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      'xpk.commands.workload.get_installed_kueue_version',
       return_value=(0, Version('0.13.0')),
   )
   _validate_sub_slicing_availability()