xpk 0.14.4__py3-none-any.whl → 0.16.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
@@ -31,9 +35,34 @@ class _Mocks:
   common_print_mock: MagicMock
   commands_print_mock: MagicMock
   commands_get_reservation_deployment_type: MagicMock
+  commands_get_pathways_machine_types: MagicMock
   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(
@@ -47,15 +76,23 @@ def mocks(mocker) -> _Mocks:
       'xpk.commands.cluster.get_reservation_deployment_type',
       return_value='DENSE',
   )
+  commands_get_pathways_machine_types = mocker.patch(
+      'xpk.commands.cluster.get_pathways_machine_types',
+      return_value=(0, []),
+  )
   return _Mocks(
       common_print_mock=common_print_mock,
       commands_get_reservation_deployment_type=commands_get_reservation_deployment_type,
       commands_print_mock=commands_print_mock,
+      commands_get_pathways_machine_types=commands_get_pathways_machine_types,
       commands_tester=CommandsTester(
           mocker,
           run_command_with_updates_path=(
               'xpk.commands.cluster.run_command_with_updates'
           ),
+          run_command_for_value_path=(
+              'xpk.commands.cluster.run_command_for_value'
+          ),
       ),
   )
 
@@ -65,6 +102,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',
@@ -72,6 +113,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       gke_version='',
       private=False,
       authorized_networks=None,
+      pathways_gce_machine_type='n2-standard-64',
       enable_pathways=False,
       enable_ray_cluster=False,
       enable_workload_identity=False,
@@ -87,11 +129,97 @@ def construct_args(**kwargs: Any) -> Namespace:
       memory_limit='100Gi',
       cpu_limit=100,
       cluster_cpu_machine_type='',
+      create_vertex_tensorboard=False,
+      enable_autoprovisioning=False,
+      sub_slicing_topology='2x2x2',
+      use_vertex_tensorboard=False,
+      env_file='',
+      env=None,
+      use_pathways=False,
+      debug_dump_gcs=False,
+      storage='',
+      restart_on_exit_codes=None,
+      ttl_seconds_after_finished=0,
+      max_restarts=1,
+      priority=0,
+      termination_grace_period_seconds=0,
+      docker_image_pull_secret='',
+      managed_mldiagnostics=False,
+      output_manifest_file='',
   )
   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 +234,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)
 
@@ -185,6 +313,64 @@ def test_validate_cluster_create_args_for_invalid_reservation(
   )
 
 
+def test_validate_cluster_create_args_for_enable_pathways_set_to_false(
+    mocks: _Mocks,
+):
+  args = construct_args(enable_pathways=False)
+  mocks.commands_get_pathways_machine_types.return_value = (1, [])
+
+  _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 0
+
+
+def test_validate_cluster_create_args_for_errored_pathways_machine_types_retrieval(
+    mocks: _Mocks,
+):
+  args = construct_args(enable_pathways=True)
+  mocks.commands_get_pathways_machine_types.return_value = (1, [])
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 1
+  assert 'Unable to retrieve' in mocks.commands_print_mock.call_args[0][0]
+
+
+def test_validate_cluster_create_args_for_invalid_pathways_machine_type(
+    mocks: _Mocks,
+):
+  args = construct_args(
+      enable_pathways=True, pathways_gce_machine_type='n2-standard-32'
+  )
+  mocks.commands_get_pathways_machine_types.return_value = (
+      0,
+      ['n2-standard-64'],
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 2
+  assert 'Available machine types' in mocks.commands_print_mock.call_args[0][0]
+
+
+def test_validate_cluster_create_args_for_valid_pathways_machine_type(
+    mocks: _Mocks,
+):
+  args = construct_args(
+      enable_pathways=True, pathways_gce_machine_type='n2-standard-32'
+  )
+  mocks.commands_get_pathways_machine_types.return_value = (
+      0,
+      ['n2-standard-32'],
+  )
+
+  _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 0
+
+
 @patch('xpk.commands.cluster.KueueManager.install_or_upgrade')
 def test_install_kueue_returns_kueue_installation_code(
     mock_kueue_manager_install: MagicMock,
@@ -209,6 +395,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 +413,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 +432,179 @@ 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(
+def test_run_gke_cluster_create_command_with_lustre_runs_correct_command(
     mocks: _Mocks,
 ):
   result = run_gke_cluster_create_command(
-      args=construct_args(gke_version=''),
+      args=construct_args(
+          enable_lustre_csi_driver=True, enable_legacy_lustre_port=False
+      ),
       gke_control_plane_version='1.2.3',
-      system=GPU_TEST_SYSTEM,
+      system=TPU_TEST_SYSTEM,
+      release_channel=ReleaseChannel.REGULAR,
+  )
+
+  assert result == 0
+  commands = mocks.commands_tester.get_matching_commands('clusters create')
+  assert len(commands) == 1
+  command = commands[0]
+  assert (
+      '--addons=LustreCsiDriver' in command
+      and '--enable-legacy-lustre-port' not in command
+  )
+
+
+def test_run_gke_cluster_create_command_with_lustre_legacy_port_adds_correct_flag(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(
+          enable_lustre_csi_driver=True, enable_legacy_lustre_port=True
+      ),
+      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',
+      '--enable-legacy-lustre-port',
+      '--addons=LustreCsiDriver',
+  )
+
+
+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,
+):
+  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'
+  )
+
+
+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 get_config as get_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])
+  get_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 = get_xpk_config().get(args.get_config_key[0])
   xpk_print(value)

xpk/commands/inspector.py

@@ -18,7 +18,7 @@ from ..core.cluster import get_cluster_credentials
 from ..core.commands import run_command_for_value
 from ..core.gcloud_context import add_zone_and_project, get_cluster_location
 from ..core.kueue_manager import CLUSTER_QUEUE_NAME, LOCAL_QUEUE_NAME
-from ..core.resources import CLUSTER_METADATA_CONFIGMAP, CLUSTER_RESOURCES_CONFIGMAP
+from ..core.resources import ConfigMapType, get_config_map_name
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import append_tmp_file, write_tmp_file
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -162,14 +162,16 @@ def inspector(args) -> None:
       (
           (
               'kubectl get configmap'
-              f' {args.cluster}-{CLUSTER_METADATA_CONFIGMAP} -o yaml'
+              f' {get_config_map_name(args.cluster, ConfigMapType.METADATA)} -o'
+              ' yaml'
           ),
           'GKE: Cluster Metadata ConfigMap Details',
       ),
       (
           (
               'kubectl get configmap'
-              f' {args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP} -o yaml'
+              f' {get_config_map_name(args.cluster, ConfigMapType.RESOURCES)} -o'
+              ' yaml'
           ),
           'GKE: Cluster Resources ConfigMap Details',
       ),

xpk/commands/kind.py

@@ -30,6 +30,7 @@ from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
     SystemCharacteristics,
     AcceleratorType,
+    DockerPlatform,
 )
 from ..utils.console import (xpk_exit, xpk_print)
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -97,9 +98,10 @@ def cluster_create(args) -> None:
       AcceleratorType.CPU,
       'kind',
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
   )
 
-  kueue_manager = KueueManager()
+  kueue_manager = KueueManager(project='', zone='')
   kueue_manager.install_or_upgrade(
       KueueConfig(
           system,