Add local tests with Kind

Author: IrvingMg

Makefile

@@ -37,6 +37,10 @@ run-unittests:
 run-integrationtests:
 	pytest src/xpk/core/tests/integration/
 
+run-kindtests:
+	chmod +x ./tools/run-kind-tests.sh
+	./tools/run-kind-tests.sh
+
 .PHONY: mkdir-bin
 mkdir-bin:
 	mkdir -p $(BIN_PATH)

README.md

@@ -1349,13 +1349,18 @@ xpk interfaces seamlessly with kind to manage Kubernetes clusters locally, facil
 
 ## Local Testing Basics
 
-Local testing is available exclusively through the `batch` and `job` commands of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally:
+Local testing is achievable through most commands of `xpk` except those that require Pathways, like `cluster create-pathways` or `workload create-pathways`. This functionality is supported by using the `--kind-cluster` flag which allows you to simulate operations locally on the `kind` tool.
+
+This example demonstrates how to run a batch job locally using the --kind-cluster flag:
 
 ```shell
 python xpk.py batch [other-options] --kind-cluster script
 ```
 
-Please note that all other xpk subcommands are intended for use with cloud systems on Google Cloud Engine (GCE) and don't support local testing. This includes commands like cluster, info, inspector, etc.
+While the `--kind-cluster` flag does extend local testing capabilities to several commands, please be aware that commands requiring specific features from Google Cloud Platform (GCP) might not yet be fully supported when tested locally. Future updates may provide enhanced support for these GCP-specific features.
+
+Currently supported local testing cases can be reviewed in the script: `tools/run-kind-tests.sh`.
+
 
 # Other advanced usage
 [Use a Jupyter notebook to interact with a Cloud TPU cluster](xpk-notebooks.md)

src/xpk/commands/cluster.py

@@ -14,6 +14,9 @@
 limitations under the License.
 """
 
+from tabulate import tabulate
+
+from ..core.cluster_private import authorize_private_cluster_access_if_necessary
 from ..core.commands import run_command_for_value, run_command_with_updates
 from ..core.core import (
     VERTEX_TENSORBOARD_FEATURE_FLAG,
@@ -26,20 +29,15 @@
     get_gke_control_plane_version,
     get_gke_node_pool_version,
     get_gke_server_config,
+    get_user_input,
     h100_device_type,
     install_nccl_on_cluster,
     run_gke_node_pool_create_command,
     set_jobset_on_cluster,
     set_up_cluster_network_for_gpu,
     zone_to_region,
-    get_user_input,
-)
-from ..core.cluster_private import authorize_private_cluster_access_if_necessary
-from ..core.kjob import (
-    verify_kjob_installed,
-    prepare_kjob,
-    apply_kjob_crds,
 )
+from ..core.kjob import apply_kjob_crds, prepare_kjob, verify_kjob_installed
 from ..core.kueue import (
     cluster_preheat_yml,
     install_kueue_crs,
@@ -55,12 +53,11 @@
     get_system_characteristics,
 )
 from ..core.workload import get_workload_list
-from ..utils.file import write_tmp_file
 from ..utils.console import xpk_exit, xpk_print
+from ..utils.file import write_tmp_file
 from . import cluster_gcluster
 from .common import set_cluster_command
-
-from tabulate import tabulate
+from .kind import set_local_cluster_command
 
 
 def cluster_create(args) -> None:
@@ -316,9 +313,13 @@ def cluster_describe(args) -> None:
     0 if successful and 1 otherwise.
   """
   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 not args.kind_cluster:
+    add_zone_and_project(args)
+    set_cluster_command_code = set_cluster_command(args)
+  else:
+    set_cluster_command_code = set_local_cluster_command(args)
+
   if set_cluster_command_code != 0:
     xpk_exit(set_cluster_command_code)
 

src/xpk/commands/info.py

@@ -23,6 +23,7 @@
 from ..core.core import (
     add_zone_and_project,
 )
+from .kind import set_local_cluster_command
 import json
 from tabulate import tabulate
 from argparse import Namespace
@@ -38,8 +39,12 @@ def info(args: Namespace) -> None:
   Returns:
     None
   """
-  add_zone_and_project(args)
-  set_cluster_command_code = set_cluster_command(args)
+  if not args.kind_cluster:
+    add_zone_and_project(args)
+    set_cluster_command_code = set_cluster_command(args)
+  else:
+    set_cluster_command_code = set_local_cluster_command(args)
+
   if set_cluster_command_code != 0:
     xpk_exit(set_cluster_command_code)
 
@@ -84,7 +89,7 @@ def get_nominal_quotas(cqs: list[dict]) -> dict[str, dict[str, str]]:
     spec = cq['spec']
     cq_name = cq['metadata']['name']
     quotas[cq_name] = {}
-    for rg in spec['resourceGroups']:
+    for rg in spec.get('resourceGroups', []):
       for flavor in rg['flavors']:
         name = flavor['name']
         for resource in flavor['resources']:

src/xpk/commands/inspector.py

@@ -26,6 +26,7 @@
 from ..utils.console import xpk_exit, xpk_print
 from .common import set_cluster_command
 from .workload import get_workload_list
+from .kind import set_local_cluster_command
 
 
 def inspector_run_command_helper(
@@ -124,52 +125,63 @@ def inspector(args) -> None:
   final_return_code = 0
   xpk_print(args)
 
-  add_zone_and_project(args)
-  set_cluster_command_code = set_cluster_command(args)
+  if not args.kind_cluster:
+    add_zone_and_project(args)
+    set_cluster_command_code = set_cluster_command(args)
+  else:
+    set_cluster_command_code = set_local_cluster_command(args)
+
   if set_cluster_command_code != 0:
     xpk_exit(set_cluster_command_code)
 
   inspector_file = write_tmp_file(
       '==================\nXPK inspector OUTPUT:\n==================\n'
   )
-  command_and_descriptions = [
-      ('gcloud version', 'Local Setup: gcloud version'),
-      (
-          (
-              'gcloud config get project; gcloud config get compute/zone;'
-              ' gcloud config get compute/region'
-          ),
-          'Local Setup: Project / Zone / Region',
-      ),
-      (
-          (
-              'gcloud beta container clusters list --project'
-              f' {args.project} --region {zone_to_region(args.zone)} | grep -e'
-              f' NAME -e {args.cluster}'
-          ),
-          'GKE: Cluster Details',
-      ),
-      (
-          (
-              'kubectl get configmap'
-              f' {args.cluster}-{CLUSTER_METADATA_CONFIGMAP} -o yaml'
-          ),
-          'GKE: Cluster Metadata ConfigMap Details',
-      ),
-      (
-          (
-              'kubectl get configmap'
-              f' {args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP} -o yaml'
-          ),
-          'GKE: Cluster Resources ConfigMap Details',
-      ),
-      (
-          (
-              f'gcloud beta container node-pools list --cluster {args.cluster} '
-              f' --project={args.project} --region={zone_to_region(args.zone)}'
-          ),
-          'GKE: Node pool Details',
-      ),
+
+  gcloud_commands_and_descriptions = []
+  if not args.kind_cluster:
+    gcloud_commands_and_descriptions = [
+        ('gcloud version', 'Local Setup: gcloud version'),
+        (
+            (
+                'gcloud config get project; gcloud config get compute/zone;'
+                ' gcloud config get compute/region'
+            ),
+            'Local Setup: Project / Zone / Region',
+        ),
+        (
+            (
+                'gcloud beta container clusters list --project'
+                f' {args.project} --region {zone_to_region(args.zone)} | grep'
+                f' -e NAME -e {args.cluster}'
+            ),
+            'GKE: Cluster Details',
+        ),
+        (
+            (
+                'kubectl get configmap'
+                f' {args.cluster}-{CLUSTER_METADATA_CONFIGMAP} -o yaml'
+            ),
+            'GKE: Cluster Metadata ConfigMap Details',
+        ),
+        (
+            (
+                'kubectl get configmap'
+                f' {args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP} -o yaml'
+            ),
+            'GKE: Cluster Resources ConfigMap Details',
+        ),
+        (
+            (
+                'gcloud beta container node-pools list --cluster'
+                f' {args.cluster} '
+                f' --project={args.project} --region={zone_to_region(args.zone)}'
+            ),
+            'GKE: Node pool Details',
+        ),
+    ]
+
+  kubectl_commands_and_descriptions = [
       (
           (
               "kubectl get node -o custom-columns='NODE_NAME:metadata.name,"
@@ -234,6 +246,9 @@ def inspector(args) -> None:
       ),
   ]
 
+  command_and_descriptions = (
+      gcloud_commands_and_descriptions + kubectl_commands_and_descriptions
+  )
   for command, description in command_and_descriptions:
     return_code = inspector_run_command_helper(
         args, command, description, inspector_file
@@ -311,45 +326,46 @@ def inspector(args) -> None:
 
   # Cloud Console Links:
   workload_links = []
-  if args.workload:
+  if args.workload and not args.kind_cluster:
     workload_links = [(
         f'Cloud Console for the workload {args.workload}',
         # pylint: disable=line-too-long
         f'https://console.cloud.google.com/kubernetes/service/{zone_to_region(args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}',
     )]
 
-  links = [
-      (
-          'Cloud Console for the GKE Cluster',
-          # pylint: disable=line-too-long
-          f'https://console.cloud.google.com/kubernetes/clusters/details/{zone_to_region(args.zone)}/{args.cluster}/details?project={args.project}',
-      ),
-      (
-          'Cloud Console for all workloads in GKE Cluster',
-          # pylint: disable=line-too-long
-          f'https://console.cloud.google.com/kubernetes/workload/overview?project={args.project}&pageState=((gke%2F{zone_to_region(args.zone)}%2F{args.cluster}))',
-      ),
-      (
-          'Cloud Console for IAM Permissions',
-          f'https://console.cloud.google.com/iam-admin/iam?project={args.project}',
-      ),
-      (
-          'Cloud Console for Quotas',
-          f'https://console.cloud.google.com/iam-admin/quotas?project={args.project}',
-      ),
-  ]
-  links.extend(workload_links)
-
-  for description, workload_link in links:
-    return_code = inspector_output_link_helper(
-        args, workload_link, description, inspector_file
-    )
-    if return_code != 0:
-      final_return_code = return_code
-      xpk_print(
-          f'inspector failed in link: {workload_link} description:'
-          f' {description} return code: {return_code}'
+  if not args.kind_cluster:
+    links = [
+        (
+            'Cloud Console for the GKE Cluster',
+            # pylint: disable=line-too-long
+            f'https://console.cloud.google.com/kubernetes/clusters/details/{zone_to_region(args.zone)}/{args.cluster}/details?project={args.project}',
+        ),
+        (
+            'Cloud Console for all workloads in GKE Cluster',
+            # pylint: disable=line-too-long
+            f'https://console.cloud.google.com/kubernetes/workload/overview?project={args.project}&pageState=((gke%2F{zone_to_region(args.zone)}%2F{args.cluster}))',
+        ),
+        (
+            'Cloud Console for IAM Permissions',
+            f'https://console.cloud.google.com/iam-admin/iam?project={args.project}',
+        ),
+        (
+            'Cloud Console for Quotas',
+            f'https://console.cloud.google.com/iam-admin/quotas?project={args.project}',
+        ),
+    ]
+    links.extend(workload_links)
+
+    for description, workload_link in links:
+      return_code = inspector_output_link_helper(
+          args, workload_link, description, inspector_file
       )
+      if return_code != 0:
+        final_return_code = return_code
+        xpk_print(
+            f'inspector failed in link: {workload_link} description:'
+            f' {description} return code: {return_code}'
+        )
 
   # Summarize inspector:
   xpk_print(f'Find xpk inspector output file: {inspector_file.name}')

src/xpk/commands/kind.py

@@ -28,8 +28,16 @@
 )
 from ..core.kueue import (
     install_kueue_on_cluster,
+    wait_for_kueue_available,
+    install_kueue_crs,
+)
+from ..core.system_characteristics import (
+    get_system_characteristics,
 )
 from ..utils.console import (xpk_exit, xpk_print)
+from ..utils.file import write_tmp_file
+
+import yaml
 
 
 def cluster_create(args) -> None:
@@ -79,6 +87,22 @@ def cluster_create(args) -> None:
   if err_code > 0:
     xpk_exit(err_code)
 
+  xpk_print('Wait for Kueue to be fully available')
+  wait_for_kueue_available_code = wait_for_kueue_available(args)
+  if wait_for_kueue_available_code != 0:
+    xpk_exit(wait_for_kueue_available_code)
+
+  args.kind_cluster = True
+  system, return_code = get_system_characteristics(args)
+  if return_code > 0:
+    xpk_print('Fetching system characteristics failed!')
+    xpk_exit(return_code)
+
+  xpk_print('Install Kueue Custom Resources')
+  enable_kueue_credentials_code = install_kueue_crs(args, system, None)
+  if enable_kueue_credentials_code != 0:
+    xpk_exit(enable_kueue_credentials_code)
+
   xpk_print('Kind commands done! Resources are created.')
   xpk_exit(0)
 
@@ -184,7 +208,24 @@ def run_kind_cluster_create_command(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
-  command = 'kind create cluster'
+  # Simulating GCP Environments in Kind
+  # Configures kind nodes to mimic GCP features such as node pools, TPU accelerators or instance types.
+  # These modifications aid in testing the compatibility and behavior of applications outside real GCP.
+  # For more details on kind configurations, visit: https://kind.sigs.k8s.io/docs/user/configuration/
+  kind_config = {
+      'kind': 'Cluster',
+      'apiVersion': 'kind.x-k8s.io/v1alpha4',
+      'nodes': [
+          {'role': 'control-plane'},
+          {
+              'role': 'worker',
+              'labels': {'cloud.google.com/gke-nodepool': 'kind-pool-0'},
+          },
+      ],
+  }
+  yaml_str = yaml.dump(kind_config)
+  tmp = write_tmp_file(yaml_str)
+  command = f'kind create cluster --config {tmp.file.name}'
 
   if args.cluster:
     command += f' --name={args.cluster}'