xpk 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

xpk/core/blueprint/blueprint_definitions.py

@@ -55,6 +55,7 @@ class Blueprint:
   """A class to represent Cluster Toolkit blueprint"""
 
   deployment_groups: list[DeploymentGroup]
+  terraform_backend_defaults: Optional[dict]
   blueprint_name: Optional[str]
   toolkit_modules_url: str
   toolkit_modules_version: str

xpk/core/blueprint/blueprint_generator.py

@@ -14,21 +14,22 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import os
 import shutil
 from typing import Optional
+
 from ruamel import yaml
-import os
 
-from .blueprint_definitions import DeploymentGroup, DeploymentModule, Blueprint
-from ..system_characteristics import get_system_characteristics_by_device_type
-from ...utils.console import xpk_print, xpk_exit
+from ...utils.console import xpk_exit, xpk_print
 from ...utils.file import ensure_directory_exists
-from ..core import CapacityType, h100_mega_device_type, h200_device_type
+from ..capacity import 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
+a3mega_device_type = H100_MEGA_DEVICE_TYPE
+a3ultra_device_type = H200_DEVICE_TYPE
 supported_device_types = {a3mega_device_type, a3ultra_device_type}
 blueprint_dependencies_dir = {
     a3mega_device_type: "src/xpk/blueprints/a3mega",
@@ -39,6 +40,16 @@ 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)
+  ]
+
+
 class BlueprintGeneratorOutput:
   """BlueprintGeneratorOutput is a class containing fields with output blueprint file path and path to blueprint dependencies.
   Atributes:
@@ -79,6 +90,7 @@ class BlueprintGenerator:
       group_placement_max_distance: int = 2,
       subnetwork_cidr_suffix: int = 24,
       reservation: str | None = None,
+      gcs_bucket: Optional[str | None] = None,
       capacity_type: CapacityType = CapacityType.ON_DEMAND,
       system_node_pool_min_node_count: int = 2,
   ) -> BlueprintGeneratorOutput:
@@ -132,6 +144,8 @@ class BlueprintGenerator:
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "enable_private_endpoint": False,
+            "enable_gcsfuse_csi": True,
+            "enable_filestore_csi": True,
             "master_authorized_networks": [{
                 "cidr_block": (
                     f"{auth_cidr}"
@@ -156,18 +170,6 @@ class BlueprintGenerator:
         },
     )
 
-    reservation_affinity = (
-        {
-            "consume_reservation_type": "NO_RESERVATION",
-            "specific_reservations": [],
-        }
-        if reservation is None
-        else {
-            "consume_reservation_type": "SPECIFIC_RESERVATION",
-            "specific_reservations": [{"name": reservation}],
-        }
-    )
-
     a3_megagpu_pool_0 = DeploymentModule(
         id="a3_megagpu_pool_0",
         source="modules/compute/gke-node-pool",
@@ -178,7 +180,9 @@ class BlueprintGenerator:
             "static_node_count": num_nodes,
             "zones": [zone],
             "host_maintenance_interval": "PERIODIC",
-            "reservation_affinity": reservation_affinity,
+            "reservation_affinity": self._getblock_reservation_affinity(
+                reservation
+            ),
             "run_workload_script": False,
             "spot": capacity_type == CapacityType.SPOT,
             "max_pods_per_node": 32,
@@ -199,6 +203,9 @@ class BlueprintGenerator:
                 "config_template_vars": {"num_chips": f"{num_chips}"},
             },
             "jobset": {"install": True, "version": "v0.7.2"},
+            "apply_manifests": [{
+                "source": f'$(ghpc_stage("{blueprint_name}"))/storage_crd.yaml'
+            }],
         },
     )
 
@@ -235,7 +242,10 @@ class BlueprintGenerator:
             workload_configmap,
         ],
     )
-    xpk_blueprint = Blueprint(
+    a3_mega_blueprint = Blueprint(
+        terraform_backend_defaults=self._getblock_terraform_backend(
+            gcs_bucket, prefix
+        ),
         blueprint_name=blueprint_name,
         toolkit_modules_url=cluster_toolkit_url,
         toolkit_modules_version=cluster_toolkit_version,
@@ -247,8 +257,9 @@ class BlueprintGenerator:
             "zone": zone,
         },
     )
+
     blueprint_file_path = self._save_blueprint_to_file(
-        blueprint_name, xpk_blueprint, prefix
+        blueprint_name, a3_mega_blueprint, prefix
     )
     blueprint_dependencies = self._get_a3_mega_blueprint_dependencies(
         blueprint_name, prefix
@@ -271,6 +282,7 @@ class BlueprintGenerator:
       region: str,
       auth_cidr: str,
       prefix: str = "",
+      gcs_bucket: Optional[str | None] = None,
   ) -> BlueprintGeneratorOutput:
     """Create a simple gke cluster
 
@@ -318,6 +330,9 @@ class BlueprintGenerator:
         modules=[network1, gke_cluster],
     )
     ml_gke = Blueprint(
+        terraform_backend_defaults=self._getblock_terraform_backend(
+            gcs_bucket, prefix
+        ),
         blueprint_name=blueprint_name,
         toolkit_modules_url=cluster_toolkit_url,
         toolkit_modules_version=cluster_toolkit_version,
@@ -328,6 +343,7 @@ class BlueprintGenerator:
             "region": region,
         },
     )
+
     blueprint_file_path = self._save_blueprint_to_file(
         blueprint_name, ml_gke, prefix
     )
@@ -337,55 +353,6 @@ class BlueprintGenerator:
         blueprint_dependencies=blueprint_dependencies,
     )
 
-  def _save_blueprint_to_file(
-      self, blueprint_name: str, xpk_blueprint: Blueprint, prefix: str = ""
-  ) -> str:
-    blueprint_path = self._get_blueprint_path(blueprint_name, prefix)
-    with open(blueprint_path, "w+", encoding="utf-8") as blueprint_file:
-      yaml.dump(xpk_blueprint, blueprint_file)
-    return blueprint_path
-
-  def _get_blueprint_path(self, blueprint_name, prefix: str = ""):
-    blueprint_path = os.path.join(
-        self._get_storage_path(prefix), f"{blueprint_name}.yaml"
-    )
-    return blueprint_path
-
-  def _get_storage_path(self, prefix):
-    storage_path_with_prefix = os.path.join(self.storage_path, prefix)
-    ensure_directory_exists(storage_path_with_prefix)
-    return storage_path_with_prefix
-
-  def blueprint_exists(self, blueprint_name, prefix: str = ""):
-    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 = ""
-  ) -> str:
-    deployment_files_path = os.path.join(
-        self._get_storage_path(prefix), blueprint_name
-    )
-    shutil.copytree(
-        blueprint_dependencies_dir[a3ultra_device_type],
-        deployment_files_path,
-        dirs_exist_ok=True,
-    )
-    return deployment_files_path
-
   def generate_a3_ultra_blueprint(
       self,
       project_id: str,
@@ -396,7 +363,9 @@ class BlueprintGenerator:
       auth_cidr: str,
       system_node_pool_machine_type: str,
       reservation: Optional[str | None] = None,
+      gcs_bucket: Optional[str | None] = None,
       num_nodes: int = 2,
+      enable_filestore_csi_driver=True,
       prefix: str = "",
       mtu_size: int = 8896,
       system_node_pool_min_node_count: int = 2,
@@ -501,6 +470,7 @@ class BlueprintGenerator:
             "system_node_pool_machine_type": system_node_pool_machine_type,
             "enable_dcgm_monitoring": True,
             "enable_gcsfuse_csi": True,
+            "enable_filestore_csi": enable_filestore_csi_driver,
             "enable_private_endpoint": False,
             "master_authorized_networks": [{
                 "cidr_block": auth_cidr,
@@ -540,6 +510,9 @@ class BlueprintGenerator:
             "zones": [zone],
             "static_node_count": num_nodes,
             "spot": capacity_type == CapacityType.SPOT,
+            "reservation_affinity": self._getblock_reservation_affinity(
+                reservation
+            ),
             "max_pods_per_node": 32,
             "guest_accelerator": [{
                 "type": "nvidia-h200-141gb",
@@ -561,11 +534,6 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
-    if reservation is not None:
-      gpu_pool.settings["reservation_affinity"] = {
-          "consume_reservation_type": "SPECIFIC_RESERVATION",
-          "specific_reservations": [{"name": reservation}],
-      }
 
     num_chips = num_nodes * system.chips_per_vm
     workload_manager_install_id = "workload-manager-install"
@@ -584,6 +552,11 @@ class BlueprintGenerator:
             "apply_manifests": [
                 {"source": nccl_installer_path},
                 {"source": mlgru_disable_path},
+                {
+                    "source": (
+                        f'$(ghpc_stage("{blueprint_name}"))/storage_crd.yaml'
+                    )
+                },
             ],
         },
     )
@@ -623,6 +596,9 @@ class BlueprintGenerator:
         ],
     )
     a3_ultra_blueprint = Blueprint(
+        terraform_backend_defaults=self._getblock_terraform_backend(
+            gcs_bucket, prefix
+        ),
         blueprint_name=blueprint_name,
         toolkit_modules_url=cluster_toolkit_url,
         toolkit_modules_version=cluster_toolkit_version,
@@ -646,6 +622,86 @@ class BlueprintGenerator:
         blueprint_dependencies=blueprint_dependencies,
     )
 
+  def _getblock_reservation_affinity(
+      self, reservation: str | None = None
+  ) -> dict:
+    return (
+        {
+            "consume_reservation_type": "NO_RESERVATION",
+            "specific_reservations": [],
+        }
+        if reservation is None
+        else {
+            "consume_reservation_type": "SPECIFIC_RESERVATION",
+            "specific_reservations": [{"name": reservation}],
+        }
+    )
+
+  def _getblock_terraform_backend(
+      self, gcs_bucket: str, prefix: str = ""
+  ) -> dict | None:
+    if gcs_bucket is None:
+      return None
+    return {
+        "type": "gcs",
+        "configuration": {
+            "bucket": gcs_bucket,
+            "prefix": self._get_terraforrm_backend_full_prefix(prefix),
+        },
+    }
+
+  def _get_terraforrm_backend_full_prefix(self, prefix: str = "") -> str:
+    return f"xpk_terraform_state/{prefix}/tfstate/"
+
+  def _save_blueprint_to_file(
+      self, blueprint_name: str, xpk_blueprint: Blueprint, prefix: str = ""
+  ) -> str:
+    blueprint_path = self._get_blueprint_path(blueprint_name, prefix)
+    with open(blueprint_path, "w+", encoding="utf-8") as blueprint_file:
+      yaml.dump(xpk_blueprint, blueprint_file)
+    return blueprint_path
+
+  def _get_blueprint_path(self, blueprint_name, prefix: str = ""):
+    blueprint_path = os.path.join(
+        self._get_storage_path(prefix), f"{blueprint_name}.yaml"
+    )
+    return blueprint_path
+
+  def _get_storage_path(self, prefix):
+    storage_path_with_prefix = os.path.join(self.storage_path, prefix)
+    ensure_directory_exists(storage_path_with_prefix)
+    return storage_path_with_prefix
+
+  def blueprint_exists(self, blueprint_name, prefix: str = ""):
+    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 = ""
+  ) -> str:
+    deployment_files_path = os.path.join(
+        self._get_storage_path(prefix), blueprint_name
+    )
+    shutil.copytree(
+        blueprint_dependencies_dir[a3ultra_device_type],
+        deployment_files_path,
+        dirs_exist_ok=True,
+    )
+    return deployment_files_path
+
 
 yaml.register_class(Blueprint)
 yaml.register_class(DeploymentGroup)

xpk/core/capacity.py

@@ -0,0 +1,185 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import enum
+
+from ..utils.console import xpk_print
+from .commands import run_command_with_updates
+
+AUTOPROVISIONING_CONFIG_VALUE = 'AUTOPROVISION'
+AUTOPROVISIONING_CONFIG_MINIMUM_KEY = 'minimum_chips'
+AUTOPROVISIONING_CONFIG_MAXIMUM_KEY = 'maximum_chips'
+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'
+RESERVATION_CONFIG_KEY = 'reservation_id'
+
+
+class CapacityType(enum.Enum):
+  ON_DEMAND = 'on_demand'
+  RESERVATION = 'reservation'
+  SPOT = 'spot'
+  UNKNOWN = 'unknown'
+
+
+def print_reservations(args) -> int:
+  """Print the reservations in the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = f'gcloud beta compute reservations list --project={args.project}'
+  return_code = run_command_with_updates(
+      command, 'Get all reservations in the project', args
+  )
+  if return_code != 0:
+    xpk_print(f'Get all reservations returned ERROR {return_code}')
+    return 1
+  return 0
+
+
+def get_capacity_type(args) -> tuple[CapacityType, int]:
+  """Determine the capacity type based on user arguments.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    Tuple with string with the system characteristics and
+    int of 0 if successful and 1 otherwise.
+  """
+  capacity_type = CapacityType.UNKNOWN
+  num_types = 0
+  return_code = 0
+
+  # Determine the capacity argument.
+  if args.on_demand:
+    capacity_type = CapacityType.ON_DEMAND
+    num_types += 1
+  if args.reservation:
+    return_code = verify_reservation_exists(args)
+    if return_code > 0:
+      return capacity_type, return_code
+    capacity_type = CapacityType.RESERVATION
+    num_types += 1
+  if args.spot:
+    capacity_type = CapacityType.SPOT
+    num_types += 1
+
+  # Check that the number of user arguments provided is valid.
+  if num_types == 0:
+    capacity_type = CapacityType.UNKNOWN
+  elif num_types != 1:
+    xpk_print(
+        'ERROR: User specified more than one of the following arguments. Please'
+        ' specify only one of `--reservation=$RESERVATION_NAME`, `--on-demand`'
+        ' or `--spot`.'
+    )
+    return_code = 1
+
+  return capacity_type, return_code
+
+
+def verify_reservation_exists(args) -> int:
+  """Verify the reservation exists.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      f'gcloud beta compute reservations describe {args.reservation}'
+      f' --project={args.project} --zone={args.zone}'
+  )
+  return_code = run_command_with_updates(command, 'Describe reservation', args)
+  if return_code != 0:
+    xpk_print(f'Describe reservation returned ERROR {return_code}')
+    xpk_print('Please confirm that your reservation name is correct.')
+    return 1
+  return 0
+
+
+def get_capacity_arguments_from_capacity_type(
+    args, capacity_type: CapacityType
+) -> tuple[str, int]:
+  """Determine the TPU Nodepool creation capacity arguments needed.
+
+  Args:
+    args: user provided arguments for running the command.
+    capacity_type: The type of capacity the user configured.
+
+  Returns:
+    Tuple with string with the capacity argument to use and
+    int of 0 if successful and 1 otherwise.
+  """
+  capacity_args = ''
+  return_code = 0
+
+  match capacity_type:
+    case CapacityType.ON_DEMAND:
+      capacity_args = ''
+    case CapacityType.SPOT:
+      capacity_args = '--spot'
+    case CapacityType.RESERVATION:
+      capacity_args = (
+          f'--reservation-affinity=specific --reservation={args.reservation}'
+      )
+    case _:
+      xpk_print(
+          f'Unknown capacity type: {capacity_type}. Unable to determine'
+          ' capacity args.'
+      )
+      return_code = 1
+  return capacity_args, return_code
+
+
+def get_capacity_node_selectors_from_capacity_type(
+    args, capacity_type: str
+) -> tuple[str, int]:
+  """Determine the node selectors for a workload to run on a specific capacity type.
+
+  Args:
+    args: user provided arguments for running the command.
+    capacity_type: The type of capacity the user configured.
+
+  Returns:
+    Tuple with string with the node selectors to use and
+    int of 0 if successful and 1 otherwise.
+  """
+  node_selector = ''
+  return_code = 0
+
+  match capacity_type:
+    case CapacityType.ON_DEMAND.name:
+      node_selector = ''
+    case CapacityType.SPOT.name:
+      node_selector = 'cloud.google.com/gke-spot="true"'
+    case CapacityType.RESERVATION.name:
+      node_selector = f'cloud.google.com/reservation-name: {args.reservation}'
+    case _:
+      xpk_print(
+          f'Unknown capacity type: {capacity_type}. Unable to determine the'
+          ' node selectors.'
+      )
+      return_code = 1
+  return node_selector, return_code