xpk 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

xpk/commands/cluster.py

@@ -19,7 +19,6 @@ from tabulate import tabulate
 from ..utils.feature_flags import FeatureFlags
 from ..utils.versions import ReleaseChannel
 from ..core.pathways import get_pathways_machine_types
-from ..core.capacity import H100_DEVICE_TYPE, get_reservation_deployment_type, parse_reservation
 from ..core.cluster import (
     get_all_clusters_programmatic,
     get_cluster_credentials,
@@ -40,7 +39,12 @@ 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.capacity import (
+    H100_DEVICE_TYPE,
+    get_capacity_type,
+    get_reservations_list,
+    get_reservation_deployment_type,
+)
 from ..core.gcloud_context import (
     add_zone_and_project,
     get_gke_control_plane_version,
@@ -240,13 +244,6 @@ def _validate_sub_slicing_reservation(args):
 
 def _validate_super_slicing_reservation(args):
   _validate_gsc_reservation(args, 'Super-slicing')
-  reservation = parse_reservation(args.reservation, args.project)
-  if reservation.block_name is None:
-    xpk_print(
-        'Error: Validation failed: Super-slicing cluster creation'
-        ' requires a block or sub-block reservation.'
-    )
-    xpk_exit(1)
 
 
 def _validate_gsc_reservation(args, creation_description: str):
@@ -257,27 +254,29 @@ def _validate_gsc_reservation(args, creation_description: str):
     )
     xpk_exit(1)
 
-  deployment_type = get_reservation_deployment_type(
-      reservation_path=args.reservation, project=args.project, zone=args.zone
-  )
-  if deployment_type != 'DENSE':
-    xpk_print(
-        'Error: Validation failed: The specified reservation'
-        f' "{args.reservation}" is not a Cluster Director reservation.'
-    )
-    xpk_print(
-        'Please provide a reservation created for Cluster Director to proceed.'
-    )
-    xpk_print('To list valid Cluster Director reservations, run:')
-    xpk_print(
-        '  gcloud compute reservations list --filter="deploymentType=DENSE"'
-    )
-    xpk_print(
-        'Refer to the documentation for more information on creating Cluster'
-        ' Director reservations:'
-        ' https://cloud.google.com/cluster-director/docs/reserve-capacity'
+  for reservation in get_reservations_list(args):
+    deployment_type = get_reservation_deployment_type(
+        reservation_path=reservation, project=args.project, zone=args.zone
     )
-    xpk_exit(1)
+    if deployment_type != 'DENSE':
+      xpk_print(
+          'Error: Validation failed: The specified reservation'
+          f' "{reservation}" is not a Cluster Director reservation.'
+      )
+      xpk_print(
+          'Please provide a reservation created for Cluster Director to'
+          ' proceed.'
+      )
+      xpk_print('To list valid Cluster Director reservations, run:')
+      xpk_print(
+          '  gcloud compute reservations list --filter="deploymentType=DENSE"'
+      )
+      xpk_print(
+          'Refer to the documentation for more information on creating Cluster'
+          ' Director reservations:'
+          ' https://cloud.google.com/cluster-director/docs/reserve-capacity'
+      )
+      xpk_exit(1)
 
 
 def _validate_num_slices_and_set_default(args):
@@ -372,7 +371,7 @@ def cluster_create(args) -> None:
 
   update_coredns_command_code = update_coredns_if_necessary()
   if update_coredns_command_code != 0:
-    xpk_exit(update_cluster_command_code)
+    xpk_exit(update_coredns_command_code)
 
   if not is_dry_run():
     k8s_client = setup_k8s_env(args)

xpk/commands/cluster_gcluster.py

@@ -32,7 +32,7 @@ from ..core.blueprint.blueprint_generator import (
     a4_device_type,
     supported_device_types,
 )
-from ..core.capacity import get_capacity_type
+from ..core.capacity import get_capacity_type, get_reservations_list
 from ..core.cluster import get_cluster_credentials
 from ..core.commands import run_command_for_value
 from ..core.docker_manager import DockerManager
@@ -304,22 +304,29 @@ def generate_blueprint(
   if args.cluster_state_gcs_bucket is not None:
     validate_state_gcs_bucket(args)
 
+  num_nodes = 2 if args.num_nodes is None else args.num_nodes
+
+  reservations = get_reservations_list(args)
+  if len(reservations) > 1:
+    xpk_print(
+        'Error: Cluster Toolkit based clusters only support a single'
+        ' reservation.'
+    )
+    xpk_exit(1)
+  reservation = reservations[0] if len(reservations) > 0 else None
+
   if args.device_type in supported_device_types:
     if args.device_type == a3mega_device_type:
-      num_nodes = args.num_nodes if not args.num_nodes is None else 2
-
       maintenance_interval = (
           get_reservation_maintenance_interval(
-              args.reservation, args.zone, args.project
+              reservation, args.zone, args.project
           )
-          if args.reservation is not None
+          if reservation is not None
           else 'PERIODIC'
       )
       placement_policy_name = (
-          get_reservation_placement_policy(
-              args.reservation, args.zone, args.project
-          )
-          if args.reservation is not None
+          get_reservation_placement_policy(reservation, args.zone, args.project)
+          if reservation is not None
           else None
       )
       placement_policy = (
@@ -342,7 +349,7 @@ def generate_blueprint(
           num_nodes=num_nodes,
           reservation_maintenance_interval=maintenance_interval,
           reservation_placement_policy=placement_policy,
-          reservation=args.reservation if args.reservation else None,
+          reservation=reservation,
           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,
@@ -351,7 +358,6 @@ def generate_blueprint(
           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
       return bpg.generate_a3_ultra_blueprint(
           blueprint_name=blueprint_name,
           prefix=prefix,
@@ -361,7 +367,7 @@ def generate_blueprint(
           zone=args.zone,
           auth_cidr=all_IPs_cidr,
           num_nodes=num_nodes,
-          reservation=args.reservation if args.reservation else None,
+          reservation=reservation,
           enable_filestore_csi_driver=args.enable_gcpfilestore_csi_driver,
           capacity_type=capacity_type,
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
@@ -371,7 +377,6 @@ def generate_blueprint(
           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
       return bpg.generate_a4_blueprint(
           blueprint_name=blueprint_name,
           prefix=prefix,
@@ -381,7 +386,7 @@ def generate_blueprint(
           zone=args.zone,
           auth_cidr=all_IPs_cidr,
           num_nodes=num_nodes,
-          reservation=args.reservation if args.reservation else None,
+          reservation=reservation,
           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,

xpk/commands/cluster_test.py

@@ -146,6 +146,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       docker_image_pull_secret='',
       managed_mldiagnostics=False,
       output_manifest_file='',
+      num_cubes=None,
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)
@@ -696,27 +697,6 @@ def test_validate_cluster_create_args_for_super_slicing_missing_reservation(
   )
 
 
-def test_validate_cluster_create_args_for_super_slicing_reservation_no_blocks(
-    mocks: _Mocks,
-):
-  FeatureFlags.SUPER_SLICING_ENABLED = True
-  args = construct_args(
-      super_slicing=True,
-      reservation='reservation',
-      num_cubes=None,
-      num_slices=None,
-  )
-
-  with pytest.raises(SystemExit):
-    _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
-
-  assert mocks.commands_print_mock.call_count == 1
-  assert (
-      'requires a block or sub-block reservation'
-      in mocks.commands_print_mock.call_args[0][0]
-  )
-
-
 def test_validate_cluster_create_args_for_super_slicing_sparse_deployment_type_reservation(
     mocks: _Mocks,
 ):

xpk/commands/common.py

@@ -14,8 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.commands import run_command_with_updates_retry
-from ..core.capacity import H100_MEGA_DEVICE_TYPE, CapacityType
+from ..core.commands import run_command_with_updates_retry, run_command_for_value
+from ..core.capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, CapacityType
 from ..core.gcloud_context import get_cluster_location
 from ..utils.console import xpk_print, xpk_exit
 from ..utils.execution_context import is_dry_run
@@ -46,9 +46,12 @@ def set_cluster_command(args) -> int:
   return return_code
 
 
-def is_TAS_possible(
+def is_GPU_TAS_possible(
     system_characteristics: SystemCharacteristics | None,
     capacity_type: CapacityType | None,
+    cluster_name: str,
+    zone: str,
+    project: str,
 ) -> bool:
   """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible"""
 
@@ -63,10 +66,40 @@ def is_TAS_possible(
     xpk_print('capacity_type data was not found in configmaps.')
     xpk_exit(1)
 
-  return (
-      system_characteristics.device_type != H100_MEGA_DEVICE_TYPE
-      or capacity_type == CapacityType.RESERVATION
+  # For A3-High and A3-Mega TAS is supported only for Flex and Reservation
+  if (
+      system_characteristics.device_type == H100_DEVICE_TYPE
+      or system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
+  ) and (
+      capacity_type != CapacityType.FLEX_START
+      and capacity_type != CapacityType.RESERVATION
+  ):
+    return False
+
+  # COMPACT placement and Flex don't work together, Flex is enough to support TAS for A3-High or newer
+  if capacity_type == CapacityType.FLEX_START:
+    return True
+
+  # For A3-Ultra or newer, all capacity types support TAS as long as COMPACT placement is used
+  command = (
+      'gcloud container node-pools list'
+      f' --cluster {cluster_name}'
+      f' --location={zone}'
+      f' --project={project}'
+      ' --filter="placementPolicy.type=COMPACT"'
+      ' --format="value(name)"'
   )
+  return_code, compact_placement_nps = run_command_for_value(
+      command=command,
+      task=(
+          'Check if there is a COMPACT placement policy nodepool in this'
+          ' cluster'
+      ),
+  )
+  if return_code != 0:
+    xpk_print('Node pool retrieval failed, assuming TAS is not possible')
+    return False
+  return compact_placement_nps.splitlines() != []
 
 
 def validate_sub_slicing_system(system: SystemCharacteristics):

xpk/commands/common_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 unittest
+from unittest.mock import MagicMock, patch
+
+from xpk.commands.common import is_GPU_TAS_possible
+from xpk.core.capacity import (
+    H100_DEVICE_TYPE,
+    H100_MEGA_DEVICE_TYPE,
+    CapacityType,
+)
+from xpk.core.system_characteristics import SystemCharacteristics
+
+
+class CommonCommandsTest(unittest.TestCase):
+
+  @patch("xpk.commands.common.run_command_for_value")
+  @patch("xpk.commands.common.xpk_exit")
+  @patch("xpk.commands.common.xpk_print")
+  @patch("xpk.commands.common.is_dry_run")
+  def test_is_GPU_TAS_possible_dry_run(
+      self, mock_is_dry_run, mock_xpk_print, mock_xpk_exit, mock_run_command
+  ):
+    """Test is_GPU_TAS_possible returns True in dry_run mode."""
+    mock_is_dry_run.return_value = True
+    self.assertTrue(
+        is_GPU_TAS_possible(None, None, "cluster", "zone", "project")
+    )
+    mock_is_dry_run.assert_called_once()
+    mock_xpk_print.assert_not_called()
+    mock_xpk_exit.assert_not_called()
+    mock_run_command.assert_not_called()
+
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  @patch("xpk.commands.common.xpk_exit")
+  @patch("xpk.commands.common.xpk_print")
+  def test_is_GPU_TAS_possible_no_system_characteristics(
+      self, mock_xpk_print, mock_xpk_exit, mock_is_dry_run
+  ):
+    """Test is_GPU_TAS_possible exits if system_characteristics is None."""
+    mock_xpk_exit.side_effect = SystemExit(1)
+    with self.assertRaises(SystemExit):
+      is_GPU_TAS_possible(None, MagicMock(), "cluster", "zone", "project")
+    mock_xpk_print.assert_called_with(
+        "system_characteristics data was not found in configmaps."
+    )
+    mock_xpk_exit.assert_called_with(1)
+
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  @patch("xpk.commands.common.xpk_exit")
+  @patch("xpk.commands.common.xpk_print")
+  def test_is_GPU_TAS_possible_no_capacity_type(
+      self, mock_xpk_print, mock_xpk_exit, mock_is_dry_run
+  ):
+    """Test is_GPU_TAS_possible exits if capacity_type is None."""
+    mock_xpk_exit.side_effect = SystemExit(1)
+    with self.assertRaises(SystemExit):
+      is_GPU_TAS_possible(MagicMock(), None, "cluster", "zone", "project")
+    mock_xpk_print.assert_called_with(
+        "capacity_type data was not found in configmaps."
+    )
+    mock_xpk_exit.assert_called_with(1)
+
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  def test_is_GPU_TAS_possible_h100_unsupported_capacity(self, mock_is_dry_run):
+    """Test is_GPU_TAS_possible for H100 with unsupported capacity type."""
+    mock_system = MagicMock(spec=SystemCharacteristics)
+    mock_system.device_type = H100_DEVICE_TYPE
+    self.assertFalse(
+        is_GPU_TAS_possible(
+            mock_system, CapacityType.ON_DEMAND, "cluster", "zone", "project"
+        )
+    )
+
+    mock_system.device_type = H100_MEGA_DEVICE_TYPE
+    self.assertFalse(
+        is_GPU_TAS_possible(
+            mock_system, CapacityType.ON_DEMAND, "cluster", "zone", "project"
+        )
+    )
+
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  def test_is_GPU_TAS_possible_flex_start_capacity(self, mock_is_dry_run):
+    """Test is_GPU_TAS_possible returns True for FLEX_START capacity."""
+    mock_system = MagicMock(spec=SystemCharacteristics)
+    mock_system.device_type = "some-device"
+    self.assertTrue(
+        is_GPU_TAS_possible(
+            mock_system, CapacityType.FLEX_START, "cluster", "zone", "project"
+        )
+    )
+
+  @patch("xpk.commands.common.run_command_for_value")
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  def test_is_GPU_TAS_possible_compact_placement_exists(
+      self, mock_is_dry_run, mock_run_command
+  ):
+    """Test is_GPU_TAS_possible with COMPACT placement returns True."""
+    mock_system = MagicMock(spec=SystemCharacteristics)
+    mock_system.device_type = "a3-ultra"
+    mock_run_command.return_value = (0, "some-nodepool\nsome-other-nodepool\n")
+    self.assertTrue(
+        is_GPU_TAS_possible(
+            mock_system,
+            CapacityType.RESERVATION,
+            "cluster",
+            "zone",
+            "project",
+        )
+    )
+
+  @patch("xpk.commands.common.run_command_for_value")
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  def test_is_GPU_TAS_possible_no_compact_placement(
+      self, mock_is_dry_run, mock_run_command
+  ):
+    """Test is_GPU_TAS_possible without COMPACT placement returns False."""
+    mock_system = MagicMock(spec=SystemCharacteristics)
+    mock_system.device_type = "a3-ultra"
+    mock_run_command.return_value = (0, "")
+    self.assertFalse(
+        is_GPU_TAS_possible(
+            mock_system,
+            CapacityType.RESERVATION,
+            "cluster",
+            "zone",
+            "project",
+        )
+    )
+
+  @patch("xpk.commands.common.xpk_print")
+  @patch("xpk.commands.common.run_command_for_value")
+  @patch("xpk.commands.common.is_dry_run", return_value=False)
+  def test_is_GPU_TAS_possible_command_fails(
+      self, mock_is_dry_run, mock_run_command, mock_xpk_print
+  ):
+    """Test is_GPU_TAS_possible when gcloud command fails."""
+    mock_system = MagicMock(spec=SystemCharacteristics)
+    mock_system.device_type = "a3-ultra"
+    mock_run_command.return_value = (1, "Error")
+    self.assertFalse(
+        is_GPU_TAS_possible(
+            mock_system,
+            CapacityType.RESERVATION,
+            "cluster",
+            "zone",
+            "project",
+        )
+    )
+    mock_xpk_print.assert_called_with(
+        "Node pool retrieval failed, assuming TAS is not possible"
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()

xpk/commands/info.py

@@ -76,7 +76,7 @@ def get_nominal_quotas(cqs: str) -> dict[str, dict[str, str]]:
   try:
     cq_list = json.loads(cqs)['items']
   except ValueError:
-    xpk_print('Incorrect respone from list clusterqueue')
+    xpk_print('Incorrect response from list clusterqueue')
     xpk_print(cqs)
     xpk_exit(1)
 
@@ -98,7 +98,7 @@ def print_formatted_cqs(cqs: str, nominalQuotas) -> None:
   try:
     cq_list = json.loads(cqs)['items']
   except ValueError:
-    xpk_print('Incorrect respone from list clusterqueue')
+    xpk_print('Incorrect response from list clusterqueue')
     xpk_print(cqs)
     xpk_exit(1)
 
@@ -114,7 +114,7 @@ def print_formatted_lqs(lqs: str, nominalQuotas) -> None:
   try:
     lq_list = json.loads(lqs)['items']
   except ValueError:
-    xpk_print('Incorrect respone from list localqueue')
+    xpk_print('Incorrect response from list localqueue')
     xpk_print(lqs)
     xpk_exit(1)
 
@@ -219,7 +219,9 @@ def run_kueuectl_list_localqueue(args: Namespace) -> str:
   command = 'kubectl kueue list localqueue -o json'
   if args.namespace != '':
     command += f' --namespace {args.namespace}'
-  return_code, val = run_command_for_value(command, 'list localqueue')
+  return_code, val = run_command_for_value(
+      command, 'list localqueue', hide_error=True
+  )
 
   if return_code != 0:
     xpk_print(f'Cluster info request returned ERROR {return_code}')
@@ -235,7 +237,9 @@ def run_kueuectl_list_clusterqueue() -> str:
   """
   command = 'kubectl kueue list clusterqueue -o json'
 
-  return_code, val = run_command_for_value(command, 'list clusterqueue')
+  return_code, val = run_command_for_value(
+      command, 'list clusterqueue', hide_error=True
+  )
 
   if return_code != 0:
     xpk_print(f'Cluster info request returned ERROR {return_code}')

xpk/commands/inspector.py

@@ -23,7 +23,7 @@ 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
 from .workload import get_workload_list
-from ..core.kueue_manager import has_sub_slicing_enabled
+from ..core.kueue_manager import has_sub_slicing_enabled, has_super_slicing_enabled
 
 
 _SPACER = '========================================================'
@@ -89,12 +89,40 @@ def inspector_run_sub_slicing_helper(args, file: str):
   if return_code != 0:
     xpk_exit(return_code)
   if result:
-    output = f'Sub-slicing topology set up.\n{_SPACER}'
+    output = f'Sub-slicing topology set up.\n{_SPACER}\n'
     append_tmp_file(output, file)
     if args.print_to_terminal:
       xpk_print(output)
 
 
+def inspector_run_slice_controller_helper(args, file: str):
+  return_code, result = has_super_slicing_enabled()
+  if return_code != 0:
+    xpk_exit(return_code)
+
+  if not result:
+    return
+
+  output = f'Super-slicing topology set up.\n{_SPACER}\n'
+  append_tmp_file(output, file)
+  if args.print_to_terminal:
+    xpk_print(output)
+
+  command = (
+      'kubectl describe deployment slice-controller-controller-manager -n'
+      ' slice-controller-system'
+  )
+  command_description = 'Slice Controller Deployment Details'
+  inspector_run_command_helper(args, command, command_description, file)
+
+  command = (
+      'kubectl logs deployment slice-controller-controller-manager -n'
+      ' slice-controller-system --tail=100 --prefix=True'
+  )
+  command_description = 'Slice Controller Logs'
+  inspector_run_command_helper(args, command, command_description, file)
+
+
 def inspector_output_link_helper(args, link, link_description, file) -> int:
   """Outputs a link for xpk inspector to the output file.
 
@@ -257,6 +285,9 @@ def inspector(args) -> None:
           f' {description} return code: {return_code}'
       )
 
+  inspector_run_sub_slicing_helper(args, inspector_file)
+  inspector_run_slice_controller_helper(args, inspector_file)
+
   # Workload list views:
   filter_by_statuses = ['EVERYTHING', 'QUEUED', 'RUNNING']
   for filter_by_status in filter_by_statuses:
@@ -321,8 +352,6 @@ def inspector(args) -> None:
           f' {command_description} return code: {return_code}'
       )
 
-  inspector_run_sub_slicing_helper(args, inspector_file)
-
   # Cloud Console Links:
   workload_links = []
   if args.workload: