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

xpk/commands/batch.py

@@ -1,144 +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 re
-from argparse import Namespace
-
-from ..core.cluster import (
-    setup_k8s_service_accounts,
-    get_cluster_credentials,
-)
-from ..core.commands import run_command_for_value
-from ..core.gcloud_context import add_zone_and_project
-from ..core.kjob import (
-    AppProfileDefaults,
-    JobTemplateDefaults,
-    get_storage_annotations,
-    prepare_kjob,
-)
-from ..core.kueue_manager import LOCAL_QUEUE_NAME
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.execution_context import is_dry_run
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from .kind import set_local_cluster_command
-from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
-
-
-def batch(args: Namespace) -> None:
-  """Run batch task.
-     This function runs passed script in non-blocking manner.
-  Args:
-    args: user provided arguments for running the command.
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  if not is_dry_run():
-    err_code = prepare_kjob(args)
-    if err_code > 0:
-      xpk_exit(err_code)
-    setup_k8s_service_accounts()
-
-  submit_job(args)
-
-
-def submit_job(args: Namespace) -> None:
-  cmd = (
-      'kubectl kjob create slurm'
-      f' --profile {AppProfileDefaults.NAME.value}'
-      f' --localqueue {LOCAL_QUEUE_NAME}'
-      f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
-      ' --first-node-ip'
-  )
-  cmd = add_gpu_networking_annotations_to_command(args, cmd)
-  cmd = add_TAS_annotations_to_command(args, cmd)
-
-  annotations = [] if is_dry_run() else get_storage_annotations(args)
-  for annotation in annotations:
-    cmd += f' --pod-template-annotation {annotation}'
-
-  if args.ignore_unknown_flags:
-    cmd += ' --ignore-unknown-flags'
-
-  cmd += f' -- {args.script} --partition {LOCAL_QUEUE_NAME}'
-
-  if args.array is not None:
-    cmd += f' --array {args.array}'
-
-  if args.cpus_per_task is not None:
-    cmd += f' --cpus-per-task {args.cpus_per_task}'
-
-  if args.gpus_per_task is not None:
-    cmd += f' --gpus-per-task {args.gpus_per_task}'
-
-  if args.mem is not None:
-    cmd += f' --mem {args.mem}'
-
-  if args.mem_per_task is not None:
-    cmd += f' --mem-per-task {args.mem_per_task}'
-
-  if args.mem_per_cpu is not None:
-    cmd += f' --mem-per-cpu {args.mem_per_cpu}'
-
-  if args.mem_per_gpu is not None:
-    cmd += f' --mem-per-gpu {args.mem_per_gpu}'
-
-  if args.nodes is not None:
-    cmd += f' --nodes {args.nodes}'
-
-  if args.ntasks is not None:
-    cmd += f' --ntasks {args.ntasks}'
-
-  if args.output is not None:
-    cmd += f' --output {args.output}'
-
-  if args.error is not None:
-    cmd += f' --error {args.error}'
-
-  if args.input is not None:
-    cmd += f' --input {args.input}'
-
-  if args.job_name is not None:
-    cmd += f' --job-name {args.job_name}'
-
-  if args.chdir is not None:
-    cmd += f' --chdir {args.chdir}'
-
-  if args.time is not None:
-    cmd += f' --time {args.time}'
-
-  return_code, return_value = run_command_for_value(cmd, 'submit job')
-
-  if return_code != 0:
-    xpk_print(f'Running batch job returned ERROR {return_code}')
-    xpk_exit(return_code)
-
-  m = re.match(r'job\.batch/([-a-z0-9]+)', return_value)
-  if m:
-    xpk_print(f'Job name: {m.group(1)}')

xpk/commands/job.py

@@ -1,244 +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 re
-import sys
-
-from ruamel.yaml import YAML
-from typing import cast
-
-from ..core.commands import run_command_for_value, run_command_with_updates
-from ..core.cluster import get_cluster_credentials
-from ..core.gcloud_context import add_zone_and_project
-from ..core.kjob import AppProfileDefaults
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from .kind import set_local_cluster_command
-
-
-JOBS_DRY_RUN_YAML = """
-items:
-- apiVersion: slurm.k8s.io/v1alpha1
-  kind: SlurmJob
-  metadata:
-    annotations:
-      kjobctl.x-k8s.io/script: echo hello
-    creationTimestamp: '2024-04-29T12:00:00Z'
-    labels:
-      kjobctl.x-k8s.io/app-profile: default
-    name: golden-job
-    namespace: default
-  spec:
-    script: echo hello
-"""
-
-PODS_DRY_RUN_RESULT = """
-foo-pod 2/2 Running 0 2d
-bar-pod 1/1 Evicted 0 1d
-"""
-
-
-def job_info(args):
-  """Run commands obtaining information about a job given by name.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  job_name = args.name
-
-  desc_command = f'kubectl-kjob describe slurm {job_name}'
-  desc_code, desc_text = run_command_for_value(desc_command, 'Getting job data')
-  if desc_code != 0:
-    xpk_print(f'Data info request returned ERROR {desc_code}')
-    xpk_exit(desc_code)
-
-  job_command = (
-      'kubectl-kjob list slurm -o yaml --field-selector'
-      f' metadata.name=={job_name}'
-  )
-  job_code, job_text = run_command_for_value(
-      job_command,
-      'Getting job info',
-      dry_run_return_val=JOBS_DRY_RUN_YAML,
-  )
-  if job_code != 0:
-    xpk_print(f'Job info request returned ERROR {job_code}')
-    xpk_exit(job_code)
-
-  pods_command = f'kubectl get pods -l=job-name={job_name} --no-headers'
-  pods_code, pods_text = run_command_for_value(
-      pods_command,
-      'Getting pods list',
-      dry_run_return_val=PODS_DRY_RUN_RESULT,
-  )
-  if pods_code != 0:
-    xpk_print(f'Pods list request returned ERROR {pods_code}')
-    xpk_exit(pods_code)
-
-  yaml = YAML(typ='safe')
-  job_yaml = yaml.load(job_text)['items'][0]
-
-  output = {
-      'Job name': job_name,
-      'Script name': get_script_name(job_yaml),
-      'Profile': get_profile(job_yaml),
-      'Labels': job_yaml.get('metadata').get('labels', []),
-      'Mounts': get_mounts(job_yaml),
-      'Pods': get_pods(pods_text),
-      'Entrypoint environment variables template': get_kjob_env_vars(desc_text),
-  }
-
-  yaml.default_flow_style = False
-  yaml.sort_base_mapping_type_on_output = False
-  yaml.dump(output, sys.stdout)
-
-
-def get_profile(job_yaml: dict) -> str:
-  containers: list[dict] = (
-      job_yaml.get('spec', {})
-      .get('template', {})
-      .get('spec', {})
-      .get('containers', [])
-  )
-  env_vars = next(iter(containers), {}).get('env', [])
-  profile = next((x['value'] for x in env_vars if x['name'] == 'PROFILE'), '')
-  return profile
-
-
-def get_mounts(job_yaml: dict) -> list[dict]:
-  containers: list[dict] = (
-      job_yaml.get('spec', {})
-      .get('template', {})
-      .get('spec', {})
-      .get('containers', [])
-  )
-  mounts: list[dict] = next(iter(containers), {}).get('volumeMounts', [])
-  return mounts
-
-
-def get_kjob_env_vars(job_desc_text: str) -> list[tuple[str, str]]:
-  regex = r'(SLURM_[A-Z_]*=.*)'
-  search_res = re.findall(regex, job_desc_text)
-  return search_res
-
-
-def get_pods(pods_text: str) -> list[dict[str, str]]:
-  pods_lines = pods_text.strip().split('\n')
-  pods_lines_tokenized = [line.split() for line in pods_lines]
-  return [
-      {
-          'Name': tokens[0],
-          'Status': tokens[2],
-      }
-      for tokens in pods_lines_tokenized
-  ]
-
-
-def get_script_name(job_yaml: dict) -> str | None:
-  return cast(
-      str | None,
-      job_yaml.get('metadata', {})
-      .get('annotations', {})
-      .get('kjobctl.x-k8s.io/script', ''),
-  )
-
-
-def job_list(args) -> None:
-  """Function around job list.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-    msg = f'Listing jobs for project {args.project} and zone {args.zone}:'
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    msg = 'Listing jobs:'
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  xpk_print(msg, flush=True)
-
-  return_code = run_slurm_job_list_command()
-  xpk_exit(return_code)
-
-
-def run_slurm_job_list_command() -> int:
-  cmd = f'kubectl-kjob list slurm  --profile {AppProfileDefaults.NAME.value}'
-
-  return_code = run_command_with_updates(cmd, 'list jobs')
-  if return_code != 0:
-    xpk_print(f'Listing jobs returned ERROR {return_code}')
-  return return_code
-
-
-def job_cancel(args) -> None:
-  """Function around job cancel.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-
-  xpk_print(f'Starting job cancel for job: {args.name}', flush=True)
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  return_code = run_slurm_job_delete_command(args)
-  xpk_exit(return_code)
-
-
-def run_slurm_job_delete_command(args) -> int:
-  list_of_jobs = ' '.join(args.name)
-  cmd = f'kubectl-kjob delete slurm {list_of_jobs}'
-
-  return_code = run_command_with_updates(cmd, 'delete job')
-  if return_code != 0:
-    xpk_print(f'Delete job request returned ERROR {return_code}')
-  return return_code

xpk/commands/kjob_common.py

@@ -1,60 +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 ..core.capacity import (
-    B200_DEVICE_TYPE,
-    H100_MEGA_DEVICE_TYPE,
-    H200_DEVICE_TYPE,
-)
-from ..core.cluster import get_gpu_type_from_cluster
-from ..core.kjob import (
-    get_a3mega_pod_template_annotations,
-    get_a3ultra_pod_template_annotations,
-    get_a4_pod_template_annotations,
-    Kueue_TAS_annotation,
-)
-from .common import is_TAS_possible
-from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
-
-
-def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
-  gpu_type = get_gpu_type_from_cluster(args)
-
-  annotations: tuple
-  if gpu_type == H100_MEGA_DEVICE_TYPE:
-    annotations = get_a3mega_pod_template_annotations()
-  elif gpu_type == H200_DEVICE_TYPE:
-    annotations = get_a3ultra_pod_template_annotations()
-  elif gpu_type == B200_DEVICE_TYPE:
-    annotations = get_a4_pod_template_annotations()
-  else:
-    annotations = tuple()
-
-  flags = [
-      f" --pod-template-annotation {annotation} " for annotation in annotations
-  ]
-  cmd += "\\\n".join(flags)
-
-  return cmd
-
-
-def add_TAS_annotations_to_command(args, cmd: str) -> str:
-  system_characteristics = get_cluster_system_characteristics(args)
-  capacity_type = get_cluster_capacity_type(args)
-  if is_TAS_possible(system_characteristics, capacity_type):
-    cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
-
-  return cmd

xpk/commands/run.py

@@ -1,140 +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 argparse import Namespace
-
-from ..core.cluster import (
-    setup_k8s_service_accounts,
-    get_cluster_credentials,
-)
-from ..core.commands import run_command_with_full_controls
-from ..core.gcloud_context import add_zone_and_project
-from ..core.kjob import (
-    AppProfileDefaults,
-    JobTemplateDefaults,
-    get_storage_annotations,
-    prepare_kjob,
-)
-from ..core.kueue_manager import LOCAL_QUEUE_NAME
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from .kind import set_local_cluster_command
-from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
-
-
-def run(args: Namespace) -> None:
-  """Run task.
-     This function runs passed script in non-blocking manner.
-  Args:
-    args: user provided arguments for running the command.
-  Returns:
-    None
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(args)
-  else:
-    set_cluster_command_code = set_local_cluster_command(args)
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
-
-  err_code = prepare_kjob(args)
-  if err_code > 0:
-    xpk_exit(err_code)
-  setup_k8s_service_accounts()
-
-  submit_job(args)
-
-
-def submit_job(args: Namespace) -> None:
-  cmd = (
-      'kubectl kjob create slurm --profile'
-      f' {AppProfileDefaults.NAME.value} '
-      f' --localqueue {LOCAL_QUEUE_NAME} '
-      f' --stream-container {JobTemplateDefaults.CONTAINER_NAME.value}'
-      f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
-      ' --wait --rm  --first-node-ip'
-  )
-  cmd = add_gpu_networking_annotations_to_command(args, cmd)
-  cmd = add_TAS_annotations_to_command(args, cmd)
-
-  for annotation in get_storage_annotations(args):
-    cmd += f' --pod-template-annotation {annotation}'
-
-  if args.timeout:
-    cmd += f' --wait-timeout {args.timeout}s'
-
-  if args.ignore_unknown_flags:
-    cmd += ' --ignore-unknown-flags'
-
-  cmd += f' -- {args.script} --partition {LOCAL_QUEUE_NAME}'
-
-  if args.array is not None:
-    cmd += f' --array {args.array}'
-
-  if args.cpus_per_task is not None:
-    cmd += f' --cpus-per-task {args.cpus_per_task}'
-
-  if args.gpus_per_task is not None:
-    cmd += f' --gpus-per-task {args.gpus_per_task}'
-
-  if args.mem is not None:
-    cmd += f' --mem {args.mem}'
-
-  if args.mem_per_task is not None:
-    cmd += f' --mem-per-task {args.mem_per_task}'
-
-  if args.mem_per_cpu is not None:
-    cmd += f' --mem-per-cpu {args.mem_per_cpu}'
-
-  if args.mem_per_gpu is not None:
-    cmd += f' --mem-per-gpu {args.mem_per_gpu}'
-
-  if args.nodes is not None:
-    cmd += f' --nodes {args.nodes}'
-
-  if args.ntasks is not None:
-    cmd += f' --ntasks {args.ntasks}'
-
-  if args.output is not None:
-    cmd += f' --output {args.output}'
-
-  if args.error is not None:
-    cmd += f' --error {args.error}'
-
-  if args.input is not None:
-    cmd += f' --input {args.input}'
-
-  if args.job_name is not None:
-    cmd += f' --job-name {args.job_name}'
-
-  if args.chdir is not None:
-    cmd += f' --chdir {args.chdir}'
-
-  if args.time is not None:
-    cmd += f' --time {args.time}'
-
-  return_code = run_command_with_full_controls(cmd, 'run task')
-
-  if return_code != 0:
-    xpk_print(f'Running task returned ERROR {return_code}')
-    xpk_exit(return_code)