xpk 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

xpk/core/cluster.py

@@ -62,8 +62,8 @@ def set_jobset_on_cluster(args) -> int:
     0 if successful and 1 otherwise.
   """
   command = (
-      'kubectl apply --server-side -f'
-      f' https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml'
+      'kubectl apply --server-side --force-conflicts'
+      f' -f https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml'
   )
   task = f'Install Jobset on {args.cluster}'
   return_code = run_command_with_updates_retry(command, task, args)
@@ -217,8 +217,8 @@ def get_cluster_nodes_info(args) -> list[dict]:
   )
   if err_code != 0:
     xpk_exit(err_code)
-  data = yaml.safe_load(val)
-  return data['items']  # pytype: disable=bad-return-type
+  data: dict[str, list[dict]] = yaml.safe_load(val)
+  return data['items']
 
 
 def count_nodes_on_cluster(args, system: SystemCharacteristics) -> int:
@@ -445,7 +445,7 @@ def setup_k8s_env(args) -> k8s_client.ApiClient:
     args.project_number = project_id_to_project_number(args.project)
 
   config.load_kube_config()
-  return k8s_client.ApiClient()  # pytype: disable=bad-return-type
+  return k8s_client.ApiClient()
 
 
 def get_gpu_type_from_cluster(args) -> str:
@@ -817,9 +817,11 @@ def update_cluster_with_clouddns_if_necessary(args) -> int:
     server_config_return_code, gke_server_config = get_gke_server_config(args)
     if server_config_return_code != 0:
       xpk_exit(server_config_return_code)
+    assert gke_server_config
+
     upgrade_master_return_code = upgrade_gke_control_plane_version(
         args,
-        gke_server_config.default_rapid_gke_version,  # pytype: disable=attribute-error
+        gke_server_config.default_rapid_gke_version,
     )
     if upgrade_master_return_code > 0:
       xpk_print("Updating GKE cluster's control plane upgrade failed!")
@@ -828,7 +830,7 @@ def update_cluster_with_clouddns_if_necessary(args) -> int:
     # Upgrade nodepools version after the master upgrade.
     node_pool_update_code = upgrade_gke_nodepools_version(
         args,
-        gke_server_config.default_rapid_gke_version,  # pytype: disable=attribute-error
+        gke_server_config.default_rapid_gke_version,
     )
     if node_pool_update_code > 0:
       xpk_print('Upgrading nodepools version failed!')

xpk/core/cluster_private.py

@@ -95,7 +95,11 @@ def add_current_machine_to_networks_if_needed(
           "Adding current machine's IP address to the authorized networks"
           ' failed!'
       )
-      return add_current_machine_to_networks_return_code, authorized_networks
+      return (
+          add_current_machine_to_networks_return_code,
+          False,
+          authorized_networks,
+      )
 
   return 0, is_current_machine_in_network, authorized_networks
 

xpk/core/commands.py

@@ -274,9 +274,9 @@ def run_command_for_value(
         else:
           if not quiet:
             xpk_print(f'Task: `{task}` terminated with code `{return_code}`')
-          out, err = child.communicate()
-          out, err = str(out, 'UTF-8'), str(err, 'UTF-8')
-          return return_code, f'{out}\n{err}'
+          out_bytes, err_bytes = child.communicate()
+          out_str, err_str = str(out_bytes, 'UTF-8'), str(err_bytes, 'UTF-8')
+          return return_code, f'{out_str}\n{err_str}'
   else:
     if not quiet:
       xpk_print(

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.11.0'
+__version__ = 'v0.12.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -71,13 +71,12 @@ class XpkConfig:
     dir_path = '/'.join(self._config.split('/')[:-1])
     file.ensure_directory_exists(dir_path)
 
-    config_yaml = {'version': 'v1', CONFIGS_KEY: {}}
     if not os.path.exists(self._config):
       return None
 
     with open(self._config, encoding='utf-8', mode='r') as stream:
       config_yaml: dict = yaml.load(stream)
-    return config_yaml
+      return config_yaml
 
   def _save_configs(self, config_yaml: dict) -> None:
     with open(self._config, encoding='utf-8', mode='w') as stream:
@@ -109,7 +108,7 @@ class XpkConfig:
 
   def get_all(
       self,
-  ) -> dict[str, dict[str, str] | str] | None:
+  ) -> dict[str, str] | None:
     config_yaml = self._open_configs()
     if config_yaml is None:
       return None

xpk/core/config_test.py

@@ -0,0 +1,71 @@
+"""
+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 xpk.core.config import XpkConfig, CFG_BUCKET_KEY, CLUSTER_NAME_KEY, PROJECT_KEY, ZONE_KEY
+
+import os
+import pytest
+
+config_tmp_path = '/tmp/config/config.yaml'
+
+
+@pytest.fixture(name='_')
+def _():
+  if os.path.exists(config_tmp_path):
+    os.remove(config_tmp_path)
+  yield
+  if os.path.exists(config_tmp_path):
+    os.remove(config_tmp_path)
+
+
+def test_config(_):
+  cfg = XpkConfig(config_tmp_path)
+  cfg.set('project-id', 'foo')
+  project_id = cfg.get('project-id')
+  assert project_id == 'foo'
+
+
+def test_config_get_all(_):
+  cfg = XpkConfig(config_tmp_path)
+  cfg.set(PROJECT_KEY, 'foo')
+  cfg.set(CLUSTER_NAME_KEY, 'bar')
+  cfg.set(ZONE_KEY, 'europe-west1-a')
+  cfg.set(CFG_BUCKET_KEY, 'cfg-bucket')
+
+  cfg_all = cfg.get_all()
+  assert cfg_all[PROJECT_KEY] == 'foo'
+  assert cfg_all[CLUSTER_NAME_KEY] == 'bar'
+  assert cfg_all[ZONE_KEY] == 'europe-west1-a'
+  assert cfg_all[CFG_BUCKET_KEY] == 'cfg-bucket'
+
+
+def test_config_get_empty(_):
+  cfg = XpkConfig(config_tmp_path)
+  val = cfg.get(PROJECT_KEY)
+  assert val is None
+
+
+def test_config_get_all_empty(_):
+  cfg = XpkConfig(config_tmp_path)
+  val = cfg.get_all()
+  assert not val
+
+
+def test_config_set_incorrect(_):
+  cfg = XpkConfig(config_tmp_path)
+  cfg.set('foo', 'bar')
+  cfg_all = cfg.get_all()
+  assert not cfg_all

xpk/core/docker_manager.py

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

xpk/core/docker_resources.py

@@ -72,7 +72,7 @@ def get_env_container(args, system: SystemCharacteristics) -> str:
   if system.accelerator_type == AcceleratorType['CPU']:
     return get_cpu_env(args, system)
 
-  return format_env_dict(args.env, system)  # pytype: disable=bad-return-type
+  return format_env_dict(args.env, system)
 
 
 def get_gpu_env(args, system) -> str:

xpk/core/filestore.py

@@ -93,11 +93,12 @@ class FilestoreClient:
 
     for instance in instancesZonal:
       if instance.name == fullname_zonal:
-        return instance  # pytype: disable=bad-return-type
+        return instance
 
     for instance in instancesRegional:
       if instance.name == fullname_regional:
-        return instance  # pytype: disable=bad-return-type
+        return instance
+    return None
 
   def check_instance_exists(self) -> bool:
     """Check if Filestore instance exists"""
@@ -111,6 +112,7 @@ class FilestoreClient:
   def get_instance_location(self) -> str:
     """Get Filestore instance's location"""
     self.load_instance()
+    assert self.instance
     return str(self.instance.name.split("/")[3])
 
   def create_instance(
@@ -192,6 +194,7 @@ class FilestoreClient:
 
   def create_sc(self, name: str, network: str) -> dict:
     """Create a yaml representing filestore StorageClass."""
+    assert self.instance
     data = templates.load(FS_SC_PATH)
     data["metadata"]["name"] = get_storage_class_name(name)
     data["parameters"]["tier"] = self.instance.tier.name
@@ -202,6 +205,7 @@ class FilestoreClient:
 
   def create_pv(self, name: str, vol: str, access_mode: str) -> dict:
     """Create a yaml representing filestore PersistentVolume."""
+    assert self.instance
     data = templates.load(FS_PV_PATH)
     data["metadata"]["name"] = get_pv_name(name)
     data["spec"]["storageClassName"] = get_storage_class_name(name)
@@ -219,6 +223,7 @@ class FilestoreClient:
 
   def create_pvc(self, name: str, access_mode: str) -> dict:
     """Create a yaml representing filestore PersistentVolumeClaim."""
+    assert self.instance
     data = templates.load(FS_PVC_PATH)
     data["metadata"]["name"] = get_pvc_name(name)
     data["spec"]["accessModes"] = [access_mode]

xpk/core/gcloud_context.py

@@ -75,7 +75,7 @@ def add_zone_and_project(args):
   xpk_print(f'Working on {args.project} and {args.zone}')
 
 
-def zone_to_region(zone) -> str:
+def zone_to_region(zone: str) -> str:
   """Helper function converts zone name to region name.
 
   Args:
@@ -85,7 +85,7 @@ def zone_to_region(zone) -> str:
      The region name.
   """
   zone_terms = zone.split('-')
-  return zone_terms[0] + '-' + zone_terms[1]  # pytype: disable=bad-return-type
+  return zone_terms[0] + '-' + zone_terms[1]
 
 
 @dataclass

xpk/core/kjob.py

@@ -277,7 +277,8 @@ def decorate_job_template_with_gpu(yml_string: str, gpu_type: str) -> str:
     job_spec = rdma_decorator.decorate_kjob_template(job_spec)
   job_template_dict = yaml.safe_load(yml_string)
   job_template_dict["template"] = job_spec
-  return yaml.dump(job_template_dict, sort_keys=False)
+  yaml_result: str = yaml.dump(job_template_dict, sort_keys=False)
+  return yaml_result
 
 
 def create_job_template_instance(

xpk/core/kueue.py

@@ -43,7 +43,7 @@ from .system_characteristics import (
 KUEUE_VERSION = 'v0.12.2'
 CLUSTER_QUEUE_NAME = 'cluster-queue'
 LOCAL_QUEUE_NAME = 'multislice-queue'
-WAIT_FOR_KUEUE_TIMEOUT = '5m'
+WAIT_FOR_KUEUE_TIMEOUT = '10m'
 MEMORY_SIZE_PER_VM = 1.2
 MIN_MEMORY_LIMIT_SIZE = 4096
 
@@ -89,6 +89,10 @@ metadata:
   name: dws-config
 spec:
   provisioningClassName: queued-provisioning.gke.io
+  podSetUpdates:
+    nodeSelector:
+    - key: autoscaling.gke.io/provisioning-request
+      valueFromProvisioningClassDetail: ResizeRequestName
   managedResources:
   - {managed_resource}
 ---
@@ -320,7 +324,7 @@ def delete_multikueueclusters_definitions(args) -> int:
   return return_code
 
 
-def get_kueue_version(args) -> (int, str):
+def get_kueue_version(args) -> tuple[int, str]:
   command = 'kubectl kueue version'
   task = 'Get kueue version on server'
   return_code, val = run_command_for_value(command, task, args)

xpk/core/nap.py

@@ -37,6 +37,7 @@ from .resources import (
 )
 from .scheduling import get_total_chips_requested_from_args
 from .system_characteristics import AcceleratorType, SystemCharacteristics
+from typing import cast
 
 AUTOPROVISIONING_CONFIG_FILE = """
 management:
@@ -269,9 +270,6 @@ def is_autoprovisioning_enabled(
     bool is true if autoprovisioning is enabled, false otherwise.
     int of 0 if successful and 1 otherwise.
   """
-  # Currently autoprovisioning is not enabled for Pathways workloads. b/360898087
-  if args.use_pathways:
-    return False, 0
 
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
   cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
@@ -339,11 +337,13 @@ def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
       )
       return node_selector_args, 1
 
-    return_code, capacity_type_str = get_value_from_map(
+    return_code, optional_capacity_type_str = get_value_from_map(
         CAPACITY_TYPE_CONFIG_KEY, cluster_config_map
     )
     if return_code != 0:
       return node_selector_args, return_code
+    # return_code==0 implies capacity_type is defined
+    capacity_type_str = cast(str, optional_capacity_type_str)
 
     if capacity_type_str == CapacityType.RESERVATION.name:
       return_code, args.reservation = get_value_from_map(

xpk/core/nodepool_test.py

@@ -0,0 +1,82 @@
+"""
+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 xpk.core.nodepool import get_desired_node_pool_names
+
+CLUSTER_NAME = "running-cucumber"
+
+
+def node_pool_name(number: int) -> str:
+  return f"{CLUSTER_NAME}-np-{number}"
+
+
+def test_compute_desired_node_pool_names_with_desired_larger_than_existing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=2,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(1)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_desired_smaller_than_existing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0), node_pool_name(1)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=1,
+  )
+
+  expected_result = [node_pool_name(0)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_consecutive_numbers_missing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0), node_pool_name(3)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=3,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(1), node_pool_name(3)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_consecutive_numbers_missing_and_desired_equal_to_existing():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[node_pool_name(0), node_pool_name(3)],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=2,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(3)]
+  assert set(result) == set(expected_result)
+
+
+def test_compute_desired_node_pool_names_with_unknown_node_pools():
+  result = get_desired_node_pool_names(
+      existing_node_pool_names=[
+          "unknown-node-pool",
+          node_pool_name(0),
+          node_pool_name(3),
+      ],
+      cluster_name=CLUSTER_NAME,
+      desired_node_pool_count=2,
+  )
+
+  expected_result = [node_pool_name(0), node_pool_name(3)]
+  assert set(result) == set(expected_result)

xpk/core/resources.py

@@ -108,13 +108,7 @@ def create_cluster_configmaps(
   device_type = system.device_type
   if system.accelerator_type == AcceleratorType['GPU']:
     resources_data = f'{device_type}: "{int(args.num_nodes)}"'
-  elif (
-      not args.enable_pathways
-      and args.enable_autoprovisioning
-      and autoprovisioning_config
-  ):
-    # Currently autoprovisioning is not supported with Pathways.
-    # Auto provisioning will have variable topologies for a gke accelerator type.
+  elif args.enable_autoprovisioning and autoprovisioning_config:
     resources_data = (
         f'{system.gke_accelerator}: {AUTOPROVISIONING_CONFIG_VALUE}'
     )

xpk/core/storage.py

@@ -17,7 +17,7 @@ limitations under the License.
 import os
 from argparse import Namespace
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
 
 import ruamel.yaml
 from google.cloud import storage as gcp_storage
@@ -95,17 +95,17 @@ class Storage:
     Args:
         data: A dictionary containing the Storage resource definition.
     """
-    metadata: k8s_client.V1ObjectMeta = data.get("metadata", {})
+    metadata = data.get("metadata", {})
     self.name = metadata.get("name")
     spec = data.get("spec", {})
-    self.type: str = spec.get("type")
-    self.auto_mount: bool = spec.get("auto_mount")
-    self.mount_point: bool = spec.get("mount_point")
-    self.readonly: bool = spec.get("readonly")
-    self.manifest: str = spec.get("manifest")
-    self.pvc: str = spec.get("pvc")
-    self.pv: str = spec.get("pv")
-    self.bucket: str = self._get_bucket()
+    self.type = spec.get("type")
+    self.auto_mount = spec.get("auto_mount")
+    self.mount_point = spec.get("mount_point")
+    self.readonly = spec.get("readonly")
+    self.manifest = spec.get("manifest")
+    self.pvc = spec.get("pvc")
+    self.pv = spec.get("pv")
+    self.bucket = self._get_bucket()
 
   def fields_as_list(self) -> list[str]:
     """
@@ -117,9 +117,9 @@ class Storage:
     return [
         self.name,
         self.type,
-        self.auto_mount,
+        str(self.auto_mount),
         self.mount_point,
-        self.readonly,
+        str(self.readonly),
         self.manifest,
     ]
 
@@ -133,7 +133,7 @@ class Storage:
     client = k8s_client.CoreV1Api()
     try:
       pv: V1PersistentVolume = client.read_persistent_volume(self.pv)
-      return pv.spec.csi.volume_handle
+      return cast(str, pv.spec.csi.volume_handle)
     except ApiException as e:
       xpk_print(
           f"Exception when calling CoreV1Api->read_persistent_volume: {e}"
@@ -150,7 +150,7 @@ class Storage:
     client = k8s_client.CoreV1Api()
     try:
       pv: V1PersistentVolume = client.read_persistent_volume(self.pv)
-      return pv.spec.mount_options
+      return cast(list[str], pv.spec.mount_options)
     except ApiException as e:
       xpk_print(
           f"Exception when calling CoreV1Api->read_persistent_volume: {e}"

xpk/core/system_characteristics.py

@@ -55,7 +55,7 @@ class SystemCharacteristics:
   gke_accelerator: str
   gce_machine_type: str
   chips_per_vm: int
-  accelerator_type: AcceleratorType  # type: ignore
+  accelerator_type: int  # TODO: use enums
   device_type: str
 
 

xpk/core/workload.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import re
 from ..utils.console import xpk_exit, xpk_print
 from .commands import run_command_for_value
 from .gcloud_context import zone_to_region
@@ -240,3 +241,13 @@ def wait_for_job_completion(args) -> int:
     xpk_print('Your workload did not complete successfully')
     return 125
   return 0
+
+
+GCP_NAME_FILTER_VALUE_REGEX = re.compile(r'[a-z0-9\-]+')
+"""Defines correct name prefix value (contains only letters, numbers and dashes) that can be used in GCP filter chips."""
+
+
+def get_jobsets_list_gcp_link(project: str) -> str:
+  """Returns a link to Cloud Console JobSets list"""
+
+  return f'https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project={project}'

xpk/core/workload_decorators/rdma_decorator.py

@@ -18,7 +18,7 @@ import yaml
 from ...utils.yaml import literal_string
 
 
-def decorate_kjob_template(job_manifest) -> str:
+def decorate_kjob_template(job_manifest: dict) -> dict:
   spec = (
       job_manifest.setdefault('spec', {})
       .setdefault('template', {})
@@ -64,7 +64,8 @@ def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
     add_tolerations(job_manifest)
     update_gpu_containers(job_manifest)
 
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_str: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_str
 
 
 def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:

xpk/core/workload_decorators/storage_decorator.py

@@ -36,7 +36,8 @@ def decorate_jobset(jobset_manifest_str, storages) -> str:
     job_manifest = job['template']
     add_annotations(job_manifest, storages)
     add_volumes(job_manifest, storage_volumes)
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_result: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_result
 
 
 def add_annotations(job_manifest, storages):

xpk/core/workload_decorators/tcpx_decorator.py

@@ -55,7 +55,8 @@ def decorate_jobset(jobset_manifest_str: str) -> str:
   for job in manifest['spec']['replicatedJobs']:
     job_manifest = job['template']
     job_manifest = decorate_job(job_manifest)
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_str: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_str
 
 
 def get_interfaces_annotation() -> dict:
@@ -131,6 +132,7 @@ def add_volumes(job_manifest: dict):
   })
   volumes.append({'name': 'sys', 'hostPath': {'path': '/sys'}})
   volumes.append({'name': 'proc-sys', 'hostPath': {'path': '/proc/sys'}})
+  volumes.append({'name': 'tcpx-socket', 'hostPath': {'path': '/run/tcpx'}})
   volumes.append(
       {'name': 'dshm', 'emptyDir': {'medium': 'Memory', 'sizeLimit': '128Gi'}}
   )
@@ -168,7 +170,7 @@ def add_tcpx_daemon_container(job_manifest):
   spec['initContainers'].append(tcpxo_daemon_container)
 
 
-def update_gpu_containers(job_manifest):
+def update_gpu_containers(job_manifest) -> None:
   for container in job_manifest['spec']['template']['spec']['containers']:
     if 'nvidia.com/gpu' in container.get('resources', {}).get('limits', {}):
       env: list = container.setdefault('env', [])