xpk 0.6.0__py3-none-any.whl → 0.7.1__py3-none-any.whl

xpk/core/config.py

@@ -0,0 +1,179 @@
+"""
+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.
+"""
+
+import os
+import re
+
+import ruamel.yaml
+
+from ..utils import file
+from ..utils.console import xpk_print
+from .system_characteristics import AcceleratorType, SystemCharacteristics
+
+# This is the version for XPK PyPI package
+__version__ = 'v0.7.1'
+XPK_CURRENT_VERSION = __version__
+XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
+
+CONFIGS_KEY = 'configs'
+CFG_BUCKET_KEY = 'cluster-state-gcs-bucket'
+CLUSTER_NAME_KEY = 'cluster-name'
+PROJECT_KEY = 'project-id'
+ZONE_KEY = 'zone'
+KJOB_BATCH_IMAGE = 'batch-image'
+KJOB_BATCH_WORKING_DIRECTORY = 'batch-working-directory'
+KJOB_SHELL_IMAGE = 'shell-image'
+KJOB_SHELL_INTERACTIVE_COMMAND = 'shell-interactive-command'
+KJOB_SHELL_WORKING_DIRECTORY = 'shell-working-directory'
+CONFIGS_KEY = 'configs'
+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,
+    ZONE_KEY,
+    DEPENDENCIES_KEY,
+    KJOB_BATCH_IMAGE,
+    KJOB_BATCH_WORKING_DIRECTORY,
+    KJOB_SHELL_IMAGE,
+    KJOB_SHELL_INTERACTIVE_COMMAND,
+    KJOB_SHELL_WORKING_DIRECTORY,
+]
+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."""
+
+  def __init__(self, custom_config_file: str = XPK_CONFIG_FILE) -> None:
+    self._config = custom_config_file
+    self._allowed_keys = DEFAULT_KEYS
+
+  def _open_configs(self) -> dict | None:
+    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
+
+  def _save_configs(self, config_yaml: dict) -> None:
+    with open(self._config, encoding='utf-8', mode='w') as stream:
+      yaml.dump(config_yaml, stream)
+
+  def set(self, key: str, value: str) -> None:
+    if key not in self._allowed_keys:
+      xpk_print(f'Key {key} is not an allowed xpk config key.')
+      return
+
+    config_yaml = self._open_configs()
+    if config_yaml is None:
+      config_yaml = {'version': 'v1', CONFIGS_KEY: {}}
+
+    config_yaml[CONFIGS_KEY][key] = value
+    self._save_configs(config_yaml)
+
+  def get(self, key: str) -> str | None:
+    if key not in self._allowed_keys:
+      xpk_print(f'Key {key} is not an allowed xpk config key.')
+      return None
+
+    config_yaml = self._open_configs()
+    if config_yaml is None:
+      return None
+
+    vals: dict[str, str] = config_yaml[CONFIGS_KEY]
+    return vals.get(key)
+
+  def get_all(
+      self,
+  ) -> dict[str, dict[str, str] | str] | None:
+    config_yaml = self._open_configs()
+    if config_yaml is None:
+      return None
+    val: dict[str, str] = config_yaml[CONFIGS_KEY]
+    return val
+
+
+def parse_env_config(args, tensorboard_config, system: SystemCharacteristics):
+  """Parses the environment configurations to the jobset config.
+
+  Args:
+    args: user provided arguments for running the command.
+    tensorboard_config: configuration of Vertex Tensorboard.
+    system: system characteristics.
+  """
+  env = {}
+
+  env_pat = re.compile(r'(^[a-zA-Z_][a-zA-Z0-9_]*?)(?:=(.*))?$', re.M)
+  if args.env_file:
+    print('Setting container environment from', args.env_file)
+    with open(file=args.env_file, mode='r', encoding='utf-8') as f:
+      for match in env_pat.finditer(f.read()):
+        variable = match.group(1)
+        if match.group(2) is not None:
+          env[variable] = match.group(2)
+        else:
+          assert variable in os.environ, (
+              f'Variable {variable} is not set in the current '
+              'environment, a value must be specified.'
+          )
+          env[variable] = os.environ[variable]
+  if args.env:
+    for var in args.env:
+      match = env_pat.match(var)
+      assert match and match.group(2) is not None, (
+          'Invalid environment variable, format must be '
+          f'`--env VARIABLE=value`: {var}'
+      )
+      variable = match.group(1)
+      env[variable] = match.group(2)
+
+  if not args.use_pathways:
+    if args.debug_dump_gcs:
+      if 'XLA_FLAGS' in env:
+        raise ValueError(
+            'Conflict: XLA_FLAGS defined in both --debug_dump_gcs '
+            'and environment file. Please choose one way to define '
+            'XLA_FLAGS.'
+        )
+      env['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_dump/'
+
+    if tensorboard_config:
+      env['UPLOAD_DATA_TO_TENSORBOARD'] = True
+      for key, value in tensorboard_config.items():
+        env[key.upper()] = value
+
+  if system.accelerator_type == AcceleratorType['GPU']:
+    # For GPUs, it has two more spaces ahead of name and value respectively
+    env_format = '''
+                  - name: {key}
+                    value: "{value}"'''
+  else:
+    env_format = '''
+                - name: {key}
+                  value: "{value}"'''
+
+  args.env = ''.join(env_format.format(key=k, value=v) for k, v in env.items())

xpk/core/docker_container.py

@@ -0,0 +1,225 @@
+"""
+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 ..utils.console import xpk_exit, xpk_print
+from .docker_image import setup_docker_image
+from .docker_resources import (
+    add_container_ports,
+    add_image_pull_policy_for_pw_or_gpu,
+    add_jax_coordinator_port,
+    get_env_container,
+    get_main_container_resources,
+    get_volume_mounts,
+)
+from .monitoring import get_gke_debugging_dashboard
+from .system_characteristics import (
+    AcceleratorType,
+    AcceleratorTypeToAcceleratorCharacteristics,
+    SystemCharacteristics,
+)
+
+
+def get_main_and_sidecar_container(args, system, docker_image) -> str:
+  """Generate yaml for main and sidecar container.
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics
+    docker_image: docker image
+
+  Returns:
+    str:
+      yaml for main and sidecar container
+  """
+  resource_type = AcceleratorTypeToAcceleratorCharacteristics[
+      system.accelerator_type
+  ].resource_type
+  main_container = get_main_container(args, system, docker_image, resource_type)
+  yaml = """- name: stacktrace-explorer
+                image: busybox:1.28
+                args: [/bin/sh, -c, "check_signal() (while [ ! -f /shared-volume/stacktrace_signal ]; do sleep 1; done; pid=$(pidof 'tail'); kill $pid;); check_signal & while [ ! -d /tmp/debugging ]; do sleep 60; done; while [ ! -e /tmp/debugging/* ]; do sleep 60; done; tail -n+1 -f /tmp/debugging/*; exit 0;"]
+                volumeMounts:
+                - name: tpu-stack-trace
+                  readOnly: true
+                  mountPath: /tmp/debugging
+                - name: shared-data
+                  mountPath: /shared-volume
+              {main_container}
+  """
+  return yaml.format(main_container=main_container)
+
+
+def get_main_container(args, system, docker_image, resource_type) -> str:
+  """Generate yaml for main container including the xpk command.
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics
+    docker_image: docker image
+    resource_type: The label to describe the resource type for TPUs/GPUs/CPUs.
+
+  Returns:
+    str:
+      yaml for main container
+  """
+
+  xpk_internal_commands = ''
+  gsutil_test_command = ''
+  if not args.use_pathways and args.debug_dump_gcs:
+    gsutil_test_command = (
+        'which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil'
+        ' is required but not installed. Aborting"; exit 24;};'
+    )
+    xpk_internal_commands += (
+        'WORKER_ID=$HOSTNAME;'
+        f'gsutil -m cp -r /tmp/xla_dump/ {args.debug_dump_gcs}/$WORKER_ID;'
+    )
+
+  command = args.command
+  if args.enable_debug_logs:
+    command = (
+        'export TPU_STDERR_LOG_LEVEL=0 &&'
+        ' export TPU_MIN_LOG_LEVEL=0 &&'
+        ' export TF_CPP_MIN_LOG_LEVEL=0 &&'
+        ' export TPU_VMODULE=real_program_continuator=1 &&'
+        f' {args.command}'
+    )
+
+  gpu_workload_terminate_command = ''
+  if system.accelerator_type == AcceleratorType['GPU']:
+    gpu_workload_terminate_command = (
+        'echo Main app is done > /usr/share/workload/workload_terminated; '
+    )
+
+  tpu_stacktrace_terminate_command = ''
+  if (
+      not args.use_pathways
+      and system.accelerator_type == AcceleratorType['TPU']
+      and args.deploy_stacktrace_sidecar
+  ):
+    tpu_stacktrace_terminate_command = (
+        'touch /shared-volume/stacktrace_signal; '
+    )
+
+  yaml = """- name: {docker_name}
+                image: {docker_image}
+                {image_pull_policy}
+                env: {env}
+                ports:
+                {container_ports}
+                {jax_coordinator_port}
+                securityContext:
+                  privileged: true
+                command:
+                - bash
+                - -c
+                - |
+                  echo XPK Start: $(date);
+                  _sigterm() (kill -SIGTERM $! 2>/dev/null;);
+                  trap _sigterm SIGTERM;
+                  {gsutil_test_command}
+                  ({command}) & PID=$!;
+                  while kill -0 $PID 2>/dev/null;
+                      do sleep 5;
+                  done;
+                  wait $PID;
+                  EXIT_CODE=$?;
+                  {xpk_internal_commands}
+                  echo XPK End: $(date);
+                  echo EXIT_CODE=$EXIT_CODE;
+                  {tpu_stacktrace_terminate_command}
+                  {gpu_workload_terminate_command}
+                  exit $EXIT_CODE
+                resources:
+                  limits:
+                    {resources}
+"""
+  volume_mounts = get_volume_mounts(args, system)
+  if volume_mounts != '':
+    yaml += """
+                volumeMounts:
+                {volume_mounts}
+"""
+  return yaml.format(
+      args=args,
+      system=system,
+      image_pull_policy=add_image_pull_policy_for_pw_or_gpu(args, system),
+      env=get_env_container(args, system),
+      container_ports=add_container_ports(args, system),
+      jax_coordinator_port=add_jax_coordinator_port(system),
+      docker_name=get_main_container_docker_image(args, system),
+      docker_image=docker_image,
+      gsutil_test_command=gsutil_test_command,
+      command=command,
+      tpu_stacktrace_terminate_command=tpu_stacktrace_terminate_command,
+      gpu_workload_terminate_command=gpu_workload_terminate_command,
+      xpk_internal_commands=xpk_internal_commands,
+      resources=get_main_container_resources(args, system, resource_type),
+      volume_mounts=volume_mounts,
+  )
+
+
+def get_user_workload_container(args, system: SystemCharacteristics):
+  """Deploy user workload container
+
+  Args:
+      args: user provided args.
+      system: system characteristics.
+
+  Returns:
+      container: main container
+      debugging_dashboard_id: id of the GKE dashboard
+  """
+
+  setup_docker_image_code, docker_image = setup_docker_image(args)
+  if setup_docker_image_code != 0:
+    xpk_exit(setup_docker_image_code)
+
+  # Determine if we deploy a sidecar and if we deploy a container.
+  debugging_dashboard_id = None
+  resource_type = AcceleratorTypeToAcceleratorCharacteristics[
+      system.accelerator_type
+  ].resource_type
+  if (
+      not args.use_pathways
+      and system.accelerator_type == AcceleratorType['TPU']
+      and args.deploy_stacktrace_sidecar
+  ):
+    xpk_print(
+        'Sidecar container to display stack traces for TPU workloads will also'
+        ' be deployed.'
+    )
+    container = get_main_and_sidecar_container(args, system, docker_image)
+    # Get GKE debugging dashboard only when sidecar container is deployed for TPU workloads
+    debugging_dashboard_id = get_gke_debugging_dashboard(args)
+  else:
+    container = get_main_container(args, system, docker_image, resource_type)
+  return container, debugging_dashboard_id
+
+
+def get_main_container_docker_image(args, system: SystemCharacteristics) -> str:
+  """Docker name for the main container.
+  Args:
+    args: user provided args.
+    system: system characteristics.
+
+  Returns:
+    str:
+      Workload docker image as a YAML string
+  """
+
+  if system.accelerator_type == AcceleratorType['GPU']:
+    return 'gpu-image'
+
+  return f'{args.docker_name}'

xpk/core/docker_image.py

@@ -0,0 +1,210 @@
+"""
+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.
+"""
+
+import datetime
+import os
+import random
+import string
+
+from ..utils.console import xpk_exit, xpk_print
+from ..utils.file import write_tmp_file
+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:
+  """Validates that the user provided docker image exists in your project.
+
+  Args:
+    docker_image: The docker image to verify.
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  project = args.project
+
+  if not any(repo in docker_image for repo in ['gcr.io', 'docker.pkg.dev']):
+    return 0
+
+  command = (
+      f'gcloud container images describe {docker_image} --project {project}'
+  )
+  return_code = run_command_with_updates(
+      command, 'Validate Docker Image', args, verbose=False
+  )
+  if return_code != 0:
+    xpk_print(
+        'Failed to validate your docker image, check that the docker image'
+        f' exists. You may be able to find the {docker_image} in {project}.'
+        ' If the docker image exists, the service account of this'
+        ' project maybe be missing the permissions to access the docker image.'
+    )
+    return return_code
+  else:
+    return 0
+
+
+def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
+  """Adds script dir to the base docker image and uploads the image.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    Tuple of:
+      0 if successful and 1 otherwise.
+      Name of the Docker image created.
+  """
+
+  # Pick a name for the docker image.
+  docker_image_prefix = os.getenv('USER', 'unknown')
+  docker_name = f'{docker_image_prefix}-runner'
+
+  script_dir_dockerfile = """FROM {base_docker_image}
+
+  # Set the working directory in the container
+  WORKDIR /app
+
+  # Copy all files from local workspace into docker container
+  COPY . .
+
+  WORKDIR /app
+  """
+
+  docker_file = script_dir_dockerfile.format(
+      base_docker_image=args.base_docker_image,
+  )
+  tmp = write_tmp_file(docker_file)
+  docker_build_command = (
+      f'docker buildx build --platform={PLATFORM} -f {str(tmp.file.name)} -t'
+      f' {docker_name} {args.script_dir}'
+  )
+  xpk_print(f'Building {args.script_dir} into docker image.')
+  return_code = run_command_with_updates(
+      docker_build_command,
+      'Building script_dir into docker image',
+      args,
+      verbose=verbose,
+  )
+  if return_code != 0:
+    xpk_print(
+        'Failed to add script_dir to docker image, check the base docker image.'
+        f' You should be able to navigate to the URL {args.base_docker_image}'
+        f' in {args.project}.'
+    )
+    xpk_exit(1)
+
+  # 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_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}')
+
+  # Tag the docker image.
+  tag_docker_image_command = f'docker tag {docker_name} {cloud_docker_image}'
+  return_code = run_command_with_updates(
+      tag_docker_image_command, 'Tag Docker Image', args, verbose=verbose
+  )
+  if return_code != 0:
+    xpk_print(
+        f'Failed to tag docker image with tag: {tag_name}.'
+        f' You should be able to navigate to the URL {cloud_docker_image} in'
+        f' {args.project}.'
+    )
+    xpk_exit(1)
+
+  # Upload image to Artifact Registry.
+  upload_docker_image_command = f'docker push {cloud_docker_image}'
+  return_code = run_command_with_updates(
+      upload_docker_image_command, 'Upload Docker Image', args, verbose=verbose
+  )
+  if return_code != 0:
+    xpk_print(
+        'Failed to upload docker image.'
+        f' You should be able to navigate to the URL {cloud_docker_image} in'
+        f' {args.project}.'
+    )
+    xpk_exit(1)
+  return return_code, cloud_docker_image
+
+
+def setup_docker_image(args) -> tuple[int, str]:
+  """Does steps to verify docker args, check image, and build image (if asked).
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    tuple:
+      0 if successful and 1 otherwise.
+      Name of the docker image to use.
+  """
+  use_base_docker_image = use_base_docker_image_or_docker_image(args)
+
+  docker_image = args.base_docker_image
+  if use_base_docker_image:
+    validate_docker_image_code = validate_docker_image(docker_image, args)
+    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
+    )
+    if build_docker_image_code != 0:
+      xpk_exit(build_docker_image_code)
+  else:
+    docker_image = args.docker_image
+    validate_docker_image_code = validate_docker_image(args.docker_image, args)
+    if validate_docker_image_code != 0:
+      xpk_exit(validate_docker_image_code)
+
+  return 0, docker_image
+
+
+def use_base_docker_image_or_docker_image(args) -> bool:
+  """Checks for correct docker image arguments.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    True if intended to use base docker image, False to use docker image.
+  """
+  use_base_docker_image = True
+  # Check if (base_docker_image and script_dir) or (docker_image) is set.
+  if args.docker_image is not None:
+    if args.script_dir is not DEFAULT_SCRIPT_DIR:
+      xpk_print(
+          '`--script-dir` and --docker-image can not be used together. Please'
+          ' see `--help` command for more details.'
+      )
+      xpk_exit(1)
+    if args.base_docker_image is not DEFAULT_DOCKER_IMAGE:
+      xpk_print(
+          '`--base-docker-image` and --docker-image can not be used together.'
+          ' Please see `--help` command for more details.'
+      )
+      xpk_exit(1)
+    use_base_docker_image = False
+  return use_base_docker_image