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

xpk/commands/batch.py

@@ -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:
@@ -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}'

xpk/commands/cluster.py

@@ -16,19 +16,22 @@ 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_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
@@ -52,8 +55,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 +77,123 @@ 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 +243,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 = {}
@@ -223,45 +313,12 @@ def cluster_create(args) -> None:
   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)
+  install_kueue(args, system, autoprovisioning_config)
 
-  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)
-
-  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 +326,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 +836,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 +869,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 +879,9 @@ def run_gke_cluster_create_command(
   if args.enable_pd_csi_driver:
     addons.append('GcePersistentDiskCsiDriver')
 
+  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 +891,87 @@ 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)
+
+
+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)
+
+
+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/common.py

@@ -15,8 +15,10 @@ 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
+from ..utils.console import xpk_print, xpk_exit
 
 
 def set_cluster_command(args) -> int:
@@ -31,6 +33,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 +45,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

xpk/commands/kjob_common.py

@@ -24,7 +24,9 @@ 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_gpu_networking_annotations_to_command(args, cmd: str) -> str:
@@ -35,7 +37,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 +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

xpk/commands/run.py

@@ -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:
@@ -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/storage.py

@@ -86,7 +86,6 @@ def storage_create(args: Namespace) -> None:
           args.vol,
           args.access_mode,
           filestore_network,
-          args.mount_options,
       )
 
     k8s_api_client = setup_k8s_env(args)
@@ -162,7 +161,6 @@ def storage_attach(args: Namespace) -> None:
           args.vol,
           args.access_mode,
           filestore_network,
-          args.mount_options,
       )
 
   elif args.type == GCS_FUSE_TYPE:
@@ -178,7 +176,11 @@ def storage_attach(args: Namespace) -> None:
         manifest = list(yaml.safe_load_all(f))
     else:
       manifest = gcsfuse.manifest(
-          args.name, args.bucket, args.size, args.mount_options
+          args.name,
+          args.bucket,
+          args.size,
+          args.mount_options,
+          args.prefetch_metadata,
       )
 
   elif args.type in [PARALLELSTORE_TYPE, GCE_PD_TYPE]:
@@ -323,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",
+  )

xpk/commands/workload.py

@@ -14,11 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.blueprint.blueprint_generator import (
-    get_subnetworks_for_a3mega,
-    get_subnetworks_for_a3ultra,
-    get_subnetworks_for_a4,
-)
 from ..core.cluster import (
     XPK_SA,
     create_xpk_k8s_service_account,
@@ -43,6 +38,7 @@ from ..core.nap import (
     get_autoprovisioning_node_selector_args,
     is_autoprovisioning_enabled,
 )
+from ..core.network import get_cluster_subnetworks
 from ..core.pathways import (
     append_custom_colocated_python_sidecar,
     append_custom_pathways_proxy_server,
@@ -93,6 +89,7 @@ from ..core.workload_decorators import (
 )
 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
@@ -217,7 +214,7 @@ spec:
               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
@@ -406,7 +403,7 @@ def workload_create(args) -> None:
           f' {parallelstore_storages}'
       )
     else:
-      xpk_print('No gcp filestore instances to add detected.')
+      xpk_print('No gcp parallelstore instances to add detected.')
 
     if len(pd_storages) > 0:
       service_account = XPK_SA
@@ -451,6 +448,13 @@ 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,
@@ -458,18 +462,16 @@ def workload_create(args) -> None:
           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)
       if args.device_type == cluster_gcluster.a3mega_device_type:
-        sub_networks = get_subnetworks_for_a3mega(args.cluster)
         yml_string = tcpxo_decorator.decorate_jobset(yml_string, sub_networks)
-
-      if args.device_type == cluster_gcluster.a3ultra_device_type:
-        sub_networks = get_subnetworks_for_a3ultra(args.cluster)
-        yml_string = rdma_decorator.decorate_jobset(yml_string, sub_networks)
-
-      if args.device_type == cluster_gcluster.a4_device_type:
-        sub_networks = get_subnetworks_for_a4()
+      elif args.device_type in [
+          cluster_gcluster.a3ultra_device_type,
+          cluster_gcluster.a4_device_type,
+      ]:
         yml_string = rdma_decorator.decorate_jobset(yml_string, sub_networks)
 
       if all_storages: