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

xpk/core/blueprint/testing/data/a3_ultra.yaml

@@ -0,0 +1,173 @@
+# 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
+#
+#     http://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.
+!Blueprint
+blueprint_name: xpk-gke-a3-ultra
+toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
+toolkit_modules_version: v1.62.2
+
+vars:
+  labels: {gke_product_type: xpk}
+
+
+terraform_backend_defaults:
+  type: gcs
+  configuration:
+    bucket: test-bucket
+    prefix: xpk_terraform_state/testdir/gke-a3-ultra/
+
+deployment_groups:
+- !DeploymentGroup
+  group: primary
+  modules:
+  - !DeploymentModule
+    id: gke-a3-ultra-net-0
+    source: modules/network/vpc
+    settings:
+      network_name: gke-a3-ultra-net-0
+      subnetworks:
+      - subnet_name: gke-a3-ultra-sub-0
+        subnet_region: us-central1
+        subnet_ip: 192.168.0.0/18
+      secondary_ranges_list:
+        - subnetwork_name : gke-a3-ultra-sub-0
+          ranges:
+          - range_name: pods
+            ip_cidr_range: 10.4.0.0/14
+          - range_name: services
+            ip_cidr_range: 10.0.32.0/20
+      firewall_rules:
+      - name: gke-a3-ultra-internal-0
+        ranges: [192.168.0.0/16]
+        allow:
+        - protocol: tcp
+          ports: ["0-65535"]
+        - protocol: udp
+          ports: ["0-65535"]
+        - protocol: icmp
+
+  - !DeploymentModule
+    id: gke-a3-ultra-net-1
+    source: modules/network/vpc
+    settings:
+      network_name: gke-a3-ultra-net-1
+      mtu: 8896
+      subnetworks:
+      - subnet_name: gke-a3-ultra-sub-1
+        subnet_region: us-central1
+        subnet_ip: 192.168.64.0/18
+      firewall_rules:
+      - name: gke-a3-ultra-internal-1
+        ranges: [192.168.0.0/16]
+        allow:
+        - protocol: tcp
+          ports: ["0-65535"]
+        - protocol: udp
+          ports: ["0-65535"]
+        - protocol: icmp
+
+  - !DeploymentModule
+    id: gke-a3-ultra-rdma-net
+    source: modules/network/gpu-rdma-vpc
+    settings:
+      network_name: gke-a3-ultra-rdma-net
+      mtu: 8896
+      network_profile: https://www.googleapis.com/compute/beta/projects/foo/global/networkProfiles/us-central1-c-vpc-roce
+      network_routing_mode: REGIONAL
+      subnetworks_template:
+        name_prefix: gke-a3-ultra-rdma-sub
+        count: 8
+        ip_range: 192.168.128.0/18
+        region: us-central1
+
+  - !DeploymentModule
+    id: gke-a3-ultra-a3-ultragpu-cluster
+    source: modules/scheduler/gke-cluster
+    use: [gke-a3-ultra-net-0]
+    settings:
+      release_channel: RAPID
+      version_prefix: '1.2'
+      min_cluster_version: 1.2.3
+      prefix_with_deployment_name: false
+      name_suffix: gke-a3-ultra
+      system_node_pool_machine_type: "e2-standard-16"
+      enable_dcgm_monitoring: true
+      enable_gcsfuse_csi: true
+      enable_filestore_csi: true
+      enable_private_endpoint: false # Allows access from authorized public IPs
+      master_authorized_networks:
+      - cidr_block: 10.0.0.0/32 # Allows your machine to run the kubectl command. Required for multi network setup.
+        display_name: "kubectl-access-network"
+      system_node_pool_node_count:
+        total_min_nodes: 2
+        total_max_nodes: 1000
+      additional_networks: $(concat([{network=gke-a3-ultra-net-1.network_name, subnetwork=gke-a3-ultra-net-1.subnetwork_name, subnetwork_project="foo", nic_type="GVNIC", queue_count=null, network_ip=null, stack_type=null, access_config=[{nat_ip=null, public_ptr_domain_name=null, network_tier=null}], ipv6_access_config=[], alias_ip_range=[]}], gke-a3-ultra-rdma-net.subnetwork_interfaces_gke))
+      k8s_network_names:
+        rdma_prefix: "gke-a3-ultra-rdma-sub-"
+        rdma_start_index: 0
+        rdma_postfix: ""
+        gvnic_prefix: "gke-a3-ultra-sub-"
+        gvnic_start_index: 1
+    outputs: [instructions]
+
+  - !DeploymentModule
+    id: gke-a3-ultra-a3u-pool
+    source: modules/compute/gke-node-pool
+    use: [gke-a3-ultra-a3-ultragpu-cluster]
+    settings:
+      machine_type: a3-ultragpu-8g
+      auto_upgrade: true
+      zones: [us-central1-c]
+      spot: false
+      max_pods_per_node: 32
+      guest_accelerator:
+      - type: nvidia-h200-141gb
+        count: 8
+        gpu_driver_installation_config:
+          gpu_driver_version: "LATEST"
+      additional_networks:
+        $(concat([{network=gke-a3-ultra-net-1.network_name, subnetwork=gke-a3-ultra-net-1.subnetwork_name, subnetwork_project="foo", nic_type="GVNIC", queue_count=null, network_ip=null, stack_type=null, access_config=[{nat_ip=null, public_ptr_domain_name=null, network_tier=null}], ipv6_access_config=[], alias_ip_range=[]}], gke-a3-ultra-rdma-net.subnetwork_interfaces_gke))
+      reservation_affinity:
+        consume_reservation_type: SPECIFIC_RESERVATION
+        specific_reservations:
+        - name: test-reservation
+      static_node_count: 2
+    outputs: [instructions]
+
+  - !DeploymentModule
+    id: workload-manager-install
+    source: modules/management/kubectl-apply
+    use: [gke-a3-ultra-a3-ultragpu-cluster]
+    settings:
+      jobset:
+        install: true
+        version: v0.7.2
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a3-ultra"))/nccl-installer.yaml
+      - source: $(ghpc_stage("xpk-gke-a3-ultra"))/mlgru-disable.yaml
+      - source: $(ghpc_stage("xpk-gke-a3-ultra"))/storage_crd.yaml
+
+  - !DeploymentModule
+    id: workload_configmap
+    source: 'modules/management/kubectl-apply'
+    use: ['gke-a3-ultra-a3-ultragpu-cluster']
+    settings:
+      apply_manifests:
+      - source: '$(ghpc_stage("xpk-gke-a3-ultra"))/config-map.yaml.tftpl'
+        template_vars: {
+          resource_config_name: "gke-a3-ultra-resources-configmap",
+          num_nodes: "2",
+          cluster_config_name: "gke-a3-ultra-metadata-configmap",
+          capacity_type: "reservation",
+          reservation: "test-reservation",
+        }

xpk/core/blueprint/testing/data/a4.yaml

@@ -0,0 +1,185 @@
+# 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
+#
+#     http://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.
+!Blueprint
+blueprint_name: xpk-gke-a4
+toolkit_modules_url: github.com/GoogleCloudPlatform/cluster-toolkit
+toolkit_modules_version: v1.62.2
+
+vars:
+  labels: {gke_product_type: xpk}
+
+
+terraform_backend_defaults:
+  type: gcs
+  configuration:
+    bucket: test-bucket
+    prefix: xpk_terraform_state/testdir/gke-a4/
+
+deployment_groups:
+- !DeploymentGroup
+  modules:
+  - !DeploymentModule
+    id: gke-a4-net-0
+    source: modules/network/vpc
+    settings:
+      network_name: gke-a4-net-0
+      mtu: 8896
+      subnetworks:
+      - subnet_name: gke-a4-sub-0
+        subnet_region: us-central1
+        subnet_ip: 192.168.0.0/18
+      secondary_ranges_list:
+      - subnetwork_name: gke-a4-sub-0
+        ranges:
+        - range_name: pods
+          ip_cidr_range: 10.4.0.0/14
+        - range_name: services
+          ip_cidr_range: 10.0.32.0/20
+      firewall_rules:
+      - name: gke-a4-internal-0
+        ranges:
+        - 192.168.0.0/16
+        allow:
+        - protocol: tcp
+          ports:
+          - 0-65535
+        - protocol: udp
+          ports:
+          - 0-65535
+        - protocol: icmp
+
+  - !DeploymentModule
+    id: gke-a4-net-1
+    source: modules/network/vpc
+    settings:
+      network_name: gke-a4-net-1
+      mtu: 8896
+      subnetworks:
+      - subnet_name: gke-a4-sub-1
+        subnet_region: us-central1
+        subnet_ip: 192.168.64.0/18
+      firewall_rules:
+      - name: gke-a4-internal-1
+        ranges:
+        - 192.168.0.0/16
+        allow:
+        - protocol: tcp
+          ports:
+          - 0-65535
+        - protocol: udp
+          ports:
+          - 0-65535
+        - protocol: icmp
+
+  - !DeploymentModule
+    id: gke-a4-rdma-net
+    source: modules/network/gpu-rdma-vpc
+    settings:
+      network_name: gke-a4-rdma-net
+      mtu: 8896
+      network_profile:
+        https://www.googleapis.com/compute/beta/projects/foo/global/networkProfiles/us-central1-c-vpc-roce
+      network_routing_mode: REGIONAL
+      subnetworks_template:
+        name_prefix: gke-a4-rdma-sub
+        count: 8
+        ip_range: 192.168.128.0/18
+        region: us-central1
+
+  - !DeploymentModule
+    id: gke-a4-a4-cluster
+    source: modules/scheduler/gke-cluster
+    outputs:
+    - instructions
+    settings:
+      system_node_pool_machine_type: e2-standard-16
+      system_node_pool_node_count:
+        total_min_nodes: 2
+        total_max_nodes: 1000
+      prefix_with_deployment_name: false
+      name_suffix: gke-a4
+      enable_dcgm_monitoring: true
+      enable_gcsfuse_csi: true
+      enable_private_endpoint: false
+      master_authorized_networks:
+      - cidr_block: 10.0.0.0/32
+        display_name: kubectl-access-network
+      additional_networks: $(concat([{network=gke-a4-net-1.network_name, subnetwork=gke-a4-net-1.subnetwork_name,
+        subnetwork_project="foo", nic_type="GVNIC", queue_count=null,
+        network_ip=null, stack_type=null, access_config=[{nat_ip=null, public_ptr_domain_name=null,
+        network_tier=null}], ipv6_access_config=[], alias_ip_range=[]}], gke-a4-rdma-net.subnetwork_interfaces_gke))
+      release_channel: RAPID
+      version_prefix: '1.2'
+      min_cluster_version: 1.2.3
+    use:
+    - gke-a4-net-0
+  - !DeploymentModule
+    id: gke-a4-a4-pool
+    source: modules/compute/gke-node-pool
+    use:
+    - gke-a4-a4-cluster
+    outputs:
+    - instructions
+    settings:
+      machine_type: a4-highgpu-8g
+      auto_upgrade: true
+      zones:
+      - us-central1-c
+      disk_type: hyperdisk-balanced
+      local_ssd_count_ephemeral_storage: 32
+      spot: false
+      reservation_affinity:
+        consume_reservation_type: SPECIFIC_RESERVATION
+        specific_reservations:
+        - name: test-reservation
+      max_pods_per_node: 32
+      guest_accelerator:
+      - type: nvidia-b200
+        count: 8
+        gpu_driver_installation_config:
+          gpu_driver_version: LATEST
+      additional_networks: $(concat([{network=gke-a4-net-1.network_name, subnetwork=gke-a4-net-1.subnetwork_name,
+        subnetwork_project="foo", nic_type="GVNIC", queue_count=null,
+        network_ip=null, stack_type=null, access_config=[{nat_ip=null, public_ptr_domain_name=null,
+        network_tier=null}], ipv6_access_config=[], alias_ip_range=[]}], gke-a4-rdma-net.subnetwork_interfaces_gke))
+      static_node_count: 2
+
+  - !DeploymentModule
+    id: workload-manager-install
+    source: modules/management/kubectl-apply
+    use:
+    - gke-a4-a4-cluster
+    settings:
+      jobset:
+        install: true
+        version: v0.7.2
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a4"))/nccl-rdma-installer-a4.yaml
+      - source: $(ghpc_stage("xpk-gke-a4"))/storage_crd.yaml
+
+  - !DeploymentModule
+    id: workload_configmap
+    source: modules/management/kubectl-apply
+    settings:
+      apply_manifests:
+      - source: $(ghpc_stage("xpk-gke-a4"))/config-map.yaml.tftpl
+        template_vars:
+          resource_config_name: gke-a4-resources-configmap
+          num_nodes: '2'
+          cluster_config_name: gke-a4-metadata-configmap
+          capacity_type: reservation
+          reservation: test-reservation
+    use:
+    - gke-a4-a4-cluster
+  group: primary

xpk/core/capacity.py

@@ -29,6 +29,8 @@ H100_DEVICE_TYPE = 'h100-80gb-8'
 H100_MEGA_DEVICE_TYPE = 'h100-mega-80gb-8'
 H200_DEVICE_TYPE = 'h200-141gb-8'
 B200_DEVICE_TYPE = 'b200-8'
+GB200_DEVICE_TYPE = 'gb200-4'
+GB200_DEVICE_TYPE_NOLSSD = 'gb200-4-no-ssd'
 RESERVATION_CONFIG_KEY = 'reservation_id'
 
 

xpk/core/cluster.py

@@ -22,7 +22,7 @@ from kubernetes import config
 from kubernetes.client.exceptions import ApiException
 
 from ..utils.console import xpk_exit, xpk_print
-from .capacity import B200_DEVICE_TYPE, H100_DEVICE_TYPE, H200_DEVICE_TYPE
+from .capacity import H200_DEVICE_TYPE
 from .commands import (
     run_command_for_value,
     run_command_with_updates,
@@ -34,16 +34,11 @@ from .gcloud_context import (
     zone_to_region,
 )
 from .resources import get_cluster_system_characteristics
-from .system_characteristics import SystemCharacteristics
+from .system_characteristics import INSTALLER_NCCL_TCPXO, SystemCharacteristics
 
 JOBSET_VERSION = 'v0.8.0'
 PATHWAYS_JOB_VERSION = 'v0.1.4'
-INSTALLER_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
-INSTALLER_NCCL_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
-INSTALLER_NCCL_RDMA = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer.yaml'
-CONFIG_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-config.yaml'
 NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml'
-MGLRU_DISABLE = 'https://raw.githubusercontent.com/GoogleCloudPlatform/cluster-toolkit/refs/heads/main/examples/gke-a3-ultragpu/mglru-disable.yaml'
 
 DEFAULT_NAMESPACE = 'default'
 XPK_SA = 'xpk-sa'
@@ -118,12 +113,12 @@ def install_nccl_on_cluster(system: SystemCharacteristics) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
-  if system.device_type == H100_DEVICE_TYPE:
-    command = f'kubectl apply -f {INSTALLER_NCCL_TCPX}'
-  elif system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
-    command = f'kubectl apply -f {INSTALLER_NCCL_RDMA}'
-  else:
-    command = f'kubectl apply -f {INSTALLER_NCCL_TCPXO}'
+  nccl_installer = (
+      system.gpu_config.nccl_installer
+      if system.gpu_config and system.gpu_config.nccl_installer
+      else INSTALLER_NCCL_TCPXO
+  )
+  command = f'kubectl apply -f {nccl_installer}'
 
   return_code = run_command_with_updates(
       command, 'Install NCCL Plugin On Cluster'
@@ -135,35 +130,6 @@ def install_nccl_on_cluster(system: SystemCharacteristics) -> int:
     )
     return 1
 
-  if system.device_type == H100_DEVICE_TYPE:
-    command = f'kubectl apply -f {CONFIG_NCCL_TCPX}'
-
-    return_code = run_command_with_updates(
-        command, 'Install NCCL Config On Cluster'
-    )
-
-    if return_code != 0:
-      xpk_print(
-          f'Install NCCL Config On Cluster request returned ERROR {return_code}'
-      )
-      return 1
-
-  return 0
-
-
-def disable_mglru_on_cluster() -> int:
-  """Disable MGLRU on the cluster.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = f'kubectl apply -f {MGLRU_DISABLE}'
-  return_code = run_command_with_updates(command, 'Disable MGLRU On Cluster')
-
-  if return_code != 0:
-    xpk_print('Disablig MGLRU On Cluster request returned ERROR')
-    return 1
-
   return 0
 
 
@@ -309,10 +275,11 @@ def update_cluster_with_lustre_driver_if_necessary(args) -> int:
   Returns:
     0 if successful and error code otherwise.
   """
-  if is_driver_enabled_on_cluster(
-      args, driver='lustreCsiDriver'
-  ) and is_driver_enabled_on_cluster(
-      args, driver='lustreCsiDriver', config_key='enableLegacyLustrePort'
+  if is_driver_enabled_on_cluster(args, driver='lustreCsiDriver') and (
+      not args.enable_legacy_lustre_port
+      or is_driver_enabled_on_cluster(
+          args, driver='lustreCsiDriver', config_key='enableLegacyLustrePort'
+      )
   ):
     return 0
   cluster_update_return_code = update_gke_cluster_with_lustre_driver_enabled(
@@ -621,9 +588,13 @@ def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --enable-legacy-lustre-port'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       ' --quiet'
   )
+  if args.enable_legacy_lustre_port:
+    command += ' --enable-legacy-lustre-port'
+  else:
+    command += ' --update-addons=LustreCsiDriver=ENABLED'
   xpk_print(
       'Updating GKE cluster to enable Lustre CSI driver, may take a while!'
   )

xpk/core/cluster_test.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 import pytest
 from .testing.commands_tester import CommandsTester
-from .cluster import get_cluster_credentials
+from .cluster import get_cluster_credentials, update_gke_cluster_with_lustre_driver_enabled, update_cluster_with_lustre_driver_if_necessary
 from pytest_mock import MockerFixture
 
 
@@ -93,3 +93,78 @@ def test_get_cluster_credentials_retries_without_dns_when_dns_retrieval_fails(
       if "dns-endpoint" not in c
   ]
   assert len(non_dns_endpoint_commands) == 1
+
+
+def test_update_cluster_with_lustre_driver_if_necessary_with_default_port_runs_correct_checks(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, "True"),
+      "gcloud container clusters describe",
+  )
+  command_args.enable_legacy_lustre_port = None
+  update_cluster_with_lustre_driver_if_necessary(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      "gcloud container clusters describe cluster --project=project"
+      " --location=us-central1"
+      ' --format="value(addonsConfig.lustreCsiDriverConfig.enabled)"'
+  ]
+
+
+def test_update_cluster_with_lustre_driver_if_necessary_with_legacy_port_runs_correct_checks(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, "True"),
+      "gcloud container clusters describe",
+  )
+  command_args.enable_legacy_lustre_port = True
+  update_cluster_with_lustre_driver_if_necessary(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      (
+          "gcloud container clusters describe cluster --project=project"
+          " --location=us-central1"
+          ' --format="value(addonsConfig.lustreCsiDriverConfig.enabled)"'
+      ),
+      (
+          "gcloud container clusters describe cluster --project=project"
+          " --location=us-central1"
+          ' --format="value(addonsConfig.lustreCsiDriverConfig.enableLegacyLustrePort)"'
+      ),
+  ]
+
+
+def test_update_gke_cluster_with_lustre_driver_enabled_default_port(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters update"
+  )
+  command_args.enable_legacy_lustre_port = None
+  update_gke_cluster_with_lustre_driver_enabled(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      "gcloud container clusters update cluster --project=project"
+      " --location=us-central1 --quiet --update-addons=LustreCsiDriver=ENABLED"
+  ]
+
+
+def test_update_gke_cluster_with_lustre_driver_enabled_legacy_port(
+    commands_tester: CommandsTester, command_args
+):
+  commands_tester.set_result_for_command(
+      (0, ""), "gcloud container clusters update"
+  )
+  command_args.enable_legacy_lustre_port = True
+  update_gke_cluster_with_lustre_driver_enabled(command_args)
+
+  executed_commands = commands_tester.get_matching_commands()
+  assert executed_commands == [
+      "gcloud container clusters update cluster --project=project"
+      " --location=us-central1 --quiet --enable-legacy-lustre-port"
+  ]

xpk/core/config.py

@@ -17,12 +17,32 @@ limitations under the License.
 import os
 
 import ruamel.yaml
-
+from abc import ABC, abstractmethod
 from ..utils import file
 from ..utils.console import xpk_print
+from setuptools_scm import get_version as setuptools_get_version
+from importlib.metadata import version, PackageNotFoundError
+
+
+def _get_version() -> str:
+  xpk_version_override = os.getenv('XPK_VERSION_OVERRIDE', '')
+  if xpk_version_override != '':
+    return xpk_version_override
+
+  try:
+    return setuptools_get_version()
+  except LookupError:
+    pass
+
+  try:
+    return version('xpk')
+  except PackageNotFoundError:
+    pass
 
-# This is the version for XPK PyPI package
-__version__ = 'v0.15.0'
+  raise LookupError('unable to determine version number')
+
+
+__version__ = _get_version()
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -41,7 +61,6 @@ KJOB_SHELL_WORKING_DIRECTORY = 'shell-working-directory'
 CONFIGS_KEY = 'configs'
 GKE_ENDPOINT_KEY = 'gke-endpoint'
 DEPENDENCIES_KEY = 'deps-verified-version'
-XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
 DEFAULT_KEYS = [
     CFG_BUCKET_KEY,
@@ -64,8 +83,28 @@ VERTEX_TENSORBOARD_FEATURE_FLAG = XPK_CURRENT_VERSION >= '0.4.0'
 yaml = ruamel.yaml.YAML()
 
 
-class XpkConfig:
-  """XpkConfig is a class for setting and getting values from .yaml config file."""
+class Config(ABC):
+  """Stores and manipulates XPK configuration."""
+
+  @abstractmethod
+  def set(self, key: str, value: str | None) -> None:
+    """Sets the config value"""
+    pass
+
+  @abstractmethod
+  def get(self, key: str) -> str | None:
+    """Reads the config value"""
+    pass
+
+  @abstractmethod
+  def get_all(
+      self,
+  ) -> dict[str, str] | None:
+    pass
+
+
+class FileSystemConfig(Config):
+  """XPK Configuration manipulation class leveraging the file system."""
 
   def __init__(self, custom_config_file: str = XPK_CONFIG_FILE) -> None:
     self._config = custom_config_file
@@ -120,4 +159,39 @@ class XpkConfig:
     return val
 
 
-xpk_config = XpkConfig()
+class InMemoryXpkConfig(Config):
+  """XPK Configuration manipulation class in memory."""
+
+  def __init__(self) -> None:
+    self._config: dict[str, str] = {}
+    self._allowed_keys = DEFAULT_KEYS
+
+  def set(self, key: str, value: str | None) -> None:
+    if key not in self._allowed_keys:
+      return
+    if value is None:
+      self._config.pop(key, None)
+    else:
+      self._config[key] = value
+
+  def get(self, key: str) -> str | None:
+    if key not in self._allowed_keys:
+      return None
+    return self._config.get(key)
+
+  def get_all(
+      self,
+  ) -> dict[str, str] | None:
+    return None if len(self._config) <= 0 else self._config
+
+
+_xpk_config: Config = InMemoryXpkConfig()
+
+
+def set_config(config: Config):
+  global _xpk_config
+  _xpk_config = config
+
+
+def get_config() -> Config:
+  return _xpk_config