xpk 0.12.0__tar.gz → 0.14.0__tar.gz

{xpk-0.12.0/src/xpk.egg-info → xpk-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.12.0
+Version: 0.14.0
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0
@@ -22,15 +22,18 @@ Requires-Dist: google-api-core==2.24.1
 Requires-Dist: packaging==24.2
 Requires-Dist: google-cloud-filestore==1.12.0
 Requires-Dist: google-cloud-storage
+Requires-Dist: Jinja2==3.1.6
 Provides-Extra: dev
 Requires-Dist: pyink==24.3.0; extra == "dev"
 Requires-Dist: pylint>=2.6.0; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-mock==3.15.1; extra == "dev"
 Requires-Dist: docker==7.1.0; extra == "dev"
 Requires-Dist: mypy~=1.17; extra == "dev"
 Requires-Dist: types-PyYAML==6.0.2; extra == "dev"
 Requires-Dist: types-docker~=7.1.0.0; extra == "dev"
+Requires-Dist: pylint-per-file-ignores==1.4.0; extra == "dev"
 Dynamic: license-file
 
 <!--
@@ -76,6 +79,7 @@ XPK supports the following TPU types:
 * v5e
 * v5p
 * Trillium (v6e)
+* Ironwood (tpu7x)
 
 and the following GPU types:
 * A100
@@ -83,6 +87,7 @@ and the following GPU types:
 * A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A4 (b200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
+* A4X (gb200)
 
 and the following CPU types:
 * n2-standard-32

{xpk-0.12.0 → xpk-0.14.0}/README.md

@@ -41,6 +41,7 @@ XPK supports the following TPU types:
 * v5e
 * v5p
 * Trillium (v6e)
+* Ironwood (tpu7x)
 
 and the following GPU types:
 * A100
@@ -48,6 +49,7 @@ and the following GPU types:
 * A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A4 (b200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
+* A4X (gb200)
 
 and the following CPU types:
 * n2-standard-32

{xpk-0.12.0 → xpk-0.14.0}/pyproject.toml

@@ -40,7 +40,8 @@ dependencies = [
   "google-api-core==2.24.1",
   "packaging==24.2",
   "google-cloud-filestore==1.12.0",
-  "google-cloud-storage"
+  "google-cloud-storage",
+  "Jinja2==3.1.6"
 ]
 
 [project.urls]
@@ -62,10 +63,12 @@ dev = [
     "pylint>=2.6.0",
     "pre-commit",
     "pytest",
+    "pytest-mock==3.15.1",
     "docker==7.1.0",
     "mypy ~= 1.17",
     "types-PyYAML == 6.0.2",
     "types-docker ~= 7.1.0.0",
+    "pylint-per-file-ignores == 1.4.0",
 ]
 
 [tool.setuptools.dynamic]

{xpk-0.12.0 → xpk-0.14.0}/src/xpk/commands/batch.py

@@ -29,8 +29,10 @@ from ..core.kjob import (
     get_storage_annotations,
     prepare_kjob,
 )
-from ..core.kueue import LOCAL_QUEUE_NAME
+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
 
@@ -43,6 +45,12 @@ def batch(args: Namespace) -> None:
   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)
@@ -51,18 +59,16 @@ def batch(args: Namespace) -> None:
     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()
+  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:
-
-  setup_k8s_service_accounts()
-
   cmd = (
       'kubectl kjob create slurm'
       f' --profile {AppProfileDefaults.NAME.value}'
@@ -73,7 +79,8 @@ def submit_job(args: Namespace) -> None:
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
   cmd = add_TAS_annotations_to_command(args, cmd)
 
-  for annotation in get_storage_annotations(args):
+  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:
@@ -126,7 +133,7 @@ def submit_job(args: Namespace) -> None:
   if args.time is not None:
     cmd += f' --time {args.time}'
 
-  return_code, return_value = run_command_for_value(cmd, 'submit job', args)
+  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}')