xpk 0.7.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

xpk/core/blueprint/blueprint_generator.py

@@ -22,22 +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"
+cluster_toolkit_version = "v1.48.0"
 
 
 def get_subnetworks_for_a3mega(cluster_name: str) -> list[str]:
@@ -50,6 +62,10 @@ def get_subnetworks_for_a3ultra(cluster_name: str) -> list[str]:
   ]
 
 
+def get_subnetworks_for_a4() -> list[str]:
+  return ["gvnic-1"] + [f"rdma-{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:
@@ -157,6 +173,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"],
     )
@@ -200,7 +221,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": [{
@@ -244,7 +265,7 @@ class BlueprintGenerator:
     )
     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'
 
 

xpk/core/cluster.py

@@ -32,7 +32,8 @@ from .gcloud_context import add_zone_and_project, get_gke_server_config, zone_to
 from .nodepool import upgrade_gke_nodepools_version
 from .system_characteristics import SystemCharacteristics
 
-JOBSET_VERSION = 'v0.7.2'
+JOBSET_VERSION = 'v0.8.0'
+PATHWAYS_JOB_VERSION = 'v0.1.0'
 INSTALLER_NCC_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
 INSTALLER_NCC_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
 
@@ -71,6 +72,35 @@ def set_jobset_on_cluster(args) -> int:
   return return_code
 
 
+def set_pathways_job_on_cluster(args) -> int:
+  """Add PathwaysJob command on server side and ask user to verify it is created.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'kubectl apply --server-side -f'
+      f' https://github.com/google/pathways-job/releases/download/{PATHWAYS_JOB_VERSION}/install.yaml'
+  )
+  task = f'Install PathwaysJob on {args.cluster}'
+  return_code = run_command_with_updates_retry(command, task, args)
+
+  if return_code != 0:
+    xpk_print(f'{task} returned with ERROR {return_code}.\n')
+    xpk_print(
+        "This LIKELY means you're missing Kubernetes Permissions, you can"
+        ' validate this by checking if the error references permission problems'
+        ' such as `requires one of ["container.*"] permission(s)`. Follow our'
+        ' readme:'
+        ' https://github.com/google/xpk/blob/main/README.md#troubleshooting for'
+        ' instructions on how to fix these permissions.'
+    )
+  return return_code
+
+
 def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
   """Install NCCL plugin on the cluster.
 
@@ -135,8 +165,48 @@ def update_cluster_with_gcpfilestore_driver_if_necessary(args) -> int:
   return 0
 
 
+def update_cluster_with_parallelstore_driver_if_necessary(args) -> int:
+  """Updates a GKE cluster to enable Parallelstore CSI driver, if not enabled already.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and error code otherwise.
+  """
+  if is_driver_enabled_on_cluster(args, driver='parallelstoreCsiDriver'):
+    return 0
+  cluster_update_return_code = update_gke_cluster_with_addon(
+      args, 'ParallelstoreCsiDriver'
+  )
+  if cluster_update_return_code > 0:
+    xpk_print('Updating GKE cluster to enable Parallelstore CSI driver failed!')
+    return cluster_update_return_code
+
+  return 0
+
+
+def update_cluster_with_pd_driver_if_necessary(args) -> int:
+  """Updates a GKE cluster to enable PersistentDisk CSI driver, if not enabled already.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and error code otherwise.
+  """
+  if is_driver_enabled_on_cluster(args, driver='gcePersistentDiskCsiDriver'):
+    return 0
+  cluster_update_return_code = update_gke_cluster_with_addon(
+      args, 'GcePersistentDiskCsiDriver'
+  )
+  if cluster_update_return_code > 0:
+    xpk_print(
+        'Updating GKE cluster to enable PersistentDisk CSI driver failed!'
+    )
+    return cluster_update_return_code
+
+  return 0
+
+
 def is_driver_enabled_on_cluster(args, driver: str) -> bool:
-  """Checks if GCSFuse CSI driver is enabled on the cluster.
+  """Checks if the CSI driver is enabled on the cluster.
   Args:
     args: user provided arguments for running the command.
     driver (str) : name of the driver
@@ -148,14 +218,14 @@ def is_driver_enabled_on_cluster(args, driver: str) -> bool:
       f' --project={args.project} --region={zone_to_region(args.zone)}'
       f' --format="value(addonsConfig.{driver}Config.enabled)"'
   )
-  return_code, gcsfuse_driver_enabled = run_command_for_value(
+  return_code, driver_enabled = run_command_for_value(
       command,
       f'Checks if {driver} driver is enabled in cluster describe.',
       args,
   )
   if return_code != 0:
     xpk_exit(return_code)
-  if gcsfuse_driver_enabled.lower() == 'true':
+  if driver_enabled.strip().lower() == 'true':
     xpk_print(f'{driver} driver is enabled on the cluster, no update needed.')
     return True
   return False
@@ -446,7 +516,7 @@ def is_gcsfuse_driver_enabled_on_cluster(args) -> bool:
   )
   if return_code != 0:
     xpk_exit(return_code)
-  if gcsfuse_driver_enabled.lower() == 'true':
+  if gcsfuse_driver_enabled.strip().lower() == 'true':
     xpk_print('GCSFuse CSI driver is enabled on the cluster, no update needed.')
     return True
   return False

xpk/core/config.py

@@ -24,7 +24,7 @@ from ..utils.console import xpk_print
 from .system_characteristics import AcceleratorType, SystemCharacteristics
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.7.1'
+__version__ = 'v0.8.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -39,6 +39,7 @@ KJOB_SHELL_IMAGE = 'shell-image'
 KJOB_SHELL_INTERACTIVE_COMMAND = 'shell-interactive-command'
 KJOB_SHELL_WORKING_DIRECTORY = 'shell-working-directory'
 CONFIGS_KEY = 'configs'
+GKE_ENDPOINT_KEY = 'gke-endpoint'
 DEPENDENCIES_KEY = 'deps-verified-version'
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -47,6 +48,7 @@ DEFAULT_KEYS = [
     CLUSTER_NAME_KEY,
     PROJECT_KEY,
     ZONE_KEY,
+    GKE_ENDPOINT_KEY,
     DEPENDENCIES_KEY,
     KJOB_BATCH_IMAGE,
     KJOB_BATCH_WORKING_DIRECTORY,

xpk/core/docker_manager.py

@@ -30,7 +30,7 @@ import time
 DockerRunCommandExitCode = 135
 dockerBuildErrorCode = 134
 ctk_dockerfile_path = "Dockerfile"
-ctk_build_ref = "v1.45.1"
+ctk_build_ref = "v1.48.0"
 ctk_docker_image = "xpk-ctk"
 ctk_container_name = "xpk-ctk-container"
 gcloud_cfg_mount_path = "/root/.config/gcloud"