xpk 0.12.0__py3-none-any.whl → 0.14.0__py3-none-any.whl

xpk/core/blueprint/blueprint_generator.py

@@ -32,7 +32,6 @@ from ..capacity import (
 )
 from ..system_characteristics import get_system_characteristics_by_device_type
 from .blueprint_definitions import Blueprint, DeploymentGroup, DeploymentModule
-from ..kueue import KUEUE_VERSION
 
 yaml_parser = yaml.YAML()
 
@@ -53,6 +52,7 @@ blueprint_dependencies_dir = {
 
 cluster_toolkit_url = "github.com/GoogleCloudPlatform/cluster-toolkit"
 cluster_toolkit_version = "v1.62.2"
+common_cluster_labels = {"gke_product_type": "xpk"}
 
 
 class BlueprintGeneratorOutput:
@@ -216,26 +216,11 @@ class BlueprintGenerator:
       a3_megagpu_pool_0.settings.update({"static_node_count": num_nodes})
 
     set_placement_policy = capacity_type != CapacityType.SPOT
-    num_chips = num_nodes * system.chips_per_vm
     workload = DeploymentModule(
         id="workload_component_install",
         source="modules/management/kubectl-apply",
         use=["gke_cluster"],
         settings={
-            "kueue": {
-                "install": True,
-                "version": KUEUE_VERSION,  # TAS feature-gates is enabled in CT
-                "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {
-                    "num_chips": num_chips,
-                    "reservation": (
-                        1 if capacity_type == CapacityType.RESERVATION else 0
-                    ),
-                    "flex_start": (
-                        1 if capacity_type == CapacityType.FLEX_START else 0
-                    ),
-                },
-            },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [{
                 "source": f'$(ghpc_stage("{blueprint_name}"))/storage_crd.yaml'
@@ -298,6 +283,7 @@ class BlueprintGenerator:
             "deployment_name": blueprint_name,
             "region": region,
             "zone": zone,
+            "labels": common_cluster_labels,
         },
     )
 
@@ -598,24 +584,12 @@ class BlueprintGenerator:
     else:
       gpu_pool.settings.update({"static_node_count": num_nodes})
 
-    num_chips = num_nodes * system.chips_per_vm
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
         id=workload_manager_install_id,
         source="modules/management/kubectl-apply",
         use=[cluster_id],
         settings={
-            "kueue": {
-                "install": True,
-                "version": KUEUE_VERSION,  # TAS feature-gates is enabled in CT
-                "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {
-                    "num_chips": num_chips,
-                    "flex_start": (
-                        1 if capacity_type == CapacityType.FLEX_START else 0
-                    ),
-                },
-            },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [
                 {"source": nccl_installer_path},
@@ -676,6 +650,7 @@ class BlueprintGenerator:
             "deployment_name": blueprint_name,
             "region": region,
             "zone": zone,
+            "labels": common_cluster_labels,
         },
     )
 
@@ -884,24 +859,12 @@ class BlueprintGenerator:
     else:
       gpu_pool.settings.update({"static_node_count": num_nodes})
 
-    num_chips = num_nodes * system.chips_per_vm
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
         id=workload_manager_install_id,
         source="modules/management/kubectl-apply",
         use=[cluster_id],
         settings={
-            "kueue": {
-                "install": True,
-                "version": KUEUE_VERSION,  # TAS feature-gates is enabled in CT
-                "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {
-                    "num_chips": num_chips,
-                    "flex_start": (
-                        1 if capacity_type == CapacityType.FLEX_START else 0
-                    ),
-                },
-            },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [
                 {"source": nccl_installer_path},
@@ -962,6 +925,7 @@ class BlueprintGenerator:
             "deployment_name": blueprint_name,
             "region": region,
             "zone": zone,
+            "labels": common_cluster_labels,
         },
     )
 

xpk/core/blueprint/blueprint_test.py

@@ -32,7 +32,6 @@ a3_spot_yaml_test_path = "src/xpk/core/blueprint/testing/data/a3_mega_spot.yaml"
 a3_ultra_yaml_test_path = "src/xpk/core/blueprint/testing/data/a3_ultra.yaml"
 a4_yaml_test_path = "src/xpk/core/blueprint/testing/data/a4.yaml"
 config_map_filename = "config-map.yaml.tftpl"
-kueue_conf_filename = "kueue-xpk-configuration.yaml.tftpl"
 tmp_test_dir = "/tmp/xpk_test"
 
 
@@ -82,11 +81,6 @@ def test_generate_a3_mega_blueprint():
               tmp_test_dir, "prefix", blueprint_name, config_map_filename
           )
       )
-      assert os.path.exists(
-          os.path.join(
-              tmp_test_dir, "prefix", blueprint_name, kueue_conf_filename
-          )
-      )
 
   shutil.rmtree(tmp_test_dir)
 

xpk/core/capacity.py

@@ -17,6 +17,7 @@ limitations under the License.
 import enum
 
 from ..utils.console import xpk_print, xpk_exit
+from ..utils.kueue import is_queued_cluster
 from .commands import run_command_with_updates, run_command_for_value
 
 AUTOPROVISIONING_CONFIG_VALUE = 'AUTOPROVISION'
@@ -50,7 +51,7 @@ def print_reservations(args) -> int:
   """
   command = f'gcloud beta compute reservations list --project={args.project}'
   return_code = run_command_with_updates(
-      command, 'Get all reservations in the project', args
+      command, 'Get all reservations in the project'
   )
   if return_code != 0:
     xpk_print(f'Get all reservations returned ERROR {return_code}')
@@ -119,7 +120,7 @@ def get_reservation_maintenance_interval(
       f' --project={project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
   )
   return_code, output = run_command_for_value(
-      command, 'Get reservation maintenance interval', None
+      command, 'Get reservation maintenance interval'
   )
   if return_code != 0:
     xpk_print(f'Get reservation maintenance interval ERROR {return_code}')
@@ -143,7 +144,7 @@ def get_reservation_placement_policy(
       f' --project={project} --zone={zone} --format="value(resourcePolicies.policy)"'
   )
   return_code, output = run_command_for_value(
-      command, 'Get reservation placement policy', None
+      command, 'Get reservation placement policy'
   )
   if return_code != 0:
     xpk_print(f'Get reservation placement policy ERROR {return_code}')
@@ -164,7 +165,7 @@ def verify_reservation_exists(args) -> int:
       f'gcloud beta compute reservations describe {args.reservation}'
       f' --project={args.project} --zone={args.zone}'
   )
-  return_code = run_command_with_updates(command, 'Describe reservation', args)
+  return_code = run_command_with_updates(command, 'Describe reservation')
   if return_code != 0:
     xpk_print(f'Describe reservation returned ERROR {return_code}')
     xpk_print('Please confirm that your reservation name is correct.')
@@ -199,7 +200,7 @@ def get_capacity_arguments_from_capacity_type(
           ' --location-policy=ANY --reservation-affinity=none'
           f' --no-enable-autorepair --max-nodes={max_nodes}'
       )
-      if args.num_slices <= 1:
+      if is_queued_cluster(args.num_slices):
         capacity_args += ' --enable-queued-provisioning'
     case CapacityType.RESERVATION:
       capacity_args = (