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

xpk/commands/batch.py

@@ -14,18 +14,26 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import re
 from argparse import Namespace
 
-from ..core.cluster import create_xpk_k8s_service_account
+from ..core.cluster import (
+    create_xpk_k8s_service_account,
+    get_cluster_credentials,
+)
 from ..core.commands import run_command_for_value
 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 .common import set_cluster_command
-from ..core.kjob import AppProfileDefaults, JobTemplateDefaults, prepare_kjob, Kueue_TAS_annotation, get_gcsfuse_annotation
-from .kjob_common import add_gpu_networking_annotations_to_command
 from .kind import set_local_cluster_command
-import re
+from .kjob_common import add_gpu_networking_annotations_to_command
 
 
 def batch(args: Namespace) -> None:
@@ -38,12 +46,11 @@ def batch(args: Namespace) -> None:
   """
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
   else:
     set_cluster_command_code = set_local_cluster_command(args)
-
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
   err_code = prepare_kjob(args)
   if err_code > 0:
@@ -66,9 +73,9 @@ def submit_job(args: Namespace) -> None:
       ' --first-node-ip'
   )
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
-  gcsfuse_annotation = get_gcsfuse_annotation(args)
-  if gcsfuse_annotation is not None:
-    cmd += f' --pod-template-annotation {gcsfuse_annotation}'
+
+  for annotation in get_storage_annotations(args):
+    cmd += f' --pod-template-annotation {annotation}'
 
   if args.ignore_unknown_flags:
     cmd += ' --ignore-unknown-flags'

xpk/commands/cluster.py

@@ -22,9 +22,13 @@ from ..core.cluster import (
     get_cluster_credentials,
     install_nccl_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,
+    update_cluster_with_gcpfilestore_driver_if_necessary,
+    update_cluster_with_parallelstore_driver_if_necessary,
+    update_cluster_with_pd_driver_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
@@ -46,7 +50,7 @@ from ..core.nap import enable_autoprovisioning_on_cluster
 from ..core.network import (
     create_cluster_network_config,
     delete_cluster_subnets,
-    set_up_cluster_network_for_gpu,
+    set_up_cluster_network_for_a3,
 )
 from ..core.nodepool import get_gke_node_pool_version, run_gke_node_pool_create_command
 from ..core.ray import install_ray_cluster
@@ -64,7 +68,6 @@ from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from . import cluster_gcluster
 from .common import set_cluster_command
-from ..core.cluster import update_cluster_with_gcpfilestore_driver_if_necessary
 
 
 def cluster_create(args) -> None:
@@ -117,11 +120,7 @@ def cluster_create(args) -> None:
 
   # 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
-      or args.enable_gcpfilestore_csi_driver
-  ):
+  if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
     update_cluster_command_code = (
         update_cluster_with_workload_identity_if_necessary(args)
     )
@@ -143,6 +142,20 @@ def cluster_create(args) -> None:
     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)
+
   # Update Pathways clusters with CloudDNS if not enabled already.
 
   get_cluster_credentials(args)
@@ -155,13 +168,12 @@ def cluster_create(args) -> None:
     if not tensorboard_config:
       xpk_exit(1)
 
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.device_type == H100_DEVICE_TYPE:
     xpk_print('Setting up Network for cluster')
-    set_up_cluster_network_code = set_up_cluster_network_for_gpu(args, system)
+    set_up_cluster_network_code = set_up_cluster_network_for_a3(args)
     if set_up_cluster_network_code != 0:
       xpk_exit(set_up_cluster_network_code)
 
-  if system.device_type == H100_DEVICE_TYPE:
     xpk_print('Creating Network Config for cluster')
     create_cluster_network_config_code = create_cluster_network_config(args)
     if create_cluster_network_config_code != 0:
@@ -207,6 +219,10 @@ def cluster_create(args) -> None:
   if set_jobset_on_cluster_code != 0:
     xpk_exit(set_jobset_on_cluster_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:
@@ -783,20 +799,21 @@ def run_gke_cluster_create_command(
   if args.enable_ray_cluster:
     command += ' --addons RayOperator'
 
-  if (
-      args.enable_workload_identity
-      or args.enable_gcsfuse_csi_driver
-      or args.enable_gcpfilestore_csi_driver
-  ):
+  if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
     command += f' --workload-pool={args.project}.svc.id.goog'
 
   addons = []
   if args.enable_gcsfuse_csi_driver:
     addons.append('GcsFuseCsiDriver')
-
   if args.enable_gcpfilestore_csi_driver:
     addons.append('GcpFilestoreCsiDriver')
 
+  if args.enable_parallelstore_csi_driver:
+    addons.append('ParallelstoreCsiDriver')
+
+  if args.enable_pd_csi_driver:
+    addons.append('GcePersistentDiskCsiDriver')
+
   if len(addons) > 0:
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'

xpk/commands/cluster_gcluster.py

@@ -16,26 +16,27 @@ limitations under the License.
 
 import os
 
-from ..core.remote_state.remote_state_client import RemoteStateClient
-from ..core.remote_state.fuse_remote_state import FuseStateClient
 from ..core.blueprint.blueprint_generator import (
     BlueprintGenerator,
     BlueprintGeneratorOutput,
     a3mega_device_type,
     a3ultra_device_type,
+    a4_device_type,
     supported_device_types,
 )
-from ..core.commands import run_command_for_value
 from ..core.capacity import get_capacity_type
+from ..core.cluster import get_cluster_credentials
+from ..core.commands import run_command_for_value
 from ..core.docker_manager import DockerManager
 from ..core.gcloud_context import zone_to_region
 from ..core.gcluster_manager import GclusterManager
+from ..core.kjob import apply_kjob_crds, prepare_kjob
+from ..core.remote_state.fuse_remote_state import FuseStateClient
+from ..core.remote_state.remote_state_client import RemoteStateClient
 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.cluster import get_cluster_credentials
-from ..core.kjob import apply_kjob_crds, prepare_kjob
 
 blueprints_path = os.path.abspath('xpkclusters/blueprints')
 gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')
@@ -266,4 +267,20 @@ def generate_blueprint(
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
           gcs_bucket=args.cluster_state_gcs_bucket,
       )
+    if args.device_type == a4_device_type:
+      num_nodes = args.num_nodes if not args.num_nodes is None else 2
+      return bpg.generate_a4_blueprint(
+          blueprint_name=blueprint_name,
+          prefix=prefix,
+          cluster_name=args.cluster,
+          region=zone_to_region(args.zone),
+          project_id=args.project,
+          zone=args.zone,
+          auth_cidr=all_IPs_cidr,
+          num_nodes=num_nodes,
+          reservation=args.reservation if args.reservation else None,
+          capacity_type=capacity_type,
+          system_node_pool_machine_type=args.default_pool_cpu_machine_type,
+          system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
+      )
   return None

xpk/commands/info.py

@@ -20,10 +20,10 @@ from argparse import Namespace
 from tabulate import tabulate
 
 from ..core.commands import run_command_for_value
+from ..core.cluster import get_cluster_credentials
 from ..core.gcloud_context import add_zone_and_project
 from ..core.kueue import verify_kueuectl
 from ..utils.console import xpk_exit, xpk_print
-from .common import set_cluster_command
 
 table_fmt = 'plain'
 
@@ -37,9 +37,7 @@ def info(args: Namespace) -> None:
     None
   """
   add_zone_and_project(args)
-  set_cluster_command_code = set_cluster_command(args)
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+  get_cluster_credentials(args)
 
   verify_kueuectl(args)
   lq, cq = bool(args.localqueue), bool(args.clusterqueue)

xpk/commands/job.py

@@ -20,10 +20,10 @@ import sys
 from ruamel.yaml import YAML
 
 from ..core.commands import run_command_for_value, run_command_with_updates
+from ..core.cluster import get_cluster_credentials
 from ..core.gcloud_context import add_zone_and_project
 from ..core.kjob import AppProfileDefaults
 from ..utils.console import xpk_exit, xpk_print
-from .common import set_cluster_command
 from .kind import set_local_cluster_command
 
 
@@ -143,14 +143,14 @@ def job_list(args) -> None:
   """
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
     msg = f'Listing jobs for project {args.project} and zone {args.zone}:'
   else:
     set_cluster_command_code = set_local_cluster_command(args)
     msg = 'Listing jobs:'
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
   xpk_print(msg, flush=True)
 
   return_code = run_slurm_job_list_command(args)
@@ -178,12 +178,11 @@ def job_cancel(args) -> None:
   xpk_print(f'Starting job cancel for job: {args.name}', flush=True)
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
   else:
     set_cluster_command_code = set_local_cluster_command(args)
-
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
   return_code = run_slurm_job_delete_command(args)
   xpk_exit(return_code)

xpk/commands/kjob_common.py

@@ -14,31 +14,34 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.kjob import get_a3mega_pod_template_annotations, get_a3ultra_pod_template_annotations
-from ..core.capacity import H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
+from ..core.capacity import (
+    B200_DEVICE_TYPE,
+    H100_MEGA_DEVICE_TYPE,
+    H200_DEVICE_TYPE,
+)
 from ..core.cluster import get_gpu_type_from_cluster
-
-
-def add_tcpxo_annotations(args, cmd: str) -> str:
-  tcpxo, interfaces, eth0 = get_a3mega_pod_template_annotations(args)
-  cmd += f" --pod-template-annotation {tcpxo} \\\n"
-  cmd += f" --pod-template-annotation {eth0} \\\n"
-  cmd += f" --pod-template-annotation {interfaces} "
-  return cmd
-
-
-def add_rdma_annotations(args, cmd) -> str:
-  eth0, interfaces = get_a3ultra_pod_template_annotations(args)
-  cmd += f" --pod-template-annotation {eth0} \\\n"
-  cmd += f" --pod-template-annotation {interfaces} \\\n"
-  return cmd
+from ..core.kjob import (
+    get_a3mega_pod_template_annotations,
+    get_a3ultra_pod_template_annotations,
+    get_a4_pod_template_annotations,
+)
 
 
 def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
   gpu_type = get_gpu_type_from_cluster(args)
 
   if gpu_type == H100_MEGA_DEVICE_TYPE:
-    return add_tcpxo_annotations(args, cmd)
-  if gpu_type == H200_DEVICE_TYPE:
-    return add_rdma_annotations(args, cmd)
+    annotations = get_a3mega_pod_template_annotations(args)
+  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()
+  else:
+    annotations = []
+
+  flags = [
+      f" --pod-template-annotation {annotation} " for annotation in annotations
+  ]
+  cmd += "\\\n".join(flags)
+
   return cmd

xpk/commands/run.py

@@ -16,15 +16,23 @@ limitations under the License.
 
 from argparse import Namespace
 
-from ..core.cluster import create_xpk_k8s_service_account
+from ..core.cluster import (
+    create_xpk_k8s_service_account,
+    get_cluster_credentials,
+)
 from ..core.commands import run_command_with_full_controls
 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 .common import set_cluster_command
-from ..core.kjob import JobTemplateDefaults, AppProfileDefaults, prepare_kjob, Kueue_TAS_annotation, get_gcsfuse_annotation
-from .kjob_common import add_gpu_networking_annotations_to_command
 from .kind import set_local_cluster_command
+from .kjob_common import add_gpu_networking_annotations_to_command
 
 
 def run(args: Namespace) -> None:
@@ -37,12 +45,11 @@ def run(args: Namespace) -> None:
   """
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
   else:
     set_cluster_command_code = set_local_cluster_command(args)
-
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
   err_code = prepare_kjob(args)
   if err_code > 0:
@@ -64,9 +71,8 @@ def submit_job(args: Namespace) -> None:
   )
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
 
-  gcsfuse_annotation = get_gcsfuse_annotation(args)
-  if gcsfuse_annotation is not None:
-    cmd += f' --pod-template-annotation {gcsfuse_annotation}'
+  for annotation in get_storage_annotations(args):
+    cmd += f' --pod-template-annotation {annotation}'
 
   if args.timeout:
     cmd += f' --wait-timeout {args.timeout}s'

xpk/commands/shell.py

@@ -20,7 +20,7 @@ from ..core.kjob import (
     AppProfileDefaults,
     prepare_kjob,
     get_pod_template_interactive_command,
-    get_gcsfuse_annotation,
+    get_storage_annotations,
 )
 
 exit_instructions = 'To exit the shell input "exit".'
@@ -89,9 +89,8 @@ def connect_to_new_interactive_shell(args: Namespace) -> int:
       f' {AppProfileDefaults.NAME.value} --pod-running-timeout 180s'
   )
 
-  gcsfuse_annotation = get_gcsfuse_annotation(args)
-  if gcsfuse_annotation is not None:
-    cmd += f' --pod-template-annotation {gcsfuse_annotation}'
+  for annotation in get_storage_annotations(args):
+    cmd += f' --pod-template-annotation {annotation}'
 
   return run_command_with_full_controls(
       command=cmd,

xpk/commands/storage.py

@@ -27,6 +27,8 @@ from ..core.cluster import (
     add_zone_and_project,
     get_cluster_network,
     setup_k8s_env,
+    update_cluster_with_parallelstore_driver_if_necessary,
+    update_cluster_with_pd_driver_if_necessary,
     update_cluster_with_gcpfilestore_driver_if_necessary,
     update_cluster_with_gcsfuse_driver_if_necessary,
     update_cluster_with_workload_identity_if_necessary,
@@ -41,6 +43,8 @@ from ..core.kjob import (
 from ..core.storage import (
     GCP_FILESTORE_TYPE,
     GCS_FUSE_TYPE,
+    GCE_PD_TYPE,
+    PARALLELSTORE_TYPE,
     STORAGE_CRD_PLURAL,
     XPK_API_GROUP_NAME,
     XPK_API_GROUP_VERSION,
@@ -78,7 +82,11 @@ def storage_create(args: Namespace) -> None:
         manifest = list(yaml.safe_load_all(f))
     else:
       manifest = filestore_client.manifest(
-          args.name, args.vol, args.access_mode, filestore_network
+          args.name,
+          args.vol,
+          args.access_mode,
+          filestore_network,
+          args.mount_options,
       )
 
     k8s_api_client = setup_k8s_env(args)
@@ -86,9 +94,10 @@ def storage_create(args: Namespace) -> None:
     create_volume_bundle_instance(
         k8s_api_client, args.name, manifest, args.readonly, args.mount_point
     )
-    return_code = update_cluster_with_workload_identity_if_necessary(args)
-    if return_code > 0:
-      xpk_exit(return_code)
+    # Not required for Filestore. Will be uncommented when adding GCSFuse create
+    # return_code = update_cluster_with_workload_identity_if_necessary(args)
+    # if return_code > 0:
+    #   xpk_exit(return_code)
     return_code = update_cluster_with_gcpfilestore_driver_if_necessary(args)
     if return_code > 0:
       xpk_exit(return_code)
@@ -131,6 +140,7 @@ def storage_delete(args: Namespace) -> None:
 
 def storage_attach(args: Namespace) -> None:
   add_zone_and_project(args)
+  manifest = [{}]
   if args.type == GCP_FILESTORE_TYPE:
     if args.instance is None:
       args.instance = args.name
@@ -148,10 +158,14 @@ def storage_attach(args: Namespace) -> None:
     else:
       filestore_network = get_cluster_network(args)
       manifest = filestore_client.manifest(
-          args.name, args.vol, args.access_mode, filestore_network
+          args.name,
+          args.vol,
+          args.access_mode,
+          filestore_network,
+          args.mount_options,
       )
 
-  else:  # args.type == GCS_FUSE_TYPE:
+  elif args.type == GCS_FUSE_TYPE:
     if args.manifest is None and args.size is None:
       xpk_print("--size is required when attaching gcsfuse storage.")
       xpk_exit(1)
@@ -164,30 +178,61 @@ def storage_attach(args: Namespace) -> None:
         manifest = list(yaml.safe_load_all(f))
     else:
       manifest = gcsfuse.manifest(
-          name=args.name, bucket=args.bucket, size=args.size
+          args.name, args.bucket, args.size, args.mount_options
+      )
+
+  elif args.type in [PARALLELSTORE_TYPE, GCE_PD_TYPE]:
+    if args.manifest is None:
+      xpk_print(
+          "Parallelstore and PersistentDisk are currently supported only with"
+          " --manifest"
       )
+      xpk_exit(1)
+
+    with open(args.manifest, "r", encoding="utf-8") as f:
+      manifest = list(yaml.safe_load_all(f))
+
+  else:
+    xpk_print(f"Storage type {args.type} is not supported.")
+    xpk_exit(1)
 
   k8s_api_client = setup_k8s_env(args)
   create_storage_crds(k8s_api_client, args, manifest)
   create_volume_bundle_instance(
       k8s_api_client, args.name, manifest, args.readonly, args.mount_point
   )
-  return_code = update_cluster_with_workload_identity_if_necessary(args)
-  if return_code > 0:
-    xpk_exit(return_code)
-
-  # args.type can have only two values after parsing
-  return_code = (
-      update_cluster_with_gcsfuse_driver_if_necessary(args)
-      if args.type == GCS_FUSE_TYPE
-      else update_cluster_with_gcpfilestore_driver_if_necessary(args)
-  )
-  if return_code > 0:
-    xpk_exit(return_code)
+
+  enable_csi_drivers_if_necessary(args)
 
   apply_kubectl_manifest(k8s_api_client, manifest)
 
 
+def enable_csi_drivers_if_necessary(args: Namespace) -> None:
+  if args.type == GCS_FUSE_TYPE:
+    return_code = update_cluster_with_workload_identity_if_necessary(args)
+    if return_code > 0:
+      xpk_exit(return_code)
+
+    return_code = update_cluster_with_gcsfuse_driver_if_necessary(args)
+    if return_code > 0:
+      xpk_exit(return_code)
+
+  if args.type == GCP_FILESTORE_TYPE:
+    return_code = update_cluster_with_gcpfilestore_driver_if_necessary(args)
+    if return_code > 0:
+      xpk_exit(return_code)
+
+  if args.type == PARALLELSTORE_TYPE:
+    return_code = update_cluster_with_parallelstore_driver_if_necessary(args)
+    if return_code > 0:
+      xpk_exit(return_code)
+
+  if args.type == GCE_PD_TYPE:
+    return_code = update_cluster_with_pd_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)
   storages = list_storages(k8s_api_client)