Introduce Storage API

pyproject.toml

@@ -30,7 +30,10 @@ keywords = []
 
 # pip dependencies installed with `pip install -e .`
 dependencies = [
- "cloud-accelerator-diagnostics"
+ "cloud-accelerator-diagnostics",
+ "kubernetes",
+ "google-cloud",
+ "google-api-core",
 ]
 
 [project.urls]
@@ -57,8 +60,9 @@ dev = [
 version = {attr = "xpk.core.core.__version__"}
 
 [tool.setuptools]
-packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands"]
+packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands", "xpk.api"]
 package-dir = {"" = "src"}
+package-data = {"xpk.api" = ["storage_crd.yaml"]}
 
 [tool.pyink]
 # Formatting configuration to follow Google style-guide.

src/xpk/api/__init__.py

@@ -0,0 +1,15 @@
+"""
+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.
+"""

src/xpk/api/storage_crd.yaml

@@ -0,0 +1,52 @@
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+ name: storages.xpk.x-k8s.io
+spec:
+ group: xpk.x-k8s.io
+ versions:
+ - name: v1
+ served: true
+ storage: true
+ schema:
+ openAPIV3Schema:
+ type: object
+ properties:
+ spec:
+ type: object
+ properties:
+ type:
+ type: string
+ cluster:
+ type: string
+ auto_mount:
+ type: boolean
+ mount_point:
+ type: string
+ readonly:
+ type: boolean
+ manifest:
+ type: string
+ pv:
+ type: string
+ pvc:
+ type: string
+ required:
+ - type
+ - cluster
+ - auto_mount
+ - mount_point
+ - readonly
+ - manifest
+ - pvc
+ - pv
+ x-kubernetes-validations:
+ - message: Value is immutable
+ rule: self == oldSelf
+ scope: Cluster
+ names:
+ plural: storages
+ singular: storage
+ kind: Storage
+ shortNames:
+ - stg

src/xpk/commands/cluster.py

@@ -14,11 +14,7 @@
 limitations under the License.
 """
 
-from ..core.commands import (
- run_command_for_value,
- run_command_with_updates,
- run_command_with_updates_retry,
-)
+from ..core.commands import run_command_for_value, run_command_with_updates
 from ..core.core import (
  VERTEX_TENSORBOARD_FEATURE_FLAG,
  add_zone_and_project,
@@ -27,6 +23,7 @@
  create_vertex_tensorboard,
  delete_cluster_subnets,
  get_all_clusters_programmatic,
+ get_cluster_credentials,
  get_gke_control_plane_version,
  get_gke_node_pool_version,
  get_gke_server_config,
@@ -113,9 +110,7 @@ def cluster_create(args) -> None:
  if update_cluster_command_code != 0:
  xpk_exit(update_cluster_command_code)
 
- set_cluster_command_code = set_cluster_command(args)
- if set_cluster_command_code != 0:
- xpk_exit(set_cluster_command_code)
+ get_cluster_credentials(args)
 
  # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
  tensorboard_config = {}
@@ -236,9 +231,7 @@ def cluster_cacheimage(args) -> None:
  )
  add_zone_and_project(args)
 
- set_cluster_command_code = set_cluster_command(args)
- if set_cluster_command_code != 0:
- xpk_exit(set_cluster_command_code)
+ get_cluster_credentials(args)
  system, return_code = get_system_characteristics(args)
 
  if return_code > 0:
@@ -287,9 +280,7 @@ def cluster_describe(args) -> None:
  xpk_print(f'Starting nodepool list for cluster: {args.cluster}', flush=True)
  add_zone_and_project(args)
 
- set_cluster_command_code = set_cluster_command(args)
- if set_cluster_command_code != 0:
- xpk_exit(set_cluster_command_code)
+ get_cluster_credentials(args)
 
  command = (
  f'gcloud container node-pools list --cluster {args.cluster} '
@@ -502,28 +493,3 @@ def run_gke_cluster_create_command(
  xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')
  return 1
  return 0
-
-
-def set_cluster_command(args) -> int:
- """Run cluster configuration command to set the kubectl config.
-
- Args:
- args: user provided arguments for running the command.
-
- Returns:
- 0 if successful and 1 otherwise.
- """
- command = (
- 'gcloud container clusters get-credentials'
- f' {args.cluster} --region={zone_to_region(args.zone)}'
- f' --project={args.project} &&'
- ' kubectl config view && kubectl config set-context --current'
- ' --namespace=default'
- )
- task = f'get-credentials to cluster {args.cluster}'
- return_code = run_command_with_updates_retry(
- command, task, args, verbose=False
- )
- if return_code != 0:
- xpk_print(f'{task} returned ERROR {return_code}')
- return return_code

src/xpk/commands/inspector.py

@@ -19,11 +19,11 @@
  CLUSTER_METADATA_CONFIGMAP,
  CLUSTER_RESOURCES_CONFIGMAP,
  add_zone_and_project,
+ get_cluster_credentials,
  zone_to_region,
 )
 from ..core.kueue import CLUSTER_QUEUE_NAME, LOCAL_QUEUE_NAME
 from ..utils import append_tmp_file, write_tmp_file, xpk_exit, xpk_print
-from .cluster import set_cluster_command
 from .workload import get_workload_list
 
 
@@ -124,9 +124,7 @@ def inspector(args) -> None:
  xpk_print(args)
 
  add_zone_and_project(args)
- set_cluster_command_code = set_cluster_command(args)
- if set_cluster_command_code != 0:
- xpk_exit(set_cluster_command_code)
+ get_cluster_credentials(args)
 
  inspector_file = write_tmp_file(
  '==================\nXPK inspector OUTPUT:\n==================\n'

src/xpk/commands/storage.py

@@ -0,0 +1,84 @@
+"""
+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 argparse import Namespace
+
+from kubernetes import client as k8s_client
+from kubernetes.client.exceptions import ApiException
+
+from ..core.core import (
+ setup_k8s_env,
+ update_cluster_with_gcsfuse_driver_if_necessary,
+ update_cluster_with_workload_identity_if_necessary,
+)
+from ..core.storage import (
+ STORAGE_CRD_KIND,
+ XPK_API_GROUP_NAME,
+ XPK_API_GROUP_VERSION,
+ create_storage_instance,
+ get_storage,
+ install_storage_crd,
+ list_storages,
+ print_storages_for_cluster,
+)
+from ..utils import apply_kubectl_manifest, xpk_exit, xpk_print
+
+
+def storage_create(args: Namespace) -> None:
+ k8s_api_client = setup_k8s_env(args)
+
+ install_storage_crd(k8s_api_client)
+ return_code = update_cluster_with_workload_identity_if_necessary(args)
+ if return_code > 0:
+ xpk_exit(return_code)
+ return_code = update_cluster_with_gcsfuse_driver_if_necessary(args)
+ if return_code > 0:
+ xpk_exit(return_code)
+
+ create_storage_instance(k8s_api_client, args)
+ apply_kubectl_manifest(k8s_api_client, args.manifest)
+
+
+def storage_list(args: Namespace) -> None:
+ k8s_api_client = setup_k8s_env(args)
+ install_storage_crd(k8s_api_client)
+ storages = list_storages(k8s_api_client)
+ print_storages_for_cluster(storages, args.cluster)
+
+
+def storage_delete(args: Namespace) -> None:
+ k8s_api_client = setup_k8s_env(args)
+ install_storage_crd(k8s_api_client)
+ api_instance = k8s_client.CustomObjectsApi(k8s_api_client)
+ core_api = k8s_client.CoreV1Api()
+ try:
+ storage = get_storage(k8s_api_client, args.name)
+
+ api_instance.delete_cluster_custom_object(
+ name=args.name,
+ group=XPK_API_GROUP_NAME,
+ version=XPK_API_GROUP_VERSION,
+ plural=STORAGE_CRD_KIND.lower() + "s",
+ )
+
+ core_api.delete_namespaced_persistent_volume_claim(storage.pvc, "default")
+ core_api.delete_persistent_volume(storage.pv)
+ except ApiException as e:
+ if e.status == 404:
+ xpk_print(f"Storage: {args.name} not found. Might be already deleted.")
+ else:
+ xpk_print(f"Encountered error during storage deletion: {e}")
+ xpk_exit(1)