Set network names based on cluster configuration when running workloads on A3 Mega, A3 Ultra and A4 (#457)

* automatically detect cluster network names when running workloads

* fix imports after merge

Author: gcie

src/xpk/commands/kjob_common.py

@@ -35,7 +35,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 = []
 

src/xpk/commands/workload.py

@@ -14,11 +14,6 @@
 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 @@
     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,
@@ -460,16 +456,13 @@ def workload_create(args) -> None:
           pod_failure_policy=pod_failure_policy,
       )
 
+      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:

src/xpk/core/blueprint/blueprint_generator.py

@@ -52,20 +52,6 @@
 cluster_toolkit_version = "v1.48.0"
 
 
-def get_subnetworks_for_a3mega(cluster_name: str) -> list[str]:
-  return [f"{cluster_name}-gpunet-{i}-subnet" for i in range(8)]
-
-
-def get_subnetworks_for_a3ultra(cluster_name: str) -> list[str]:
-  return [f"{cluster_name}-sub-1"] + [
-      f"{cluster_name}-rdma-sub-{i}" for i in range(8)
-  ]
-
-
-def get_subnetworks_for_a4() -> list[str]:
-  return ["gvnic-1"] + [f"rdma-{i}" for i in range(8)]
-
-
 class BlueprintGeneratorOutput:
   """BlueprintGeneratorOutput is a class containing fields with output blueprint file path and path to blueprint dependencies.
   Atributes:

src/xpk/core/kjob.py

@@ -22,24 +22,9 @@
 from kubernetes.client import ApiClient
 from kubernetes.client.rest import ApiException
 
-from ..core.blueprint.blueprint_generator import (
-    get_subnetworks_for_a3mega,
-    get_subnetworks_for_a3ultra,
-    get_subnetworks_for_a4,
-)
-from ..core.capacity import (
-    H100_DEVICE_TYPE,
-    H100_MEGA_DEVICE_TYPE,
-    H200_DEVICE_TYPE,
-)
-from ..core.storage import GCS_FUSE_ANNOTATIONS, PARALLELSTORE_ANNOTATIONS
-from ..core.workload_decorators import (
-    rdma_decorator,
-    tcpx_decorator,
-    tcpxo_decorator,
-)
 from ..utils import templates
 from ..utils.console import xpk_exit, xpk_print
+from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .cluster import DEFAULT_NAMESPACE, XPK_SA, setup_k8s_env
 from .commands import (
     run_command_for_value,
@@ -54,16 +39,24 @@
     KJOB_SHELL_WORKING_DIRECTORY,
     XpkConfig,
 )
+from .network import get_cluster_subnetworks
 from .resources import (
     AcceleratorType,
     SystemCharacteristics,
     get_cluster_system_characteristics,
 )
 from .storage import (
+    GCS_FUSE_ANNOTATIONS,
+    PARALLELSTORE_ANNOTATIONS,
     get_auto_mount_gcsfuse_storages,
     get_auto_mount_parallelstore_storages,
     get_auto_mount_storages,
 )
+from .workload_decorators import (
+    rdma_decorator,
+    tcpx_decorator,
+    tcpxo_decorator,
+)
 from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
 
 KJOB_API_GROUP_NAME = "kjobctl.x-k8s.io"
@@ -176,8 +169,8 @@ class PodTemplateDefaults(Enum):
 default_interface_annotation = "networking.gke.io/default-interface=eth0"
 
 
-def get_a4_pod_template_annotations() -> tuple[str, str]:
-  sub_networks = get_subnetworks_for_a4()
+def get_a4_pod_template_annotations(args) -> tuple[str, str]:
+  sub_networks = get_cluster_subnetworks(args)
   interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
       sub_networks
   )
@@ -189,7 +182,7 @@ def get_a4_pod_template_annotations() -> tuple[str, str]:
 
 
 def get_a3ultra_pod_template_annotations(args: Namespace) -> tuple[str, str]:
-  sub_networks = get_subnetworks_for_a3ultra(args.cluster)
+  sub_networks = get_cluster_subnetworks(args)
   interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
       sub_networks
   )
@@ -204,7 +197,7 @@ def get_a3mega_pod_template_annotations(
     args: Namespace,
 ) -> tuple[str, str, str]:
   """Adds or updates annotations in the Pod template."""
-  sub_networks = get_subnetworks_for_a3mega(args.cluster)
+  sub_networks = get_cluster_subnetworks(args)
   tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
   interfaces_key, interfaces_value = tcpxo_decorator.get_interfaces_entry(
       sub_networks

src/xpk/core/network.py

@@ -14,7 +14,7 @@
 limitations under the License.
 """
 
-from ..utils.console import xpk_print
+from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from .commands import run_command_for_value, run_command_with_updates
 from .gcloud_context import zone_to_region
@@ -235,6 +235,28 @@ def create_cluster_network_config(args) -> int:
   return 0
 
 
+def get_cluster_subnetworks(args) -> list[str]:
+  """Gets the list of cluster networks.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    list[str]: list of cluster networks
+  """
+  command = 'kubectl get GKENetworkParamSet'
+  return_code, stdout = run_command_for_value(
+      command, 'Get Cluster Networks', args
+  )
+  if return_code != 0:
+    xpk_print('GKE Cluster Get NetworkParamSet failed')
+    xpk_exit(return_code)
+
+  networks = [line.split()[0] for line in stdout.splitlines()][1:]
+
+  return networks
+
+
 def set_up_cluster_network_for_a3(args) -> int:
   """Set up GKE Cluster networks, subnets and firewall rules for A3.
   Note: there are 4 NICs for GPU-GPU bw and 1 NIC for host in an A3 node.

src/xpk/core/workload_decorators/rdma_decorator.py

@@ -68,16 +68,12 @@ def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
 
 
 def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:
-  interfaces = [
-      '[',
-      '    {"interfaceName":"eth0","network":"default"},',
-      *[
-          f'    {{"interfaceName":"eth{i + 1}","network":"{sub_networks[i]}"}}{"," if i<8 else ""}'
-          for i in range(9)
-      ],
-      ']',
-  ]
-  return 'networking.gke.io/interfaces', literal_string('\n'.join(interfaces))
+  entries = ',\n'.join([
+      f'    {{"interfaceName":"eth{i}","network":"{network}"}}'
+      for i, network in enumerate(sub_networks)
+  ])
+  interfaces = f'[\n{entries}\n]'
+  return 'networking.gke.io/interfaces', literal_string(interfaces)
 
 
 def add_annotations(job_manifest: dict, sub_networks: list[str]):

src/xpk/core/workload_decorators/tcpxo_decorator.py

@@ -77,16 +77,12 @@ def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
 
 
 def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:
-  interfaces = [
-      '[',
-      '    {"interfaceName":"eth0","network":"default"},',
-      *[
-          f'    {{"interfaceName":"eth{i + 1}","network":"{sub_networks[i]}"}}{"," if i<7 else ""}'
-          for i in range(8)
-      ],
-      ']',
-  ]
-  return 'networking.gke.io/interfaces', literal_string('\n'.join(interfaces))
+  entries = ',\n'.join([
+      f'    {{"interfaceName":"eth{i}","network":"{network}"}}'
+      for i, network in enumerate(sub_networks)
+  ])
+  interfaces = f'[\n{entries}\n]'
+  return 'networking.gke.io/interfaces', literal_string(interfaces)
 
 
 def get_tcpxo_deamon_entry() -> tuple[str, str]: