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

xpk/commands/kjob_common.py

@@ -14,31 +14,43 @@ 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
+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
 
 
-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_gpu_networking_annotations_to_command(args, cmd: str) -> str:
+  gpu_type = get_gpu_type_from_cluster(args)
 
+  if gpu_type == H100_MEGA_DEVICE_TYPE:
+    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(args)
+  else:
+    annotations = []
+
+  flags = [
+      f" --pod-template-annotation {annotation} " for annotation in annotations
+  ]
+  cmd += "\\\n".join(flags)
 
-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
 
 
-def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
-  gpu_type = get_gpu_type_from_cluster(args)
+def add_TAS_annotations_to_command(args, cmd: str) -> str:
+  if is_TAS_possible(args):
+    cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
 
-  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)
   return cmd

xpk/commands/run.py

@@ -16,15 +16,22 @@ 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,
+    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, add_TAS_annotations_to_command
 
 
 def run(args: Namespace) -> None:
@@ -37,12 +44,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:
@@ -57,16 +63,15 @@ 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)
 
-  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,10 @@ 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,
       )
 
     k8s_api_client = setup_k8s_env(args)
@@ -86,9 +93,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 +139,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 +157,13 @@ 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,
       )
 
-  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 +176,65 @@ 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,
+          args.prefetch_metadata,
       )
 
+  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)
@@ -278,3 +325,12 @@ def delete_storage_resources(k8s_api_client: ApiClient, storage: Storage):
       storage.name,
       "Storage",
   )
+
+  # remove kubernetes.io/pvc-protection
+  delete_resource(
+      lambda name: core_api.patch_namespaced_persistent_volume_claim(
+          name, "default", {"metadata": {"finalizers": None}}
+      ),
+      storage.pvc,
+      "Persistent Volume Claim finalizers",
+  )