xpk 0.0.1__py3-none-any.whl

xpk/core/filestore.py

@@ -0,0 +1,251 @@
+"""
+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.
+"""
+
+from enum import Enum
+
+from google.cloud import filestore_v1
+from google.cloud.exceptions import GoogleCloudError
+from google.cloud.filestore_v1.types import (
+    FileShareConfig,
+    Instance,
+    NetworkConfig,
+)
+
+from ..utils import templates
+from ..utils.console import xpk_exit, xpk_print
+from .cluster import zone_to_region
+
+FS_PV_PATH = "/../templates/filestore-pv.yaml"
+FS_PVC_PATH = "/../templates/filestore-pvc.yaml"
+FS_SC_PATH = "/../templates/filestore-sc.yaml"
+
+
+class Availability(Enum):
+  ZONAL = "Zonal"
+  REGIONAL = "Regional"
+
+
+TIERS = {
+    "BASIC_HDD": Availability.ZONAL,
+    "BASIC_SSD": Availability.ZONAL,
+    "ZONAL": Availability.ZONAL,
+    "REGIONAL": Availability.REGIONAL,
+    "ENTERPRISE": Availability.REGIONAL,
+}
+
+
+def get_storage_class_name(storage_name: str) -> str:
+  return f"{storage_name}-sc"
+
+
+def get_pv_name(storage_name: str) -> str:
+  return f"{storage_name}-pv"
+
+
+def get_pvc_name(storage_name: str) -> str:
+  return f"{storage_name}-pvc"
+
+
+class FilestoreClient:
+  """FilestoreClient is a class for interacting with GCP filestore instances."""
+
+  def __init__(
+      self,
+      zone: str,
+      name: str,
+      project: str,
+  ) -> None:
+    self.zone = zone
+    self.region = zone_to_region(zone)
+    self.name = name
+    self.project = project
+    self._client = filestore_v1.CloudFilestoreManagerClient()
+    self.instance: Instance | None = None
+
+  def get_instance(self) -> Instance | None:
+    """Get existing Filestore instance"""
+    parentZonal = self.get_parent(self.zone)
+    parentRegional = self.get_parent(self.region)
+    reqZonal = filestore_v1.ListInstancesRequest(parent=parentZonal)
+    reqRegional = filestore_v1.ListInstancesRequest(parent=parentRegional)
+    try:
+      instancesZonal = self._client.list_instances(reqZonal)
+      instancesRegional = self._client.list_instances(reqRegional)
+    except GoogleCloudError as e:
+      xpk_print(f"Exception while trying to list instances {e}")
+      xpk_exit(1)
+
+    fullname_zonal = self.get_instance_fullname(self.zone)
+    fullname_regional = self.get_instance_fullname(self.region)
+
+    for instance in instancesZonal:
+      if instance.name == fullname_zonal:
+        return instance  # pytype: disable=bad-return-type
+
+    for instance in instancesRegional:
+      if instance.name == fullname_regional:
+        return instance  # pytype: disable=bad-return-type
+
+  def check_instance_exists(self) -> bool:
+    """Check if Filestore instance exists"""
+    instance = self.get_instance()
+    return instance is not None
+
+  def load_instance(self) -> None:
+    if self.instance is None:
+      self.instance = self.get_instance()
+
+  def get_instance_location(self) -> str:
+    """Get Filestore instance's location"""
+    self.load_instance()
+    return str(self.instance.name.split("/")[3])
+
+  def create_instance(
+      self,
+      vol: str,
+      size: int,
+      tier: str,
+      connect_mode=None,
+      reserved_ip_range=None,
+      network: str = "default",
+      description: str = "XPK created filestore instance",
+      kms_key_name=None,
+      source_backup=None,
+      nfs_export_options=None,
+      modes=None,
+  ) -> None:
+    """Create new Filestore instance"""
+
+    location = (
+        self.zone
+        if TIERS[tier].value == Availability.ZONAL.value
+        else self.region
+    )
+
+    file_shares = [
+        FileShareConfig(
+            name=vol,
+            capacity_gb=size,
+            source_backup=source_backup,
+            nfs_export_options=nfs_export_options,
+        )
+    ]
+    networks = [
+        NetworkConfig(
+            network=network,
+            modes=modes,
+            reserved_ip_range=reserved_ip_range,
+            connect_mode=connect_mode,
+        )
+    ]
+    request = filestore_v1.CreateInstanceRequest(
+        parent=self.get_parent(location),
+        instance_id=self.name,
+        instance=Instance(
+            description=description,
+            tier=tier,
+            kms_key_name=kms_key_name,
+            file_shares=file_shares,
+            networks=networks,
+        ),
+    )
+    # Make the request
+    operation = self._client.create_instance(request=request)
+    xpk_print("Waiting for filestore creation to complete...")
+    self.instance = None
+    try:
+      self.instance = operation.result()
+    except GoogleCloudError as e:
+      xpk_print(f"Error while creating Filestore instance: {e}")
+      xpk_exit(1)
+    xpk_print(
+        f"Filestore instance {self.get_instance_fullname(location)} created"
+    )
+
+  def delete_filestore_instance(self):
+    # Initialize request
+    name = self.get_instance_fullname()
+    request = filestore_v1.DeleteInstanceRequest(name=name)
+
+    # Make the request
+    operation = self._client.delete_instance(request)
+    xpk_print("Waiting for filestore deletion to complete...")
+    try:
+      operation.result()
+    except GoogleCloudError as e:
+      xpk_print(f"Error while deleting Filestore instance: {e}")
+      xpk_exit(1)
+    xpk_print(f"Filestore instance {name} deleted")
+
+  def create_sc(self, name: str, network: str) -> dict:
+    """Create a yaml representing filestore StorageClass."""
+    data = templates.load(FS_SC_PATH)
+    data["metadata"]["name"] = get_storage_class_name(name)
+    data["parameters"]["tier"] = self.instance.tier.name
+    data["parameters"][
+        "network"
+    ] = f"projects/{self.project}/global/networks/{network}"
+    return data
+
+  def create_pv(self, name: str, vol: str, access_mode: str) -> dict:
+    """Create a yaml representing filestore PersistentVolume."""
+    data = templates.load(FS_PV_PATH)
+    data["metadata"]["name"] = get_pv_name(name)
+    data["spec"]["storageClassName"] = get_storage_class_name(name)
+    data["spec"]["capacity"]["storage"] = self.instance.file_shares[
+        0
+    ].capacity_gb
+    data["spec"]["accessModes"] = [access_mode]
+    volumeHandle = f"{self.get_instance_fullname()}/volumes/{vol}"
+    data["spec"]["csi"]["volumeHandle"] = volumeHandle
+    data["spec"]["csi"]["volumeAttributes"]["ip"] = self.instance.networks[
+        0
+    ].ip_addresses[0]
+    data["spec"]["csi"]["volumeAttributes"]["volume"] = vol
+    return data
+
+  def create_pvc(self, name: str, access_mode: str) -> dict:
+    """Create a yaml representing filestore PersistentVolumeClaim."""
+    data = templates.load(FS_PVC_PATH)
+    data["metadata"]["name"] = get_pvc_name(name)
+    data["spec"]["accessModes"] = [access_mode]
+    data["spec"]["storageClassName"] = get_storage_class_name(name)
+    data["spec"]["volumeName"] = get_pv_name(name)
+    data["spec"]["resources"]["requests"]["storage"] = (
+        self.instance.file_shares[0].capacity_gb
+    )
+    return data
+
+  def manifest(
+      self, name: str, vol: str, access_mode: str, network: str
+  ) -> list[dict]:
+    self.load_instance()
+    pv = self.create_pv(name, vol, access_mode)
+    pvc = self.create_pvc(name, access_mode)
+    sc = self.create_sc(name, network)
+    return [pv, pvc, sc]
+
+  def get_parent(self, location: str | None = None) -> str:
+    """Get the Filestore's parent's name"""
+    if location is None:
+      location = self.get_instance_location()
+    return f"projects/{self.project}/locations/{location}"
+
+  def get_instance_fullname(self, location: str | None = None) -> str:
+    """Get the Filestore's full name"""
+    if location is None:
+      location = self.get_instance_location()
+    return f"projects/{self.project}/locations/{location}/instances/{self.name}"

xpk/core/gcloud_context.py

@@ -0,0 +1,196 @@
+"""
+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 subprocess
+import sys
+from dataclasses import dataclass
+
+from ..utils.console import xpk_print
+from .commands import run_command_for_value
+
+
+def get_project():
+  """Get GCE project from `gcloud config get project`.
+
+  Returns:
+     The project name.
+  """
+  completed_command = subprocess.run(
+      ['gcloud', 'config', 'get', 'project'], check=True, capture_output=True
+  )
+  project_outputs = completed_command.stdout.decode().strip().split('\n')
+  if len(project_outputs) < 1 or project_outputs[-1] == '':
+    sys.exit(
+        'You must specify the project in the project flag or set it with'
+        " 'gcloud config set project <project>'"
+    )
+  return project_outputs[
+      -1
+  ]  # The project name lives on the last line of the output
+
+
+def get_zone():
+  """Get GCE zone from `gcloud config get compute/zone`.
+
+  Returns:
+     The zone name.
+  """
+  completed_command = subprocess.run(
+      ['gcloud', 'config', 'get', 'compute/zone'],
+      check=True,
+      capture_output=True,
+  )
+  zone_outputs = completed_command.stdout.decode().strip().split('\n')
+  if len(zone_outputs) < 1 or zone_outputs[-1] == '':
+    sys.exit(
+        "You must specify the zone in the zone flag or set it with 'gcloud"
+        " config set compute/zone <zone>'"
+    )
+  return zone_outputs[-1]  # The zone name lives on the last line of the output
+
+
+def add_zone_and_project(args):
+  """Obtains the zone and project names from gcloud configs if not defined.
+
+  Args:
+    args: user provided arguments for running the command.
+  """
+  if not args.project:
+    args.project = get_project()
+  if not args.zone:
+    args.zone = get_zone()
+  xpk_print(f'Working on {args.project} and {args.zone}')
+
+
+def zone_to_region(zone) -> str:
+  """Helper function converts zone name to region name.
+
+  Args:
+    zone: zone name.
+
+  Returns:
+     The region name.
+  """
+  zone_terms = zone.split('-')
+  return zone_terms[0] + '-' + zone_terms[1]  # pytype: disable=bad-return-type
+
+
+@dataclass
+class GkeServerConfig:
+  """Stores the valid gke versions based on gcloud recommendations."""
+
+  default_rapid_gke_version: str
+  valid_versions: set[str]
+
+
+def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
+  """Determine the GKE versions supported by gcloud currently.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    Tuple of
+    int: 0 if successful and 1 otherwise.
+    GkeServerConfig: stores valid gke version to use in node pool and cluster.
+  """
+  base_command = (
+      'gcloud container get-server-config'
+      f' --project={args.project} --region={zone_to_region(args.zone)}'
+  )
+  default_rapid_gke_version_cmd = (
+      base_command
+      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      ' --format="value(channels.defaultVersion)"'
+  )
+  valid_versions_cmd = (
+      base_command
+      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      ' --format="value(channels.validVersions)"'
+  )
+  base_command_description = 'Determine server supported GKE versions for '
+
+  server_config_commands_and_descriptions = [
+      (
+          default_rapid_gke_version_cmd,
+          base_command_description + 'default rapid gke version',
+      ),
+      (
+          valid_versions_cmd,
+          base_command_description + 'valid versions',
+      ),
+  ]
+  command_outputs = []
+
+  for command, command_description in server_config_commands_and_descriptions:
+    return_code, cmd_output = run_command_for_value(
+        command,
+        command_description,
+        args,
+        hide_error=True,
+    )
+    if return_code != 0:
+      xpk_print(f'Unable to get server config for {command_description}.')
+      return return_code, None
+    command_outputs.append(cmd_output)
+
+  return 0, GkeServerConfig(
+      default_rapid_gke_version=command_outputs[0].strip(),
+      valid_versions=set(command_outputs[1].split(';')),
+  )
+
+
+def get_gke_control_plane_version(
+    args, gke_server_config: GkeServerConfig
+) -> tuple[int, str | None]:
+  """Determine gke control plane version for cluster creation.
+
+  Args:
+    args: user provided arguments for running the command.
+    gke_server_config: holds valid gke versions and recommended default version.
+
+  Returns:
+    Tuple of
+    int: 0 if successful and 1 otherwise.
+    str: gke control plane version to use.
+  """
+
+  # Override with user provide gke version if specified.
+  if args.gke_version is not None:
+    master_gke_version = args.gke_version
+  else:
+    master_gke_version = gke_server_config.default_rapid_gke_version
+
+  is_valid_version = master_gke_version in gke_server_config.valid_versions
+
+  if not is_valid_version:
+    xpk_print(
+        f'Planned GKE Version: {master_gke_version}\n Valid Versions:'
+        f'\n{gke_server_config.valid_versions}\nRecommended / Default GKE'
+        f' Version: {gke_server_config.default_rapid_gke_version}'
+    )
+    xpk_print(
+        f'Error: Planned GKE Version {master_gke_version} is not valid.'
+        f'Checks failed: Is Version Valid: {is_valid_version}'
+    )
+    xpk_print(
+        'Please select a gke version from the above list using --gke-version=x'
+        ' argument or rely on the default gke version:'
+        f' {gke_server_config.default_rapid_gke_version}'
+    )
+    return 1, None
+
+  return 0, master_gke_version

xpk/core/gcluster_manager.py

@@ -0,0 +1,176 @@
+"""
+Copyright 2024 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.
+"""
+
+from .docker_manager import CommandRunner
+from ..utils.console import xpk_exit, xpk_print
+from .remote_state.remote_state_client import RemoteStateClient
+
+xpk_gcloud_cfg_path = '~/gcloud/cfg'
+xpk_deployment_dir = '/deployment'
+gcluster_deploy_command = 'gcluster deploy'
+gcluster_create_command = 'gcluster create'
+gcluster_destroy_command = 'gcluster destroy'
+blueprint_file_name = 'xpk_blueprint.yaml'
+deployment_module = '/out/xpk-deployment'
+a3_utils_dir_name = 'a3-mega-xpk'
+config_map_repo_path = 'src/xpk/blueprints/a3-mega-xpk/config-map.yaml.tftpl'
+kueue_config_repo_path = (
+    'src/xpk/blueprints/a3-mega-xpk/kueue-xpk-configuration.yaml.tftpl'
+)
+
+
+class GclusterManager:
+  """Manager is a class responsible for running cluster toolkit commands.
+  Attributes:
+    - gcluster_command_runner (CommandRunner) : instance of class implementing CommandRunner abstract methods.
+  Methods:
+    - deploy : run a deployment process of cluster toolkit. This method will invoke gcluster create and than gcluster deploy commands.
+    - destroy_deployment : run gcluster command to destroy existing deployment.
+  """
+
+  def __init__(
+      self,
+      gcluster_command_runner: CommandRunner,
+      remote_state_client: RemoteStateClient | None,
+  ) -> None:
+    self.gcluster_command_runner = gcluster_command_runner
+    self.remote_state_client = remote_state_client
+
+  def _run_create_deployment_cmd(
+      self, blueprint_container_path: str, prefix: str = ''
+  ):
+    xpk_print('Creating deployment resources...')
+    cluster_create_cmd = (
+        f'{gcluster_create_command} -o {self._get_deployment_path(prefix)}'
+        f' {blueprint_container_path} -w --force'
+    )
+    self.gcluster_command_runner.run_command(cluster_create_cmd)
+    xpk_print('Creating deployment resources completed.')
+
+  def _run_deploy_cmd(
+      self,
+      deployment_name: str,
+      auto_approve: bool,
+      dry_run: bool,
+      prefix: str = '',
+  ):
+    xpk_print('Deploying resources...')
+    deploy_cmd = (
+        f'{gcluster_deploy_command} {self._get_deployment_path(prefix)}/{deployment_name}'
+    )
+    if auto_approve is True:
+      deploy_cmd += ' --auto-approve'
+    if dry_run is True:
+      return
+    self.gcluster_command_runner.run_command(deploy_cmd)
+    xpk_print('Deployment completed.')
+
+  def deploy(
+      self,
+      blueprint_path: str,
+      deployment_name: str,
+      prefix: str = '',
+      auto_approve: bool = True,
+      dry_run: bool = False,
+  ) -> None:
+    """ "deploy method provisions a new cluster using Cluster Toolkit.
+    It will invoke gcluster create and then gcluster deploy commands.
+    The files staged or created during running gcluster command will be managed by gcluster_command_runner in its working directory."
+
+    Args:
+        blueprint_path (str): path pointing to blueprint which will be deployed.
+        deployment_name (str): name of the deployment.
+        auto_approve (bool, optional): If set to true deployment command will be auto approved. Currently only True is supported. Defaults to True.
+        dry_run (bool, optional): If set to True gcluster will not deploy. Defaults to False.
+    Returns:
+      None
+    """
+    xpk_print(f'Deploying blueprint from path {blueprint_path} ...')
+    self._run_create_deployment_cmd(
+        blueprint_container_path=blueprint_path, prefix=prefix
+    )
+    self._run_deploy_cmd(
+        deployment_name=deployment_name,
+        prefix=prefix,
+        auto_approve=auto_approve,
+        dry_run=dry_run,
+    )
+    xpk_print('Deploying blueprint completed!')
+
+  def _run_destroy_command(
+      self,
+      deployment_name: str,
+      prefix: str = '',
+      auto_approve: bool = True,
+      dry_run: bool = False,
+  ):
+    destroy_cmd = (
+        f'{gcluster_destroy_command} {self._get_deployment_path(prefix)}/{deployment_name}'
+    )
+    if auto_approve is True:
+      destroy_cmd += ' --auto-approve'
+    if dry_run is True:
+      xpk_print(f'executing command {destroy_cmd}')
+      return
+    self.gcluster_command_runner.run_command(destroy_cmd)
+
+  def _get_deployment_path(self, prefix: str = '') -> str:
+    prefix = f'/{prefix}' if prefix != '' else ''
+    return f'deployments{prefix}'
+
+  def destroy_deployment(self, deployment_name: str, prefix: str = '') -> None:
+    """Destroy deployment.
+
+    Args:
+        deployment_name (str): name of deployment to destroy.
+    """
+    xpk_print(f'Destroying {deployment_name} started...')
+    self._run_destroy_command(deployment_name, prefix=prefix)
+    xpk_print(f'Destroying {deployment_name} completed!')
+
+  def stage_files(
+      self, blueprint_file: str, blueprint_dependencies: str, prefix: str = ''
+  ) -> str:
+    """Uploads blueprint file and directory to gcluster working directory."""
+    xpk_print(
+        "Staging (sending) blueprint file to gcluster's working directory..."
+    )
+    staged_blueprint = self.gcluster_command_runner.upload_file_to_working_dir(
+        blueprint_file, prefix
+    )
+    if len(blueprint_dependencies) > 0:
+      self.gcluster_command_runner.upload_directory_to_working_dir(
+          blueprint_dependencies, prefix
+      )
+    xpk_print('Staging blueprint completed!')
+    xpk_print(f"File path in gcluster's working directory: {staged_blueprint}")
+    return staged_blueprint
+
+  def upload_state(self) -> None:
+    xpk_print('Uploading state.')
+    if self.remote_state_client is None:
+      xpk_print('No remote state defined')
+      xpk_exit(1)
+    self.remote_state_client.upload_state()
+
+  def download_state(self) -> None:
+    if self.remote_state_client is None:
+      xpk_print('No remote state defined')
+      xpk_exit(1)
+
+    if self.remote_state_client.check_remote_state_exists():
+      self.remote_state_client.download_state()
+    xpk_print('Remote state not found.')

xpk/core/gcsfuse.py

@@ -0,0 +1,50 @@
+"""
+Copyright 2024 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.
+"""
+
+from ..utils import templates
+
+FUSE_PV_PATH = "/../templates/fuse-pv.yaml"
+FUSE_PVC_PATH = "/../templates/fuse-pvc.yaml"
+
+
+def create_pv(name: str, size: int, bucket: str) -> dict:
+  data = templates.load(FUSE_PV_PATH)
+  data["metadata"]["name"] = f"{name}-pv"
+  data["spec"]["capacity"]["storage"] = f"{size}Gi"
+  data["spec"]["csi"]["volumeHandle"] = bucket
+  return data
+
+
+def create_pvc(name: str, size: int) -> dict:
+  data = templates.load(FUSE_PVC_PATH)
+  data["metadata"]["name"] = f"{name}-pvc"
+  data["spec"]["resources"]["requests"]["storage"] = f"{size}Gi"
+  data["spec"]["volumeName"] = f"{name}-pv"
+  return data
+
+
+def manifest(name: str, bucket: str, size: int) -> list[dict]:
+  """Creates GCS FUSE manifest file.
+
+  Args:
+      path (str): path to the file where the manifest will be created
+      name (str): base name of the volumes
+      bucket (str): name of the storage bucket
+      size (str): size of the storage
+  """
+  pv = create_pv(name, size, bucket)
+  pvc = create_pvc(name, size)
+  return [pv, pvc]