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

xpk/core/blueprint/blueprint_test.py

@@ -0,0 +1,218 @@
+"""
+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.
+"""
+
+import os
+import shutil
+
+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
+
+yaml = ruamel.yaml.YAML()
+
+yaml.register_class(Blueprint)
+
+a3_yaml_test_path = "src/xpk/core/blueprint/testing/data/a3_mega.yaml"
+a3_spot_yaml_test_path = "src/xpk/core/blueprint/testing/data/a3_mega_spot.yaml"
+a3_ultra_yaml_test_path = "src/xpk/core/blueprint/testing/data/a3_ultra.yaml"
+a4_yaml_test_path = "src/xpk/core/blueprint/testing/data/a4.yaml"
+config_map_filename = "config-map.yaml.tftpl"
+kueue_conf_filename = "kueue-xpk-configuration.yaml.tftpl"
+tmp_test_dir = "/tmp/xpk_test"
+
+
+def prepare_test():
+  if os.path.exists(tmp_test_dir):
+    shutil.rmtree(tmp_test_dir)
+  os.mkdir(tmp_test_dir)
+
+
+def test_generate_a3_mega_blueprint():
+  prepare_test()
+  blueprint_name = "xpk-gke-a3-megagpu"
+  bp_generator = BlueprintGenerator(tmp_test_dir)
+  bp = bp_generator.generate_a3_mega_blueprint(
+      project_id="foo",
+      cluster_name="bar",
+      blueprint_name=blueprint_name,
+      prefix="prefix",
+      region="us-central1",
+      zone="us-central1-c",
+      auth_cidr="10.0.0.0/32",
+      reservation_placement_policy={
+          "type": "COMPACT",
+          "name": "test-reservation-placement",
+      },
+      reservation="test-reservation",
+      capacity_type=CapacityType.RESERVATION,
+      system_node_pool_min_node_count=5,
+  )
+
+  assert bp.blueprint_file.endswith("/prefix/xpk-gke-a3-megagpu.yaml")
+
+  with open(a3_yaml_test_path, encoding="utf-8") as stream:
+    ctk_yaml = yaml.load(stream)
+    with open(bp.blueprint_file, encoding="utf-8") as generated_blueprint:
+      ctk_test = yaml.load(generated_blueprint)
+      assert ctk_yaml.blueprint_name == ctk_test.blueprint_name
+      assert ctk_test.terraform_backend_defaults is None
+      assert ctk_yaml.toolkit_modules_url == ctk_test.toolkit_modules_url
+      assert (
+          ctk_yaml.toolkit_modules_version == ctk_test.toolkit_modules_version
+      )
+      assert ctk_yaml.vars == ctk_test.vars
+      assert ctk_test.deployment_groups == ctk_yaml.deployment_groups
+      assert os.path.exists(
+          os.path.join(
+              tmp_test_dir, "prefix", blueprint_name, config_map_filename
+          )
+      )
+      assert os.path.exists(
+          os.path.join(
+              tmp_test_dir, "prefix", blueprint_name, kueue_conf_filename
+          )
+      )
+
+  shutil.rmtree(tmp_test_dir)
+
+
+def test_generate_a3_mega_spot_blueprint():
+  prepare_test()
+  blueprint_name = "xpk-gke-a3-megagpu"
+  bp_generator = BlueprintGenerator(tmp_test_dir)
+  bp = bp_generator.generate_a3_mega_blueprint(
+      project_id="foo",
+      cluster_name="bar",
+      blueprint_name=blueprint_name,
+      prefix="prefix",
+      region="us-central1",
+      zone="us-central1-c",
+      auth_cidr="10.0.0.0/32",
+      capacity_type=CapacityType.SPOT,
+      system_node_pool_min_node_count=5,
+  )
+
+  assert bp.blueprint_file.endswith("/prefix/xpk-gke-a3-megagpu.yaml")
+
+  with open(a3_spot_yaml_test_path, encoding="utf-8") as stream:
+    ctk_yaml = yaml.load(stream)
+    with open(bp.blueprint_file, encoding="utf-8") as generated_blueprint:
+      ctk_test = yaml.load(generated_blueprint)
+      assert ctk_yaml.blueprint_name == ctk_test.blueprint_name
+      assert ctk_test.terraform_backend_defaults is None
+      assert ctk_yaml.toolkit_modules_url == ctk_test.toolkit_modules_url
+      assert (
+          ctk_yaml.toolkit_modules_version == ctk_test.toolkit_modules_version
+      )
+      assert ctk_yaml.vars == ctk_test.vars
+      assert ctk_test.deployment_groups == ctk_yaml.deployment_groups
+
+  shutil.rmtree(tmp_test_dir)
+
+
+def test_generate_a3_ultra_blueprint():
+  prepare_test()
+  blueprint_name = "xpk-gke-a3-ultra"
+  bp_generator = BlueprintGenerator(tmp_test_dir)
+  bp = bp_generator.generate_a3_ultra_blueprint(
+      project_id="foo",
+      cluster_name="gke-a3-ultra",
+      blueprint_name=blueprint_name,
+      region="us-central1",
+      zone="us-central1-c",
+      auth_cidr="10.0.0.0/32",
+      reservation="test-reservation",
+      system_node_pool_machine_type="e2-standard-16",
+      capacity_type=CapacityType.RESERVATION,
+      gcs_bucket="test-bucket",
+      prefix="testdir",
+  )
+  with open(a3_ultra_yaml_test_path, encoding="utf-8") as stream:
+    ctk_yaml = yaml.load(stream)
+    with open(bp.blueprint_file, encoding="utf-8") as generated_blueprint:
+      ctk_test = yaml.load(generated_blueprint)
+      assert ctk_yaml.blueprint_name == ctk_test.blueprint_name
+      assert (
+          ctk_yaml.terraform_backend_defaults
+          == ctk_test.terraform_backend_defaults
+      )
+      assert ctk_yaml.toolkit_modules_url == ctk_test.toolkit_modules_url
+      assert (
+          ctk_yaml.toolkit_modules_version == ctk_test.toolkit_modules_version
+      )
+      assert ctk_test.deployment_groups == ctk_yaml.deployment_groups
+      assert os.path.exists(
+          os.path.join(
+              tmp_test_dir, "testdir", blueprint_name, "mlgru-disable.yaml"
+          )
+      )
+      assert os.path.exists(
+          os.path.join(
+              tmp_test_dir, "testdir", blueprint_name, "nccl-installer.yaml"
+          )
+      )
+
+  shutil.rmtree(tmp_test_dir)
+
+
+def test_generate_a4_blueprint():
+  prepare_test()
+  blueprint_name = "xpk-gke-a4"
+  bp_generator = BlueprintGenerator(tmp_test_dir)
+  bp = bp_generator.generate_a4_blueprint(
+      project_id="foo",
+      cluster_name="gke-a4",
+      blueprint_name=blueprint_name,
+      region="us-central1",
+      zone="us-central1-c",
+      auth_cidr="10.0.0.0/32",
+      reservation="test-reservation",
+      system_node_pool_machine_type="e2-standard-16",
+      capacity_type=CapacityType.RESERVATION,
+      gcs_bucket="test-bucket",
+      prefix="testdir",
+  )
+  with open(a4_yaml_test_path, encoding="utf-8") as stream:
+    ctk_yaml = yaml.load(stream)
+    with open(bp.blueprint_file, encoding="utf-8") as generated_blueprint:
+      ctk_test = yaml.load(generated_blueprint)
+      assert ctk_yaml.blueprint_name == ctk_test.blueprint_name
+      assert (
+          ctk_yaml.terraform_backend_defaults
+          == ctk_test.terraform_backend_defaults
+      )
+      assert ctk_yaml.toolkit_modules_url == ctk_test.toolkit_modules_url
+      assert (
+          ctk_yaml.toolkit_modules_version == ctk_test.toolkit_modules_version
+      )
+      assert ctk_test.deployment_groups == ctk_yaml.deployment_groups
+      assert os.path.exists(
+          os.path.join(
+              tmp_test_dir, "testdir", blueprint_name, "storage_crd.yaml"
+          )
+      )
+      assert os.path.exists(
+          os.path.join(
+              tmp_test_dir,
+              "testdir",
+              blueprint_name,
+              "nccl-rdma-installer-a4.yaml",
+          )
+      )
+
+  shutil.rmtree(tmp_test_dir)

xpk/core/capacity.py

@@ -195,10 +195,12 @@ def get_capacity_arguments_from_capacity_type(
       capacity_args = '--spot'
     case CapacityType.FLEX_START:
       capacity_args = (
-          ' --flex-start --enable-queued-provisioning --enable-autoscaling'
+          ' --flex-start --enable-autoscaling'
           ' --location-policy=ANY --reservation-affinity=none'
           f' --no-enable-autorepair --max-nodes={max_nodes}'
       )
+      if args.num_slices <= 1:
+        capacity_args += ' --enable-queued-provisioning'
     case CapacityType.RESERVATION:
       capacity_args = (
           f'--reservation-affinity=specific --reservation={args.reservation}'

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:
@@ -442,10 +442,14 @@ def setup_k8s_env(args) -> k8s_client.ApiClient:
   if not getattr(args, 'kind_cluster', False):
     add_zone_and_project(args)
     get_cluster_credentials(args)
-    args.project_number = project_id_to_project_number(args.project)
+    args.project_number = (
+        project_id_to_project_number(args.project)
+        if not args.dry_run
+        else abs(hash(args.project) % (10**12))  # 12 digit hash
+    )
 
   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 +821,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 +834,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

@@ -19,6 +19,7 @@ from ..utils.network import (
     add_current_machine_to_networks,
     is_current_machine_in_any_network,
 )
+from ..utils.execution_context import is_dry_run
 from ..utils.objects import is_text_true
 from .commands import run_command_for_value, run_command_with_updates
 from .gcloud_context import zone_to_region
@@ -37,7 +38,7 @@ def authorize_private_cluster_access_if_necessary(args) -> int:
     if not args.private and args.authorized_networks is None:
       xpk_print('Cluster is public and no need to authorize networks.')
       return 0
-    else:
+    elif not is_dry_run():
       xpk_print(
           'Cannot convert an existing public cluster to private. The arguments'
           ' --private and --authorized-networks are not acceptable for public'
@@ -95,7 +96,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
 
@@ -160,6 +165,7 @@ def get_cluster_authorized_networks(args) -> list[str]:
       command,
       'Fetching the list of authorized network from cluster describe.',
       args,
+      dry_run_return_val='127.0.0.1/32',
   )
 
   if return_code != 0:

xpk/core/commands.py

@@ -78,14 +78,13 @@ def run_command_batch(commands, jobname, per_command_name, output_logs):
     The max return code and a list of all the return codes.
   """
 
+  files = [open(f, 'w', encoding='utf-8') for f in output_logs]
   children = []
   start_time = datetime.datetime.now()
-  for i, command in enumerate(commands):
+  for command, file in zip(commands, files):
     children.append(
         # subprocess managed by list pylint: disable=consider-using-with
-        subprocess.Popen(
-            command, stdout=output_logs[i], stderr=output_logs[i], shell=True
-        )
+        subprocess.Popen(command, stdout=file, stderr=file, shell=True)
     )
 
   while True:
@@ -99,7 +98,7 @@ def run_command_batch(commands, jobname, per_command_name, output_logs):
       slow_worker_text = per_command_name[slow_worker_index]
       slow_str = (
           f', task {slow_worker_text} still working, logfile'
-          f' {output_logs[slow_worker_index].name}'
+          f' {output_logs[slow_worker_index]}'
       )
     else:
       slow_str = ''
@@ -116,7 +115,7 @@ def run_command_batch(commands, jobname, per_command_name, output_logs):
       )
       xpk_print(
           f'Failure is {per_command_name[failing_index]}'
-          f' and logfile {output_logs[failing_index].name}'
+          f' and logfile {output_logs[failing_index]}'
       )
       for child in children:
         child.terminate()
@@ -126,6 +125,10 @@ def run_command_batch(commands, jobname, per_command_name, output_logs):
       break
 
     time.sleep(1)
+
+  for file in files:
+    file.close()
+
   return max_returncode, returncodes
 
 
@@ -274,9 +277,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(
@@ -351,6 +354,6 @@ def run_command_with_full_controls(
 
 def run_kubectl_apply(yml_string: str, task: str, args: Namespace) -> int:
   tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp.file.name)}'
+  command = f'kubectl apply -f {str(tmp)}'
   err_code = run_command_with_updates(command, task, args)
   return err_code

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.13.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_image.py

@@ -21,6 +21,7 @@ import string
 
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
+from ..utils.execution_context import is_dry_run
 from .commands import run_command_with_updates
 
 DEFAULT_DOCKER_IMAGE = 'python:3.10'
@@ -75,7 +76,9 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   """
 
   # Pick a name for the docker image.
-  docker_image_prefix = os.getenv('USER', 'unknown')
+  docker_image_prefix = (
+      'dry-run' if is_dry_run() else os.getenv('USER', 'unknown')
+  )
   docker_name = f'{docker_image_prefix}-runner'
 
   script_dir_dockerfile = """FROM {base_docker_image}
@@ -94,7 +97,7 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   )
   tmp = write_tmp_file(docker_file)
   docker_build_command = (
-      f'docker buildx build --platform={PLATFORM} -f {str(tmp.file.name)} -t'
+      f'docker buildx build --platform={PLATFORM} -f {str(tmp)} -t'
       f' {docker_name} {args.script_dir}'
   )
   xpk_print(f'Building {args.script_dir} into docker image.')
@@ -114,10 +117,16 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
 
   # Pick a randomly generated `tag_length` character docker tag.
   tag_length = 4
-  tag_random_prefix = ''.join(
-      random.choices(string.ascii_lowercase, k=tag_length)
+  tag_random_prefix = (
+      'prefix'
+      if is_dry_run()
+      else ''.join(random.choices(string.ascii_lowercase, k=tag_length))
+  )
+  tag_datetime = (
+      'current'
+      if is_dry_run()
+      else datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
   )
-  tag_datetime = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
   tag_name = f'{tag_random_prefix}-{tag_datetime}'
   cloud_docker_image = f'gcr.io/{args.project}/{docker_name}:{tag_name}'
   xpk_print(f'Adding Docker Image: {cloud_docker_image} to {args.project}')

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

@@ -20,6 +20,7 @@ from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .cluster import setup_k8s_env
 from .storage import GCS_FUSE_TYPE, GCP_FILESTORE_TYPE, PARALLELSTORE_TYPE, GCE_PD_TYPE, LUSTRE_TYPE, Storage, get_storages_to_mount
 from .system_characteristics import AcceleratorType, SystemCharacteristics
+from ..utils.execution_context import is_dry_run
 
 
 def get_main_container_resources(
@@ -72,7 +73,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:
@@ -272,8 +273,10 @@ def get_volumes(args, system: SystemCharacteristics) -> str:
               - name: shared-data
               """
 
-  storages: list[Storage] = get_storages_to_mount(
-      setup_k8s_env(args), args.storage
+  storages: list[Storage] = (
+      []
+      if is_dry_run()
+      else get_storages_to_mount(setup_k8s_env(args), args.storage)
   )
   for storage in storages:
     if storage.type in {
@@ -325,8 +328,10 @@ def get_volume_mounts(args, system: SystemCharacteristics) -> str:
   elif system.accelerator_type == AcceleratorType['GPU']:
     volume_mount_yaml = ''
 
-  storages: list[Storage] = get_storages_to_mount(
-      setup_k8s_env(args), args.storage
+  storages: list[Storage] = (
+      []
+      if is_dry_run()
+      else get_storages_to_mount(setup_k8s_env(args), args.storage)
   )
   for storage in storages:
     if storage.type in {

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/jobset.py

@@ -134,7 +134,7 @@ def update_jobset_resources_if_necessary(args):
       memory_limit_size=new_memory_limit,
   )
   tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp.file.name)}'
+  command = f'kubectl apply -f {str(tmp)}'
 
   task = 'Updating jobset Controller Manager resources'
   return_code = run_command_with_updates_retry(command, task, args)