xpk 0.9.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
@@ -54,14 +54,14 @@ 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'

xpk/commands/cluster.py

@@ -31,6 +31,7 @@ from ..core.cluster import (
     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
@@ -42,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 (
@@ -170,7 +173,6 @@ def cluster_adapt(args) -> None:
   install_kueue(args, system, autoprovisioning_config)
 
   install_kjob(args)
-
   if system.accelerator_type == AcceleratorType['GPU']:
     prepare_gpus(args, system)
 
@@ -308,6 +310,9 @@ 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:
@@ -879,6 +884,10 @@ 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')
 
@@ -922,6 +931,13 @@ def install_storage_csis(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')
@@ -957,6 +973,11 @@ def install_kueue(args, system: SystemCharacteristics, 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')

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,10 +15,12 @@ limitations under the License.
 """
 
 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, xpk_exit
+from ..core.system_characteristics import (
+    SystemCharacteristics,
+)
 
 
 def set_cluster_command(args) -> int:
@@ -47,7 +49,11 @@ def set_cluster_command(args) -> int:
   return return_code
 
 
-def is_TAS_possible(args) -> bool:
+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:
@@ -56,8 +62,6 @@ def is_TAS_possible(args) -> bool:
   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.')
@@ -67,9 +71,12 @@ def is_TAS_possible(args) -> bool:
     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.SPOT
+      and capacity_type != CapacityType.RESERVATION
   ):
     return False
 

xpk/commands/kjob_common.py

@@ -27,6 +27,7 @@ from ..core.kjob import (
     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:
@@ -50,7 +51,9 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
 
 
 def add_TAS_annotations_to_command(args, cmd: str) -> str:
-  if is_TAS_possible(args):
+  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
@@ -53,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)
 

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'

xpk/commands/storage.py

@@ -29,6 +29,7 @@ from ..core.cluster import (
     setup_k8s_env,
     update_cluster_with_parallelstore_driver_if_necessary,
     update_cluster_with_pd_driver_if_necessary,
+    update_cluster_with_lustre_driver_if_necessary,
     update_cluster_with_gcpfilestore_driver_if_necessary,
     update_cluster_with_gcsfuse_driver_if_necessary,
     update_cluster_with_workload_identity_if_necessary,
@@ -45,6 +46,7 @@ from ..core.storage import (
     GCS_FUSE_TYPE,
     GCE_PD_TYPE,
     PARALLELSTORE_TYPE,
+    LUSTRE_TYPE,
     STORAGE_CRD_PLURAL,
     XPK_API_GROUP_NAME,
     XPK_API_GROUP_VERSION,
@@ -183,11 +185,11 @@ def storage_attach(args: Namespace) -> None:
           args.prefetch_metadata,
       )
 
-  elif args.type in [PARALLELSTORE_TYPE, GCE_PD_TYPE]:
+  elif args.type in [PARALLELSTORE_TYPE, GCE_PD_TYPE, LUSTRE_TYPE]:
     if args.manifest is None:
       xpk_print(
-          "Parallelstore and PersistentDisk are currently supported only with"
-          " --manifest"
+          "Parallelstore, PersistentDisk, and Lustre are currently supported"
+          " only with --manifest"
       )
       xpk_exit(1)
 
@@ -234,6 +236,11 @@ def enable_csi_drivers_if_necessary(args: Namespace) -> None:
     if return_code > 0:
       xpk_exit(return_code)
 
+  if args.type == LUSTRE_TYPE:
+    return_code = update_cluster_with_lustre_driver_if_necessary(args)
+    if return_code > 0:
+      xpk_exit(return_code)
+
 
 def storage_list(args: Namespace) -> None:
   k8s_api_client = setup_k8s_env(args)

xpk/commands/workload.py

@@ -14,23 +14,25 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from ..core.blueprint.blueprint_generator import (
+    a3high_device_type,
+    a3mega_device_type,
+    a3ultra_device_type,
+    a4_device_type,
+)
 from ..core.cluster import (
     XPK_SA,
-    create_xpk_k8s_service_account,
+    setup_k8s_service_accounts,
     get_cluster_credentials,
     setup_k8s_env,
 )
 from ..core.commands import run_command_with_updates, run_commands
-from ..core.config import (
-    VERTEX_TENSORBOARD_FEATURE_FLAG,
-    XPK_CURRENT_VERSION,
-    parse_env_config,
-)
+from ..core.config import (VERTEX_TENSORBOARD_FEATURE_FLAG, XPK_CURRENT_VERSION)
 from ..core.docker_container import (
     get_main_container_docker_image,
     get_user_workload_container,
 )
-from ..core.docker_resources import get_volumes
+from ..core.docker_resources import get_volumes, parse_env_config
 from ..core.gcloud_context import add_zone_and_project
 from ..core.kueue import LOCAL_QUEUE_NAME
 from ..core.monitoring import get_gke_outlier_dashboard
@@ -50,6 +52,10 @@ from ..core.pathways import (
     get_user_workload_for_pathways,
     try_to_delete_pathwaysjob_first,
 )
+from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
+from ..core.capacity import (
+    CapacityType,
+)
 from ..core.resources import CLUSTER_METADATA_CONFIGMAP, get_cluster_configmap
 from ..core.scheduling import (
     check_if_workload_can_schedule,
@@ -65,6 +71,7 @@ from ..core.storage import (
     GCP_FILESTORE_TYPE,
     GCS_FUSE_TYPE,
     PARALLELSTORE_TYPE,
+    LUSTRE_TYPE,
     Storage,
     add_bucket_iam_members,
     get_storage_annotations,
@@ -76,7 +83,6 @@ from ..core.system_characteristics import (
 )
 from ..core.vertex import create_vertex_experiment
 from ..core.workload import (
-    add_gpu_rxdm_container,
     check_if_workload_exists,
     get_workload_list,
     wait_for_job_completion,
@@ -85,12 +91,13 @@ from ..core.workload import (
 from ..core.workload_decorators import (
     rdma_decorator,
     storage_decorator,
+    tcpx_decorator,
     tcpxo_decorator,
 )
 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
+from .common import is_TAS_possible
 
 WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2
 kind: JobSet
@@ -123,6 +130,8 @@ spec:
                 {storage_annotations}
             spec:
               schedulerName: {args.scheduler}
+              imagePullSecrets:
+              - name: {args.docker_image_pull_secret}
               restartPolicy: Never
               {affinity}
               nodeSelector:
@@ -136,6 +145,8 @@ spec:
               containers:
               {container}
               serviceAccountName: {service_account}
+              tolerations:
+              {tpu_toleration}
               volumes:
               {volumes}
 """
@@ -175,6 +186,8 @@ spec:
               {gpu_scheduler}
               priorityClassName: {args.priority}
               restartPolicy: Never
+              imagePullSecrets:
+              - name: {args.docker_image_pull_secret}
               hostNetwork: true
               dnsPolicy: ClusterFirstWithHostNet
               terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
@@ -213,11 +226,12 @@ spec:
             metadata:
               labels:
                 xpk.google.com/workload: {args.workload}
-              annotations:
-                {kueue_TAS_annotation}
+              annotations: {annotations}
             spec:
               priorityClassName: {args.priority}
               restartPolicy: Never
+              imagePullSecrets:
+              - name: {args.docker_image_pull_secret}
               dnsPolicy: ClusterFirstWithHostNet
               terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
               serviceAccountName: {service_account}
@@ -291,7 +305,7 @@ def workload_create(args) -> None:
     0 if successful and 1 otherwise.
   """
   k8s_api_client = setup_k8s_env(args)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   workload_exists = check_if_workload_exists(args)
 
@@ -347,7 +361,7 @@ def workload_create(args) -> None:
     if not tensorboard_config:
       xpk_exit(1)
 
-  parse_env_config(args, tensorboard_config, system)
+  parse_env_config(args, tensorboard_config)
 
   autoprovisioning_args = ''
   autoprovisioning_enabled, return_code = is_autoprovisioning_enabled(
@@ -382,6 +396,9 @@ def workload_create(args) -> None:
     pd_storages: list[Storage] = list(
         filter(lambda storage: storage.type == GCE_PD_TYPE, storages)
     )
+    lustre_storages: list[Storage] = list(
+        filter(lambda storage: storage.type == LUSTRE_TYPE, storages)
+    )
     if len(gcs_fuse_storages) > 0:
       service_account = XPK_SA
       xpk_print(f'Detected gcsfuse Storages to add: {gcs_fuse_storages}')
@@ -411,11 +428,18 @@ def workload_create(args) -> None:
     else:
       xpk_print('No gce persistent disk instances to add detected.')
 
+    if len(lustre_storages) > 0:
+      service_account = XPK_SA
+      xpk_print(f'Detected managed lustre instances to add: {lustre_storages}')
+    else:
+      xpk_print('No managed lustre instances to add detected.')
+
     all_storages = (
         gcs_fuse_storages
         + gcpfilestore_storages
         + parallelstore_storages
         + pd_storages
+        + lustre_storages
     )
 
   # Currently failure policy rules are supported for Pathways workloads. b/408465881
@@ -447,31 +471,41 @@ 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"'
+    system_characteristics = get_cluster_system_characteristics(args)
+    capacity_type = get_cluster_capacity_type(args)
+
+    annotations = (
+        ''
+        if not is_TAS_possible(
+            system_characteristics,
+            capacity_type,
+            flex=True if capacity_type == CapacityType.FLEX_START else False,
+        )
+        else (
+            '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:
+    if (
+        system.device_type in cluster_gcluster.supported_device_types
+        or system.device_type == a3high_device_type
+    ):
       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,
+          annotations=annotations,
       )
 
       sub_networks = get_cluster_subnetworks(args)
-      if args.device_type == cluster_gcluster.a3mega_device_type:
+      if args.device_type == a3high_device_type:
+        yml_string = tcpx_decorator.decorate_jobset(yml_string)
+      elif args.device_type == a3mega_device_type:
         yml_string = tcpxo_decorator.decorate_jobset(yml_string, sub_networks)
-      elif args.device_type in [
-          cluster_gcluster.a3ultra_device_type,
-          cluster_gcluster.a4_device_type,
-      ]:
+      elif args.device_type in [a3ultra_device_type, a4_device_type]:
         yml_string = rdma_decorator.decorate_jobset(yml_string, sub_networks)
 
       if all_storages:
@@ -489,7 +523,6 @@ def workload_create(args) -> None:
           failure_policy_rules=failure_policy_rules,
           pod_failure_policy=pod_failure_policy,
       )
-      yml_string = add_gpu_rxdm_container(yml_string, system, all_storages)
 
   elif args.use_pathways and ensure_pathways_workload_prerequisites(
       args, system
@@ -526,6 +559,10 @@ def workload_create(args) -> None:
             get_storage_annotations(all_storages)
         ),
         service_account=service_account,
+        tpu_toleration="""
+              - operator: "Exists"
+                key: google.com/tpu
+        """ if system.accelerator_type == AcceleratorType['TPU'] else '',
         failure_policy_rules=failure_policy_rules,
         pod_failure_policy=pod_failure_policy,
     )