xpk 0.14.4__py3-none-any.whl → 0.15.0__py3-none-any.whl

xpk/core/blueprint/blueprint_generator.py

@@ -19,10 +19,13 @@ import shutil
 from typing import Optional
 
 from ruamel import yaml
+from packaging.version import parse
 
 from ...utils.console import xpk_exit, xpk_print
+from ...utils.versions import ReleaseChannel
 from ...utils.file import ensure_directory_exists
 
+
 from ..capacity import (
     H100_DEVICE_TYPE,
     B200_DEVICE_TYPE,
@@ -84,6 +87,8 @@ class BlueprintGenerator:
       region: str,
       zone: str,
       auth_cidr: str,
+      cluster_version: str,
+      release_channel: ReleaseChannel,
       prefix: str = "",
       num_nodes: int = 2,
       pods_ip_cidr_range: str = "10.4.0.0/14",
@@ -142,11 +147,17 @@ class BlueprintGenerator:
         },
     )
 
+    sanitized_version = cluster_version.replace("-", "+", 1)
+    version = parse(sanitized_version)
+    version_prefix = f"{version.major}.{version.minor}"
     gke_cluster = DeploymentModule(
         id="gke_cluster",
         source="modules/scheduler/gke-cluster",
         use=[primary_vpc_name, gpu_subnets_name],
         settings={
+            "release_channel": release_channel.value,
+            "version_prefix": version_prefix,
+            "min_master_version": cluster_version,
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "enable_private_endpoint": False,
@@ -171,6 +182,16 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+    if release_channel != ReleaseChannel.RAPID:
+      gke_cluster.set_setting(
+          "maintenance_exclusions",
+          [{
+              "name": "no-minor-or-node-upgrades-indefinite",
+              "start_time": "2024-12-01T00:00:00Z",
+              "end_time": "2026-01-16T00:00:00Z",
+              "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+          }],
+      )
 
     group_placement_0 = DeploymentModule(
         id="group_placement_0",
@@ -215,6 +236,9 @@ class BlueprintGenerator:
     else:
       a3_megagpu_pool_0.update_settings({"static_node_count": num_nodes})
 
+    if release_channel == ReleaseChannel.RAPID:
+      a3_megagpu_pool_0.set_setting("auto_upgrade", True)
+
     set_placement_policy = capacity_type != CapacityType.SPOT
     workload = DeploymentModule(
         id="workload_component_install",
@@ -391,6 +415,8 @@ class BlueprintGenerator:
       zone: str,
       auth_cidr: str,
       system_node_pool_machine_type: str,
+      cluster_version: str,
+      release_channel: ReleaseChannel,
       reservation: Optional[str | None] = None,
       gcs_bucket: Optional[str | None] = None,
       num_nodes: int = 2,
@@ -480,28 +506,19 @@ class BlueprintGenerator:
             },
         },
     )
+
+    sanitized_version = cluster_version.replace("-", "+", 1)
+    version = parse(sanitized_version)
+    version_prefix = f"{version.major}.{version.minor}"
     cluster_id = f"{cluster_name}-a3-ultragpu-cluster"
     a3_ultra_cluster = DeploymentModule(
         id=cluster_id,
         source="modules/scheduler/gke-cluster",
         use=[net_0_id],
         settings={
-            "release_channel": (
-                "UNSPECIFIED"
-                if capacity_type == CapacityType.FLEX_START
-                else "RAPID"
-            ),
-            "version_prefix": "1.32.",
-            "maintenance_exclusions": (
-                []
-                if capacity_type == CapacityType.FLEX_START
-                else [{
-                    "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",
-                }]
-            ),
+            "release_channel": release_channel.value,
+            "version_prefix": version_prefix,
+            "min_cluster_version": cluster_version,
             "prefix_with_deployment_name": False,
             "name_suffix": cluster_name,
             "system_node_pool_machine_type": system_node_pool_machine_type,
@@ -537,6 +554,17 @@ class BlueprintGenerator:
         },
         outputs=["instructions"],
     )
+    if release_channel != ReleaseChannel.RAPID:
+      a3_ultra_cluster.set_setting(
+          "maintenance_exclusions",
+          [{
+              "name": "no-minor-or-node-upgrades-indefinite",
+              "start_time": "2024-12-01T00:00:00Z",
+              "end_time": "2026-01-16T00:00:00Z",
+              "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+          }],
+      )
+
     system, _ = get_system_characteristics_by_device_type(a3ultra_device_type)
     if system is None:
       xpk_print(
@@ -584,6 +612,9 @@ class BlueprintGenerator:
     else:
       gpu_pool.update_settings({"static_node_count": num_nodes})
 
+    if release_channel == ReleaseChannel.RAPID:
+      gpu_pool.set_setting("auto_upgrade", True)
+
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
         id=workload_manager_install_id,
@@ -674,6 +705,8 @@ class BlueprintGenerator:
       zone: str,
       auth_cidr: str,
       system_node_pool_machine_type: str,
+      cluster_version: str,
+      release_channel: ReleaseChannel,
       reservation: Optional[str | None] = None,
       gcs_bucket: Optional[str | None] = None,
       num_nodes: int = 2,
@@ -761,12 +794,19 @@ class BlueprintGenerator:
             },
         },
     )
+
+    sanitized_version = cluster_version.replace("-", "+", 1)
+    version = parse(sanitized_version)
+    version_prefix = f"{version.major}.{version.minor}"
     cluster_id = f"{cluster_name}-a4-cluster"
     a4_cluster = DeploymentModule(
         id=cluster_id,
         source="modules/scheduler/gke-cluster",
         use=[net_0_id],
         settings={
+            "release_channel": release_channel.value,
+            "version_prefix": version_prefix,
+            "min_cluster_version": cluster_version,
             "system_node_pool_machine_type": system_node_pool_machine_type,
             "system_node_pool_node_count": {
                 "total_min_nodes": system_node_pool_min_node_count,
@@ -791,25 +831,20 @@ class BlueprintGenerator:
                 " alias_ip_range=[]}],"
                 f" {cluster_name}-rdma-net.subnetwork_interfaces_gke))"
             ),
-            "version_prefix": "1.32.",
-            "release_channel": (
-                "UNSPECIFIED"
-                if capacity_type == CapacityType.FLEX_START
-                else "RAPID"
-            ),
-            "maintenance_exclusions": (
-                []
-                if capacity_type == CapacityType.FLEX_START
-                else [{
-                    "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"],
     )
+    if release_channel != ReleaseChannel.RAPID:
+      a4_cluster.set_setting(
+          "maintenance_exclusions",
+          [{
+              "name": "no-minor-or-node-upgrades-indefinite",
+              "start_time": "2024-12-01T00:00:00Z",
+              "end_time": "2026-01-16T00:00:00Z",
+              "exclusion_scope": "NO_MINOR_OR_NODE_UPGRADES",
+          }],
+      )
+
     system, _ = get_system_characteristics_by_device_type(a4_device_type)
     if system is None:
       xpk_print(
@@ -859,6 +894,9 @@ class BlueprintGenerator:
     else:
       gpu_pool.update_settings({"static_node_count": num_nodes})
 
+    if release_channel == ReleaseChannel.RAPID:
+      gpu_pool.set_setting("auto_upgrade", True)
+
     workload_manager_install_id = "workload-manager-install"
     workload_manager_install = DeploymentModule(
         id=workload_manager_install_id,
@@ -1019,7 +1057,6 @@ class BlueprintGenerator:
         "enable_flex_start": True,
         "enable_queued_provisioning": True,
         "autoscaling_total_min_nodes": 0,
-        "release_channel": "UNSPECIFIED",
         "auto_repair": False,
         "auto_upgrade": False,
     }

xpk/core/blueprint/blueprint_test.py

@@ -22,6 +22,7 @@ import ruamel.yaml
 from xpk.core.blueprint.blueprint_definitions import Blueprint
 from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
 from xpk.core.capacity import CapacityType
+from xpk.utils.versions import ReleaseChannel
 
 yaml = ruamel.yaml.YAML()
 
@@ -60,6 +61,8 @@ def test_generate_a3_mega_blueprint():
       reservation="test-reservation",
       capacity_type=CapacityType.RESERVATION,
       system_node_pool_min_node_count=5,
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
 
   assert bp.blueprint_file.endswith("/prefix/xpk-gke-a3-megagpu.yaml")
@@ -99,6 +102,8 @@ def test_generate_a3_mega_spot_blueprint():
       auth_cidr="10.0.0.0/32",
       capacity_type=CapacityType.SPOT,
       system_node_pool_min_node_count=5,
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
 
   assert bp.blueprint_file.endswith("/prefix/xpk-gke-a3-megagpu.yaml")
@@ -135,6 +140,8 @@ def test_generate_a3_ultra_blueprint():
       capacity_type=CapacityType.RESERVATION,
       gcs_bucket="test-bucket",
       prefix="testdir",
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
   with open(a3_ultra_yaml_test_path, encoding="utf-8") as stream:
     ctk_yaml = yaml.load(stream)
@@ -180,6 +187,8 @@ def test_generate_a4_blueprint():
       capacity_type=CapacityType.RESERVATION,
       gcs_bucket="test-bucket",
       prefix="testdir",
+      release_channel=ReleaseChannel.RAPID,
+      cluster_version="1.2.3",
   )
   with open(a4_yaml_test_path, encoding="utf-8") as stream:
     ctk_yaml = yaml.load(stream)

xpk/core/capacity.py

@@ -115,9 +115,12 @@ def get_reservation_maintenance_interval(
   Returns:
     0 if successful and 1 otherwise.
   """
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      reservation, project
+  )
   command = (
-      f'gcloud beta compute reservations describe {reservation}'
-      f' --project={project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
   )
   return_code, output = run_command_for_value(
       command, 'Get reservation maintenance interval'
@@ -139,9 +142,12 @@ def get_reservation_placement_policy(
   Returns:
     0 if successful and 1 otherwise.
   """
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      reservation, project
+  )
   command = (
-      f'gcloud beta compute reservations describe {reservation}'
-      f' --project={project} --zone={zone} --format="value(resourcePolicies.policy)"'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={zone} --format="value(resourcePolicies.policy)"'
   )
   return_code, output = run_command_for_value(
       command, 'Get reservation placement policy'
@@ -156,9 +162,12 @@ def get_reservation_deployment_type(
     reservation: str, zone: str, project: str
 ) -> str:
   """Get reservation deployment type."""
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      reservation, project
+  )
   command = (
-      f'gcloud beta compute reservations describe {reservation}'
-      f' --project={project} --zone={zone} --format="value(deploymentType)"'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={zone} --format="value(deploymentType)"'
   )
   return_code, output = run_command_for_value(
       command, 'Get reservation deployment type', dry_run_return_val='DENSE'
@@ -178,9 +187,12 @@ def verify_reservation_exists(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
+  reservation_project, reservation_name = get_reservation_project_and_name(
+      args.reservation, args.project
+  )
   command = (
-      f'gcloud beta compute reservations describe {args.reservation}'
-      f' --project={args.project} --zone={args.zone}'
+      f'gcloud beta compute reservations describe {reservation_name}'
+      f' --project={reservation_project} --zone={args.zone}'
   )
   return_code = run_command_with_updates(command, 'Describe reservation')
   if return_code != 0:
@@ -264,3 +276,29 @@ def get_capacity_node_selectors_from_capacity_type(
       )
       return_code = 1
   return node_selector, return_code
+
+
+def get_reservation_project_and_name(
+    reservation_name_or_path: str, cluster_project: str
+) -> tuple[str, str]:
+  """Get the reservation project and name.
+
+  Args:
+    reservation_name_or_path: either reservation name or reservation path in format
+      projects/RESERVATION_PROJECT_ID/reservations/RESERVATION_NAME
+    cluster_project: the cluster project
+
+  Returns:
+    Tuple with reservation project and reservation name.
+  """
+  if '/' not in reservation_name_or_path:
+    return cluster_project, reservation_name_or_path
+  reservation_parts = reservation_name_or_path.split('/')
+  if (
+      len(reservation_parts) != 4
+      or reservation_parts[0] != 'projects'
+      or reservation_parts[2] != 'reservations'
+  ):
+    xpk_print('Unable to parse reservation: ', reservation_name_or_path)
+    xpk_exit(1)
+  return reservation_parts[1], reservation_parts[3]

xpk/core/capacity_test.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 import pytest
 from unittest.mock import MagicMock, patch
-from .capacity import get_reservation_deployment_type
+from .capacity import get_reservation_deployment_type, get_reservation_project_and_name
 
 
 @patch('xpk.core.capacity.xpk_print')
@@ -48,3 +48,34 @@ def test_get_reservation_deployment_type_returns_deployment_type_when_command_su
       reservation='reservation', zone='zone', project='project'
   )
   assert result == 'DENSE'
+
+
+def test_get_reservation_project_and_name_parses_local_reservation():
+  project, name = get_reservation_project_and_name(
+      'test-reservation', 'cluster-project'
+  )
+
+  assert project == 'cluster-project'
+  assert name == 'test-reservation'
+
+
+def test_get_reservation_project_and_name_parses_shared_reservation():
+  project, name = get_reservation_project_and_name(
+      'projects/reservation-project/reservations/test-reservation',
+      'cluster-project',
+  )
+
+  assert project == 'reservation-project'
+  assert name == 'test-reservation'
+
+
+@patch('xpk.core.capacity.xpk_print')
+def test_get_reservation_project_and_name_fails_for_invalid_reservation(
+    xpk_print: MagicMock, mocker
+):
+  with pytest.raises(SystemExit):
+    get_reservation_project_and_name(
+        'invalid/reservation',
+        'cluster-project',
+    )
+  assert 'Unable to parse reservation' in xpk_print.mock_calls[0].args[0]

xpk/core/cluster.py

@@ -729,78 +729,58 @@ def update_cluster_with_gcsfuse_driver_if_necessary(args) -> int:
   return 0
 
 
-def test_and_retry_credentials_with_dns_logic(args) -> int:
-  """Tests kubectl credentials and retries with default settings if a DNS error is found.
+def get_cluster_credentials(args) -> int:
+  """Run cluster configuration command to set the kubectl config.
 
   Args:
     args: user provided arguments for running the command.
 
   Returns:
-    0 if credentials are valid after retrying, 1 otherwise.
+    0 if successful and 1 otherwise.
   """
-
-  xpk_print('Testing credentials with kubectl...')
-  kubectl_command = 'kubectl get pods'
-  kubectl_return_code, kubectl_output = run_command_for_value(
-      kubectl_command, 'kubectl get pods'
-  )
-  if kubectl_return_code == 0:
-    xpk_print('Credentials test succeeded.')
-    return 0
-
-  dns_endpoint_error = (
-      'control_plane_endpoints_config.dns_endpoint_config.allow_external_traffic'
-      ' is disabled'
-  )
-  if dns_endpoint_error not in kubectl_output:
-    xpk_print(f'kubectl failed with an unhandled error: {kubectl_output}')
-    xpk_exit(kubectl_return_code)
-  xpk_print(
-      'Detected DNS endpoint-related error. Retrying without --dns-endpoint'
-      ' flag...'
-  )
-
   location = get_cluster_location(args.project, args.cluster, args.zone)
-  without_dns_command = (
-      'gcloud container clusters get-credentials'
-      f' {args.cluster} --location={location}'
-      f' --project={args.project} &&'
-      ' kubectl config view && kubectl config set-context --current'
-      ' --namespace=default'
-  )
-  return_code = run_command_with_updates(
-      without_dns_command, 'get-credentials to cluster', verbose=False
+
+  return_code = _get_credentials(
+      project=args.project,
+      cluster=args.cluster,
+      location=location,
+      dns_endpoint=True,
   )
   if return_code != 0:
-    xpk_print('Failed to get credentials even without --dns-endpoint. Exiting.')
-    xpk_exit(return_code)
-  return 0
-
+    return return_code
+
+  if not _are_credentials_valid():
+    xpk_print('Detected error. Retrying without --dns-endpoint flag...')
+    return_code = _get_credentials(
+        project=args.project,
+        cluster=args.cluster,
+        location=location,
+        dns_endpoint=False,
+    )
+    if return_code != 0:
+      return return_code
 
-def get_cluster_credentials(args) -> int:
-  """Run cluster configuration command to set the kubectl config.
+  xpk_print('Finished get-credentials and kubectl setup.')
+  return 0
 
-  Args:
-    args: user provided arguments for running the command.
 
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  location = get_cluster_location(args.project, args.cluster, args.zone)
+def _get_credentials(
+    project: str, cluster: str, location: str, dns_endpoint: bool
+) -> int:
+  dns_endpoint_arg = '--dns-endpoint' if dns_endpoint else ''
   command = (
       'gcloud container clusters get-credentials'
-      f' {args.cluster} --location={location} --dns-endpoint'
-      f' --project={args.project} && kubectl config view && kubectl config'
+      f' {cluster} --location={location} {dns_endpoint_arg}'
+      f' --project={project} && kubectl config view && kubectl config'
       ' set-context --current --namespace=default'
   )
-  task = f'get-credentials-dns-endpoint to cluster {args.cluster}'
-  return_code = run_command_with_updates_retry(command, task, verbose=False)
+  task = f'get-credentials-dns-endpoint to cluster {cluster}'
+  return run_command_with_updates(command, task, verbose=False)
 
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-    xpk_exit(return_code)
 
-  return_code = test_and_retry_credentials_with_dns_logic(args)
-  xpk_print('Finished get-credentials and kubectl setup.')
-
-  return return_code
+def _are_credentials_valid() -> bool:
+  kubectl_command = 'kubectl get pods'
+  kubectl_return_code = run_command_with_updates(
+      kubectl_command, 'Test kubectl credentials'
+  )
+  return kubectl_return_code == 0

xpk/core/cluster_test.py

@@ -0,0 +1,95 @@
+"""
+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 pytest
+from .testing.commands_tester import CommandsTester
+from .cluster import get_cluster_credentials
+from pytest_mock import MockerFixture
+
+
+@pytest.fixture(autouse=True)
+def commands_tester(mocker: MockerFixture) -> CommandsTester:
+  return CommandsTester(
+      mocker=mocker,
+      run_command_for_value_path="xpk.core.cluster.run_command_for_value",
+      run_command_with_updates_path="xpk.core.cluster.run_command_with_updates",
+  )
+
+
+@pytest.fixture(autouse=True)
+def mock_location(mocker: MockerFixture):
+  mocker.patch(
+      "xpk.core.cluster.get_cluster_location", return_value="us-central1"
+  )
+
+
+@pytest.fixture(autouse=True)
+def command_args(mocker: MockerFixture):
+  return mocker.Mock(cluster="cluster", project="project", zone="zone")
+
+
+def test_get_cluster_credentials_returns_1_when_retrieval_command_fails(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (1, ""), "gcloud container clusters get-credentials"
+  )
+  assert get_cluster_credentials(command_args) == 1
+
+
+def test_get_cluster_credentials_returns_0_when_retrieval_succeeds(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters get-credentials"
+  )
+  assert get_cluster_credentials(command_args) == 0
+
+
+def test_get_cluster_credentials_does_not_retry_with_dns_when_retrieval_succeeds(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters get-credentials --dns-endpoint"
+  )
+  commands_tester.set_result_for_command((0, ""), "kubectl get pods")
+  get_cluster_credentials(command_args)
+  non_dns_endpoint_commands = [
+      c
+      for c in commands_tester.get_matching_commands(
+          "gcloud container clusters get-credentials"
+      )
+      if "dns-endpoint" not in c
+  ]
+  assert len(non_dns_endpoint_commands) == 0
+
+
+def test_get_cluster_credentials_retries_without_dns_when_dns_retrieval_fails(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters get-credentials --dns-endpoint"
+  )
+  commands_tester.set_result_for_command((1, ""), "kubectl get pods")
+  get_cluster_credentials(command_args)
+  non_dns_endpoint_commands = [
+      c
+      for c in commands_tester.get_matching_commands(
+          "gcloud container clusters get-credentials"
+      )
+      if "dns-endpoint" not in c
+  ]
+  assert len(non_dns_endpoint_commands) == 1

xpk/core/commands.py

@@ -195,16 +195,13 @@ def run_command_with_updates(command, task, verbose=True) -> int:
         return_code = child.poll()
         if return_code is None:
           xpk_print(f'Waiting for `{task}`, for {i} seconds...', end='\r')
-          time.sleep(1)
-          i += 1
+          time.sleep(10)
+          i += 10
         else:
           xpk_print(f'Task: `{task}` terminated with code `{return_code}`')
           return return_code
   else:
-    xpk_print(
-        f'Task: `{task}` is implemented by `{command}`, hiding output unless'
-        ' there is an error.'
-    )
+    xpk_print(f'Task: `{task}` is implemented by `{command}`')
     try:
       subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT)
     except subprocess.CalledProcessError as e:
@@ -277,10 +274,7 @@ def run_command_for_value(
           return return_code, f'{out_str}\n{err_str}'
   else:
     if not quiet:
-      xpk_print(
-          f'Task: `{task}` is implemented by `{command}`, hiding output unless'
-          ' there is an error.'
-      )
+      xpk_print(f'Task: `{task}` is implemented by `{command}`')
     try:
       output = subprocess.check_output(
           command,

xpk/core/config.py

@@ -22,7 +22,7 @@ from ..utils import file
 from ..utils.console import xpk_print
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.14.4'
+__version__ = 'v0.15.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -30,6 +30,8 @@ CONFIGS_KEY = 'configs'
 CFG_BUCKET_KEY = 'cluster-state-gcs-bucket'
 CLUSTER_NAME_KEY = 'cluster-name'
 PROJECT_KEY = 'project-id'
+CLIENT_ID_KEY = 'client-id'
+SEND_TELEMETRY_KEY = 'send-telemetry'
 ZONE_KEY = 'zone'
 KJOB_BATCH_IMAGE = 'batch-image'
 KJOB_BATCH_WORKING_DIRECTORY = 'batch-working-directory'
@@ -45,6 +47,8 @@ DEFAULT_KEYS = [
     CFG_BUCKET_KEY,
     CLUSTER_NAME_KEY,
     PROJECT_KEY,
+    CLIENT_ID_KEY,
+    SEND_TELEMETRY_KEY,
     ZONE_KEY,
     GKE_ENDPOINT_KEY,
     DEPENDENCIES_KEY,
@@ -82,7 +86,7 @@ class XpkConfig:
     with open(self._config, encoding='utf-8', mode='w') as stream:
       yaml.dump(config_yaml, stream)
 
-  def set(self, key: str, value: str) -> None:
+  def set(self, key: str, value: str | None) -> None:
     if key not in self._allowed_keys:
       xpk_print(f'Key {key} is not an allowed xpk config key.')
       return
@@ -114,3 +118,6 @@ class XpkConfig:
       return None
     val: dict[str, str] = config_yaml[CONFIGS_KEY]
     return val
+
+
+xpk_config = XpkConfig()