xpk 0.15.0__py3-none-any.whl → 0.16.1__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.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.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.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

@@ -31,6 +31,8 @@ from ..capacity import (
     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
@@ -42,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,

xpk/core/blueprint/testing/data/a3_mega.yaml

@@ -0,0 +1,129 @@
+# Copyright 2024 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
+#
+#     http://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.
+
+---
+!Blueprint
+blueprint_name: xpk-gke-a3-megagpu
+toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
+toolkit_modules_version: v1.62.2
+
+vars:
+  project_id: "foo"
+  deployment_name: xpk-gke-a3-megagpu
+  region: us-central1
+  zone: us-central1-c
+  labels: {gke_product_type: xpk}
+
+deployment_groups:
+- !DeploymentGroup
+  group: primary
+  modules:
+  - !DeploymentModule
+    id: network1
+    source: modules/network/vpc
+    settings:
+      subnetwork_name: bar-xpk-gke-a3-megagpu-subnet
+      secondary_ranges:
+        bar-xpk-gke-a3-megagpu-subnet:
+        - range_name: pods
+          ip_cidr_range: 10.4.0.0/14
+        - range_name: services
+          ip_cidr_range: 10.0.32.0/20
+  - !DeploymentModule
+    id: gpunets
+    source: modules/network/multivpc
+    settings:
+      network_name_prefix: bar-gpunet
+      global_ip_address_range: 192.169.0.0/16
+      network_count: 8
+      subnetwork_cidr_suffix: 24
+  - !DeploymentModule
+    id: gke_cluster
+    source: modules/scheduler/gke-cluster
+    use: [network1, gpunets]
+    settings:
+      release_channel: RAPID
+      version_prefix: '1.2'
+      min_master_version: 1.2.3
+      prefix_with_deployment_name: false
+      name_suffix: bar
+      enable_private_endpoint: false
+      enable_gcsfuse_csi: true
+      enable_filestore_csi: true
+      master_authorized_networks:
+      - cidr_block: 10.0.0.0/32  # Allows your machine run kubectl command. It's required for the multi-network setup.
+        display_name: "kubectl-access-network"
+      system_node_pool_machine_type: "e2-standard-32"
+      system_node_pool_node_count:
+        total_min_nodes: 5
+        total_max_nodes: 1000
+      k8s_network_names:
+        gvnic_prefix: "bar-gpunet-"
+        gvnic_postfix: "-subnet"
+        gvnic_start_index: 0
+    outputs: [instructions]
+
+  - !DeploymentModule
+    id: a3_megagpu_pool_0
+    source: modules/compute/gke-node-pool
+    use: [gke_cluster, gpunets]
+    settings:
+      name: bar-a3-megagpu-pool-0
+      machine_type: a3-megagpu-8g
+      zones: [us-central1-c]
+      host_maintenance_interval: 'PERIODIC'
+      reservation_affinity:
+        consume_reservation_type: SPECIFIC_RESERVATION
+        specific_reservations:
+        - name: test-reservation
+      run_workload_script: false
+      spot: false
+      max_pods_per_node: 32
+      guest_accelerator:
+      - type: nvidia-h100-mega-80gb
+        count: 8
+        gpu_driver_installation_config:
+          gpu_driver_version: "LATEST"
+      auto_upgrade: true
+      static_node_count: 2
+      placement_policy:
+        type: COMPACT
+        name: test-reservation-placement
+    outputs: [instructions]
+
+  - !DeploymentModule
+    id: workload_component_install
+    source: modules/management/kubectl-apply
+    use: [gke_cluster]
+    settings:
+      jobset:
+        install: true
+        version: v0.7.2
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a3-megagpu"))/storage_crd.yaml
+
+  - !DeploymentModule
+    id: workload_configmap
+    source: modules/management/kubectl-apply
+    use: [gke_cluster]
+    settings:
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a3-megagpu"))/config-map.yaml.tftpl
+        template_vars: {
+          resource_config_name: "bar-resources-configmap",
+          num_nodes: "2",
+          cluster_config_name: "bar-metadata-configmap",
+          capacity_type: "reservation",
+          reservation: "test-reservation",
+          }

xpk/core/blueprint/testing/data/a3_mega_spot.yaml

@@ -0,0 +1,125 @@
+# Copyright 2024 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
+#
+#     http://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.
+
+---
+!Blueprint
+blueprint_name: xpk-gke-a3-megagpu
+toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
+toolkit_modules_version: v1.62.2
+
+vars:
+  project_id: "foo"
+  deployment_name: xpk-gke-a3-megagpu
+  region: us-central1
+  zone: us-central1-c
+  labels: {gke_product_type: xpk}
+
+deployment_groups:
+- !DeploymentGroup
+  group: primary
+  modules:
+  - !DeploymentModule
+    id: network1
+    source: modules/network/vpc
+    settings:
+      subnetwork_name: bar-xpk-gke-a3-megagpu-subnet
+      secondary_ranges:
+        bar-xpk-gke-a3-megagpu-subnet:
+        - range_name: pods
+          ip_cidr_range: 10.4.0.0/14
+        - range_name: services
+          ip_cidr_range: 10.0.32.0/20
+  - !DeploymentModule
+    id: gpunets
+    source: modules/network/multivpc
+    settings:
+      network_name_prefix: bar-gpunet
+      global_ip_address_range: 192.169.0.0/16
+      network_count: 8
+      subnetwork_cidr_suffix: 24
+  - !DeploymentModule
+    id: gke_cluster
+    source: modules/scheduler/gke-cluster
+    use: [network1, gpunets]
+    settings:
+      release_channel: RAPID
+      version_prefix: '1.2'
+      min_master_version: 1.2.3
+      prefix_with_deployment_name: false
+      name_suffix: bar
+      enable_private_endpoint: false
+      enable_gcsfuse_csi: true
+      enable_filestore_csi: true
+      master_authorized_networks:
+      - cidr_block: 10.0.0.0/32  # Allows your machine run kubectl command. It's required for the multi-network setup.
+        display_name: "kubectl-access-network"
+      system_node_pool_machine_type: "e2-standard-32"
+      system_node_pool_node_count:
+        total_min_nodes: 5
+        total_max_nodes: 1000
+      k8s_network_names:
+        gvnic_prefix: "bar-gpunet-"
+        gvnic_postfix: "-subnet"
+        gvnic_start_index: 0
+    outputs: [instructions]
+
+  - !DeploymentModule
+    id: a3_megagpu_pool_0
+    source: modules/compute/gke-node-pool
+    use: [gke_cluster, gpunets]
+    settings:
+      name: bar-a3-megagpu-pool-0
+      machine_type: a3-megagpu-8g
+      zones: [us-central1-c]
+      host_maintenance_interval: PERIODIC
+      reservation_affinity:
+        consume_reservation_type: NO_RESERVATION
+        specific_reservations: []
+      run_workload_script: false
+      max_pods_per_node: 32
+      spot: true
+      guest_accelerator:
+      - type: nvidia-h100-mega-80gb
+        count: 8
+        gpu_driver_installation_config:
+          gpu_driver_version: "LATEST"
+      auto_upgrade: true
+      static_node_count: 2
+    outputs: [instructions]
+
+  - !DeploymentModule
+    id: workload_component_install
+    source: modules/management/kubectl-apply
+    use: [gke_cluster]
+    settings:
+      jobset:
+        install: true
+        version: v0.7.2
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a3-megagpu"))/storage_crd.yaml
+
+  - !DeploymentModule
+    id: workload_configmap
+    source: modules/management/kubectl-apply
+    use: [gke_cluster]
+    settings:
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a3-megagpu"))/config-map.yaml.tftpl
+        template_vars: {
+          resource_config_name: "bar-resources-configmap",
+          num_nodes: "2",
+          cluster_config_name: "bar-metadata-configmap",
+          capacity_type: "spot",
+          reservation: "None",
+          }