Placement Policy removed from A3-Mega blueprints with --spot (#478)

* Placement Policy removed from A3-Mega blueprints with --spot

* Kueue config updated to support optional topologyName

* a3_mega_spot.yaml file

* tcpx(o) container moved to initContainers

Author: sharabiani

src/xpk/blueprints/a3mega/kueue-xpk-configuration.yaml.tftpl

@@ -16,7 +16,7 @@ metadata:
 spec:
   nodeLabels:
     cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
-  topologyName: "gke-default"
+  ${tas_name}
 ---
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: ClusterQueue

src/xpk/commands/batch.py

@@ -26,14 +26,13 @@
 from ..core.kjob import (
     AppProfileDefaults,
     JobTemplateDefaults,
-    Kueue_TAS_annotation,
     get_storage_annotations,
     prepare_kjob,
 )
 from ..core.kueue import LOCAL_QUEUE_NAME
 from ..utils.console import xpk_exit, xpk_print
 from .kind import set_local_cluster_command
-from .kjob_common import add_gpu_networking_annotations_to_command
+from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
 
 
 def batch(args: Namespace) -> None:
@@ -68,11 +67,11 @@ def submit_job(args: Namespace) -> None:
       'kubectl kjob create slurm'
       f' --profile {AppProfileDefaults.NAME.value}'
       f' --localqueue {LOCAL_QUEUE_NAME}'
-      f' --pod-template-annotation {Kueue_TAS_annotation}'
       f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
       ' --first-node-ip'
   )
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
+  cmd = add_TAS_annotations_to_command(args, cmd)
 
   for annotation in get_storage_annotations(args):
     cmd += f' --pod-template-annotation {annotation}'

src/xpk/commands/common.py

@@ -15,8 +15,10 @@
 """
 
 from ..core.commands import run_command_with_updates_retry
+from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
+from ..core.capacity import H100_MEGA_DEVICE_TYPE, CapacityType
 from ..core.gcloud_context import zone_to_region
-from ..utils.console import xpk_print
+from ..utils.console import xpk_print, xpk_exit
 
 
 def set_cluster_command(args) -> int:
@@ -42,3 +44,32 @@ def set_cluster_command(args) -> int:
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')
   return return_code
+
+
+def is_TAS_possible(args) -> bool:
+  """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    True if possible and False otherwise.
+  """
+  system_characteristics = get_cluster_system_characteristics(args)
+  capacity_type = get_cluster_capacity_type(args)
+
+  if system_characteristics is None:
+    xpk_print('system_characteristics data was not found in configmaps.')
+    xpk_exit(1)
+
+  if capacity_type is None:
+    xpk_print('capacity_type data was not found in configmaps.')
+    xpk_exit(1)
+
+  if (
+      system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
+      and capacity_type == CapacityType.SPOT
+  ):
+    return False
+
+  return True

src/xpk/commands/kjob_common.py

@@ -24,7 +24,9 @@
     get_a3mega_pod_template_annotations,
     get_a3ultra_pod_template_annotations,
     get_a4_pod_template_annotations,
+    Kueue_TAS_annotation,
 )
+from .common import is_TAS_possible
 
 
 def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
@@ -45,3 +47,10 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
   cmd += "\\\n".join(flags)
 
   return cmd
+
+
+def add_TAS_annotations_to_command(args, cmd: str) -> str:
+  if is_TAS_possible(args):
+    cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
+
+  return cmd

src/xpk/commands/run.py

@@ -25,14 +25,13 @@
 from ..core.kjob import (
     AppProfileDefaults,
     JobTemplateDefaults,
-    Kueue_TAS_annotation,
     get_storage_annotations,
     prepare_kjob,
 )
 from ..core.kueue import LOCAL_QUEUE_NAME
 from ..utils.console import xpk_exit, xpk_print
 from .kind import set_local_cluster_command
-from .kjob_common import add_gpu_networking_annotations_to_command
+from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
 
 
 def run(args: Namespace) -> None:
@@ -64,12 +63,12 @@ def submit_job(args: Namespace) -> None:
       'kubectl kjob create slurm --profile'
       f' {AppProfileDefaults.NAME.value} '
       f' --localqueue {LOCAL_QUEUE_NAME} '
-      f" --pod-template-annotation '{Kueue_TAS_annotation}'"
       f' --stream-container {JobTemplateDefaults.CONTAINER_NAME.value}'
       f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
       ' --wait --rm  --first-node-ip'
   )
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
+  cmd = add_TAS_annotations_to_command(args, cmd)
 
   for annotation in get_storage_annotations(args):
     cmd += f' --pod-template-annotation {annotation}'

src/xpk/commands/workload.py

@@ -89,6 +89,7 @@
 )
 from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
+from .common import is_TAS_possible
 from . import cluster_gcluster
 
 WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2
@@ -213,7 +214,7 @@
               labels:
                 xpk.google.com/workload: {args.workload}
               annotations:
-                kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
+                {kueue_TAS_annotation}
             spec:
               priorityClassName: {args.priority}
               restartPolicy: Never
@@ -447,13 +448,21 @@ def workload_create(args) -> None:
     if return_code != 0:
       xpk_exit(return_code)
 
+    kueue_TAS_annotation = (
+        'kueue.x-k8s.io/podset-preferred-topology:'
+        ' "cloud.google.com/gce-topology-host"'
+    )
+    if not is_TAS_possible(args):
+      kueue_TAS_annotation = ''
+
     if system.device_type in cluster_gcluster.supported_device_types:
       yml_string = A3_GPU_WORKLOAD_CREATE_YAML.format(
           args=args,
           container=container,
           service_account=XPK_SA,
           failure_policy_rules=failure_policy_rules,
           pod_failure_policy=pod_failure_policy,
+          kueue_TAS_annotation=kueue_TAS_annotation,
       )
 
       sub_networks = get_cluster_subnetworks(args)

src/xpk/core/blueprint/blueprint_generator.py

@@ -180,7 +180,7 @@ def generate_a3_mega_blueprint(
     a3_megagpu_pool_0 = DeploymentModule(
         id="a3_megagpu_pool_0",
         source="modules/compute/gke-node-pool",
-        use=["gke_cluster", gpu_subnets_name, "group_placement_0"],
+        use=["gke_cluster", gpu_subnets_name],
         settings={
             "name": f"{cluster_name}-a3-megagpu-pool-0",
             "machine_type": system.gce_machine_type,
@@ -197,6 +197,9 @@ def generate_a3_mega_blueprint(
         },
         outputs=["instructions"],
     )
+
+    set_placement_policy = capacity_type != CapacityType.SPOT
+    tas_name = "topologyName: 'gke-default'" if set_placement_policy else ""
     num_chips = num_nodes * system.chips_per_vm
     workload = DeploymentModule(
         id="workload_component_install",
@@ -207,7 +210,10 @@ def generate_a3_mega_blueprint(
                 "install": True,
                 "version": "v0.10.0",  # 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},
+                "config_template_vars": {
+                    "num_chips": num_chips,
+                    "tas_name": tas_name,
+                },
             },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [{
@@ -243,12 +249,16 @@ def generate_a3_mega_blueprint(
             primary_vpc,
             gpunets,
             gke_cluster,
-            group_placement_0,
             a3_megagpu_pool_0,
             workload,
             workload_configmap,
         ],
     )
+
+    if set_placement_policy:
+      a3_megagpu_pool_0.use.append(group_placement_0.id)
+      primary_group.modules.append(group_placement_0)
+
     a3_mega_blueprint = Blueprint(
         terraform_backend_defaults=self._getblock_terraform_backend(
             gcs_bucket, cluster_name, prefix

src/xpk/core/resources.py

@@ -236,3 +236,24 @@ def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
       return system
 
   return None
+
+
+def get_cluster_capacity_type(args) -> CapacityType | None:
+  """Get systemCharcteristics based on the cluster resources configMap
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    returns system characteristics
+  """
+  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+
+  if cluster_config_map is None:
+    return None
+
+  capacityValue = cluster_config_map.get('capacity_type')
+  if capacityValue is not None:
+    return CapacityType[capacityValue.upper()]
+
+  return None

src/xpk/core/tests/data/a3_mega.yaml

@@ -70,12 +70,6 @@ deployment_groups:
         gvnic_postfix: "-subnet"
         gvnic_start_index: 0
     outputs: [instructions]
-  - !DeploymentModule
-    id: group_placement_0
-    source: modules/compute/resource-policy
-    settings:
-      name: bar-gp-np-0
-      group_placement_max_distance: 2
 
   - !DeploymentModule
     id: a3_megagpu_pool_0
@@ -108,6 +102,7 @@ deployment_groups:
         config_path: $(ghpc_stage("xpk-gke-a3-megagpu"))/kueue-xpk-configuration.yaml.tftpl
         config_template_vars:
           num_chips: 16
+          tas_name: "topologyName: 'gke-default'"
       jobset:
         install: true
         version: v0.7.2
@@ -127,4 +122,11 @@ deployment_groups:
           cluster_config_name: "bar-metadata-configmap",
           capacity_type: "reservation",
           reservation: "test-reservation",
-          }
+          }
+
+  - !DeploymentModule
+    id: group_placement_0
+    source: modules/compute/resource-policy
+    settings:
+      name: bar-gp-np-0
+      group_placement_max_distance: 2