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

xpk/parser/workload.py

@@ -20,9 +20,9 @@ from ..commands.workload import (
     workload_delete,
     workload_list,
 )
-from ..core.core import default_docker_image, default_script_dir
-from .validators import directory_path_type, workload_name_type
+from ..core.docker_image import DEFAULT_DOCKER_IMAGE, DEFAULT_SCRIPT_DIR
 from .common import add_shared_arguments
+from .validators import directory_path_type, name_type
 
 
 def set_workload_parsers(workload_parser):
@@ -67,11 +67,7 @@ def set_workload_parsers(workload_parser):
           'Arguments for configuring autoprovisioning.',
       )
   )
-  workload_pathways_workload_arguments = workload_create_parser.add_argument_group(
-      'Pathways Image Arguments',
-      'If --use-pathways is provided, user wants to set up a'
-      'Pathways workload on xpk.',
-  )
+
   workload_vertex_tensorboard_arguments = (
       workload_create_parser.add_argument_group(
           'Vertex Tensorboard Arguments',
@@ -114,6 +110,12 @@ def set_workload_parsers(workload_parser):
       ),
   )
 
+  workload_create_parser_optional_arguments.add_argument(
+      '--storage',
+      action='append',
+      default=[],
+      help='Names of storages the workload uses',
+  )
   workload_create_parser_optional_arguments.add_argument(
       '--num-nodes',
       type=int,
@@ -151,6 +153,15 @@ def set_workload_parsers(workload_parser):
       ),
   )
 
+  workload_create_parser_optional_arguments.add_argument(
+      '--use-pathways',
+      action='store_true',
+      help=(
+          'Please use `xpk workload create-pathways` instead to'
+          ' create Pathways workloads.'
+      ),
+  )
+
   # Autoprovisioning workload arguments
   workload_create_autoprovisioning_arguments.add_argument(
       '--on-demand',
@@ -178,16 +189,6 @@ def set_workload_parsers(workload_parser):
       ),
   )
 
-  # Pathways workload arguments
-  workload_pathways_workload_arguments.add_argument(
-      '--use-pathways',
-      action='store_true',
-      help=(
-          'DECRATING SOON!!! Please use `xpk workload create-pathways` instead.'
-          ' Provide this argument to create Pathways workloads.'
-      ),
-  )
-
   # "workload create-pathways" command parser.
   workload_create_pathways_parser = workload_subcommands.add_parser(
       'create-pathways', help='Create a new job.'
@@ -230,6 +231,45 @@ def set_workload_parsers(workload_parser):
       help='The tpu type to use, v5litepod-16, etc.',
   )
 
+  ### "workload create-pathways" Optional arguments, specific to Pathways
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--headless',
+      action='store_true',
+      help=(
+          'Please provide this argument to create Pathways workloads in'
+          ' headless mode. This arg can only be used in `xpk workload'
+          ' create-pathways`.'
+      ),
+  )
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--proxy-server-image',
+      type=str,
+      default=(
+          'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest'
+      ),
+      help=(
+          'Please provide the proxy server image for Pathways. This arg can'
+          ' only be used in `xpk workload create-pathways`.'
+      ),
+  )
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--server-image',
+      type=str,
+      default='us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest',
+      help=(
+          'Please provide the server image for Pathways. This arg can only be'
+          ' used in `xpk workload create-pathways`.'
+      ),
+  )
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--pathways-gcs-location',
+      type=str,
+      default='gs://cloud-pathways-staging/tmp',
+      help=(
+          'Please provide the GCS location to store Pathways artifacts. This'
+          ' arg can only be used in `xpk workload create-pathways`.'
+      ),
+  )
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--command',
       type=str,
@@ -243,6 +283,45 @@ def set_workload_parsers(workload_parser):
       ),
       required=False,
   )
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--storage',
+      action='append',
+      default=[],
+      help='Names of storages the workload uses',
+  )
+
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--custom-pathways-server-args',
+      type=str,
+      default=None,
+      help=(
+          'Provide custom Pathways server args as follows -'
+          " --custom-pathways-server-args='--arg_1=xxx --arg2=yyy'"
+      ),
+      required=False,
+  )
+
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--custom-pathways-proxy-server-args',
+      type=str,
+      default=None,
+      help=(
+          'Provide custom Pathways proxy server args as follows -'
+          " --custom-pathways-proxy-server-args='--arg_1=xxx --arg2=yyy'"
+      ),
+      required=False,
+  )
+
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--custom-pathways-worker-args',
+      type=str,
+      default=None,
+      help=(
+          'Provide custom Pathways worker args as follows -'
+          " --custom-pathways-worker-args='--arg_1=xxx --arg2=yyy'"
+      ),
+      required=False,
+  )
 
   add_shared_workload_create_required_arguments([
       workload_create_parser_required_arguments,
@@ -293,7 +372,7 @@ def set_workload_parsers(workload_parser):
   ### "workload delete" Required arguments
   workload_delete_parser_required_arguments.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='The name of the cluster to delete the job on.',
       required=True,
@@ -301,7 +380,7 @@ def set_workload_parsers(workload_parser):
   ### "workload delete" Optional arguments
   workload_delete_parser_optional_arguments.add_argument(
       '--workload',
-      type=workload_name_type,
+      type=name_type,
       default=None,
       help=(
           'The name of the workload to delete. If the workload is not'
@@ -352,7 +431,7 @@ def set_workload_parsers(workload_parser):
 
   workload_list_parser.add_argument(
       '--cluster',
-      type=str,
+      type=name_type,
       default=None,
       help='The name of the cluster to list jobs on.',
       required=True,
@@ -428,14 +507,14 @@ def add_shared_workload_create_required_arguments(args_parsers):
   for custom_parser in args_parsers:
     custom_parser.add_argument(
         '--workload',
-        type=workload_name_type,
+        type=name_type,
         default=None,
         help='The name of the workload to run.',
         required=True,
     )
     custom_parser.add_argument(
         '--cluster',
-        type=str,
+        type=name_type,
         default=None,
         help='The name of the cluster to run the job on.',
         required=True,
@@ -503,6 +582,12 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' event or deletion request.Defaults to 30 seconds.'
         ),
     )
+    custom_parser.add_argument(
+        '--remote-python-sidecar-image',
+        type=str,
+        default=None,
+        help='Remote Python sidecar server image.',
+    )
     custom_parser.add_argument(
         '--enable-debug-logs',
         action='store_true',
@@ -512,59 +597,16 @@ def add_shared_workload_create_optional_arguments(args_parsers):
         ),
     )
     custom_parser.add_argument(
-        '--restart-on-user-code-failure',
-        action='store_true',
-        help=(
-            'Adding this argument will return user failures back to the jobset'
-            ' manager allowing restarts on user code when --max-restarts is set'
-            ' greater than 0. By default, this is not enabled, and workloads'
-            ' will not restart from user code failures. This is enabled by'
-            ' default on Pathways workloads.'
-        ),
-    )
-    custom_parser.add_argument(
-        '--headless',
-        action='store_true',
-        help=(
-            'Please provide this argument to create Pathways workloads in'
-            ' headless mode. This arg can only be used in `xpk workload'
-            ' create-pathways`(preferred) or `xpk workload create'
-            ' --use-pathways.` (--use-pathways will be deprecated soon).'
-        ),
-    )
-    custom_parser.add_argument(
-        '--proxy-server-image',
+        '--restart-on-exit-codes',
         type=str,
-        default=(
-            'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest'
-        ),
-        help=(
-            'Please provide the proxy server image for Pathways. This arg can'
-            ' only be used in `xpk workload create-pathways`(preferred) or `xpk'
-            ' workload create --use-pathways.` (--use-pathways will be'
-            ' deprecated soon).'
-        ),
-    )
-    custom_parser.add_argument(
-        '--server-image',
-        type=str,
-        default='us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest',
-        help=(
-            'Please provide the server image for Pathways. This arg can only be'
-            ' used in `xpk workload create-pathways`(preferred) or `xpk'
-            ' workload create --use-pathways.` (--use-pathways will be'
-            ' deprecated soon).'
-        ),
-    )
-    custom_parser.add_argument(
-        '--pathways-gcs-location',
-        type=str,
-        default='gs://cloud-pathways-staging/tmp',
+        default=None,
         help=(
-            'Please provide the GCS location to store Pathways artifacts. This'
-            ' arg can only be used in `xpk workload create-pathways`(preferred)'
-            ' or `xpk workload create --use-pathways.` (--use-pathways will be'
-            ' deprecated soon).'
+            'Adding this argument specifies additional user-defined exit codes'
+            ' that allow restarting the workload when --max-restarts is set to'
+            ' a value greater than 0. By default, workloads restart on exit'
+            ' codes 42 and 127-255. Any exit codes provided through this flag'
+            ' will be included alongside the default codes for restarting'
+            ' conditions.'
         ),
     )
     custom_parser.add_argument(
@@ -618,10 +660,10 @@ def add_shared_workload_base_docker_image_arguments(args_parsers):
     custom_parser.add_argument(
         '--base-docker-image',
         type=str,
-        default=default_docker_image,
+        default=DEFAULT_DOCKER_IMAGE,
         help=(
             'The base docker-image to use, default'
-            f' {default_docker_image}. If using a custom docker image it'
+            f' {DEFAULT_DOCKER_IMAGE}. If using a custom docker image it'
             ' is typically addressed as gcr.io/${PROJECT}/${NAME}:latest.'
             ' This docker image will be used as a base image by default and'
             ' the `--script-dir` by default will be added to the image.'
@@ -630,7 +672,7 @@ def add_shared_workload_base_docker_image_arguments(args_parsers):
     custom_parser.add_argument(
         '--script-dir',
         type=directory_path_type,
-        default=default_script_dir,
+        default=DEFAULT_SCRIPT_DIR,
         help=(
             'The local location of the directory to copy to the docker image'
             ' and run the main command from. Defaults to current working'

xpk/templates/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright 2023 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.
+"""

xpk/templates/storage.yaml

@@ -0,0 +1,13 @@
+apiVersion: xpk.x-k8s.io/v1
+kind: Storage
+metadata:
+  name: $NAME
+spec:
+  auto_mount:
+  cluster:
+  manifest:
+  mount_point:
+  readonly:
+  type: $NAME
+  pvc:
+  pv:

xpk/utils/gcs_utils.py

@@ -0,0 +1,125 @@
+"""
+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 pathlib import Path
+
+from google.cloud.storage import transfer_manager, Client
+from .console import xpk_print
+
+
+def upload_file_to_gcs(
+    storage_client: Client, bucket_name: str, bucket_path: str, file: str
+):
+  bucket = storage_client.bucket(bucket_name)
+  blob = bucket.blob(bucket_path)
+  blob.upload_from_filename(file)
+
+
+def upload_directory_to_gcs(
+    storage_client: Client,
+    bucket_name: str,
+    bucket_path: str,
+    source_directory: str,
+    workers: int = 8,
+):
+  """Upload every file in a directory, including all files in subdirectories.
+
+  Each blob name is derived from the filename, not including the `directory`
+  parameter itself. For complete control of the blob name for each file (and
+  other aspects of individual blob metadata), use
+  transfer_manager.upload_many() instead.
+  """
+  xpk_print(f"Uploading directory {source_directory} to bucket {bucket_name}")
+  bucket = storage_client.bucket(bucket_name)
+
+  directory_as_path_obj = Path(source_directory)
+  paths = directory_as_path_obj.rglob("*")
+
+  # Filter so the list only includes files, not directories themselves.
+  file_paths = [path for path in paths if path.is_file()]
+
+  # These paths are relative to the current working directory. Next, make them
+  # relative to `directory`
+  relative_paths = [path.relative_to(source_directory) for path in file_paths]
+
+  # Finally, convert them all to strings.
+  string_paths = [str(path) for path in relative_paths]
+
+  xpk_print(f"Found {len(string_paths)} files.")
+  # Start the upload.
+  results = transfer_manager.upload_many_from_filenames(
+      bucket=bucket,
+      filenames=string_paths,
+      source_directory=source_directory,
+      max_workers=workers,
+      blob_name_prefix=bucket_path,
+  )
+
+  for name, result in zip(string_paths, results):
+    # The results list is either `None` or an exception for each filename in
+    # the input list, in order.
+
+    if isinstance(result, Exception):
+      xpk_print(f"Failed to upload {name} due to exception: {result}")
+    else:
+      xpk_print(f"Uploaded {name} to {bucket.name}.")
+
+
+def check_file_exists(
+    storage_client: Client, bucket_name: str, filename: str
+) -> bool:
+  xpk_print(f"Checking if file {filename} exists in bucket: {bucket_name}")
+  bucket = storage_client.get_bucket(bucket_name)
+  is_file: bool = bucket.blob(filename).exists()
+  return is_file
+
+
+def download_bucket_to_dir(
+    storage_client: Client,
+    bucket_name: str,
+    bucket_path: str,
+    destination_directory: str = "",
+    workers: int = 8,
+    max_results: int = 1000,
+):
+  """Download all of the blobs in a bucket, concurrently in a process pool.
+
+  The filename of each blob once downloaded is derived from the blob name and
+  the `destination_directory `parameter. For complete control of the filename
+  of each blob, use transfer_manager.download_many() instead.
+
+  Directories will be created automatically as needed, for instance to
+  accommodate blob names that include slashes.
+  """
+  bucket = storage_client.bucket(bucket_name)
+
+  blob_names = [
+      blob.name
+      for blob in bucket.list_blobs(max_results=max_results, prefix=bucket_path)
+  ]
+
+  results = transfer_manager.download_many_to_path(
+      bucket,
+      blob_names,
+      destination_directory=destination_directory,
+      max_workers=workers,
+  )
+
+  for name, result in zip(blob_names, results):
+    if isinstance(result, Exception):
+      xpk_print(f"Failed to download {name} due to exception: {result}")
+    else:
+      xpk_print(f"Downloaded {name} to {destination_directory + name}.")

xpk/utils/kubectl.py

@@ -0,0 +1,57 @@
+"""
+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 kubernetes.client.exceptions import ApiException
+from kubernetes.dynamic import DynamicClient
+
+from .console import xpk_print
+
+
+def apply_kubectl_manifest(client, manifest):
+  xpk_print('Applying manifest')
+  dynamic_client = DynamicClient(client)
+
+  for obj in manifest:
+    api_version = obj['apiVersion']
+    kind = obj['kind']
+    namespace = obj.get('metadata', {}).get('namespace', 'default')
+
+    api_resource = dynamic_client.resources.get(
+        api_version=api_version, kind=kind
+    )
+
+    try:
+      api_resource.get(name=obj['metadata']['name'], namespace=namespace)
+      api_resource.patch(
+          body=obj,
+          namespace=namespace,
+          name=obj['metadata']['name'],
+          content_type='application/merge-patch+json',
+      )
+      xpk_print(
+          f"Updated {kind} '{obj['metadata']['name']}' in namespace"
+          f" '{namespace}'"
+      )
+
+    except ApiException as e:
+      if e.status == 404:
+        api_resource.create(body=obj, namespace=namespace)
+        xpk_print(
+            f"Applied {kind} '{obj['metadata']['name']}' in namespace"
+            f" '{namespace}'"
+        )
+      else:
+        xpk_print(f'Error applying {kind}: {e}')

xpk/utils/objects.py

@@ -31,7 +31,9 @@ def chunks(lst: list, n: int):
   return [lst[i : i + n] for i in range(0, len(lst), n)]
 
 
-def get_value_from_map(key: str, map_to_search: dict) -> tuple[int, str | None]:
+def get_value_from_map(
+    key: str, map_to_search: dict, verbose: bool = True
+) -> tuple[int, str | None]:
   """Helper function to get value from a map if the key exists.
 
   Args:
@@ -47,10 +49,11 @@ def get_value_from_map(key: str, map_to_search: dict) -> tuple[int, str | None]:
   if value:
     return 0, value
   else:
-    xpk_print(
-        f'Unable to find key: {key} in map: {map_to_search}.'
-        f'The map has the following keys: {map_to_search.keys()}'
-    )
+    if verbose:
+      xpk_print(
+          f'Unable to find key: {key} in map: {map_to_search}.'
+          f'The map has the following keys: {map_to_search.keys()}'
+      )
     return 1, value
 
 

xpk/utils/templates.py

@@ -0,0 +1,28 @@
+"""
+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 ruamel.yaml
+
+yaml = ruamel.yaml.YAML()
+
+
+def load(path: str) -> dict:
+  template_path = os.path.dirname(__file__) + path
+  with open(template_path, "r", encoding="utf-8") as file:
+    data: dict = yaml.load(file)
+  return data

xpk/utils/validation.py

@@ -0,0 +1,80 @@
+"""
+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.commands import run_command_for_value
+from .console import xpk_exit, xpk_print
+from ..commands.config import xpk_cfg
+from ..core.config import DEPENDENCIES_KEY
+from ..commands.version import get_xpk_version
+
+
+validation_commands = {
+    'kubectl': {
+        'command': 'kubectl --help',
+        'message': (
+            '`kubectl` not installed. Please follow'
+            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+            ' to install xpk prerequisites.'
+        ),
+    },
+    'kjob': {
+        'command': 'kubectl kjob --help',
+        'message': (
+            '`kjobctl` not installed. Please follow'
+            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+            ' to install xpk prerequisites.'
+        ),
+    },
+    'gcloud': {
+        'command': 'gcloud version',
+        'message': (
+            '`gcloud not installed. Please follow'
+            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+            ' to install xpk prerequisites.'
+        ),
+    },
+    'docker': {
+        'command': 'docker version',
+        'message': (
+            '`docker` not installed. Please follow'
+            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+            ' to install xpk prerequisites.'
+        ),
+    },
+    'kueuectl': {
+        'command': 'kubectl kueue --help',
+        'message': (
+            '`kueuectl` not installed. Please follow'
+            ' https://github.com/AI-Hypercomputer/xpk?tab=readme-ov-file#prerequisites'
+            ' to install xpk prerequisites.'
+        ),
+    },
+}
+
+
+def validate_dependencies():
+  deps_version = xpk_cfg.get(DEPENDENCIES_KEY)
+  xpk_version = get_xpk_version()
+  if deps_version is None or deps_version != xpk_version:
+    for name, check in validation_commands.items():
+      cmd, message = check['command'], check['message']
+      code, _ = run_command_for_value(
+          cmd, f'Validate {name} installation.', None
+      )
+      if code != 0:
+        xpk_print(message)
+        xpk_exit(code)
+    xpk_cfg.set(DEPENDENCIES_KEY, get_xpk_version())