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

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.14.4'
+  raise LookupError('unable to determine version number')
+
+
+__version__ = _get_version()
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -30,6 +50,8 @@ CONFIGS_KEY = 'configs'
 CFG_BUCKET_KEY = 'cluster-state-gcs-bucket'
 CLUSTER_NAME_KEY = 'cluster-name'
 PROJECT_KEY = 'project-id'
+CLIENT_ID_KEY = 'client-id'
+SEND_TELEMETRY_KEY = 'send-telemetry'
 ZONE_KEY = 'zone'
 KJOB_BATCH_IMAGE = 'batch-image'
 KJOB_BATCH_WORKING_DIRECTORY = 'batch-working-directory'
@@ -39,12 +61,13 @@ 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,
     CLUSTER_NAME_KEY,
     PROJECT_KEY,
+    CLIENT_ID_KEY,
+    SEND_TELEMETRY_KEY,
     ZONE_KEY,
     GKE_ENDPOINT_KEY,
     DEPENDENCIES_KEY,
@@ -60,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
@@ -82,7 +125,7 @@ class XpkConfig:
     with open(self._config, encoding='utf-8', mode='w') as stream:
       yaml.dump(config_yaml, stream)
 
-  def set(self, key: str, value: str) -> None:
+  def set(self, key: str, value: str | None) -> None:
     if key not in self._allowed_keys:
       xpk_print(f'Key {key} is not an allowed xpk config key.')
       return
@@ -114,3 +157,41 @@ class XpkConfig:
       return None
     val: dict[str, str] = config_yaml[CONFIGS_KEY]
     return val
+
+
+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

xpk/core/config_test.py

@@ -14,7 +14,9 @@ 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
+from xpk.core.config import FileSystemConfig, InMemoryXpkConfig, CFG_BUCKET_KEY, CLUSTER_NAME_KEY, PROJECT_KEY, ZONE_KEY, _get_version
+from unittest.mock import patch
+from importlib.metadata import PackageNotFoundError
 
 import os
 import pytest
@@ -31,15 +33,60 @@ def _():
     os.remove(config_tmp_path)
 
 
-def test_config(_):
-  cfg = XpkConfig(config_tmp_path)
+@patch('os.getenv', return_value='10.0.0')
+def test_get_version_returns_overriden_value_when_it_is_overriden(_):
+  assert _get_version() == '10.0.0'
+
+
+@patch('os.getenv', return_value='')
+@patch('xpk.core.config.setuptools_get_version', return_value='10.0.0')
+def test_get_version_returns_value_from_setuptools_scm_when_there_is_no_override(
+    *_,
+):
+  assert _get_version() == '10.0.0'
+
+
+@patch('os.getenv', return_value='')
+@patch(
+    'xpk.core.config.setuptools_get_version',
+    side_effect=LookupError('unable to find git version'),
+)
+@patch('xpk.core.config.version', return_value='10.0.0')
+def test_get_version_returns_value_from_pip_when_there_is_no_setuptools_could_be_resolved(
+    *_,
+):
+  assert _get_version() == '10.0.0'
+
+
+@patch('os.getenv', return_value='')
+@patch(
+    'xpk.core.config.setuptools_get_version',
+    side_effect=LookupError('unable to find git version'),
+)
+@patch(
+    'xpk.core.config.version',
+    side_effect=PackageNotFoundError('unable to locate package'),
+)
+def test_get_version_returns_none_when_no_version_could_be_resolved(*_):
+  with pytest.raises(LookupError):
+    _get_version()
+
+
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config(_, cfg):
   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)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_get_all(_, cfg):
   cfg.set(PROJECT_KEY, 'foo')
   cfg.set(CLUSTER_NAME_KEY, 'bar')
   cfg.set(ZONE_KEY, 'europe-west1-a')
@@ -52,20 +99,29 @@ def test_config_get_all(_):
   assert cfg_all[CFG_BUCKET_KEY] == 'cfg-bucket'
 
 
-def test_config_get_empty(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_get_empty(_, cfg):
   val = cfg.get(PROJECT_KEY)
   assert val is None
 
 
-def test_config_get_all_empty(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_get_all_empty(_, cfg):
   val = cfg.get_all()
   assert not val
 
 
-def test_config_set_incorrect(_):
-  cfg = XpkConfig(config_tmp_path)
+@pytest.mark.parametrize(
+    argnames='cfg',
+    argvalues=[(FileSystemConfig(config_tmp_path)), (InMemoryXpkConfig())],
+)
+def test_config_set_incorrect(cfg, _):
   cfg.set('foo', 'bar')
   cfg_all = cfg.get_all()
   assert not cfg_all

xpk/core/docker_container.py

@@ -182,7 +182,9 @@ def get_user_workload_container(args, system: SystemCharacteristics):
       debugging_dashboard_id: id of the GKE dashboard
   """
 
-  setup_docker_image_code, docker_image = setup_docker_image(args)
+  setup_docker_image_code, docker_image = setup_docker_image(
+      args, system.docker_platform
+  )
   if setup_docker_image_code != 0:
     xpk_exit(setup_docker_image_code)
 

xpk/core/docker_image.py

@@ -19,6 +19,7 @@ import os
 import random
 import string
 
+from .system_characteristics import DockerPlatform
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
@@ -26,7 +27,6 @@ from .commands import run_command_with_updates
 
 DEFAULT_DOCKER_IMAGE = 'python:3.10'
 DEFAULT_SCRIPT_DIR = os.getcwd()
-PLATFORM = 'linux/amd64'
 
 
 def validate_docker_image(docker_image, args) -> int:
@@ -63,7 +63,9 @@ def validate_docker_image(docker_image, args) -> int:
     return 0
 
 
-def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
+def build_docker_image_from_base_image(
+    args, docker_platform: DockerPlatform, verbose=True
+) -> tuple[int, str]:
   """Adds script dir to the base docker image and uploads the image.
 
   Args:
@@ -97,8 +99,8 @@ 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)} -t'
-      f' {docker_name} {args.script_dir}'
+      f'docker buildx build --platform={docker_platform.value} -f'
+      f' {str(tmp)} -t {docker_name} {args.script_dir}'
   )
   xpk_print(f'Building {args.script_dir} into docker image.')
   return_code = run_command_with_updates(
@@ -158,7 +160,9 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   return return_code, cloud_docker_image
 
 
-def setup_docker_image(args) -> tuple[int, str]:
+def setup_docker_image(
+    args, docker_platform: DockerPlatform
+) -> tuple[int, str]:
   """Does steps to verify docker args, check image, and build image (if asked).
 
   Args:
@@ -177,7 +181,7 @@ def setup_docker_image(args) -> tuple[int, str]:
     if validate_docker_image_code != 0:
       xpk_exit(validate_docker_image_code)
     build_docker_image_code, docker_image = build_docker_image_from_base_image(
-        args
+        args, docker_platform
     )
     if build_docker_image_code != 0:
       xpk_exit(build_docker_image_code)

xpk/core/docker_resources.py

@@ -16,7 +16,6 @@ limitations under the License.
 
 import os
 import re
-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
@@ -109,14 +108,6 @@ def get_gpu_env(args, system) -> str:
                     value: "{args.command}"
                   {custom_envs}"""
 
-  gpu_direct_name = 'fastrak'
-  if args.device_type == H100_DEVICE_TYPE:
-    gpu_direct_name = 'tcpx'
-  elif args.device_type == H100_MEGA_DEVICE_TYPE:
-    gpu_direct_name = 'tcpxo'
-  elif args.device_type == H200_DEVICE_TYPE:
-    gpu_direct_name = 'rdma'
-
   gpu_env_dic = {
       'JAX_COORDINATOR_PORT': '6002',
       'JAX_COORDINATOR_ADDRESS': (
@@ -129,7 +120,7 @@ def get_gpu_env(args, system) -> str:
   return gpu_env_yaml.format(
       args=args,
       chips_per_vm=system.chips_per_vm,
-      gpu_direct_name=gpu_direct_name,
+      gpu_direct_name=system.gpu_config.gpu_direct_name,
       custom_envs=format_env_dict(args.env, system),
   )
 

xpk/core/gcloud_context.py

@@ -19,6 +19,7 @@ import sys
 from dataclasses import dataclass
 
 from ..utils.console import xpk_print, xpk_exit
+from ..utils.versions import ReleaseChannel
 from .commands import run_command_for_value
 from functools import lru_cache
 
@@ -117,15 +118,18 @@ def get_cluster_location(project: str, name: str, zone: str) -> str:
 class GkeServerConfig:
   """Stores the valid gke versions based on gcloud recommendations."""
 
-  default_rapid_gke_version: str
+  default_gke_version: str
   valid_versions: set[str]
 
 
-def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
+def get_gke_server_config(
+    args, release_channel: ReleaseChannel
+) -> tuple[int, GkeServerConfig | None]:
   """Determine the GKE versions supported by gcloud currently.
 
   Args:
     args: user provided arguments for running the command.
+    release_channel: the release channel to use.
 
   Returns:
     Tuple of
@@ -136,22 +140,24 @@ def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
       'gcloud container get-server-config'
       f' --project={args.project} --region={zone_to_region(args.zone)}'
   )
-  default_rapid_gke_version_cmd = (
+  default_gke_version_cmd = (
       base_command
-      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      + ' --flatten="channels"'
+      f' --filter="channels.channel={release_channel.value}"'
       ' --format="value(channels.defaultVersion)"'
   )
   valid_versions_cmd = (
       base_command
-      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      + ' --flatten="channels"'
+      f' --filter="channels.channel={release_channel.value}"'
       ' --format="value(channels.validVersions)"'
   )
   base_command_description = 'Determine server supported GKE versions for '
 
   server_config_commands_and_descriptions = [
       (
-          default_rapid_gke_version_cmd,
-          base_command_description + 'default rapid gke version',
+          default_gke_version_cmd,
+          base_command_description + 'default gke version',
       ),
       (
           valid_versions_cmd,
@@ -172,8 +178,8 @@ def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
     command_outputs.append(cmd_output)
 
   return 0, GkeServerConfig(
-      default_rapid_gke_version=command_outputs[0].strip(),
-      valid_versions=set(command_outputs[1].split(';')),
+      default_gke_version=command_outputs[0].strip(),
+      valid_versions=set([s.strip() for s in command_outputs[1].split(';')]),
   )
 
 
@@ -196,7 +202,7 @@ def get_gke_control_plane_version(
   if args.gke_version is not None:
     master_gke_version = args.gke_version
   else:
-    master_gke_version = gke_server_config.default_rapid_gke_version
+    master_gke_version = gke_server_config.default_gke_version
 
   is_valid_version = master_gke_version in gke_server_config.valid_versions
 
@@ -204,7 +210,7 @@ def get_gke_control_plane_version(
     xpk_print(
         f'Planned GKE Version: {master_gke_version}\n Valid Versions:'
         f'\n{gke_server_config.valid_versions}\nRecommended / Default GKE'
-        f' Version: {gke_server_config.default_rapid_gke_version}'
+        f' Version: {gke_server_config.default_gke_version}'
     )
     xpk_print(
         f'Error: Planned GKE Version {master_gke_version} is not valid.'
@@ -213,7 +219,7 @@ def get_gke_control_plane_version(
     xpk_print(
         'Please select a gke version from the above list using --gke-version=x'
         ' argument or rely on the default gke version:'
-        f' {gke_server_config.default_rapid_gke_version}'
+        f' {gke_server_config.default_gke_version}'
     )
     return 1, None
 

xpk/core/gcloud_context_test.py

@@ -15,7 +15,20 @@ limitations under the License.
 """
 
 import pytest
-from .gcloud_context import get_cluster_location, zone_to_region
+from unittest.mock import MagicMock
+from .gcloud_context import (
+    get_cluster_location,
+    get_gke_control_plane_version,
+    get_gke_server_config,
+    GkeServerConfig,
+    zone_to_region,
+)
+from ..utils.versions import ReleaseChannel
+
+
+@pytest.fixture(autouse=True)
+def xpk_print(mocker):
+  return mocker.patch("xpk.core.gcloud_context.xpk_print")
 
 
 def test_zone_to_region_raises_when_zone_is_invalid():
@@ -94,3 +107,100 @@ def test_get_cluster_location_invokes_command_for_different_input_args(mocker):
   get_cluster_location(project="project6", name="name6", zone="us-central1-a")
 
   assert mock.call_count == 2
+
+
+def test_get_gke_server_config_success(mocker):
+  mock_run_command = mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      side_effect=[
+          (0, "1.2.3"),
+          (0, "1.2.3;1.2.4;1.3.0"),
+      ],
+  )
+  args = mocker.Mock(project="test-project", zone="us-central1")
+
+  return_code, config = get_gke_server_config(args, ReleaseChannel.STABLE)
+
+  assert return_code == 0
+  assert isinstance(config, GkeServerConfig)
+  assert config.default_gke_version == "1.2.3"
+  assert config.valid_versions == {"1.2.3", "1.2.4", "1.3.0"}
+  assert mock_run_command.call_count == 2
+
+
+def test_get_gke_server_config_fails_on_default_version_command(mocker):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      return_value=(1, "error"),
+  )
+  args = mocker.Mock(project="test-project", zone="us-central1")
+
+  return_code, config = get_gke_server_config(args, ReleaseChannel.STABLE)
+
+  assert return_code == 1
+  assert config is None
+
+
+def test_get_gke_server_config_fails_on_valid_versions_command(mocker):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      side_effect=[(0, "1.2.3"), (1, "error")],
+  )
+  args = mocker.Mock(project="test-project", zone="us-central1")
+
+  return_code, config = get_gke_server_config(args, ReleaseChannel.STABLE)
+
+  assert return_code == 1
+  assert config is None
+
+
+def test_get_gke_control_plane_version_uses_default_when_not_specified(mocker):
+  args = mocker.Mock(gke_version=None)
+  gke_server_config = GkeServerConfig(
+      default_gke_version="1.2.3", valid_versions={"1.2.3", "1.2.4"}
+  )
+
+  return_code, version = get_gke_control_plane_version(args, gke_server_config)
+
+  assert return_code == 0
+  assert version == "1.2.3"
+
+
+def test_get_gke_control_plane_version_uses_user_version_when_valid(mocker):
+  args = mocker.Mock(gke_version="1.2.4")
+  gke_server_config = GkeServerConfig(
+      default_gke_version="1.2.3", valid_versions={"1.2.3", "1.2.4"}
+  )
+
+  return_code, version = get_gke_control_plane_version(args, gke_server_config)
+
+  assert return_code == 0
+  assert version == "1.2.4"
+
+
+def test_get_gke_control_plane_version_fails_for_invalid_user_version(
+    mocker, xpk_print: MagicMock
+):
+  args = mocker.Mock(gke_version="1.2.5")
+  gke_server_config = GkeServerConfig(
+      default_gke_version="1.2.3", valid_versions={"1.2.3", "1.2.4"}
+  )
+
+  return_code, version = get_gke_control_plane_version(args, gke_server_config)
+
+  assert return_code == 1
+  assert version is None
+  assert "Planned GKE Version: 1.2.5" in xpk_print.mock_calls[0].args[0]
+  assert (
+      "Recommended / Default GKE Version: 1.2.3"
+      in xpk_print.mock_calls[0].args[0]
+  )
+  assert (
+      "Error: Planned GKE Version 1.2.5 is not valid."
+      in xpk_print.mock_calls[1].args[0]
+  )
+  assert (
+      "Please select a gke version from the above list using --gke-version=x"
+      " argument or rely on the default gke version: 1.2.3"
+      in xpk_print.mock_calls[2].args[0]
+  )

xpk/core/kjob.py

@@ -25,7 +25,6 @@ from kubernetes.client.rest import ApiException
 from ..utils import templates
 from ..utils.execution_context import is_dry_run
 from ..utils.console import xpk_exit, xpk_print
-from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .cluster import DEFAULT_NAMESPACE, XPK_SA, setup_k8s_env
 from .commands import (
     run_command_for_value,
@@ -38,7 +37,7 @@ from .config import (
     KJOB_SHELL_IMAGE,
     KJOB_SHELL_INTERACTIVE_COMMAND,
     KJOB_SHELL_WORKING_DIRECTORY,
-    XpkConfig,
+    get_config,
 )
 from .network import get_cluster_subnetworks
 from .system_characteristics import AcceleratorType, SystemCharacteristics
@@ -52,7 +51,6 @@ from .storage import (
 )
 from .workload_decorators import (
     rdma_decorator,
-    tcpx_decorator,
     tcpxo_decorator,
 )
 from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
@@ -234,8 +232,7 @@ def get_pod_template_interactive_command() -> str:
   Returns:
     str - PodTemplate's interactive command
   """
-  config = XpkConfig()
-  pod_command = config.get(KJOB_SHELL_INTERACTIVE_COMMAND)
+  pod_command = get_config().get(KJOB_SHELL_INTERACTIVE_COMMAND)
   if pod_command is None or len(pod_command) == 0:
     pod_command = PodTemplateDefaults.INTERACTIVE_COMMAND.value
 
@@ -261,14 +258,17 @@ def create_app_profile_instance(volume_bundles: list[str]) -> int:
   )
 
 
-def decorate_job_template_with_gpu(yml_string: str, gpu_type: str) -> str:
+def decorate_job_template_with_gpu(
+    yml_string: str, system: SystemCharacteristics
+) -> str:
   job_spec = yaml.safe_load(yml_string)["template"]
-  if gpu_type == H100_DEVICE_TYPE:
-    job_spec = tcpx_decorator.decorate_kjob_template(job_spec)
-  if gpu_type == H100_MEGA_DEVICE_TYPE:
-    job_spec = tcpxo_decorator.decorate_kjob_template(job_spec)
-  if gpu_type == H200_DEVICE_TYPE:
-    job_spec = rdma_decorator.decorate_kjob_template(job_spec)
+  kjob_decorator = (
+      system.gpu_config.kjob_decorator_fn
+      if system.gpu_config and system.gpu_config.kjob_decorator_fn
+      else None
+  )
+  if kjob_decorator:
+    job_spec = kjob_decorator(job_spec)
   job_template_dict = yaml.safe_load(yml_string)
   job_template_dict["template"] = job_spec
   yaml_result: str = yaml.dump(job_template_dict, sort_keys=False)
@@ -287,11 +287,10 @@ def create_job_template_instance(
   Returns:
     exit_code > 0 if creating JobTemplate fails, 0 otherwise
   """
-  config = XpkConfig()
-  job_image = config.get(KJOB_BATCH_IMAGE)
+  job_image = get_config().get(KJOB_BATCH_IMAGE)
   if job_image is None or len(job_image) == 0:
     job_image = JobTemplateDefaults.IMAGE.value
-  working_directory = config.get(KJOB_BATCH_WORKING_DIRECTORY)
+  working_directory = get_config().get(KJOB_BATCH_WORKING_DIRECTORY)
   if working_directory is None or len(working_directory) == 0:
     working_directory = JobTemplateDefaults.WORKING_DIRECTORY.value
   resources = (
@@ -318,7 +317,7 @@ def create_job_template_instance(
       service_account=service_account,
   )
   if system is not None and system.accelerator_type == AcceleratorType.GPU:
-    yml_string = decorate_job_template_with_gpu(yml_string, system.device_type)
+    yml_string = decorate_job_template_with_gpu(yml_string, system)
 
   return run_kubectl_apply(
       yml_string,
@@ -332,11 +331,10 @@ def create_pod_template_instance(service_account: str) -> int:
   Returns:
     exit_code > 0 if creating PodTemplate fails, 0 otherwise
   """
-  config = XpkConfig()
-  pod_image = config.get(KJOB_SHELL_IMAGE)
+  pod_image = get_config().get(KJOB_SHELL_IMAGE)
   if pod_image is None or len(pod_image) == 0:
     pod_image = PodTemplateDefaults.IMAGE.value
-  working_directory = config.get(KJOB_SHELL_WORKING_DIRECTORY)
+  working_directory = get_config().get(KJOB_SHELL_WORKING_DIRECTORY)
   if working_directory is None or len(working_directory) == 0:
     working_directory = PodTemplateDefaults.WORKING_DIRECTORY.value