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

xpk/core/capacity.py

@@ -29,6 +29,8 @@ H100_DEVICE_TYPE = 'h100-80gb-8'
 H100_MEGA_DEVICE_TYPE = 'h100-mega-80gb-8'
 H200_DEVICE_TYPE = 'h200-141gb-8'
 B200_DEVICE_TYPE = 'b200-8'
+GB200_DEVICE_TYPE = 'gb200-4'
+GB200_DEVICE_TYPE_NOLSSD = 'gb200-4-no-ssd'
 RESERVATION_CONFIG_KEY = 'reservation_id'
 
 
@@ -115,9 +117,12 @@ def get_reservation_maintenance_interval(
   Returns:
     0 if successful and 1 otherwise.
   """
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      reservation, project
+  )
   command = (
-      f'gcloud beta compute reservations describe {reservation}'
-      f' --project={project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
   )
   return_code, output = run_command_for_value(
       command, 'Get reservation maintenance interval'
@@ -139,9 +144,12 @@ def get_reservation_placement_policy(
   Returns:
     0 if successful and 1 otherwise.
   """
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      reservation, project
+  )
   command = (
-      f'gcloud beta compute reservations describe {reservation}'
-      f' --project={project} --zone={zone} --format="value(resourcePolicies.policy)"'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={zone} --format="value(resourcePolicies.policy)"'
   )
   return_code, output = run_command_for_value(
       command, 'Get reservation placement policy'
@@ -156,9 +164,12 @@ def get_reservation_deployment_type(
     reservation: str, zone: str, project: str
 ) -> str:
   """Get reservation deployment type."""
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      reservation, project
+  )
   command = (
-      f'gcloud beta compute reservations describe {reservation}'
-      f' --project={project} --zone={zone} --format="value(deploymentType)"'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={zone} --format="value(deploymentType)"'
   )
   return_code, output = run_command_for_value(
       command, 'Get reservation deployment type', dry_run_return_val='DENSE'
@@ -178,9 +189,12 @@ def verify_reservation_exists(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      args.reservation, args.project
+  )
   command = (
-      f'gcloud beta compute reservations describe {args.reservation}'
-      f' --project={args.project} --zone={args.zone}'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={args.zone}'
   )
   return_code = run_command_with_updates(command, 'Describe reservation')
   if return_code != 0:
@@ -264,3 +278,29 @@ def get_capacity_node_selectors_from_capacity_type(
       )
       return_code = 1
   return node_selector, return_code
+
+
+def get_reservation_project_and_name(
+    reservation_name_or_path: str, cluster_project: str
+) -> tuple[str, str]:
+  """Get the reservation project and name.
+
+  Args:
+    reservation_name_or_path: either reservation name or reservation path in format
+      projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME
+    cluster_project: the cluster project
+
+  Returns:
+    Tuple with reservation project and reservation name.
+  """
+  if '/' not in reservation_name_or_path:
+    return cluster_project, reservation_name_or_path
+  reservation_parts = reservation_name_or_path.split('/')
+  if (
+      len(reservation_parts) != 4
+      or reservation_parts[0] != 'projects'
+      or reservation_parts[2] != 'reservations'
+  ):
+    xpk_print('Unable to parse reservation: ', reservation_name_or_path)
+    xpk_exit(1)
+  return reservation_parts[1], reservation_parts[3]

xpk/core/capacity_test.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 import pytest
 from unittest.mock import MagicMock, patch
-from .capacity import get_reservation_deployment_type
+from .capacity import get_reservation_deployment_type, get_reservation_project_and_name
 
 
 @patch('xpk.core.capacity.xpk_print')
@@ -48,3 +48,34 @@ def test_get_reservation_deployment_type_returns_deployment_type_when_command_su
       reservation='reservation', zone='zone', project='project'
   )
   assert result == 'DENSE'
+
+
+def test_get_reservation_project_and_name_parses_local_reservation():
+  project, name = get_reservation_project_and_name(
+      'test-reservation', 'cluster-project'
+  )
+
+  assert project == 'cluster-project'
+  assert name == 'test-reservation'
+
+
+def test_get_reservation_project_and_name_parses_shared_reservation():
+  project, name = get_reservation_project_and_name(
+      'projects/reservation-project/reservations/test-reservation',
+      'cluster-project',
+  )
+
+  assert project == 'reservation-project'
+  assert name == 'test-reservation'
+
+
+@patch('xpk.core.capacity.xpk_print')
+def test_get_reservation_project_and_name_fails_for_invalid_reservation(
+    xpk_print: MagicMock, mocker
+):
+  with pytest.raises(SystemExit):
+    get_reservation_project_and_name(
+        'invalid/reservation',
+        'cluster-project',
+    )
+  assert 'Unable to parse reservation' in xpk_print.mock_calls[0].args[0]

xpk/core/cluster.py

@@ -22,7 +22,7 @@ from kubernetes import config
 from kubernetes.client.exceptions import ApiException
 
 from ..utils.console import xpk_exit, xpk_print
-from .capacity import B200_DEVICE_TYPE, H100_DEVICE_TYPE, H200_DEVICE_TYPE
+from .capacity import H200_DEVICE_TYPE
 from .commands import (
     run_command_for_value,
     run_command_with_updates,
@@ -34,16 +34,11 @@ from .gcloud_context import (
     zone_to_region,
 )
 from .resources import get_cluster_system_characteristics
-from .system_characteristics import SystemCharacteristics
+from .system_characteristics import INSTALLER_NCCL_TCPXO, SystemCharacteristics
 
 JOBSET_VERSION = 'v0.8.0'
 PATHWAYS_JOB_VERSION = 'v0.1.4'
-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'
-CONFIG_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-config.yaml'
 NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml'
-MGLRU_DISABLE = 'https://raw.githubusercontent.com/GoogleCloudPlatform/cluster-toolkit/refs/heads/main/examples/gke-a3-ultragpu/mglru-disable.yaml'
 
 DEFAULT_NAMESPACE = 'default'
 XPK_SA = 'xpk-sa'
@@ -118,12 +113,12 @@ def install_nccl_on_cluster(system: SystemCharacteristics) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
-  if system.device_type == H100_DEVICE_TYPE:
-    command = f'kubectl apply -f {INSTALLER_NCCL_TCPX}'
-  elif system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
-    command = f'kubectl apply -f {INSTALLER_NCCL_RDMA}'
-  else:
-    command = f'kubectl apply -f {INSTALLER_NCCL_TCPXO}'
+  nccl_installer = (
+      system.gpu_config.nccl_installer
+      if system.gpu_config and system.gpu_config.nccl_installer
+      else INSTALLER_NCCL_TCPXO
+  )
+  command = f'kubectl apply -f {nccl_installer}'
 
   return_code = run_command_with_updates(
       command, 'Install NCCL Plugin On Cluster'
@@ -135,35 +130,6 @@ def install_nccl_on_cluster(system: SystemCharacteristics) -> int:
     )
     return 1
 
-  if system.device_type == H100_DEVICE_TYPE:
-    command = f'kubectl apply -f {CONFIG_NCCL_TCPX}'
-
-    return_code = run_command_with_updates(
-        command, 'Install NCCL Config On Cluster'
-    )
-
-    if return_code != 0:
-      xpk_print(
-          f'Install NCCL Config On Cluster request returned ERROR {return_code}'
-      )
-      return 1
-
-  return 0
-
-
-def disable_mglru_on_cluster() -> int:
-  """Disable MGLRU on the cluster.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = f'kubectl apply -f {MGLRU_DISABLE}'
-  return_code = run_command_with_updates(command, 'Disable MGLRU On Cluster')
-
-  if return_code != 0:
-    xpk_print('Disablig MGLRU On Cluster request returned ERROR')
-    return 1
-
   return 0
 
 
@@ -309,10 +275,11 @@ def update_cluster_with_lustre_driver_if_necessary(args) -> int:
   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'
+  if is_driver_enabled_on_cluster(args, driver='lustreCsiDriver') and (
+      not args.enable_legacy_lustre_port
+      or is_driver_enabled_on_cluster(
+          args, driver='lustreCsiDriver', config_key='enableLegacyLustrePort'
+      )
   ):
     return 0
   cluster_update_return_code = update_gke_cluster_with_lustre_driver_enabled(
@@ -621,9 +588,13 @@ def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --enable-legacy-lustre-port'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       ' --quiet'
   )
+  if args.enable_legacy_lustre_port:
+    command += ' --enable-legacy-lustre-port'
+  else:
+    command += ' --update-addons=LustreCsiDriver=ENABLED'
   xpk_print(
       'Updating GKE cluster to enable Lustre CSI driver, may take a while!'
   )
@@ -729,78 +700,58 @@ def update_cluster_with_gcsfuse_driver_if_necessary(args) -> int:
   return 0
 
 
-def test_and_retry_credentials_with_dns_logic(args) -> int:
-  """Tests kubectl credentials and retries with default settings if a DNS error is found.
+def get_cluster_credentials(args) -> int:
+  """Run cluster configuration command to set the kubectl config.
 
   Args:
     args: user provided arguments for running the command.
 
   Returns:
-    0 if credentials are valid after retrying, 1 otherwise.
+    0 if successful and 1 otherwise.
   """
-
-  xpk_print('Testing credentials with kubectl...')
-  kubectl_command = 'kubectl get pods'
-  kubectl_return_code, kubectl_output = run_command_for_value(
-      kubectl_command, 'kubectl get pods'
-  )
-  if kubectl_return_code == 0:
-    xpk_print('Credentials test succeeded.')
-    return 0
-
-  dns_endpoint_error = (
-      'control_plane_endpoints_config.dns_endpoint_config.allow_external_traffic'
-      ' is disabled'
-  )
-  if dns_endpoint_error not in kubectl_output:
-    xpk_print(f'kubectl failed with an unhandled error: {kubectl_output}')
-    xpk_exit(kubectl_return_code)
-  xpk_print(
-      'Detected DNS endpoint-related error. Retrying without --dns-endpoint'
-      ' flag...'
-  )
-
   location = get_cluster_location(args.project, args.cluster, args.zone)
-  without_dns_command = (
-      'gcloud container clusters get-credentials'
-      f' {args.cluster} --location={location}'
-      f' --project={args.project} &&'
-      ' kubectl config view && kubectl config set-context --current'
-      ' --namespace=default'
-  )
-  return_code = run_command_with_updates(
-      without_dns_command, 'get-credentials to cluster', verbose=False
+
+  return_code = _get_credentials(
+      project=args.project,
+      cluster=args.cluster,
+      location=location,
+      dns_endpoint=True,
   )
   if return_code != 0:
-    xpk_print('Failed to get credentials even without --dns-endpoint. Exiting.')
-    xpk_exit(return_code)
-  return 0
-
+    return return_code
+
+  if not _are_credentials_valid():
+    xpk_print('Detected error. Retrying without --dns-endpoint flag...')
+    return_code = _get_credentials(
+        project=args.project,
+        cluster=args.cluster,
+        location=location,
+        dns_endpoint=False,
+    )
+    if return_code != 0:
+      return return_code
 
-def get_cluster_credentials(args) -> int:
-  """Run cluster configuration command to set the kubectl config.
+  xpk_print('Finished get-credentials and kubectl setup.')
+  return 0
 
-  Args:
-    args: user provided arguments for running the command.
 
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  location = get_cluster_location(args.project, args.cluster, args.zone)
+def _get_credentials(
+    project: str, cluster: str, location: str, dns_endpoint: bool
+) -> int:
+  dns_endpoint_arg = '--dns-endpoint' if dns_endpoint else ''
   command = (
       'gcloud container clusters get-credentials'
-      f' {args.cluster} --location={location} --dns-endpoint'
-      f' --project={args.project} && kubectl config view && kubectl config'
+      f' {cluster} --location={location} {dns_endpoint_arg}'
+      f' --project={project} && kubectl config view && kubectl config'
       ' set-context --current --namespace=default'
   )
-  task = f'get-credentials-dns-endpoint to cluster {args.cluster}'
-  return_code = run_command_with_updates_retry(command, task, verbose=False)
+  task = f'get-credentials-dns-endpoint to cluster {cluster}'
+  return run_command_with_updates(command, task, verbose=False)
 
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-    xpk_exit(return_code)
-
-  return_code = test_and_retry_credentials_with_dns_logic(args)
-  xpk_print('Finished get-credentials and kubectl setup.')
 
-  return return_code
+def _are_credentials_valid() -> bool:
+  kubectl_command = 'kubectl get pods'
+  kubectl_return_code = run_command_with_updates(
+      kubectl_command, 'Test kubectl credentials'
+  )
+  return kubectl_return_code == 0

xpk/core/cluster_test.py

@@ -0,0 +1,170 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import pytest
+from .testing.commands_tester import CommandsTester
+from .cluster import get_cluster_credentials, update_gke_cluster_with_lustre_driver_enabled, update_cluster_with_lustre_driver_if_necessary
+from pytest_mock import MockerFixture
+
+
+@pytest.fixture(autouse=True)
+def commands_tester(mocker: MockerFixture) -> CommandsTester:
+  return CommandsTester(
+      mocker=mocker,
+      run_command_for_value_path="xpk.core.cluster.run_command_for_value",
+      run_command_with_updates_path="xpk.core.cluster.run_command_with_updates",
+  )
+
+
+@pytest.fixture(autouse=True)
+def mock_location(mocker: MockerFixture):
+  mocker.patch(
+      "xpk.core.cluster.get_cluster_location", return_value="us-central1"
+  )
+
+
+@pytest.fixture(autouse=True)
+def command_args(mocker: MockerFixture):
+  return mocker.Mock(cluster="cluster", project="project", zone="zone")
+
+
+def test_get_cluster_credentials_returns_1_when_retrieval_command_fails(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (1, ""), "gcloud container clusters get-credentials"
+  )
+  assert get_cluster_credentials(command_args) == 1
+
+
+def test_get_cluster_credentials_returns_0_when_retrieval_succeeds(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters get-credentials"
+  )
+  assert get_cluster_credentials(command_args) == 0
+
+
+def test_get_cluster_credentials_does_not_retry_with_dns_when_retrieval_succeeds(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters get-credentials --dns-endpoint"
+  )
+  commands_tester.set_result_for_command((0, ""), "kubectl get pods")
+  get_cluster_credentials(command_args)
+  non_dns_endpoint_commands = [
+      c
+      for c in commands_tester.get_matching_commands(
+          "gcloud container clusters get-credentials"
+      )
+      if "dns-endpoint" not in c
+  ]
+  assert len(non_dns_endpoint_commands) == 0
+
+
+def test_get_cluster_credentials_retries_without_dns_when_dns_retrieval_fails(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters get-credentials --dns-endpoint"
+  )
+  commands_tester.set_result_for_command((1, ""), "kubectl get pods")
+  get_cluster_credentials(command_args)
+  non_dns_endpoint_commands = [
+      c
+      for c in commands_tester.get_matching_commands(
+          "gcloud container clusters get-credentials"
+      )
+      if "dns-endpoint" not in c
+  ]
+  assert len(non_dns_endpoint_commands) == 1
+
+
+def test_update_cluster_with_lustre_driver_if_necessary_with_default_port_runs_correct_checks(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, "True"),
+      "gcloud container clusters describe",
+  )
+  command_args.enable_legacy_lustre_port = None
+  update_cluster_with_lustre_driver_if_necessary(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      "gcloud container clusters describe cluster --project=project"
+      " --location=us-central1"
+      ' --format="value(addonsConfig.lustreCsiDriverConfig.enabled)"'
+  ]
+
+
+def test_update_cluster_with_lustre_driver_if_necessary_with_legacy_port_runs_correct_checks(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, "True"),
+      "gcloud container clusters describe",
+  )
+  command_args.enable_legacy_lustre_port = True
+  update_cluster_with_lustre_driver_if_necessary(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      (
+          "gcloud container clusters describe cluster --project=project"
+          " --location=us-central1"
+          ' --format="value(addonsConfig.lustreCsiDriverConfig.enabled)"'
+      ),
+      (
+          "gcloud container clusters describe cluster --project=project"
+          " --location=us-central1"
+          ' --format="value(addonsConfig.lustreCsiDriverConfig.enableLegacyLustrePort)"'
+      ),
+  ]
+
+
+def test_update_gke_cluster_with_lustre_driver_enabled_default_port(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters update"
+  )
+  command_args.enable_legacy_lustre_port = None
+  update_gke_cluster_with_lustre_driver_enabled(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      "gcloud container clusters update cluster --project=project"
+      " --location=us-central1 --quiet --update-addons=LustreCsiDriver=ENABLED"
+  ]
+
+
+def test_update_gke_cluster_with_lustre_driver_enabled_legacy_port(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters update"
+  )
+  command_args.enable_legacy_lustre_port = True
+  update_gke_cluster_with_lustre_driver_enabled(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      "gcloud container clusters update cluster --project=project"
+      " --location=us-central1 --quiet --enable-legacy-lustre-port"
+  ]

xpk/core/commands.py

@@ -195,16 +195,13 @@ def run_command_with_updates(command, task, verbose=True) -> int:
         return_code = child.poll()
         if return_code is None:
           xpk_print(f'Waiting for `{task}`, for {i} seconds...', end='\r')
-          time.sleep(1)
-          i += 1
+          time.sleep(10)
+          i += 10
         else:
           xpk_print(f'Task: `{task}` terminated with code `{return_code}`')
           return return_code
   else:
-    xpk_print(
-        f'Task: `{task}` is implemented by `{command}`, hiding output unless'
-        ' there is an error.'
-    )
+    xpk_print(f'Task: `{task}` is implemented by `{command}`')
     try:
       subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT)
     except subprocess.CalledProcessError as e:
@@ -277,10 +274,7 @@ def run_command_for_value(
           return return_code, f'{out_str}\n{err_str}'
   else:
     if not quiet:
-      xpk_print(
-          f'Task: `{task}` is implemented by `{command}`, hiding output unless'
-          ' there is an error.'
-      )
+      xpk_print(f'Task: `{task}` is implemented by `{command}`')
     try:
       output = subprocess.check_output(
           command,