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

xpk/core/blueprint/blueprint_generator.py

@@ -22,32 +22,34 @@ from ruamel import yaml
 
 from ...utils.console import xpk_exit, xpk_print
 from ...utils.file import ensure_directory_exists
-from ..capacity import H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE, CapacityType
+from ..capacity import (
+    B200_DEVICE_TYPE,
+    H100_MEGA_DEVICE_TYPE,
+    H200_DEVICE_TYPE,
+    CapacityType,
+)
 from ..system_characteristics import get_system_characteristics_by_device_type
 from .blueprint_definitions import Blueprint, DeploymentGroup, DeploymentModule
 
+
 yaml = yaml.YAML()
 
 a3mega_device_type = H100_MEGA_DEVICE_TYPE
 a3ultra_device_type = H200_DEVICE_TYPE
-supported_device_types = {a3mega_device_type, a3ultra_device_type}
+a4_device_type = B200_DEVICE_TYPE
+supported_device_types = {
+    a3mega_device_type,
+    a3ultra_device_type,
+    a4_device_type,
+}
 blueprint_dependencies_dir = {
     a3mega_device_type: "src/xpk/blueprints/a3mega",
     a3ultra_device_type: "src/xpk/blueprints/a3ultra",
+    a4_device_type: "src/xpk/blueprints/a4",
 }
 
 cluster_toolkit_url = "github.com/GoogleCloudPlatform/cluster-toolkit"
-cluster_toolkit_version = "v1.45.1"
-
-
-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)
-  ]
+cluster_toolkit_version = "v1.48.0"
 
 
 class BlueprintGeneratorOutput:
@@ -157,6 +159,11 @@ class BlueprintGenerator:
                 "total_min_nodes": system_node_pool_min_node_count,
                 "total_max_nodes": 1000,
             },
+            "k8s_network_names": {
+                "gvnic_prefix": f"{cluster_name}-gpunet-",
+                "gvnic_postfix": "-subnet",
+                "gvnic_start_index": 0,
+            },
         },
         outputs=["instructions"],
     )
@@ -173,13 +180,17 @@ class BlueprintGenerator:
     a3_megagpu_pool_0 = DeploymentModule(
         id="a3_megagpu_pool_0",
         source="modules/compute/gke-node-pool",
-        use=["gke_cluster", gpu_subnets_name, "group_placement_0"],
+        use=["gke_cluster", gpu_subnets_name],
         settings={
             "name": f"{cluster_name}-a3-megagpu-pool-0",
             "machine_type": system.gce_machine_type,
             "static_node_count": num_nodes,
             "zones": [zone],
-            "host_maintenance_interval": "PERIODIC",
+            "host_maintenance_interval": (
+                None
+                if capacity_type == CapacityType.RESERVATION
+                else "PERIODIC"
+            ),
             "reservation_affinity": self._getblock_reservation_affinity(
                 reservation
             ),
@@ -190,6 +201,9 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+
+    set_placement_policy = capacity_type != CapacityType.SPOT
+    tas_name = "topologyName: 'gke-default'" if set_placement_policy else ""
     num_chips = num_nodes * system.chips_per_vm
     workload = DeploymentModule(
         id="workload_component_install",
@@ -200,7 +214,10 @@ class BlueprintGenerator:
                 "install": True,
                 "version": "v0.10.0",  # TAS feature-gates is enabled in CT
                 "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {"num_chips": f"{num_chips}"},
+                "config_template_vars": {
+                    "num_chips": num_chips,
+                    "tas_name": tas_name,
+                },
             },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [{
@@ -236,15 +253,19 @@ class BlueprintGenerator:
             primary_vpc,
             gpunets,
             gke_cluster,
-            group_placement_0,
             a3_megagpu_pool_0,
             workload,
             workload_configmap,
         ],
     )
+
+    if set_placement_policy:
+      a3_megagpu_pool_0.use.append(group_placement_0.id)
+      primary_group.modules.append(group_placement_0)
+
     a3_mega_blueprint = Blueprint(
         terraform_backend_defaults=self._getblock_terraform_backend(
-            gcs_bucket, prefix
+            gcs_bucket, cluster_name, prefix
         ),
         blueprint_name=blueprint_name,
         toolkit_modules_url=cluster_toolkit_url,
@@ -261,8 +282,8 @@ class BlueprintGenerator:
     blueprint_file_path = self._save_blueprint_to_file(
         blueprint_name, a3_mega_blueprint, prefix
     )
-    blueprint_dependencies = self._get_a3_mega_blueprint_dependencies(
-        blueprint_name, prefix
+    blueprint_dependencies = self._get_blueprint_dependencies(
+        a3mega_device_type, blueprint_name, prefix
     )
     xpk_print(f"Blueprint file path: {blueprint_file_path}")
     xpk_print(
@@ -331,7 +352,7 @@ class BlueprintGenerator:
     )
     ml_gke = Blueprint(
         terraform_backend_defaults=self._getblock_terraform_backend(
-            gcs_bucket, prefix
+            gcs_bucket, cluster_name, prefix
         ),
         blueprint_name=blueprint_name,
         toolkit_modules_url=cluster_toolkit_url,
@@ -490,6 +511,13 @@ class BlueprintGenerator:
                 " alias_ip_range=[]}],"
                 f" {cluster_name}-rdma-net.subnetwork_interfaces_gke))"
             ),
+            "k8s_network_names": {
+                "rdma_prefix": f"{cluster_name}-rdma-sub-",
+                "rdma_start_index": 0,
+                "rdma_postfix": "",
+                "gvnic_prefix": f"{cluster_name}-sub-",
+                "gvnic_start_index": 1,
+            },
         },
         outputs=["instructions"],
     )
@@ -546,7 +574,7 @@ class BlueprintGenerator:
                 "install": True,
                 "version": "v0.10.0",  # TAS feature-gates is enabled in CT
                 "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
-                "config_template_vars": {"num_chips": f"{num_chips}"},
+                "config_template_vars": {"num_chips": num_chips},
             },
             "jobset": {"install": True, "version": "v0.7.2"},
             "apply_manifests": [
@@ -597,7 +625,7 @@ class BlueprintGenerator:
     )
     a3_ultra_blueprint = Blueprint(
         terraform_backend_defaults=self._getblock_terraform_backend(
-            gcs_bucket, prefix
+            gcs_bucket, cluster_name, prefix
         ),
         blueprint_name=blueprint_name,
         toolkit_modules_url=cluster_toolkit_url,
@@ -614,8 +642,276 @@ class BlueprintGenerator:
     blueprint_file_path = self._save_blueprint_to_file(
         blueprint_name, a3_ultra_blueprint, prefix
     )
-    blueprint_dependencies = self._get_a3_ultra_blueprint_dependencies(
-        blueprint_name, prefix
+    blueprint_dependencies = self._get_blueprint_dependencies(
+        a3ultra_device_type, blueprint_name, prefix
+    )
+    return BlueprintGeneratorOutput(
+        blueprint_file=blueprint_file_path,
+        blueprint_dependencies=blueprint_dependencies,
+    )
+
+  def generate_a4_blueprint(
+      self,
+      project_id: str,
+      cluster_name: str,
+      blueprint_name: str,
+      region: str,
+      zone: str,
+      auth_cidr: str,
+      system_node_pool_machine_type: str,
+      reservation: Optional[str | None] = None,
+      gcs_bucket: Optional[str | None] = None,
+      num_nodes: int = 2,
+      prefix: str = "",
+      system_node_pool_min_node_count: int = 2,
+      capacity_type: CapacityType = CapacityType.ON_DEMAND,
+  ) -> BlueprintGeneratorOutput:
+    """Create A4 blueprint.
+
+    Args:
+    Returns:
+      - Blueprint representing cluster toolkit blueprint
+    """
+    nccl_installer_path = (
+        f'$(ghpc_stage("{blueprint_name}"))/nccl-rdma-installer-a4.yaml'
+    )
+
+    net_0_id = f"{cluster_name}-net-0"
+    gpu_net_0 = DeploymentModule(
+        id=net_0_id,
+        source="modules/network/vpc",
+        settings={
+            "network_name": f"{cluster_name}-net-0",
+            "mtu": 8896,
+            "subnetworks": [{
+                "subnet_name": f"{cluster_name}-sub-0",
+                "subnet_region": region,
+                "subnet_ip": "192.168.0.0/18",
+            }],
+            "secondary_ranges_list": [{
+                "subnetwork_name": f"{cluster_name}-sub-0",
+                "ranges": [
+                    {"range_name": "pods", "ip_cidr_range": "10.4.0.0/14"},
+                    {"range_name": "services", "ip_cidr_range": "10.0.32.0/20"},
+                ],
+            }],
+            "firewall_rules": [{
+                "name": f"{cluster_name}-internal-0",
+                "ranges": ["192.168.0.0/16"],
+                "allow": [
+                    {"protocol": "tcp", "ports": ["0-65535"]},
+                    {"protocol": "udp", "ports": ["0-65535"]},
+                    {"protocol": "icmp"},
+                ],
+            }],
+        },
+    )
+    net_1_id = f"{cluster_name}-net-1"
+    gpu_net_1 = DeploymentModule(
+        id=net_1_id,
+        source="modules/network/vpc",
+        settings={
+            "network_name": f"{cluster_name}-net-1",
+            "mtu": 8896,
+            "subnetworks": [{
+                "subnet_name": f"{cluster_name}-sub-1",
+                "subnet_region": region,
+                "subnet_ip": "192.168.64.0/18",
+            }],
+            "firewall_rules": [{
+                "name": f"{cluster_name}-internal-1",
+                "ranges": ["192.168.0.0/16"],
+                "allow": [
+                    {"protocol": "tcp", "ports": ["0-65535"]},
+                    {"protocol": "udp", "ports": ["0-65535"]},
+                    {"protocol": "icmp"},
+                ],
+            }],
+        },
+    )
+    rma_net_id = f"{cluster_name}-rdma-net"
+    rma_net = DeploymentModule(
+        id=rma_net_id,
+        source="modules/network/gpu-rdma-vpc",
+        settings={
+            "network_name": f"{cluster_name}-rdma-net",
+            "mtu": 8896,
+            "network_profile": f"https://www.googleapis.com/compute/beta/projects/{project_id}/global/networkProfiles/{zone}-vpc-roce",
+            "network_routing_mode": "REGIONAL",
+            "subnetworks_template": {
+                "name_prefix": f"{cluster_name}-rdma-sub",
+                "count": 8,
+                "ip_range": "192.168.128.0/18",
+                "region": region,
+            },
+        },
+    )
+    cluster_id = f"{cluster_name}-a4-cluster"
+    a4_cluster = DeploymentModule(
+        id=cluster_id,
+        source="modules/scheduler/gke-cluster",
+        use=[net_0_id],
+        settings={
+            "system_node_pool_machine_type": system_node_pool_machine_type,
+            "system_node_pool_node_count": {
+                "total_min_nodes": system_node_pool_min_node_count,
+                "total_max_nodes": 1000,
+            },
+            "prefix_with_deployment_name": False,
+            "name_suffix": cluster_name,
+            "enable_dcgm_monitoring": True,
+            "enable_gcsfuse_csi": True,
+            "enable_private_endpoint": False,
+            "master_authorized_networks": [{
+                "cidr_block": auth_cidr,
+                "display_name": "kubectl-access-network",
+            }],
+            "additional_networks": (
+                f"$(concat([{{network={cluster_name}-net-1.network_name,"
+                f" subnetwork={cluster_name}-net-1.subnetwork_name,"
+                f' subnetwork_project="{project_id}", nic_type="GVNIC",'
+                " queue_count=null, network_ip=null, stack_type=null,"
+                " access_config=[{nat_ip=null, public_ptr_domain_name=null,"
+                " network_tier=null}], ipv6_access_config=[],"
+                " alias_ip_range=[]}],"
+                f" {cluster_name}-rdma-net.subnetwork_interfaces_gke))"
+            ),
+            "version_prefix": "1.32.",
+            "release_channel": "RAPID",
+            "maintenance_exclusions": [{
+                "name": "no-minor-or-node-upgrades-indefinite",
+                "start_time": "2024-12-01T00:00:00Z",
+                "end_time": "2025-12-22T00:00:00Z",
+                "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+            }],
+        },
+        outputs=["instructions"],
+    )
+    system, _ = get_system_characteristics_by_device_type(a4_device_type)
+    if system is None:
+      xpk_print(
+          "Error: Could not retrieve system characteristics for"
+          f" {a4_device_type} device_type."
+      )
+      xpk_exit(1)
+    gpu_pool = DeploymentModule(
+        id=f"{cluster_name}-a4-pool",
+        source="modules/compute/gke-node-pool",
+        use=[cluster_id],
+        settings={
+            "machine_type": system.gce_machine_type,
+            "auto_upgrade": True,
+            "zones": [zone],
+            "disk_type": "hyperdisk-balanced",
+            "static_node_count": num_nodes,
+            "local_ssd_count_ephemeral_storage": 32,
+            "spot": capacity_type == CapacityType.SPOT,
+            "reservation_affinity": self._getblock_reservation_affinity(
+                reservation
+            ),
+            "max_pods_per_node": 32,
+            "guest_accelerator": [{
+                "type": system.gke_accelerator,
+                "count": 8,
+                "gpu_driver_installation_config": {
+                    "gpu_driver_version": "LATEST"
+                },
+            }],
+            "additional_networks": (
+                f"$(concat([{{network={cluster_name}-net-1.network_name,"
+                f" subnetwork={cluster_name}-net-1.subnetwork_name,"
+                f' subnetwork_project="{project_id}", nic_type="GVNIC",'
+                " queue_count=null, network_ip=null, stack_type=null,"
+                " access_config=[{nat_ip=null, public_ptr_domain_name=null,"
+                " network_tier=null}], ipv6_access_config=[],"
+                " alias_ip_range=[]}],"
+                f" {cluster_name}-rdma-net.subnetwork_interfaces_gke))"
+            ),
+        },
+        outputs=["instructions"],
+    )
+
+    num_chips = num_nodes * system.chips_per_vm
+    workload_manager_install_id = "workload-manager-install"
+    workload_manager_install = DeploymentModule(
+        id=workload_manager_install_id,
+        source="modules/management/kubectl-apply",
+        use=[cluster_id],
+        settings={
+            "kueue": {
+                "install": True,
+                "version": "v0.10.0",  # TAS feature-gates is enabled in CT
+                "config_path": f'$(ghpc_stage("{blueprint_name}"))/kueue-xpk-configuration.yaml.tftpl',
+                "config_template_vars": {"num_chips": num_chips},
+            },
+            "jobset": {"install": True, "version": "v0.7.2"},
+            "apply_manifests": [
+                {"source": nccl_installer_path},
+                {
+                    "source": (
+                        f'$(ghpc_stage("{blueprint_name}"))/storage_crd.yaml'
+                    )
+                },
+            ],
+        },
+    )
+
+    workload_configmap = DeploymentModule(
+        id="workload_configmap",
+        source="modules/management/kubectl-apply",
+        use=[cluster_id],
+        settings={
+            "apply_manifests": [{
+                "source": (
+                    f'$(ghpc_stage("{blueprint_name}"))/config-map.yaml.tftpl'
+                ),
+                "template_vars": {
+                    "resource_config_name": (
+                        f"{cluster_name}-resources-configmap"
+                    ),
+                    "num_nodes": f"{num_nodes}",
+                    "cluster_config_name": f"{cluster_name}-metadata-configmap",
+                    "capacity_type": f"{capacity_type.value}",
+                    "reservation": f"{reservation}",
+                },
+            }]
+        },
+    )
+
+    primary_group = DeploymentGroup(
+        group="primary",
+        modules=[
+            gpu_net_0,
+            gpu_net_1,
+            rma_net,
+            a4_cluster,
+            gpu_pool,
+            workload_manager_install,
+            workload_configmap,
+        ],
+    )
+
+    a4_blueprint = Blueprint(
+        terraform_backend_defaults=self._getblock_terraform_backend(
+            gcs_bucket, cluster_name, prefix
+        ),
+        blueprint_name=blueprint_name,
+        toolkit_modules_url=cluster_toolkit_url,
+        toolkit_modules_version=cluster_toolkit_version,
+        deployment_groups=[primary_group],
+        vars={
+            "project_id": project_id,
+            "deployment_name": blueprint_name,
+            "region": region,
+            "zone": zone,
+        },
+    )
+
+    blueprint_file_path = self._save_blueprint_to_file(
+        blueprint_name, a4_blueprint, prefix
+    )
+    blueprint_dependencies = self._get_blueprint_dependencies(
+        a4_device_type, blueprint_name, prefix
     )
     return BlueprintGeneratorOutput(
         blueprint_file=blueprint_file_path,
@@ -638,7 +934,7 @@ class BlueprintGenerator:
     )
 
   def _getblock_terraform_backend(
-      self, gcs_bucket: str, prefix: str = ""
+      self, gcs_bucket: str, cluster_name: str, prefix: str = ""
   ) -> dict | None:
     if gcs_bucket is None:
       return None
@@ -646,12 +942,19 @@ class BlueprintGenerator:
         "type": "gcs",
         "configuration": {
             "bucket": gcs_bucket,
-            "prefix": self._get_terraforrm_backend_full_prefix(prefix),
+            "prefix": self._get_terraforrm_backend_full_prefix(
+                cluster_name, prefix
+            ),
         },
     }
 
-  def _get_terraforrm_backend_full_prefix(self, prefix: str = "") -> str:
-    return f"xpk_terraform_state/{prefix}/tfstate/"
+  def _get_terraforrm_backend_full_prefix(
+      self, cluster_name: str, prefix: str = ""
+  ) -> str:
+    full_prefix = "xpk_terraform_state"
+    if prefix:
+      full_prefix += f"/{prefix}"
+    return f"{full_prefix}/{cluster_name}/"
 
   def _save_blueprint_to_file(
       self, blueprint_name: str, xpk_blueprint: Blueprint, prefix: str = ""
@@ -676,27 +979,14 @@ class BlueprintGenerator:
     blueprint_path = self._get_blueprint_path(blueprint_name, prefix)
     return os.path.exists(blueprint_path)
 
-  def _get_a3_mega_blueprint_dependencies(
-      self, blueprint_name: str, prefix: str = ""
-  ) -> str:
-    deployment_files_path = os.path.join(
-        self._get_storage_path(prefix), blueprint_name
-    )
-    shutil.copytree(
-        blueprint_dependencies_dir[a3mega_device_type],
-        deployment_files_path,
-        dirs_exist_ok=True,
-    )
-    return deployment_files_path
-
-  def _get_a3_ultra_blueprint_dependencies(
-      self, blueprint_name: str, prefix: str = ""
+  def _get_blueprint_dependencies(
+      self, device_type: str, blueprint_name: str, prefix: str = ""
   ) -> str:
     deployment_files_path = os.path.join(
         self._get_storage_path(prefix), blueprint_name
     )
     shutil.copytree(
-        blueprint_dependencies_dir[a3ultra_device_type],
+        blueprint_dependencies_dir[device_type],
         deployment_files_path,
         dirs_exist_ok=True,
     )

xpk/core/capacity.py

@@ -27,6 +27,7 @@ CAPACITY_TYPE_CONFIG_KEY = 'capacity_type'
 H100_DEVICE_TYPE = 'h100-80gb-8'
 H100_MEGA_DEVICE_TYPE = 'h100-mega-80gb-8'
 H200_DEVICE_TYPE = 'h200-141gb-8'
+B200_DEVICE_TYPE = 'b200-8'
 RESERVATION_CONFIG_KEY = 'reservation_id'