xpk 0.17.1__py3-none-any.whl → 0.17.3__py3-none-any.whl

xpk/core/kjob.py

@@ -1,473 +0,0 @@
-"""
-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.
-"""
-
-from argparse import Namespace
-from enum import Enum
-
-import yaml
-from kubernetes import client as k8s_client
-from kubernetes.client import ApiClient
-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 .cluster import DEFAULT_NAMESPACE, XPK_SA, setup_k8s_env
-from .commands import (
-    run_command_for_value,
-    run_command_with_updates,
-    run_kubectl_apply,
-)
-from .config import (
-    KJOB_BATCH_IMAGE,
-    KJOB_BATCH_WORKING_DIRECTORY,
-    KJOB_SHELL_IMAGE,
-    KJOB_SHELL_INTERACTIVE_COMMAND,
-    KJOB_SHELL_WORKING_DIRECTORY,
-    get_config,
-)
-from .network import get_cluster_subnetworks
-from .system_characteristics import AcceleratorType, SystemCharacteristics
-from .resources import get_cluster_system_characteristics
-from .storage import (
-    GCS_FUSE_ANNOTATIONS,
-    PARALLELSTORE_ANNOTATIONS,
-    get_auto_mount_gcsfuse_storages,
-    get_auto_mount_parallelstore_storages,
-    get_auto_mount_storages,
-)
-from .workload_decorators import (
-    rdma_decorator,
-    tcpxo_decorator,
-)
-from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
-
-KJOB_API_GROUP_NAME = "kjobctl.x-k8s.io"
-KJOB_API_GROUP_VERSION = "v1alpha1"
-KJOB_API_VOLUME_BUNDLE_PLURAL = "volumebundles"
-VOLUME_BUNDLE_TEMPLATE_PATH = "/../templates/volume_bundle.yaml"
-
-
-class AppProfileDefaults(Enum):
-  NAME = "xpk-def-app-profile"
-
-
-class JobTemplateDefaults(Enum):
-  NAME = "xpk-def-batch"
-  PARALLELISM = 1
-  COMPLETIONS = 1
-  CONTAINER_NAME = "xpk-batch-container"
-  IMAGE = "ubuntu:22.04"
-  WORKING_DIRECTORY = "/"
-
-
-class PodTemplateDefaults(Enum):
-  NAME = "xpk-def-pod"
-  CONTAINER_NAME = "xpk-interactive-container"
-  IMAGE = "busybox:1.28"
-  WORKING_DIRECTORY = "/"
-  INTERACTIVE_COMMAND = "/bin/sh"
-
-
-job_template_yaml = """
-  apiVersion: kjobctl.x-k8s.io/v1alpha1
-  kind: JobTemplate
-  metadata:
-    name: {name}
-    namespace: default
-  template:
-    spec:
-      parallelism: {parallelism}
-      completions: {completions}
-      completionMode: Indexed
-      template:
-        spec:
-          dnsPolicy: ClusterFirstWithHostNet
-          tolerations:
-            - operator: "Exists"
-              key: nvidia.com/gpu
-          containers:
-            - name: {container_name}
-              image: {image}
-              workingDir: {working_directory}
-              {resources}
-          {node_selector}
-          priorityClassName: {priority}
-          restartPolicy: OnFailure
-          serviceAccountName: {service_account}
-"""
-job_node_selector_template = """
-          nodeSelector:
-            cloud.google.com/gke-accelerator: {gpu_name}
-"""
-job_resources_template = """
-              resources:
-                limits:
-                  nvidia.com/gpu: {gpu_per_node}
-"""
-
-app_profile_yaml = """
-apiVersion: kjobctl.x-k8s.io/v1alpha1
-kind: ApplicationProfile
-metadata:
-  name: {name}
-  namespace: default
-spec:
-  supportedModes:
-    - name: Slurm
-      template: {batch_template}
-      requiredFlags: []
-    - name: Interactive
-      template: {interactive_template}
-  volumeBundles: {volume_bundles}
-"""
-
-pod_template_yaml = """
-apiVersion: v1
-kind: PodTemplate
-metadata:
-  name: {name}
-  namespace: default
-template:
-  spec:
-    tolerations:
-      - effect: NoSchedule
-        key: components.gke.io/gke-managed-components
-        operator: Equal
-        value: "true"
-    containers:
-      - name: {container_name}
-        image: {image}
-        command: [{interactive_command}]
-        workingDir: {working_directory}
-    initContainers:
-      - name: init
-        image: {image}
-        command: ['/bin/mkdir', '-p', '{working_directory}']
-    serviceAccountName: {service_account}
-"""
-
-Kueue_TAS_annotation = "kueue.x-k8s.io/podset-preferred-topology=cloud.google.com/gce-topology-host"
-
-default_interface_annotation = "networking.gke.io/default-interface=eth0"
-
-
-def get_a4_pod_template_annotations() -> tuple[str, str]:
-  sub_networks = get_cluster_subnetworks()
-  interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
-      sub_networks
-  )
-
-  return (
-      default_interface_annotation,
-      f"{interfaces_key}=$'{interfaces_value}'",
-  )
-
-
-def get_a3ultra_pod_template_annotations() -> tuple[str, str]:
-  sub_networks = get_cluster_subnetworks()
-  interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
-      sub_networks
-  )
-
-  return (
-      default_interface_annotation,
-      f"{interfaces_key}=$'{interfaces_value}'",
-  )
-
-
-def get_a3mega_pod_template_annotations() -> tuple[str, str, str]:
-  """Adds or updates annotations in the Pod template."""
-  sub_networks = get_cluster_subnetworks()
-  tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
-  interfaces_key, interfaces_value = tcpxo_decorator.get_interfaces_entry(
-      sub_networks
-  )
-  tcpxo = f"{tcpxo_deamon_key}=$'{tcpxo_deamon_paths}'"
-  interfaces = f"{interfaces_key}=$'{interfaces_value}'"
-  return tcpxo, interfaces, default_interface_annotation
-
-
-def verify_kjob_installed() -> int:
-  """Check if kjob is installed. If not provide user with proper communicate and exit.
-  Returns:
-    error code > if kjob not installed, otherwise 0
-  """
-  command = "kubectl-kjob help"
-  task = "Verify kjob installation "
-  verify_kjob_installed_code, _ = run_command_for_value(command, task)
-
-  if verify_kjob_installed_code == 0:
-    xpk_print("kjob found")
-    return 0
-
-  if verify_kjob_installed_code != 0:
-    xpk_print(
-        " kjob not found. Please follow"
-        " https://github.com/kubernetes-sigs/kjob/blob/main/docs/installation.md"
-        " to install kjob."
-    )
-    return verify_kjob_installed_code
-  return 0
-
-
-def get_pod_template_interactive_command() -> str:
-  """Gets the interactive command for PodTemplate from config otherwise the default value.
-
-  Args:
-    args - user provided arguments
-  Returns:
-    str - PodTemplate's 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
-
-  return pod_command
-
-
-def create_app_profile_instance(volume_bundles: list[str]) -> int:
-  """Create new AppProfile instance on cluster with default settings.
-
-  Args:
-    args - user provided arguments
-  Returns:
-    exit_code > 0 if creating AppProfile fails, 0 otherwise
-  """
-  return run_kubectl_apply(
-      yml_string=app_profile_yaml.format(
-          name=AppProfileDefaults.NAME.value,
-          batch_template=JobTemplateDefaults.NAME.value,
-          interactive_template=PodTemplateDefaults.NAME.value,
-          volume_bundles=volume_bundles,
-      ),
-      task="Creating AppProfile",
-  )
-
-
-def decorate_job_template_with_gpu(
-    yml_string: str, system: SystemCharacteristics
-) -> str:
-  job_spec = yaml.safe_load(yml_string)["template"]
-  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)
-  return yaml_result
-
-
-def create_job_template_instance(
-    args: Namespace,
-    system: SystemCharacteristics | None,
-    service_account: str,
-) -> int:
-  """Create new JobTemplate instance on cluster with default settings.
-
-  Args:
-    args - user provided arguments
-  Returns:
-    exit_code > 0 if creating JobTemplate fails, 0 otherwise
-  """
-  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 = get_config().get(KJOB_BATCH_WORKING_DIRECTORY)
-  if working_directory is None or len(working_directory) == 0:
-    working_directory = JobTemplateDefaults.WORKING_DIRECTORY.value
-  resources = (
-      job_resources_template.format(gpu_per_node=system.chips_per_vm)
-      if system is not None and system.accelerator_type == AcceleratorType.GPU
-      else ""
-  )
-
-  node_selector = (
-      job_node_selector_template.format(gpu_name=system.gke_accelerator)
-      if system is not None and system.accelerator_type == AcceleratorType.GPU
-      else ""
-  )
-  yml_string = job_template_yaml.format(
-      name=JobTemplateDefaults.NAME.value,
-      parallelism=JobTemplateDefaults.PARALLELISM.value,
-      completions=JobTemplateDefaults.COMPLETIONS.value,
-      container_name=JobTemplateDefaults.CONTAINER_NAME.value,
-      image=job_image,
-      working_directory=working_directory,
-      resources=resources,
-      node_selector=node_selector,
-      priority=args.priority if hasattr(args, "priority") else "medium",
-      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)
-
-  return run_kubectl_apply(
-      yml_string,
-      task="Creating JobTemplate",
-  )
-
-
-def create_pod_template_instance(service_account: str) -> int:
-  """Create new PodTemplate instance on cluster with default settings.
-
-  Returns:
-    exit_code > 0 if creating PodTemplate fails, 0 otherwise
-  """
-  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 = get_config().get(KJOB_SHELL_WORKING_DIRECTORY)
-  if working_directory is None or len(working_directory) == 0:
-    working_directory = PodTemplateDefaults.WORKING_DIRECTORY.value
-
-  return run_kubectl_apply(
-      yml_string=pod_template_yaml.format(
-          name=PodTemplateDefaults.NAME.value,
-          container_name=PodTemplateDefaults.CONTAINER_NAME.value,
-          image=pod_image,
-          working_directory=working_directory,
-          interactive_command=get_pod_template_interactive_command(),
-          service_account=service_account,
-      ),
-      task="Creating PodTemplate",
-  )
-
-
-def prepare_kjob(args: Namespace) -> int:
-  system = get_cluster_system_characteristics(args)
-
-  storages = []
-  if not is_dry_run():
-    k8s_api_client = setup_k8s_env(args)
-    storages = get_auto_mount_storages(k8s_api_client)
-
-  service_account = ""
-  if len(storages) > 0:
-    service_account = XPK_SA
-
-  job_err_code = create_job_template_instance(args, system, service_account)
-  if job_err_code > 0:
-    return job_err_code
-  pod_err_code = create_pod_template_instance(service_account)
-  if pod_err_code > 0:
-    return pod_err_code
-
-  volume_bundles = [item.name for item in storages]
-
-  return create_app_profile_instance(volume_bundles)
-
-
-def apply_kjob_crds() -> int:
-  """Apply kjob CRDs on cluster.
-
-  This function install kjob CRDs files from kjobctl printcrds.
-  It creates all neccessary kjob CRDs.
-
-  Returns:
-    None
-  """
-  command = "kubectl kjob printcrds | kubectl apply --server-side -f -"
-  task = "Create kjob CRDs on cluster"
-  return_code = run_command_with_updates(command, task)
-  if return_code != 0:
-    xpk_print(f"{task} returned ERROR {return_code}")
-    return return_code
-  xpk_print("Creating kjob CRDs succeeded")
-  return 0
-
-
-def create_volume_bundle_instance(
-    k8s_api_client: ApiClient,
-    name: str,
-    manifest: list[dict],
-    readonly: bool,
-    mount_point: str,
-) -> None:
-  """
-  Creates a new VolumeBundle resource in the Kubernetes cluster.
-
-  This function reads a VolumeBundle template from a YAML file, populates it with
-  values from the provided arguments, and then creates the VolumeBundle object
-  in the cluster.
-
-  Args:
-      k8s_api_client: An ApiClient object for interacting with the Kubernetes API.
-      args: An argparse Namespace object containing the arguments for creating
-            the Storage resource.
-  """
-  data = templates.load(VOLUME_BUNDLE_TEMPLATE_PATH)
-  data["metadata"]["name"] = name
-  spec = data["spec"]
-  spec["volumes"] = []
-  spec["containerVolumeMounts"] = []
-
-  for obj in manifest:
-    if obj["kind"] == "PersistentVolumeClaim":
-      spec["volumes"].append({
-          "name": obj["metadata"]["name"],
-          "persistentVolumeClaim": {
-              "claimName": obj["metadata"]["name"],
-              "readOnly": readonly,
-          },
-      })
-      spec["containerVolumeMounts"].append({
-          "name": obj["metadata"]["name"],
-          "mountPath": mount_point,
-      })
-
-  data["spec"] = spec
-
-  api_instance = k8s_client.CustomObjectsApi(k8s_api_client)
-  try:
-    api_instance.create_namespaced_custom_object(
-        namespace=DEFAULT_NAMESPACE,
-        group=KJOB_API_GROUP_NAME,
-        version=KJOB_API_GROUP_VERSION,
-        plural=KJOB_API_VOLUME_BUNDLE_PLURAL,
-        body=data,
-    )
-    xpk_print(
-        f"Created {KJOB_API_VOLUME_BUNDLE_PLURAL}.{KJOB_API_GROUP_NAME} object:"
-        f" {data['metadata']['name']}"
-    )
-  except ApiException as e:
-    if e.status == 409:
-      xpk_print(f"VolumeBundle: {name} already exists. Skipping its creation")
-    else:
-      xpk_print(f"Encountered error during VolumeBundle creation: {e}")
-      xpk_exit(1)
-
-
-def get_storage_annotations(args: Namespace) -> list[str]:
-  annotations = []
-  k8s_api_client = setup_k8s_env(args)
-
-  gcsfuse_storages = get_auto_mount_gcsfuse_storages(k8s_api_client)
-  if len(gcsfuse_storages) > 0:
-    for key, value in GCS_FUSE_ANNOTATIONS.items():
-      annotations.append(f"{key}={value}")
-
-  parallelstore_storages = get_auto_mount_parallelstore_storages(k8s_api_client)
-  if len(parallelstore_storages) > 0:
-    for key, value in PARALLELSTORE_ANNOTATIONS.items():
-      annotations.append(f"{key}={value}")
-
-  return annotations

xpk/parser/batch.py

@@ -1,43 +0,0 @@
-"""
-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.
-"""
-
-from .common import (
-    add_shared_arguments,
-    add_slurm_arguments,
-    add_cluster_arguments,
-    add_kind_cluster_arguments,
-)
-from ..commands.batch import batch
-
-
-def set_batch_parser(batch_parser):
-  batch_required_arguments = batch_parser.add_argument_group(
-      'batch Built-in Arguments', 'Arguments required for `batch`.'
-  )
-  batch_optional_arguments = batch_parser.add_argument_group(
-      'Optional Arguments', 'Arguments optional for `batch`.'
-  )
-
-  ### "batch" Required arguments
-  batch_required_arguments.add_argument(
-      'script', help='script with batch task to run'
-  )
-
-  add_cluster_arguments(batch_optional_arguments)
-  add_kind_cluster_arguments(batch_optional_arguments)
-  add_shared_arguments(batch_optional_arguments)
-  add_slurm_arguments(batch_optional_arguments)
-  batch_parser.set_defaults(func=batch)

xpk/parser/job.py

@@ -1,147 +0,0 @@
-"""
-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 argparse
-from ..commands.job import job_info, job_list, job_cancel
-
-from .common import add_shared_arguments
-from .validators import name_type
-
-
-def set_job_parser(job_parser: argparse.ArgumentParser):
-  job_subcommands = job_parser.add_subparsers(
-      title='job subcommands',
-      dest='xpk_job_subcommands',
-      help=(
-          'These are commands related to job management. Look at help for'
-          ' specific subcommands for more details.'
-      ),
-  )
-  set_job_info_parser(
-      job_info_parser=job_subcommands.add_parser(
-          'info', help='Show information about specified job.'
-      )
-  )
-  set_job_list_parser(
-      job_list_parser=job_subcommands.add_parser('ls', help='List jobs.')
-  )
-  set_job_cancel_parser(
-      job_cancel_parser=job_subcommands.add_parser(
-          'cancel', help='Cancel job execution.'
-      )
-  )
-
-
-def set_job_info_parser(job_info_parser: argparse.ArgumentParser):
-  job_info_required_arguments = job_info_parser.add_argument_group(
-      'Required arguments',
-      'The basic information required to identify the job.',
-  )
-  job_info_optional_arguments = job_info_parser.add_argument_group(
-      'Optional Arguments', 'Arguments optional for job info.'
-  )
-
-  ### Required arguments
-  job_info_required_arguments.add_argument(
-      '--cluster',
-      type=name_type,
-      default=None,
-      help='The name of the cluster to info jobs on.',
-      required=True,
-  )
-
-  job_info_optional_arguments.add_argument(
-      '--kind-cluster',
-      type=bool,
-      action=argparse.BooleanOptionalAction,
-      default=False,
-      help='Apply command to a local test cluster.',
-  )
-  job_info_required_arguments.add_argument(
-      'name',
-      type=str,
-      default=None,
-      help='Name of the job.',
-  )
-  job_info_parser.set_defaults(func=job_info)
-  add_shared_arguments(job_info_parser)
-
-
-def set_job_list_parser(job_list_parser: argparse.ArgumentParser):
-  job_list_required_arguments = job_list_parser.add_argument_group(
-      'Required Arguments',
-      'Arguments required for job list.',
-  )
-  job_list_optional_arguments = job_list_parser.add_argument_group(
-      'Optional Arguments', 'Arguments optional for job list.'
-  )
-
-  ### Required arguments
-  job_list_required_arguments.add_argument(
-      '--cluster',
-      type=name_type,
-      default=None,
-      help='The name of the cluster to list jobs on.',
-      required=True,
-  )
-
-  job_list_optional_arguments.add_argument(
-      '--kind-cluster',
-      type=bool,
-      action=argparse.BooleanOptionalAction,
-      default=False,
-      help='Apply command to a local test cluster.',
-  )
-
-  job_list_parser.set_defaults(func=job_list)
-  add_shared_arguments(job_list_optional_arguments)
-
-
-def set_job_cancel_parser(job_cancel_parser: argparse.ArgumentParser):
-  job_cancel_required_arguments = job_cancel_parser.add_argument_group(
-      'Required Arguments',
-      'Arguments required for job cancel.',
-  )
-  job_cancel_optional_arguments = job_cancel_parser.add_argument_group(
-      'Optional Arguments', 'Arguments optional for job cancel.'
-  )
-
-  job_cancel_required_arguments.add_argument(
-      'name',
-      type=str,
-      default=None,
-      help='The name of the job to be cancelled.',
-      nargs='+',
-  )
-
-  job_cancel_required_arguments.add_argument(
-      '--cluster',
-      type=name_type,
-      default=None,
-      help='The name of the cluster to delete the job on.',
-      required=True,
-  )
-
-  job_cancel_optional_arguments.add_argument(
-      '--kind-cluster',
-      type=bool,
-      action=argparse.BooleanOptionalAction,
-      default=False,
-      help='Apply command to a local test cluster.',
-  )
-
-  job_cancel_parser.set_defaults(func=job_cancel)
-  add_shared_arguments(job_cancel_optional_arguments)

xpk/parser/run.py

@@ -1,47 +0,0 @@
-"""
-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 ..commands.run import run
-from .common import (
-    add_shared_arguments,
-    add_slurm_arguments,
-    add_cluster_arguments,
-    add_kind_cluster_arguments,
-)
-
-
-def set_run_parser(run_parser):
-  run_required_arguments = run_parser.add_argument_group(
-      'Required Arguments', 'Arguments required for `run`.'
-  )
-  run_optional_arguments = run_parser.add_argument_group(
-      'Optional Arguments', 'Arguments optional for `run`.'
-  )
-
-  run_required_arguments.add_argument('script', help='script with task to run')
-  run_optional_arguments.add_argument(
-      '--timeout',
-      type=int,
-      default=None,
-      help='Amount of time to wait for job in seconds',
-      required=False,
-  )
-
-  add_cluster_arguments(run_optional_arguments)
-  add_kind_cluster_arguments(run_optional_arguments)
-  add_slurm_arguments(run_optional_arguments)
-  add_shared_arguments(run_parser)
-  run_parser.set_defaults(func=run)