xpk 0.16.1__py3-none-any.whl → 0.17.0__py3-none-any.whl

xpk/core/kueue_manager_test.py

@@ -21,7 +21,7 @@ from pytest_mock import MockerFixture
 import yaml
 from unittest.mock import MagicMock, patch
 
-from xpk.core.kueue_manager import KueueConfig, KueueManager, has_sub_slicing_enabled
+from xpk.core.kueue_manager import KueueConfig, KueueManager, has_sub_slicing_enabled, has_super_slicing_enabled
 from xpk.core.system_characteristics import GpuConfig, DockerPlatform, AcceleratorType, SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
 from packaging.version import Version
@@ -35,6 +35,7 @@ TPU_SYSTEM: SystemCharacteristics = SystemCharacteristics(
     accelerator_type=AcceleratorType.TPU,
     device_type="v5p-8",
     supports_sub_slicing=False,
+    supports_super_slicing=False,
     docker_platform=DockerPlatform.ARM,
 )
 
@@ -44,6 +45,7 @@ KUEUE_CONFIG: KueueConfig = KueueConfig(
     cpu_limit=100,
     memory_limit="100Gi",
     configure_sub_slicing=False,
+    configure_super_slicing=False,
 )
 
 
@@ -370,6 +372,7 @@ def test_configure_generates_manifest_with_admission_checks_for_flex_single_slic
       KUEUE_CONFIG,
       num_slices=1,
       flex=True,
+      system=UserFacingNameToSystemCharacteristics["l4-1"],
   )
 
   kueue_manager.install_or_upgrade(kueue_config)
@@ -382,7 +385,7 @@ def test_configure_generates_manifest_with_admission_checks_for_flex_single_slic
   )
   assert (
       cluster_queue["spec"]["resourceGroups"][0]["flavors"][0]["name"]
-      == "1xv5p-8"
+      == "1xl4-1"
   )
   assert cluster_queue["spec"]["admissionChecks"][0] == "dws-prov"
 
@@ -406,6 +409,7 @@ def test_configure_generates_correct_manifest_with_gke_default_topology(
           accelerator_type=AcceleratorType.GPU,
           device_type="h100-mega-80gb-8",
           supports_sub_slicing=False,
+          supports_super_slicing=False,
           docker_platform=DockerPlatform.ARM,
           gpu_config=GpuConfig(requires_topology=True),
       ),
@@ -462,6 +466,47 @@ def test_configure_generates_correct_manifest_with_sub_slicing(
   assert actual_levels == expected_levels
 
 
+@patch("xpk.core.kueue_manager.write_tmp_file")
+def test_configure_generates_correct_manifest_with_super_slicing(
+    write_tmp_file_mock: MagicMock,
+    mock_commands: CommandsTester,
+    kueue_manager: KueueManager,
+):
+  """Test that __configure generates correct manifest with super-slicing topology."""
+  set_installed_kueue_version(mock_commands, None)
+  kueue_config = dataclasses.replace(
+      KUEUE_CONFIG,
+      configure_super_slicing=True,
+      system=UserFacingNameToSystemCharacteristics["tpu7x-4x4x4"],
+  )
+
+  kueue_manager.install_or_upgrade(kueue_config)
+
+  rendered_manifest: str = write_tmp_file_mock.call_args[0][0]
+  manifest_docs = list(yaml.safe_load_all(rendered_manifest))
+  resource_flavor = _first(
+      doc for doc in manifest_docs if doc["kind"] == "ResourceFlavor"
+  )
+  assert resource_flavor["spec"]["topologyName"] == "super-slice-topology"
+  assert resource_flavor["spec"]["nodeLabels"] == {
+      "cloud.google.com/gke-tpu-accelerator": "tpu7x",
+      "cloud.google.com/gke-tpu-partition-4x4x4-state": "HEALTHY",
+  }
+  topology = _first(doc for doc in manifest_docs if doc["kind"] == "Topology")
+  assert topology["metadata"]["name"] == "super-slice-topology"
+  expected_levels = [
+      "cloud.google.com/gce-topology-block",
+      "cloud.google.com/gke-tpu-partition-4x4x4-id",
+      "kubernetes.io/hostname",
+  ]
+  actual_levels = [level["nodeLabel"] for level in topology["spec"]["levels"]]
+  assert actual_levels == expected_levels
+  cluster_queue = _first(
+      doc for doc in manifest_docs if doc["kind"] == "ClusterQueue"
+  )
+  assert cluster_queue["spec"]["admissionChecks"][0] == "ss-kueue-operator"
+
+
 @patch("xpk.core.kueue_manager.write_tmp_file")
 def test_configure_generates_correct_manifest_with_pathways(
     write_tmp_file_mock: MagicMock,
@@ -549,7 +594,7 @@ def test_has_sub_slicing_enabled_returns_false_when_sub_slicing_topology_is_not_
   assert result is False
 
 
-def test_has_sub_slicing_enabled_returns_true_when_sub_slicing_topology_is_not_present(
+def test_has_sub_slicing_enabled_returns_true_when_sub_slicing_topology_is_present(
     mock_commands: CommandsTester,
 ):
   mock_commands.set_result_for_command(
@@ -562,6 +607,41 @@ def test_has_sub_slicing_enabled_returns_true_when_sub_slicing_topology_is_not_p
   assert result is True
 
 
+def test_has_super_slicing_enabled_returns_exit_code_when_command_fails(
+    mock_commands: CommandsTester,
+):
+  mock_commands.set_result_for_command((1, ""), "kubectl get topology")
+
+  return_code, result = has_super_slicing_enabled()
+
+  assert return_code == 1
+  assert result is None
+
+
+def test_has_super_slicing_enabled_returns_false_when_super_slicing_topology_is_not_present(
+    mock_commands: CommandsTester,
+):
+  mock_commands.set_result_for_command((0, ""), "kubectl get topology")
+
+  return_code, result = has_super_slicing_enabled()
+
+  assert return_code == 0
+  assert result is False
+
+
+def test_has_super_slicing_enabled_returns_true_when_super_slicing_topology_is_present(
+    mock_commands: CommandsTester,
+):
+  mock_commands.set_result_for_command(
+      (0, "super-slice-topology"), "kubectl get topology"
+  )
+
+  return_code, result = has_super_slicing_enabled()
+
+  assert return_code == 0
+  assert result is True
+
+
 T = TypeVar("T")
 
 

xpk/core/nap.py

@@ -147,17 +147,18 @@ def enable_autoprovisioning_on_cluster(
 
   for i, command in enumerate(commands):
     xpk_print(f'To complete {task_names[i]} we are executing {command}')
-  max_return_code = run_commands(
+  maybe_failure = run_commands(
       commands,
       'Update node pools with autoprovisioning support',
       task_names,
   )
-  if max_return_code != 0:
+  if maybe_failure is not None:
     xpk_print(
         'Update node pools with autoprovisioning support returned ERROR:'
-        f' {max_return_code}'
+        f' {maybe_failure.return_code}'
     )
-    return None, max_return_code
+    return None, maybe_failure.return_code
+
   return autoprovisioning_config, return_code
 
 

xpk/core/nodepool.py

@@ -15,6 +15,8 @@ limitations under the License.
 """
 
 from typing import List
+
+from ..utils.feature_flags import FeatureFlags
 from ..utils.console import ask_for_user_consent, xpk_print
 from .scheduling import get_placement_policy_name, is_placement_policy_supported
 from .capacity import (
@@ -25,14 +27,14 @@ from .capacity import (
     get_capacity_type,
     print_reservations,
 )
-from .commands import run_command_for_value, run_commands
+from .commands import run_command_for_value, run_commands, FailedCommand
 from .gcloud_context import GkeServerConfig, get_cluster_location, zone_to_region
 from .resources import (
     ConfigMapType,
     check_cluster_resources,
     update_cluster_configmap,
 )
-from .system_characteristics import AcceleratorType
+from .system_characteristics import AcceleratorType, SystemCharacteristics
 
 
 CLOUD_PLATFORM_AUTH_SCOPE_URL = (
@@ -43,7 +45,7 @@ OLDER_PATHWAYS_CPU_NP_TO_DELETE = ['cpu-rm-np', 'cpu-proxy-np', 'cpu-user-np']
 
 
 def run_gke_node_pool_create_command(
-    args, system, gke_node_pool_version
+    args, system: SystemCharacteristics, gke_node_pool_version: str
 ) -> int:
   """Run the Create GKE Node Pool request.
 
@@ -84,7 +86,7 @@ def run_gke_node_pool_create_command(
   else:
     max_nodes = 1000
   capacity_args, return_code = get_capacity_arguments_from_capacity_type(
-      args, capacity_type, max_nodes
+      args, capacity_type, max_nodes, system.accelerator_type
   )
   if return_code > 0:
     xpk_print('Parsing capacity arguments failed!')
@@ -200,13 +202,13 @@ def run_gke_node_pool_create_command(
       xpk_print(
           f'To complete {delete_task_names[i]} we are executing {command}'
       )
-    max_return_code = run_commands(
+    maybe_failure = run_commands(
         delete_commands,
         'Delete Nodepools',
         delete_task_names,
     )
-    if max_return_code != 0:
-      xpk_print(f'Delete Nodepools returned ERROR {max_return_code}')
+    if maybe_failure is not None:
+      xpk_print(f'Delete Nodepools returned ERROR {maybe_failure.return_code}')
       return 1
 
   # Enable Workload Identity on existing Nodepools
@@ -224,15 +226,15 @@ def run_gke_node_pool_create_command(
         xpk_print(
             f'To complete {update_WI_task_names[i]} we are executing {command}'
         )
-      max_return_code = run_commands(
+      maybe_failure = run_commands(
           update_WI_commands,
           'Enable Workload Identity on existing Nodepools',
           update_WI_task_names,
       )
-      if max_return_code != 0:
+      if maybe_failure is not None:
         xpk_print(
             'Enable Workload Identity on existing Nodepools returned ERROR'
-            f' {max_return_code}'
+            f' {maybe_failure.return_code}'
         )
         return 1
 
@@ -256,12 +258,17 @@ def run_gke_node_pool_create_command(
 
   placement_args = ''
   if is_placement_policy_supported(system):
-    placement_policy = get_placement_policy_name(system)
+    super_slicing = FeatureFlags.SUPER_SLICING_ENABLED and args.super_slicing
+    placement_policy = get_placement_policy_name(
+        system,
+        super_slicing,
+    )
     ensure_resource_policy_exists(
         resource_policy_name=placement_policy,
         project=args.project,
         zone=args.zone,
         topology=system.topology,
+        super_slicing=super_slicing,
     )
     placement_args = f' --placement-policy={placement_policy}'
 
@@ -358,19 +365,41 @@ def run_gke_node_pool_create_command(
 
   for i, command in enumerate(create_commands):
     xpk_print(f'To complete {create_task_names[i]} we are executing {command}')
-  max_return_code = run_commands(
+  maybe_failure = run_commands(
       create_commands,
       'Create Nodepools',
       create_task_names,
   )
-  if max_return_code != 0:
-    xpk_print(f'Create Nodepools returned ERROR {max_return_code}')
+  if maybe_failure is not None:
+    display_nodepool_creation_error(maybe_failure)
     return 1
 
   xpk_print('Create or delete node pool request complete.')
   return 0
 
 
+def display_nodepool_creation_error(maybe_failure: FailedCommand) -> None:
+  """Display nodepool creation errors to the user."""
+
+  xpk_print(f'Create Nodepools returned ERROR {maybe_failure.return_code}')
+  try:
+    with open(maybe_failure.logfile, 'r', encoding='utf-8') as f:
+      contents = f.read()
+    error_marker = 'finished with error:'
+    error = contents[contents.index(error_marker) + len(error_marker) :].strip()
+    # the longest error we're expecting to see is 256 characters + np name
+    max_error_display_length = 400
+    xpk_print(f'Nodepool creation error: {error[:max_error_display_length]}')
+    if (
+        error.find('lack of capacity') != -1
+        or error.find('Requested resource is exhausted') != -1
+    ):
+      xpk_print('NOTE: this error might be caused by a stockout')
+  except (FileNotFoundError, IOError, ValueError):
+    # silently ignore any log parsing errors
+    pass
+
+
 def get_node_pools_to_delete(
     args, system, existing_node_pool_names, desired_node_pool_names
 ) -> list:
@@ -585,18 +614,22 @@ def get_desired_node_pool_names(
   while len(result) < desired_node_pool_count:
     result.add(f'{cluster_name}-np-{i}')
     i += 1
-  return list(result)
+  return list(sorted(result))
 
 
 def ensure_resource_policy_exists(
-    resource_policy_name: str, project: str, zone: str, topology: str
+    resource_policy_name: str,
+    project: str,
+    zone: str,
+    topology: str,
+    super_slicing: bool,
 ) -> None:
   return_code, _ = run_command_for_value(
       (
           'gcloud compute resource-policies describe'
-          f' {resource_policy_name} '
-          f'--project={project} '
-          f'--region={zone_to_region(zone)}'
+          f' {resource_policy_name}'
+          f' --project={project}'
+          f' --region={zone_to_region(zone)}'
       ),
       'Retrieve resource policy',
   )
@@ -604,11 +637,15 @@ def ensure_resource_policy_exists(
   if return_code == 0:
     return
 
+  # TODO: b/465696970 - Verify the flag below before launching SUPER_SLICING:
+  accelerator_topology_mode = (
+      ' --accelerator-topology-mode=PROVISION_ONLY' if super_slicing else ''
+  )
   return_code, _ = run_command_for_value(
       (
           'gcloud compute resource-policies create workload-policy'
           f' {resource_policy_name} --project={project} --region={zone_to_region(zone)} --type=HIGH_THROUGHPUT'
-          f' --accelerator-topology={topology}'
+          f' --accelerator-topology={topology}{accelerator_topology_mode}'
       ),
       'Create resource policy',
   )

xpk/core/nodepool_test.py

@@ -16,13 +16,23 @@ limitations under the License.
 
 import pytest
 from xpk.core.nodepool import (
+    display_nodepool_creation_error,
     ensure_resource_policy_exists,
     get_desired_node_pool_names,
     run_gke_node_pool_create_command,
 )
 from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics, DockerPlatform, GpuConfig
+from xpk.core.commands import FailedCommand
+from xpk.core.testing.commands_tester import CommandsTester
+
 
 CLUSTER_NAME = "running-cucumber"
+maybe_failure = FailedCommand(
+    return_code=1,
+    name="create-nodepool",
+    command="test-command",
+    logfile="logfile_path",
+)
 
 
 def node_pool_name(number: int) -> str:
@@ -88,61 +98,107 @@ def test_compute_desired_node_pool_names_with_unknown_node_pools():
   assert set(result) == set(expected_result)
 
 
+@pytest.fixture
+def commands_tester(mocker):
+  return CommandsTester(
+      mocker,
+      run_command_for_value_path="xpk.core.nodepool.run_command_for_value",
+  )
+
+
 def test_ensure_resource_policy_exists_with_existing_policy_retrieves_existing_policy(
-    mocker,
+    commands_tester: CommandsTester,
 ):
-  args = mocker.Mock(project="test-project", zone="us-central1-a")
-  mocker.patch("xpk.core.nodepool.get_cluster_location", return_value=args.zone)
-  mock = mocker.patch(
-      "xpk.core.nodepool.run_command_for_value", return_value=(0, "")
-  )
   ensure_resource_policy_exists(
       resource_policy_name="resource-policy",
       project="test-project",
       zone="us-central1-a",
       topology="2x2x1",
+      super_slicing=False,
+  )
+
+  assert len(commands_tester.commands_history) == 1
+  commands_tester.assert_command_run(
+      "gcloud compute resource-policies describe resource-policy",
+      "--project=test-project",
+      "--region=us-central1",
   )
-  mock.assert_called_once()
 
 
 def test_ensure_resource_policy_exists_without_existing_policy_creates_policy(
-    mocker,
+    commands_tester: CommandsTester,
 ):
-  args = mocker.Mock(project="test-project", zone="us-central1-a")
-  mocker.patch("xpk.core.nodepool.get_cluster_location", return_value=args.zone)
-  mock = mocker.patch(
-      "xpk.core.nodepool.run_command_for_value", side_effect=[(1, ""), (0, "")]
+  commands_tester.set_result_for_command(
+      (1, ""), "gcloud compute resource-policies describe"
   )
+
   ensure_resource_policy_exists(
       resource_policy_name="resource-policy",
       project="test-project",
       zone="us-central1-a",
       topology="2x2x1",
+      super_slicing=False,
+  )
+
+  assert len(commands_tester.commands_history) == 2
+  commands_tester.assert_command_run(
+      "gcloud compute resource-policies describe"
+  )
+  commands_tester.assert_command_run(
+      "gcloud compute resource-policies create workload-policy resource-policy",
+      "--project=test-project",
+      "--region=us-central1",
+      "--accelerator-topology=2x2x1",
+  )
+  commands_tester.assert_command_not_run(
+      "gcloud compute resource-policies create workload-policy",
+      "--accelerator-topology-mode",
+  )
+
+
+def test_ensure_resource_policy_exists_without_existing_policy_creates_policy_for_super_slicing(
+    commands_tester: CommandsTester,
+):
+  commands_tester.set_result_for_command(
+      (1, ""), "gcloud compute resource-policies describe"
+  )
+
+  ensure_resource_policy_exists(
+      resource_policy_name="ss-resource-policy",
+      project="test-project",
+      zone="us-central1-a",
+      topology="2x2x1",
+      super_slicing=True,
+  )
+
+  commands_tester.assert_command_run(
+      "gcloud compute resource-policies create workload-policy",
+      "--accelerator-topology-mode",
   )
-  assert mock.call_count == 2
-  assert mock.call_args_list[0].args[1] == "Retrieve resource policy"
 
 
 def test_ensure_resource_policy_exits_without_existing_policy_throws_when_creation_fails(
-    mocker,
+    commands_tester: CommandsTester,
 ):
   with pytest.raises(RuntimeError):
-    args = mocker.Mock(project="test-project", zone="us-central1-a")
-    mocker.patch(
-        "xpk.core.nodepool.get_cluster_location", return_value=args.zone
-    )
-    mocker.patch(
-        "xpk.core.nodepool.run_command_for_value",
-        side_effect=[(1, ""), (1, "")],
+    commands_tester.set_result_for_command(
+        (1, ""), "gcloud compute resource-policies"
     )
+
     ensure_resource_policy_exists(
         resource_policy_name="resource-policy",
         project="test-project",
         zone="us-central1-a",
         topology="2x2x1",
+        super_slicing=False,
     )
 
 
+@pytest.fixture
+def mock_xpk_print(mocker):
+  return mocker.patch("xpk.core.nodepool.xpk_print")
+
+
 @pytest.fixture
 def mock_nodepool_dependencies(mocker):
   """Mocks dependencies for run_gke_node_pool_create_command."""
@@ -159,7 +215,7 @@ def mock_nodepool_dependencies(mocker):
   mocker.patch(
       "xpk.core.nodepool.get_cluster_location", return_value="us-central1"
   )
-  mocker.patch("xpk.core.nodepool.run_commands", return_value=0)
+  mocker.patch("xpk.core.nodepool.run_commands", return_value=None)
   mocker.patch("xpk.core.nodepool.ask_for_user_consent", return_value=True)
   mock_is_placement_policy_supported = mocker.patch(
       "xpk.core.nodepool.is_placement_policy_supported"
@@ -194,6 +250,7 @@ def test_placement_policy_created_for_gpu_with_valid_topology(
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-80gb-8",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
       gpu_config=GpuConfig(requires_topology=True),
   )
@@ -226,6 +283,7 @@ def test_placement_policy_not_created_for_gpu_with_invalid_topology(
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-80gb-8",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
       gpu_config=GpuConfig(requires_topology=True),
   )
@@ -261,6 +319,7 @@ def test_placement_policy_created_for_tpu7x_with_valid_topology(
       device_type="tpu7x-8",
       requires_workload_policy=True,
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
 
@@ -294,9 +353,79 @@ def test_placement_policy_not_created_for_non7x_tpu(
       accelerator_type=AcceleratorType.TPU,
       device_type="v6e-4",
       supports_sub_slicing=True,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
 
   run_gke_node_pool_create_command(args, system, "1.2.3")
 
   mock_ensure_resource_policy.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    argnames="error_message,is_stockout",
+    argvalues=[
+        (
+            (
+                "Requested resource is exhausted: Zone 'us-central1-c' is not"
+                " available. Please try another zone."
+            ),
+            True,
+        ),
+        (
+            (
+                "TPU: the nodes (in pool test-pool) cannot be created now due"
+                " to lack of capacity in your reservation. They will be created"
+                " asynchronously once capacity is available. You can either"
+                " wait for the nodes to be up, or delete the node pool and try"
+                " re-creating it again later"
+            ),
+            True,
+        ),
+        ("Generic error message", False),
+    ],
+)
+def test_display_nodepool_creation_error_handles_error_messages(
+    mocker, mock_xpk_print, error_message, is_stockout
+):
+  """Tests that display_nodepool_creation_error surfaces errors and detects stockouts."""
+
+  log_contents = """Operation [
+  ...
+  ] finished with error: """ + error_message + "\n"
+  mocker.patch("builtins.open", mocker.mock_open(read_data=log_contents))
+  display_nodepool_creation_error(maybe_failure)
+
+  assert mock_xpk_print.call_count == 3 if is_stockout else 2
+  assert (
+      mock_xpk_print.call_args_list[0].args[0]
+      == "Create Nodepools returned ERROR 1"
+  )
+  assert (
+      mock_xpk_print.call_args_list[1].args[0]
+      == "Nodepool creation error: " + error_message
+  )
+  assert (
+      not is_stockout
+      or mock_xpk_print.call_args_list[2].args[0]
+      == "NOTE: this error might be caused by a stockout"
+  )
+
+
+def test_display_nodepool_creation_ignores_logs_without_errors(
+    mocker,
+    mock_xpk_print,
+):
+  """Tests that display_nodepool_creation_error ignores log files with no errors."""
+
+  log_contents = """Operation [
+  ...
+  ] succeeded!"""
+  mocker.patch("builtins.open", mocker.mock_open(read_data=log_contents))
+  display_nodepool_creation_error(maybe_failure)
+
+  assert mock_xpk_print.call_count == 1
+  assert (
+      mock_xpk_print.call_args_list[0].args[0]
+      == "Create Nodepools returned ERROR 1"
+  )

xpk/core/pathways.py

@@ -325,9 +325,10 @@ def try_to_delete_pathwaysjob_first(args, workloads) -> bool:
   if len(workloads) == 1:
     return_code = run_command_with_updates(commands[0], 'Delete Workload')
   else:
-    return_code = run_commands(
+    maybe_failure = run_commands(
         commands, 'Delete Workload', task_names, batch=100
     )
+    return_code = 0 if maybe_failure is None else maybe_failure.return_code
 
   if return_code != 0:
     xpk_print(f'Delete Workload request returned ERROR {return_code}')

xpk/core/resources.py

@@ -201,15 +201,15 @@ def _create_or_update_cluster_configmap(configmap_yml: dict[str, str]) -> int:
     task_name = f'ConfigMap CreateOrUpdate-{configmap_name}'
     task_names.append(task_name)
 
-  return_code = run_commands(
+  maybe_failure = run_commands(
       commands,
       'GKE Cluster CreateOrUpdate ConfigMap(s)',
       task_names,
   )
-  if return_code != 0:
+  if maybe_failure is not None:
     xpk_print(
         'GKE Cluster Create/Update ConfigMap(s) request returned ERROR'
-        f' {return_code}'
+        f' {maybe_failure.return_code}'
     )
     return 1
   return 0