xpk 0.10.1__py3-none-any.whl → 0.12.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}'
@@ -232,9 +234,9 @@ def get_capacity_node_selectors_from_capacity_type(
     case CapacityType.ON_DEMAND.name:
       node_selector = ''
     case CapacityType.FLEX_START.name:
-      node_selector = 'cloud.google.com/gke-queued="true"'
+      node_selector = 'cloud.google.com/gke-queued: "true"'
     case CapacityType.SPOT.name:
-      node_selector = 'cloud.google.com/gke-spot="true"'
+      node_selector = 'cloud.google.com/gke-spot: "true"'
     case CapacityType.RESERVATION.name:
       node_selector = f'cloud.google.com/reservation-name: {args.reservation}'
     case _:

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.10.1'
+__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/jobset.py

@@ -81,7 +81,7 @@ spec:
           limits:
             memory: {memory_limit_size}
           requests:
-            cpu: 500m
+            cpu: 1000m
             memory: 128Mi
         securityContext:
           allowPrivilegeEscalation: false

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}
 ---
@@ -244,14 +248,16 @@ spec:
           periodSeconds: 10
         resources:
           limits:
-            cpu: 500m
+            cpu: 1000m
             memory: {memory_limit_size}
           requests:
-            cpu: 500m
+            cpu: 1000m
             memory: 512Mi
         securityContext:
           allowPrivilegeEscalation: false
         volumeMounts:
+        - mountPath: /visibility
+          name: visibility
         - mountPath: /tmp/k8s-webhook-server/serving-certs
           name: cert
           readOnly: true
@@ -263,6 +269,8 @@ spec:
       serviceAccountName: kueue-controller-manager
       terminationGracePeriodSeconds: 10
       volumes:
+      - name: visibility
+        emptyDir: {{}}
       - name: cert
         secret:
           defaultMode: 420
@@ -316,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,11 +37,14 @@ 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:
   autoRepair: true
   autoUpgrade: true
+scopes:
+  - "https://www.googleapis.com/auth/devstorage.read_write"
 autoprovisioningLocations:
   {zones}
 {resource_limits}
@@ -106,6 +109,18 @@ def enable_autoprovisioning_on_cluster(
     xpk_print(f'{task} request returned ERROR {return_code}')
     return autoprovisioning_config, return_code
 
+  command = (
+      'gcloud container clusters update'
+      f' {args.cluster} --project={args.project}'
+      f' --region={zone_to_region(args.zone)}'
+      ' --autoscaling-profile=optimize-utilization'
+  )
+  task = 'Update cluster with autoscaling-profile'
+  return_code = run_command_with_updates(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} request returned ERROR {return_code}')
+    return autoprovisioning_config, return_code
+
   # Update created accelerator node pools to support autoprovisioning.
   existing_node_pool_names, return_code = get_all_nodepools_programmatic(args)
   if return_code != 0:
@@ -171,11 +186,11 @@ def create_autoprovisioning_config(
   # is not controlled by NAP.
   cpu_limits = """
   minimum: 1
-  maximum: 10000
+  maximum: 1000000
   """
   memory_limits = """
   minimum: 1
-  maximum: 10000
+  maximum: 10000000
   """
 
   # By default, the maximum chips is set to be the current number of resources used
@@ -255,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)
@@ -325,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(