xpk 0.8.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

xpk/commands/batch.py

@@ -18,7 +18,7 @@ import re
 from argparse import Namespace
 
 from ..core.cluster import (
-    create_xpk_k8s_service_account,
+    setup_k8s_service_accounts,
     get_cluster_credentials,
 )
 from ..core.commands import run_command_for_value
@@ -26,14 +26,13 @@ from ..core.gcloud_context import add_zone_and_project
 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:
@@ -55,24 +54,24 @@ def batch(args: Namespace) -> None:
   err_code = prepare_kjob(args)
   if err_code > 0:
     xpk_exit(err_code)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   submit_job(args)
 
 
 def submit_job(args: Namespace) -> None:
 
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   cmd = (
       '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}'

xpk/commands/cluster.py

@@ -16,19 +16,23 @@ limitations under the License.
 
 from tabulate import tabulate
 
-from ..core.capacity import H100_DEVICE_TYPE
+from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE
 from ..core.cluster import (
     get_all_clusters_programmatic,
     get_cluster_credentials,
     install_nccl_on_cluster,
+    install_nri_on_cluster,
     set_jobset_on_cluster,
     set_pathways_job_on_cluster,
     setup_k8s_env,
-    update_cluster_with_gcsfuse_driver_if_necessary,
-    update_cluster_with_workload_identity_if_necessary,
+    disable_mglru_on_cluster,
+    count_nodes_on_cluster,
     update_cluster_with_gcpfilestore_driver_if_necessary,
+    update_cluster_with_gcsfuse_driver_if_necessary,
     update_cluster_with_parallelstore_driver_if_necessary,
     update_cluster_with_pd_driver_if_necessary,
+    update_cluster_with_lustre_driver_if_necessary,
+    update_cluster_with_workload_identity_if_necessary,
 )
 from ..core.cluster_private import authorize_private_cluster_access_if_necessary
 from ..core.commands import run_command_for_value, run_command_with_updates
@@ -39,12 +43,14 @@ from ..core.gcloud_context import (
     get_gke_server_config,
     zone_to_region,
 )
+from ..core.jobset import update_jobset_resources_if_necessary
 from ..core.kjob import apply_kjob_crds, prepare_kjob, verify_kjob_installed
 from ..core.kueue import (
     cluster_preheat_yml,
     install_kueue_crs,
     install_kueue_on_cluster,
     wait_for_kueue_available,
+    update_kueue_resources_if_necessary,
 )
 from ..core.nap import enable_autoprovisioning_on_cluster
 from ..core.network import (
@@ -52,8 +58,12 @@ from ..core.network import (
     delete_cluster_subnets,
     set_up_cluster_network_for_a3,
 )
-from ..core.nodepool import get_gke_node_pool_version, run_gke_node_pool_create_command
+from ..core.nodepool import (
+    get_gke_node_pool_version,
+    run_gke_node_pool_create_command,
+)
 from ..core.ray import install_ray_cluster
+from ..core.mtc import install_mtc_on_cluster
 from ..core.resources import create_cluster_configmaps
 from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
@@ -70,14 +80,122 @@ from . import cluster_gcluster
 from .common import set_cluster_command
 
 
+def cluster_adapt(args) -> None:
+  """Function that performs cluster adaptation.
+
+  Args:
+    args: user provided arguments for running the command.
+  """
+  args.enable_pathways = False
+
+  system, return_code = get_system_characteristics(args)
+
+  if return_code > 0:
+    xpk_print('Fetching system characteristics failed!')
+    xpk_exit(return_code)
+
+  xpk_print(
+      f'Starting cluster adaptation for cluster {args.cluster}:', flush=True
+  )
+  add_zone_and_project(args)
+
+  if system.accelerator_type == AcceleratorType['GPU'] and not getattr(
+      args, 'num_nodes'
+  ):
+    xpk_print(
+        'Argument --num-nodes was not provided, trying to determine number of'
+        ' nodes based on the available nodes in the cluster...'
+    )
+    args.num_nodes = count_nodes_on_cluster(args, system)
+    if args.num_nodes == 0:
+      xpk_print(
+          'Found unexpected number of nodes. Is the --device-type correct?'
+      )
+      xpk_exit(1)
+    else:
+      xpk_print(f'Using {args.num_nodes} nodes.')
+
+  # ToDo(roshanin@) - Re-enable CloudDNS on Pathways clusters conditionally.
+  # Enable WorkloadIdentity if not enabled already.
+  if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_workload_identity_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  get_cluster_credentials(args)
+
+  k8s_client = setup_k8s_env(args)
+
+  install_storage_crd(k8s_client)
+  install_storage_csis(args)
+
+  # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
+  tensorboard_config = {}
+  if VERTEX_TENSORBOARD_FEATURE_FLAG and args.create_vertex_tensorboard:
+    tensorboard_config = create_vertex_tensorboard(args)
+    # exit if failed to create Tensorboard in Vertex AI
+    if not tensorboard_config:
+      xpk_exit(1)
+
+  # Provision node pools dynamically based on incoming workloads:
+  # Currently autoprovisioning is not supported with Pathways.
+  autoprovisioning_config = None
+  if args.enable_autoprovisioning:
+    xpk_print('Enabling Autoprovisioning')
+    autoprovisioning_config, return_code = enable_autoprovisioning_on_cluster(
+        args, system
+    )
+    if return_code != 0:
+      xpk_exit(return_code)
+
+  xpk_print('Creating ConfigMap for cluster')
+  create_cluster_configmaps_code = create_cluster_configmaps(
+      args, system, tensorboard_config, autoprovisioning_config
+  )
+  if create_cluster_configmaps_code != 0:
+    xpk_exit(create_cluster_configmaps_code)
+
+  xpk_print(
+      'Enabling the jobset API on our cluster, to be deprecated when Jobset is'
+      ' globally available'
+  )
+  set_jobset_on_cluster_code = set_jobset_on_cluster(args)
+  if set_jobset_on_cluster_code != 0:
+    xpk_exit(set_jobset_on_cluster_code)
+
+  # TODO: Uncomment when cluster_adapt will support TPU cluters
+  # set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args)
+  # if set_pathways_job_on_cluster_code != 0:
+  #   xpk_exit(set_pathways_job_on_cluster_code)
+
+  install_kueue(args, system, autoprovisioning_config)
+
+  install_kjob(args)
+  if system.accelerator_type == AcceleratorType['GPU']:
+    prepare_gpus(args, system)
+
+  if args.enable_ray_cluster:
+    return_code = install_ray_cluster(args, system)
+    if return_code != 0:
+      xpk_print('Installation of RayCluster failed.')
+      xpk_exit(return_code)
+
+  xpk_print('GKE commands done! Resources are created.')
+  xpk_print(
+      'See your GKE Cluster here:'
+      # pylint: disable=line-too-long
+      f' https://console.cloud.google.com/kubernetes/clusters/details/{zone_to_region(args.zone)}/{args.cluster}/details?project={args.project}'
+  )
+  xpk_exit(0)
+
+
 def cluster_create(args) -> None:
   """Function around cluster creation.
 
   Args:
     args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
   """
   system, return_code = get_system_characteristics(args)
 
@@ -127,38 +245,12 @@ def cluster_create(args) -> None:
     if update_cluster_command_code != 0:
       xpk_exit(update_cluster_command_code)
 
-  # Enable GCSFuse CSI Driver if not enabled already.
-  if args.enable_gcsfuse_csi_driver:
-    update_cluster_command_code = (
-        update_cluster_with_gcsfuse_driver_if_necessary(args)
-    )
-    if update_cluster_command_code != 0:
-      xpk_exit(update_cluster_command_code)
-
-  if args.enable_gcpfilestore_csi_driver:
-    update_cluster_command_code = (
-        update_cluster_with_gcpfilestore_driver_if_necessary(args)
-    )
-    if update_cluster_command_code != 0:
-      xpk_exit(update_cluster_command_code)
-
-  if args.enable_parallelstore_csi_driver:
-    update_cluster_command_code = (
-        update_cluster_with_parallelstore_driver_if_necessary(args)
-    )
-    if update_cluster_command_code != 0:
-      xpk_exit(update_cluster_command_code)
-
-  if args.enable_pd_csi_driver:
-    update_cluster_command_code = update_cluster_with_pd_driver_if_necessary(
-        args
-    )
-    if update_cluster_command_code != 0:
-      xpk_exit(update_cluster_command_code)
+  get_cluster_credentials(args)
 
-  # Update Pathways clusters with CloudDNS if not enabled already.
+  k8s_client = setup_k8s_env(args)
 
-  get_cluster_credentials(args)
+  install_storage_crd(k8s_client)
+  install_storage_csis(args)
 
   # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
   tensorboard_config = {}
@@ -218,50 +310,20 @@ def cluster_create(args) -> None:
   set_jobset_on_cluster_code = set_jobset_on_cluster(args)
   if set_jobset_on_cluster_code != 0:
     xpk_exit(set_jobset_on_cluster_code)
+  update_jobset_resources_code = update_jobset_resources_if_necessary(args)
+  if update_jobset_resources_code != 0:
+    xpk_exit(update_jobset_resources_code)
 
   set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args)
   if set_pathways_job_on_cluster_code != 0:
     xpk_exit(set_pathways_job_on_cluster_code)
 
-  xpk_print('Enabling Kueue on the cluster')
-  install_kueue_on_cluster_code = install_kueue_on_cluster(args)
-  if install_kueue_on_cluster_code != 0:
-    xpk_exit(install_kueue_on_cluster_code)
-
-  xpk_print('Verifying kjob installation')
-  err_code = verify_kjob_installed(args)
-  if err_code > 0:
-    xpk_exit(err_code)
-
-  xpk_print('Applying kjob CDRs')
-  err_code = apply_kjob_crds(args)
-  if err_code > 0:
-    xpk_exit(err_code)
+  install_kueue(args, system, autoprovisioning_config)
 
-  err_code = prepare_kjob(args)
-  if err_code > 0:
-    xpk_exit(err_code)
-
-  k8s_client = setup_k8s_env(args)
-  install_storage_crd(k8s_client)
-
-  xpk_print('Wait for Kueue to be fully available')
-  wait_for_kueue_available_code = wait_for_kueue_available(args)
-  if wait_for_kueue_available_code != 0:
-    xpk_exit(wait_for_kueue_available_code)
-
-  xpk_print('Install Kueue Custom Resources')
-  enable_kueue_credentials_code = install_kueue_crs(
-      args, system, autoprovisioning_config
-  )
-  if enable_kueue_credentials_code != 0:
-    xpk_exit(enable_kueue_credentials_code)
+  install_kjob(args)
 
   if system.accelerator_type == AcceleratorType['GPU']:
-    xpk_print('Installing NCCL Plugin for cluster')
-    install_nccl_code = install_nccl_on_cluster(args, system)
-    if install_nccl_code != 0:
-      xpk_exit(install_nccl_code)
+    prepare_gpus(args, system)
 
   if args.enable_ray_cluster:
     return_code = install_ray_cluster(args, system)
@@ -269,6 +331,12 @@ def cluster_create(args) -> None:
       xpk_print('Installation of RayCluster failed.')
       xpk_exit(return_code)
 
+  if hasattr(args, 'enable_mtc') and args.enable_mtc:
+    return_code = install_mtc_on_cluster(args, system)
+    if return_code != 0:
+      xpk_print('Installation of MTC failed.')
+      xpk_exit(return_code)
+
   xpk_print('GKE commands done! Resources are created.')
   xpk_print(
       'See your GKE Cluster here:'
@@ -773,6 +841,7 @@ def run_gke_cluster_create_command(
       f' --num-nodes {args.default_pool_cpu_num_nodes}'
       f' {args.custom_cluster_arguments}'
       f' {rapid_release_cmd}'
+      ' --enable-dns-access'
   )
 
   enable_ip_alias = False
@@ -805,6 +874,7 @@ def run_gke_cluster_create_command(
   addons = []
   if args.enable_gcsfuse_csi_driver:
     addons.append('GcsFuseCsiDriver')
+
   if args.enable_gcpfilestore_csi_driver:
     addons.append('GcpFilestoreCsiDriver')
 
@@ -814,6 +884,13 @@ def run_gke_cluster_create_command(
   if args.enable_pd_csi_driver:
     addons.append('GcePersistentDiskCsiDriver')
 
+  if args.enable_lustre_csi_driver:
+    addons.append('LustreCsiDriver')
+    command += ' --enable-legacy-lustre-port'
+
+  if hasattr(args, 'enable_mtc') and args.enable_mtc:
+    addons.append('HighScaleCheckpointing')
+
   if len(addons) > 0:
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'
@@ -823,3 +900,99 @@ def run_gke_cluster_create_command(
     xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')
     return 1
   return 0
+
+
+def install_storage_csis(args):
+  if args.enable_gcsfuse_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_gcsfuse_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_gcpfilestore_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_gcpfilestore_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_parallelstore_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_parallelstore_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_pd_csi_driver:
+    update_cluster_command_code = update_cluster_with_pd_driver_if_necessary(
+        args
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_lustre_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_lustre_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+
+def install_kjob(args):
+  xpk_print('Verifying kjob installation')
+  err_code = verify_kjob_installed(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  xpk_print('Applying kjob CDRs')
+  err_code = apply_kjob_crds(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+
+def install_kueue(args, system: SystemCharacteristics, autoprovisioning_config):
+  xpk_print('Enabling Kueue on the cluster')
+  install_kueue_on_cluster_code = install_kueue_on_cluster(args)
+  if install_kueue_on_cluster_code != 0:
+    xpk_exit(install_kueue_on_cluster_code)
+
+  xpk_print('Wait for Kueue to be fully available')
+  wait_for_kueue_available_code = wait_for_kueue_available(args)
+  if wait_for_kueue_available_code != 0:
+    xpk_exit(wait_for_kueue_available_code)
+
+  xpk_print('Install Kueue Custom Resources')
+  enable_kueue_credentials_code = install_kueue_crs(
+      args, system, autoprovisioning_config
+  )
+  if enable_kueue_credentials_code != 0:
+    xpk_exit(enable_kueue_credentials_code)
+
+  xpk_print('Update Kueue Controller Manager resources')
+  update_kueue_resources_code = update_kueue_resources_if_necessary(args)
+  if update_kueue_resources_code != 0:
+    xpk_exit(update_kueue_resources_code)
+
+
+def prepare_gpus(args, system: SystemCharacteristics):
+  xpk_print('Installing NCCL Plugin for cluster')
+  install_nccl_code = install_nccl_on_cluster(args, system)
+  if install_nccl_code != 0:
+    xpk_exit(install_nccl_code)
+
+  if system.device_type == H100_DEVICE_TYPE:
+    xpk_print('Installing NRI device injector for cluster')
+    install_nri_code = install_nri_on_cluster(args)
+    if install_nri_code != 0:
+      xpk_exit(install_nri_code)
+
+  if system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
+    xpk_print('Disabling MGLRU')
+    err_code = disable_mglru_on_cluster(args)
+    if err_code > 0:
+      xpk_exit(err_code)

xpk/commands/cluster_gcluster.py

@@ -37,6 +37,7 @@ from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import ensure_directory_exists
 from ..utils.network import all_IPs_cidr
 from ..utils.objects import hash_string
+from ..core.capacity import get_reservation_maintenance_interval, get_reservation_placement_policy
 
 blueprints_path = os.path.abspath('xpkclusters/blueprints')
 gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')
@@ -234,6 +235,30 @@ def generate_blueprint(
   if args.device_type in supported_device_types:
     if args.device_type == a3mega_device_type:
       num_nodes = args.num_nodes if not args.num_nodes is None else 2
+
+      maintenance_interval = (
+          get_reservation_maintenance_interval(
+              args.reservation, args.zone, args.project
+          )
+          if args.reservation is not None
+          else 'PERIODIC'
+      )
+      placement_policy_name = (
+          get_reservation_placement_policy(
+              args.reservation, args.zone, args.project
+          )
+          if args.reservation is not None
+          else None
+      )
+      placement_policy = (
+          {
+              'type': 'COMPACT',
+              'name': placement_policy_name.split('/')[-1],
+          }
+          if placement_policy_name is not None
+          and len(placement_policy_name) > 0
+          else None
+      )
       return bpg.generate_a3_mega_blueprint(
           blueprint_name=blueprint_name,
           prefix=prefix,
@@ -243,6 +268,8 @@ def generate_blueprint(
           zone=args.zone,
           auth_cidr=all_IPs_cidr,
           num_nodes=num_nodes,
+          reservation_maintenance_interval=maintenance_interval,
+          reservation_placement_policy=placement_policy,
           reservation=args.reservation if args.reservation else None,
           capacity_type=capacity_type,
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,

xpk/commands/common.py

@@ -15,8 +15,12 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_with_updates_retry
+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
+from ..core.system_characteristics import (
+    SystemCharacteristics,
+)
 
 
 def set_cluster_command(args) -> int:
@@ -31,6 +35,7 @@ def set_cluster_command(args) -> int:
   command = (
       'gcloud container clusters get-credentials'
       f' {args.cluster} --region={zone_to_region(args.zone)}'
+      ' --dns-endpoint'
       f' --project={args.project} &&'
       ' kubectl config view && kubectl config set-context --current'
       ' --namespace=default'
@@ -42,3 +47,37 @@ 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(
+    system_characteristics: SystemCharacteristics,
+    capacity_type: CapacityType,
+    flex: bool,
+) -> 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.
+  """
+
+  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 flex:
+    return False
+
+  if (
+      system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
+      and capacity_type != CapacityType.RESERVATION
+  ):
+    return False
+
+  return True

xpk/commands/kjob_common.py

@@ -24,7 +24,10 @@ from ..core.kjob import (
     get_a3mega_pod_template_annotations,
     get_a3ultra_pod_template_annotations,
     get_a4_pod_template_annotations,
+    Kueue_TAS_annotation,
 )
+from .common import is_TAS_possible
+from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
 
 
 def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
@@ -35,7 +38,7 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
   elif gpu_type == H200_DEVICE_TYPE:
     annotations = get_a3ultra_pod_template_annotations(args)
   elif gpu_type == B200_DEVICE_TYPE:
-    annotations = get_a4_pod_template_annotations()
+    annotations = get_a4_pod_template_annotations(args)
   else:
     annotations = []
 
@@ -45,3 +48,12 @@ 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:
+  system_characteristics = get_cluster_system_characteristics(args)
+  capacity_type = get_cluster_capacity_type(args)
+  if is_TAS_possible(system_characteristics, capacity_type, flex=False):
+    cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
+
+  return cmd

xpk/commands/run.py

@@ -17,7 +17,7 @@ limitations under the License.
 from argparse import Namespace
 
 from ..core.cluster import (
-    create_xpk_k8s_service_account,
+    setup_k8s_service_accounts,
     get_cluster_credentials,
 )
 from ..core.commands import run_command_with_full_controls
@@ -25,14 +25,13 @@ from ..core.gcloud_context import add_zone_and_project
 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:
@@ -54,7 +53,7 @@ def run(args: Namespace) -> None:
   err_code = prepare_kjob(args)
   if err_code > 0:
     xpk_exit(err_code)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   submit_job(args)
 
@@ -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}'

xpk/commands/shell.py

@@ -12,7 +12,7 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_with_full_controls, run_command_for_value, run_command_with_updates
-from ..core.cluster import get_cluster_credentials, add_zone_and_project, create_xpk_k8s_service_account
+from ..core.cluster import get_cluster_credentials, add_zone_and_project, setup_k8s_service_accounts
 from ..utils.console import xpk_exit, xpk_print
 from argparse import Namespace
 
@@ -82,7 +82,7 @@ def connect_to_new_interactive_shell(args: Namespace) -> int:
   err_code = prepare_kjob(args)
   if err_code > 0:
     xpk_exit(err_code)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   cmd = (
       'kubectl-kjob create interactive --profile'