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

xpk/commands/workload_test.py

@@ -16,10 +16,13 @@ limitations under the License.
 
 import dataclasses
 from unittest.mock import MagicMock, patch
+import yaml
 import pytest
-from ..core.system_characteristics import SystemCharacteristics, AcceleratorType
-from .workload import _validate_sub_slicing_topology, _validate_sub_slicing_availability
-from packaging.version import Version
+
+from ..core.scheduling import WorkloadScheduling
+from ..core.system_characteristics import DockerPlatform, SystemCharacteristics, AcceleratorType, UserFacingNameToSystemCharacteristics, GpuConfig
+from .workload import workload_create
+from .cluster_test import construct_args
 
 
 SYSTEM_CHARACTERISTICS = SystemCharacteristics(
@@ -32,133 +35,172 @@ SYSTEM_CHARACTERISTICS = SystemCharacteristics(
     device_type='l4-1',
     supports_sub_slicing=True,
     requires_workload_policy=False,
+    docker_platform=DockerPlatform.AMD,
 )
 
 
-@pytest.fixture(autouse=True)
+@dataclasses.dataclass
+class _WorkloadCreateMocks:
+  """Holds all the mocked dependencies for the workload_create function."""
+
+  get_user_workload_container: MagicMock
+  get_gpu_scheduler: MagicMock
+  get_storages_to_mount: MagicMock
+  add_bucket_iam_members: MagicMock
+  get_gke_outlier_dashboard: MagicMock
+  check_if_workload_exists: MagicMock
+  get_cluster_configmap: MagicMock
+  check_if_workload_can_schedule: MagicMock
+  setup_k8s_env: MagicMock
+  setup_k8s_service_accounts: MagicMock
+  validate_dependencies_list: MagicMock
+  write_tmp_file: MagicMock
+  get_cluster_capacity_type: MagicMock
+  is_TAS_possible: MagicMock
+  get_cluster_location: MagicMock
+  xpk_exit: MagicMock
+  run_command_with_updates: MagicMock
+  ensure_resource_policy_exists: MagicMock
+  get_cluster_subnetworks: MagicMock
+
+
+@pytest.fixture
 def xpk_print(mocker):
   return mocker.patch('xpk.commands.workload.xpk_print')
 
 
-def test_validate_sub_slicing_topology_exits_for_unsupported_topology(
-    xpk_print: MagicMock,
-):
-  with pytest.raises(SystemExit):
-    _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '2x1')
-
-  assert (
-      'shape is invalid. It has to be one of' in xpk_print.mock_calls[0].args[0]
-  )
-
-
-def test_validate_sub_slicing_topology_exits_for_too_large_topology(
-    xpk_print: MagicMock,
-):
-  with pytest.raises(SystemExit):
-    _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '16x16')
-
-  assert (
-      'shape is too large. The shape cannot be'
-      in xpk_print.mock_calls[0].args[0]
-  )
-
-
-def test_validate_sub_slicing_topology_does_nothing_for_supported_topology():
-  _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '4x4')
-
-
-def test_validate_sub_slicing_availability_exits_when_getting_topologies_fails(
-    xpk_print: MagicMock, mocker
-):
-  mocker.patch(
-      'xpk.commands.workload.has_sub_slicing_enabled',
-      return_value=(1, None),
-  )
-  with pytest.raises(SystemExit):
-    _validate_sub_slicing_availability()
-
-  assert (
-      'Unable to validate sub-slicing support'
-      in xpk_print.mock_calls[0].args[0]
-  )
-
-
-def test_validate_sub_slicing_availability_exits_when_subslicing_topology_is_not_defined(
-    xpk_print: MagicMock, mocker
-):
-  mocker.patch(
-      'xpk.commands.workload.has_sub_slicing_enabled',
-      return_value=(0, False),
-  )
-  with pytest.raises(SystemExit):
-    _validate_sub_slicing_availability()
-
-  assert (
-      'Cluster has not been not set up for Sub-slicing.'
-      in xpk_print.mock_calls[0].args[0]
+@pytest.fixture
+def workload_create_mocks(mocker) -> _WorkloadCreateMocks:
+  """Mocks all dependencies for the workload_create function."""
+  return _WorkloadCreateMocks(
+      get_user_workload_container=mocker.patch(
+          'xpk.commands.workload.get_user_workload_container',
+          return_value=('', None),
+      ),
+      get_gpu_scheduler=mocker.patch(
+          'xpk.commands.workload.get_gpu_scheduler', return_value=('', 0)
+      ),
+      get_storages_to_mount=mocker.patch(
+          'xpk.commands.workload.get_storages_to_mount', return_value=[]
+      ),
+      add_bucket_iam_members=mocker.patch(
+          'xpk.commands.workload.add_bucket_iam_members'
+      ),
+      get_gke_outlier_dashboard=mocker.patch(
+          'xpk.commands.workload.get_gke_outlier_dashboard'
+      ),
+      check_if_workload_exists=mocker.patch(
+          'xpk.commands.workload.check_if_workload_exists', return_value=False
+      ),
+      get_cluster_configmap=mocker.patch(
+          'xpk.commands.workload.get_cluster_configmap', return_value={}
+      ),
+      check_if_workload_can_schedule=mocker.patch(
+          'xpk.commands.workload.check_if_workload_can_schedule',
+          return_value=WorkloadScheduling.AVAILABLE,
+      ),
+      setup_k8s_env=mocker.patch('xpk.commands.workload.setup_k8s_env'),
+      setup_k8s_service_accounts=mocker.patch(
+          'xpk.commands.workload.setup_k8s_service_accounts'
+      ),
+      validate_dependencies_list=mocker.patch(
+          'xpk.commands.workload.validate_dependencies_list'
+      ),
+      write_tmp_file=mocker.patch('xpk.commands.workload.write_tmp_file'),
+      get_cluster_capacity_type=mocker.patch(
+          'xpk.commands.workload.get_cluster_capacity_type',
+          return_value='on-demand',
+      ),
+      is_TAS_possible=mocker.patch(
+          'xpk.commands.workload.is_TAS_possible', return_value=False
+      ),
+      get_cluster_location=mocker.patch(
+          'xpk.commands.workload.get_cluster_location',
+          return_value='us-central1',
+      ),
+      xpk_exit=mocker.patch('xpk.commands.workload.xpk_exit'),
+      run_command_with_updates=mocker.patch(
+          'xpk.commands.workload.run_command_with_updates', return_value=0
+      ),
+      ensure_resource_policy_exists=mocker.patch(
+          'xpk.commands.workload.ensure_resource_policy_exists'
+      ),
+      get_cluster_subnetworks=mocker.patch(
+          'xpk.commands.workload.get_cluster_subnetworks', return_value=[]
+      ),
   )
 
 
-def test_validate_sub_slicing_availability_exits_when_kueue_version_cannot_be_determined(
-    xpk_print: MagicMock, mocker
+def test_workload_create_for_a4x_has_arm_toleration(
+    workload_create_mocks: _WorkloadCreateMocks,
 ):
-  mocker.patch(
-      'xpk.commands.workload.has_sub_slicing_enabled',
-      return_value=(0, True),
+  """Tests that the generated YAML for an A4X workload has arm64 toleration."""
+  # Copy and overwrite the decorator with a no-op lambda.
+  gb200_system_chars = UserFacingNameToSystemCharacteristics['gb200-4']
+  gb200_system_chars_no_decorator = dataclasses.replace(
+      gb200_system_chars,
+      gpu_config=GpuConfig(
+          requires_topology=False, jobset_decorator_fn=lambda yml, *_: yml
+      ),
   )
-  mocker.patch(
-      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
-      return_value=(1, None),
-  )
-  with pytest.raises(SystemExit):
-    _validate_sub_slicing_availability()
-
-  assert 'Unable to validate sub-slicing' in xpk_print.mock_calls[0].args[0]
+  # Patch the function that returns the system characteristics
+  # to return our modified object.
+  with patch(
+      'xpk.commands.workload.get_system_characteristics',
+      return_value=(gb200_system_chars_no_decorator, 0),
+  ):
+    args = construct_args(
+        device_type='gb200-4',
+        workload='test-workload',
+        command='echo hello',
+        num_nodes=1,
+        restart_on_exit_codes=None,
+    )
+    workload_create(args)
+
+  assert workload_create_mocks.write_tmp_file.called
+  yaml_content = workload_create_mocks.write_tmp_file.call_args[0][0]
+  jobset = yaml.safe_load(yaml_content)
+
+  tolerations = jobset['spec']['replicatedJobs'][0]['template']['spec'][
+      'template'
+  ]['spec']['tolerations']
+  assert {
+      'key': 'kubernetes.io/arch',
+      'operator': 'Equal',
+      'value': 'arm64',
+      'effect': 'NoSchedule',
+  } in tolerations
+
+
+def test_workload_create_dry_run_with_output_file(mocker):
+  args = MagicMock()
+  args.workload = 'test-workload'
+  args.output_manifest_file = 'manifest.yaml'
+  args.use_pathways = False
+  args.use_vertex_tensorboard = False
+  args.project = 'test-project'
+  args.cluster = 'test-cluster'
+  args.zone = 'test-zone'
+  args.sub_slicing_topology = None
+
+  # Mock dependencies to avoid external calls and simulate state
+  mocker.patch('xpk.utils.execution_context.dry_run', True)
+  mocks = {
+      'get_system_characteristics': (SYSTEM_CHARACTERISTICS, 0),
+      'get_user_workload_container': ('container_yaml', None),
+      'write_tmp_file': 'tmp_file',
+      'parse_env_config': None,
+  }
+  for name, return_value in mocks.items():
+    mocker.patch(f'xpk.commands.workload.{name}', return_value=return_value)
+
+  mock_open = mocker.patch('builtins.open', mocker.mock_open())
 
-
-def test_validate_sub_slicing_availability_exits_when_kueue_version_does_not_meet_minimum_requirements(
-    xpk_print: MagicMock, mocker
-):
-  mocker.patch(
-      'xpk.commands.workload.has_sub_slicing_enabled',
-      return_value=(0, True),
-  )
-  mocker.patch(
-      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
-      return_value=(0, Version('0.0.0')),
-  )
   with pytest.raises(SystemExit):
-    _validate_sub_slicing_availability()
-
-  assert 'The minimal required version is' in xpk_print.mock_calls[0].args[0]
-
+    workload_create(args)
 
-def test_validate_sub_slicing_availability_does_nothing_when_cluster_is_correctly_configured_for_subslicing(
-    mocker,
-):
-  mocker.patch(
-      'xpk.commands.workload.has_sub_slicing_enabled',
-      return_value=(0, True),
-  )
-  mocker.patch(
-      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
-      return_value=(0, Version('0.13.0')),
-  )
-  _validate_sub_slicing_availability()
-
-
-@patch('xpk.commands.common.xpk_print')
-def test_validate_sub_slicing_topology_fails_for_unsupported_system(
-    common_xpk_print: MagicMock,
-):
-  unsupported_system = dataclasses.replace(
-      SYSTEM_CHARACTERISTICS, supports_sub_slicing=False
-  )
-
-  with pytest.raises(SystemExit):
-    _validate_sub_slicing_topology(unsupported_system, '4x4')
-
-  assert (
-      'l4-1 does not support Sub-slicing.'
-      in common_xpk_print.mock_calls[0].args[0]
-  )
+  mock_open.assert_called_once_with('manifest.yaml', 'w', encoding='utf-8')
+  written_content = mock_open.return_value.write.call_args[0][0]
+  assert 'test-workload' in written_content
+  assert 'cloud.google.com/gke-tpu-topology: 8x8' in written_content

xpk/core/blueprint/blueprint_generator.py

@@ -19,15 +19,20 @@ import shutil
 from typing import Optional
 
 from ruamel import yaml
+from packaging.version import parse
 
 from ...utils.console import xpk_exit, xpk_print
+from ...utils.versions import ReleaseChannel
 from ...utils.file import ensure_directory_exists
 
+
 from ..capacity import (
     H100_DEVICE_TYPE,
     B200_DEVICE_TYPE,
     H100_MEGA_DEVICE_TYPE,
     H200_DEVICE_TYPE,
+    GB200_DEVICE_TYPE,
+    GB200_DEVICE_TYPE_NOLSSD,
     CapacityType,
 )
 from ..system_characteristics import get_system_characteristics_by_device_type
@@ -39,6 +44,7 @@ a3high_device_type = H100_DEVICE_TYPE
 a3mega_device_type = H100_MEGA_DEVICE_TYPE
 a3ultra_device_type = H200_DEVICE_TYPE
 a4_device_type = B200_DEVICE_TYPE
+a4x_device_types = (GB200_DEVICE_TYPE, GB200_DEVICE_TYPE_NOLSSD)
 supported_device_types = {
     a3mega_device_type,
     a3ultra_device_type,
@@ -84,6 +90,8 @@ class BlueprintGenerator:
       region: str,
       zone: str,
       auth_cidr: str,
+      cluster_version: str,
+      release_channel: ReleaseChannel,
       prefix: str = "",
       num_nodes: int = 2,
       pods_ip_cidr_range: str = "10.4.0.0/14",
@@ -142,11 +150,17 @@ class BlueprintGenerator:
         },
     )
 
+    sanitized_version = cluster_version.replace("-", "+", 1)
+    version = parse(sanitized_version)
+    version_prefix = f"{version.major}.{version.minor}"
     gke_cluster = DeploymentModule(
         id="gke_cluster",
         source="modules/scheduler/gke-cluster",
         use=[primary_vpc_name, gpu_subnets_name],
         settings={
+            "release_channel": release_channel.value,
+            "version_prefix": version_prefix,
+            "min_master_version": cluster_version,
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "enable_private_endpoint": False,
@@ -171,6 +185,16 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+    if release_channel != ReleaseChannel.RAPID:
+      gke_cluster.set_setting(
+          "maintenance_exclusions",
+          [{
+              "name": "no-minor-or-node-upgrades-indefinite",
+              "start_time": "2024-12-01T00:00:00Z",
+              "end_time": "2026-01-16T00:00:00Z",
+              "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+          }],
+      )
 
     group_placement_0 = DeploymentModule(
         id="group_placement_0",
@@ -215,6 +239,9 @@ class BlueprintGenerator:
     else:
       a3_megagpu_pool_0.update_settings({"static_node_count": num_nodes})
 
+    if release_channel == ReleaseChannel.RAPID:
+      a3_megagpu_pool_0.set_setting("auto_upgrade", True)
+
     set_placement_policy = capacity_type != CapacityType.SPOT
     workload = DeploymentModule(
         id="workload_component_install",
@@ -391,6 +418,8 @@ class BlueprintGenerator:
       zone: str,
       auth_cidr: str,
       system_node_pool_machine_type: str,
+      cluster_version: str,
+      release_channel: ReleaseChannel,
       reservation: Optional[str | None] = None,
       gcs_bucket: Optional[str | None] = None,
       num_nodes: int = 2,
@@ -480,28 +509,19 @@ class BlueprintGenerator:
             },
         },
     )
+
+    sanitized_version = cluster_version.replace("-", "+", 1)
+    version = parse(sanitized_version)
+    version_prefix = f"{version.major}.{version.minor}"
     cluster_id = f"{cluster_name}-a3-ultragpu-cluster"
     a3_ultra_cluster = DeploymentModule(
         id=cluster_id,
         source="modules/scheduler/gke-cluster",
         use=[net_0_id],
         settings={
-            "release_channel": (
-                "UNSPECIFIED"
-                if capacity_type == CapacityType.FLEX_START
-                else "RAPID"
-            ),
-            "version_prefix": "1.32.",
-            "maintenance_exclusions": (
-                []
-                if capacity_type == CapacityType.FLEX_START
-                else [{
-                    "name": "no-minor-or-node-upgrades-indefinite",
-                    "start_time": "2024-12-01T00:00:00Z",
-                    "end_time": "2025-12-22T00:00:00Z",
-                    "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
-                }]
-            ),
+            "release_channel": release_channel.value,
+            "version_prefix": version_prefix,
+            "min_cluster_version": cluster_version,
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "system_node_pool_machine_type": system_node_pool_machine_type,
@@ -537,6 +557,17 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+    if release_channel != ReleaseChannel.RAPID:
+      a3_ultra_cluster.set_setting(
+          "maintenance_exclusions",
+          [{
+              "name": "no-minor-or-node-upgrades-indefinite",
+              "start_time": "2024-12-01T00:00:00Z",
+              "end_time": "2026-01-16T00:00:00Z",
+              "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+          }],
+      )
+
     system, _ = get_system_characteristics_by_device_type(a3ultra_device_type)
     if system is None:
       xpk_print(
@@ -584,6 +615,9 @@ class BlueprintGenerator:
     else:
       gpu_pool.update_settings({"static_node_count": num_nodes})
 
+    if release_channel == ReleaseChannel.RAPID:
+      gpu_pool.set_setting("auto_upgrade", True)
+
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
         id=workload_manager_install_id,
@@ -674,6 +708,8 @@ class BlueprintGenerator:
       zone: str,
       auth_cidr: str,
       system_node_pool_machine_type: str,
+      cluster_version: str,
+      release_channel: ReleaseChannel,
       reservation: Optional[str | None] = None,
       gcs_bucket: Optional[str | None] = None,
       num_nodes: int = 2,
@@ -761,12 +797,19 @@ class BlueprintGenerator:
             },
         },
     )
+
+    sanitized_version = cluster_version.replace("-", "+", 1)
+    version = parse(sanitized_version)
+    version_prefix = f"{version.major}.{version.minor}"
     cluster_id = f"{cluster_name}-a4-cluster"
     a4_cluster = DeploymentModule(
         id=cluster_id,
         source="modules/scheduler/gke-cluster",
         use=[net_0_id],
         settings={
+            "release_channel": release_channel.value,
+            "version_prefix": version_prefix,
+            "min_cluster_version": cluster_version,
             "system_node_pool_machine_type": system_node_pool_machine_type,
             "system_node_pool_node_count": {
                 "total_min_nodes": system_node_pool_min_node_count,
@@ -791,25 +834,20 @@ class BlueprintGenerator:
                 " alias_ip_range=[]}],"
                 f" {cluster_name}-rdma-net.subnetwork_interfaces_gke))"
             ),
-            "version_prefix": "1.32.",
-            "release_channel": (
-                "UNSPECIFIED"
-                if capacity_type == CapacityType.FLEX_START
-                else "RAPID"
-            ),
-            "maintenance_exclusions": (
-                []
-                if capacity_type == CapacityType.FLEX_START
-                else [{
-                    "name": "no-minor-or-node-upgrades-indefinite",
-                    "start_time": "2024-12-01T00:00:00Z",
-                    "end_time": "2025-12-22T00:00:00Z",
-                    "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
-                }]
-            ),
         },
         outputs=["instructions"],
     )
+    if release_channel != ReleaseChannel.RAPID:
+      a4_cluster.set_setting(
+          "maintenance_exclusions",
+          [{
+              "name": "no-minor-or-node-upgrades-indefinite",
+              "start_time": "2024-12-01T00:00:00Z",
+              "end_time": "2026-01-16T00:00:00Z",
+              "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+          }],
+      )
+
     system, _ = get_system_characteristics_by_device_type(a4_device_type)
     if system is None:
       xpk_print(
@@ -859,6 +897,9 @@ class BlueprintGenerator:
     else:
       gpu_pool.update_settings({"static_node_count": num_nodes})
 
+    if release_channel == ReleaseChannel.RAPID:
+      gpu_pool.set_setting("auto_upgrade", True)
+
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
         id=workload_manager_install_id,
@@ -1019,7 +1060,6 @@ class BlueprintGenerator:
         "enable_flex_start": True,
         "enable_queued_provisioning": True,
         "autoscaling_total_min_nodes": 0,
-        "release_channel": "UNSPECIFIED",
         "auto_repair": False,
         "auto_upgrade": False,
     }

xpk/core/blueprint/blueprint_test.py

@@ -22,6 +22,7 @@ import ruamel.yaml
 from xpk.core.blueprint.blueprint_definitions import Blueprint
 from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
 from xpk.core.capacity import CapacityType
+from xpk.utils.versions import ReleaseChannel
 
 yaml = ruamel.yaml.YAML()
 
@@ -60,6 +61,8 @@ def test_generate_a3_mega_blueprint():
       reservation="test-reservation",
       capacity_type=CapacityType.RESERVATION,
       system_node_pool_min_node_count=5,
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
 
   assert bp.blueprint_file.endswith("/prefix/xpk-gke-a3-megagpu.yaml")
@@ -99,6 +102,8 @@ def test_generate_a3_mega_spot_blueprint():
       auth_cidr="10.0.0.0/32",
       capacity_type=CapacityType.SPOT,
       system_node_pool_min_node_count=5,
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
 
   assert bp.blueprint_file.endswith("/prefix/xpk-gke-a3-megagpu.yaml")
@@ -135,6 +140,8 @@ def test_generate_a3_ultra_blueprint():
       capacity_type=CapacityType.RESERVATION,
       gcs_bucket="test-bucket",
       prefix="testdir",
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
   with open(a3_ultra_yaml_test_path, encoding="utf-8") as stream:
     ctk_yaml = yaml.load(stream)
@@ -180,6 +187,8 @@ def test_generate_a4_blueprint():
       capacity_type=CapacityType.RESERVATION,
       gcs_bucket="test-bucket",
       prefix="testdir",
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
   with open(a4_yaml_test_path, encoding="utf-8") as stream:
     ctk_yaml = yaml.load(stream)