Update Docker Resource Limits for Pathways Main Workload

Author: SujeethJinesh

src/xpk/core/cluster.py

@@ -617,14 +617,14 @@ def update_cluster_with_workload_identity_if_necessary(args) -> int:
 
   if is_workload_identity_enabled_on_cluster(args):
     return 0
-  cluster_update_return_code = (
-      update_gke_cluster_with_workload_identity_enabled(args)
-  )
-  if cluster_update_return_code > 0:
-    xpk_print(
-        'Updating GKE cluster to enable Workload Identity Federation failed!'
-    )
-    return cluster_update_return_code
+  # cluster_update_return_code = (
+  #     update_gke_cluster_with_workload_identity_enabled(args)
+  # )
+  # if cluster_update_return_code > 0:
+  #   xpk_print(
+  #       'Updating GKE cluster to enable Workload Identity Federation failed!'
+  #   )
+  #   return cluster_update_return_code
 
   return 0
 

src/xpk/core/docker_resources.py

@@ -34,7 +34,7 @@ def get_main_container_resources(
       Workload resources port as a YAML string
   """
   # Resources requirements for Pathways workload containers are known.
-  resources_yaml = """cpu: "24"
+  resources_yaml = f"""cpu: "24"
                     memory: 100G"""
   if args.use_pathways:
     return resources_yaml

src/xpk/core/nodepool.py

@@ -310,18 +310,26 @@ def run_gke_node_pool_create_command(
     create_commands.append(command)
     create_task_names.append(task)
 
-  desired_pw_cpu_node_pools = ['cpu-np']
+  desired_pw_cpu_node_pools = ['cpu-np'] #, 'highmem-cpu-np']
   if args.enable_pathways:
     # Pathways needs CPU nodepools in addition to TPU nodepools
     for node_pool_name in desired_pw_cpu_node_pools:
       if node_pool_name in existing_node_pool_names:
         continue
+      # if node_pool_name == 'cpu-np':
       command = (
           'gcloud beta container node-pools create'
           f' {node_pool_name} --node-version={gke_node_pool_version} --cluster={args.cluster} --project={args.project} --node-locations={args.zone} --region={zone_to_region(args.zone)} --num-nodes=1'
           f' --machine-type={args.pathways_gce_machine_type} --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL} --enable-autoscaling'
           ' --min-nodes=1 --max-nodes=20'
       )
+      # else:
+      #   command = (
+      #       'gcloud beta container node-pools create'
+      #       f' {node_pool_name} --node-version={gke_node_pool_version} --cluster={args.cluster} --project={args.project} --node-locations={args.zone} --region={zone_to_region(args.zone)} --num-nodes=1'
+      #       f' --machine-type={args.pathways_highmem_gce_machine_type} --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL} --enable-autoscaling'
+      #       ' --min-nodes=1 --max-nodes=20'
+      #   )
       task = f'NodepoolCreate-{node_pool_name}'
       create_commands.append(command)
       create_task_names.append(task)

src/xpk/core/pathways.py

@@ -33,6 +33,14 @@ def add_pw_resource_flavors(args):
   nodeLabels:
     cloud.google.com/gke-nodepool: cpu-np
 ---"""
+# apiVersion: kueue.x-k8s.io/v1beta1
+# kind: ResourceFlavor
+# metadata:
+#   name: highmem-cpu-user
+# spec:
+#   nodeLabels:
+#     cloud.google.com/gke-nodepool: highmem-cpu-np
+# ---"""
   if args.enable_pathways:
     return resource_flavor_yaml
   return ''
@@ -48,6 +56,12 @@ def add_pw_resources_to_kueue(args):
         nominalQuota: 480
       - name: "memory"
         nominalQuota: 2000G"""
+    # - name: highmem-cpu-user
+    #   resources:
+    #   - name: "cpu"
+    #     nominalQuota: 480
+    #   - name: "memory"
+    #     nominalQuota: 4000G"""
   if args.enable_pathways:
     return resources_yaml
   return ''
@@ -79,7 +93,7 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool:
 
   # Ensure the cluster and CPU nodepools were created with create-pathways
   all_node_pools = get_all_nodepools_programmatic(args)
-  desired_pw_cpu_node_pools = {'cpu-np'}
+  desired_pw_cpu_node_pools = {'cpu-np'} #, 'highmem-cpu-np'}
   if not desired_pw_cpu_node_pools.issubset(set(all_node_pools[0])):
     xpk_print(
         'Cluster needs to be created with `xpk create-pathways` to run'
@@ -263,8 +277,6 @@ def get_user_workload_for_pathways(
           spec:
             containers:
               {container}
-            nodeSelector:
-              cloud.google.com/gke-nodepool: cpu-np
             hostNetwork: true
             dnsPolicy: ClusterFirstWithHostNet
             restartPolicy: Never

src/xpk/parser/cluster.py

@@ -528,6 +528,12 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
       default='n2-standard-64',
       help='The CPU type for Pathways CPU nodepools',
   )
+#   parser.add_argument(
+#       '--pathways-highmem-gce-machine-type',
+#       type=str,
+#       default='c4-highmem-192',
+#       help='The highmem CPU type for Pathways CPU nodepools',
+#   )
   parser.add_argument(
       '--default-pool-cpu-machine-type',
       type=str,