xpk 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

xpk/core/workload_decorators/tcpx_decorator_test.py

@@ -0,0 +1,267 @@
+"""
+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 copy
+
+import yaml
+
+from xpk.core.workload_decorators import tcpx_decorator
+from xpk.utils.yaml import literal_string
+
+# Minimal JobSet manifest for testing
+BASE_JOBSET_MANIFEST_STR = """
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: test-jobset
+spec:
+  replicatedJobs:
+    - name: slice-job
+      template:
+        spec:
+          template:
+            metadata:
+              annotations:
+                existing-annotation: "true"
+            spec:
+              containers:
+              - name: main-gpu-container
+                image: my-gpu-image
+                resources:
+                  limits:
+                    nvidia.com/gpu: 8
+              - name: sidecar-container
+                image: my-sidecar-image
+"""
+
+# Minimal kjob template for testing
+BASE_KJOB_TEMPLATE = {
+    "spec": {
+        "template": {
+            "spec": {
+                "containers": [
+                    {
+                        "name": "main-gpu-container",
+                        "image": "my-gpu-image",
+                        "resources": {"limits": {"nvidia.com/gpu": 8}},
+                    },
+                    {"name": "sidecar-container", "image": "my-sidecar-image"},
+                ]
+            }
+        }
+    }
+}
+
+# Minimal job manifest for testing
+BASE_JOB_MANIFEST = {
+    "spec": {
+        "template": {
+            "metadata": {"annotations": {"existing-annotation": "true"}},
+            "spec": {
+                "containers": [
+                    {
+                        "name": "main-gpu-container",
+                        "image": "my-gpu-image",
+                        "resources": {"limits": {"nvidia.com/gpu": 8}},
+                    },
+                    {"name": "sidecar-container", "image": "my-sidecar-image"},
+                ]
+            },
+        }
+    }
+}
+
+
+def test_get_interfaces_annotation():
+  """Tests get_interfaces_annotation."""
+  annotation = tcpx_decorator.get_interfaces_annotation()
+  assert "networking.gke.io/interfaces" in annotation
+  assert isinstance(annotation["networking.gke.io/interfaces"], literal_string)
+  expected_value = (
+      "[\n"
+      '    {"interfaceName":"eth0","network":"default"},\n'
+      '    {"interfaceName":"eth1","network":"vpc1"},\n'
+      '    {"interfaceName":"eth2","network":"vpc2"},\n'
+      '    {"interfaceName":"eth3","network":"vpc3"},\n'
+      '    {"interfaceName":"eth4","network":"vpc4"}\n'
+      "]"
+  )
+  assert str(annotation["networking.gke.io/interfaces"]) == expected_value
+
+
+def test_get_tcpx_deamon_annotation():
+  """Tests get_tcpx_deamon_annotation."""
+  annotation = tcpx_decorator.get_tcpx_deamon_annotation()
+  assert "devices.gke.io/container.tcpx-daemon" in annotation
+  assert isinstance(
+      annotation["devices.gke.io/container.tcpx-daemon"], literal_string
+  )
+  expected_value = (
+      "- path: /dev/nvidia0\n"
+      "- path: /dev/nvidia1\n"
+      "- path: /dev/nvidia2\n"
+      "- path: /dev/nvidia3\n"
+      "- path: /dev/nvidia4\n"
+      "- path: /dev/nvidia5\n"
+      "- path: /dev/nvidia6\n"
+      "- path: /dev/nvidia7\n"
+      "- path: /dev/nvidiactl\n"
+      "- path: /dev/nvidia-uvm\n"
+  )
+  assert (
+      str(annotation["devices.gke.io/container.tcpx-daemon"]) == expected_value
+  )
+
+
+def test_decorate_jobset():
+  """Tests decorate_jobset."""
+  decorated_str = tcpx_decorator.decorate_jobset(BASE_JOBSET_MANIFEST_STR)
+  manifest = yaml.safe_load(decorated_str)
+
+  pod_template_spec = manifest["spec"]["replicatedJobs"][0]["template"]["spec"][
+      "template"
+  ]["spec"]
+  pod_template_metadata = manifest["spec"]["replicatedJobs"][0]["template"][
+      "spec"
+  ]["template"]["metadata"]
+
+  # Check annotations
+  annotations = pod_template_metadata["annotations"]
+  assert "existing-annotation" in annotations
+  assert "devices.gke.io/container.tcpx-daemon" in annotations
+  assert "networking.gke.io/default-interface" in annotations
+  assert "networking.gke.io/interfaces" in annotations
+
+  # Check tolerations
+  tolerations = pod_template_spec["tolerations"]
+  assert {
+      "key": "user-workload",
+      "operator": "Equal",
+      "value": "true",
+      "effect": "NoSchedule",
+  } in tolerations
+
+  # Check volumes
+  volumes = pod_template_spec["volumes"]
+  volume_names = {v["name"] for v in volumes}
+  assert "libraries" in volume_names
+  assert "sys" in volume_names
+  assert "proc-sys" in volume_names
+  assert "tcpx-socket" in volume_names
+  assert "dshm" in volume_names
+
+  # Check init container
+  init_containers = pod_template_spec["initContainers"]
+  assert len(init_containers) == 1
+  tcpx_daemon = init_containers[0]
+  assert tcpx_daemon["name"] == "tcpx-daemon"
+  assert tcpx_daemon["image"].endswith(f":{tcpx_decorator.tcpx}")
+
+  # Check GPU container update
+  gpu_container = pod_template_spec["containers"][0]
+  assert gpu_container["name"] == "main-gpu-container"
+
+  # Check env
+  env_vars = {e["name"]: e["value"] for e in gpu_container["env"]}
+  assert env_vars["LD_LIBRARY_PATH"] == "/usr/local/nvidia/lib64"
+
+  # Check volume mounts
+  volume_mounts = {
+      vm["name"]: vm["mountPath"] for vm in gpu_container["volumeMounts"]
+  }
+  assert volume_mounts["tcpx-socket"] == "/tmp"
+  assert volume_mounts["libraries"] == "/usr/local/nvidia/lib64"
+  assert volume_mounts["dshm"] == "/dev/shm"
+
+  # Check non-GPU container is not updated
+  sidecar_container = pod_template_spec["containers"][1]
+  assert "env" not in sidecar_container
+  assert "volumeMounts" not in sidecar_container
+
+
+def test_decorate_job():
+  """Tests decorate_job."""
+  job_manifest = copy.deepcopy(BASE_JOB_MANIFEST)
+
+  decorated_manifest = tcpx_decorator.decorate_job(job_manifest)
+  pod_template_metadata = decorated_manifest["spec"]["template"]["metadata"]
+
+  # Check annotations
+  annotations = pod_template_metadata["annotations"]
+  assert "existing-annotation" in annotations
+  assert "devices.gke.io/container.tcpx-daemon" in annotations
+  assert "networking.gke.io/default-interface" in annotations
+  assert "networking.gke.io/interfaces" in annotations
+
+
+def test_decorate_kjob_template():
+  """Tests decorate_kjob_template."""
+  kjob_template = copy.deepcopy(BASE_KJOB_TEMPLATE)
+
+  decorated_manifest = tcpx_decorator.decorate_kjob_template(kjob_template)
+
+  pod_template_spec = decorated_manifest["spec"]["template"]["spec"]
+
+  # Check annotations are NOT added
+  assert "annotations" not in decorated_manifest["spec"]["template"].get(
+      "metadata", {}
+  )
+
+  # Check tolerations
+  tolerations = pod_template_spec["tolerations"]
+  assert {
+      "key": "user-workload",
+      "operator": "Equal",
+      "value": "true",
+      "effect": "NoSchedule",
+  } in tolerations
+
+  # Check volumes
+  volumes = pod_template_spec["volumes"]
+  volume_names = {v["name"] for v in volumes}
+  assert "libraries" in volume_names
+  assert "sys" in volume_names
+  assert "proc-sys" in volume_names
+  assert "tcpx-socket" in volume_names
+  assert "dshm" in volume_names
+
+  # Check init container
+  init_containers = pod_template_spec["initContainers"]
+  assert len(init_containers) == 1
+  tcpx_daemon = init_containers[0]
+  assert tcpx_daemon["name"] == "tcpx-daemon"
+  assert tcpx_daemon["image"].endswith(f":{tcpx_decorator.tcpx}")
+
+  # Check GPU container update
+  gpu_container = pod_template_spec["containers"][0]
+  assert gpu_container["name"] == "main-gpu-container"
+
+  # Check env
+  env_vars = {e["name"]: e["value"] for e in gpu_container["env"]}
+  assert env_vars["LD_LIBRARY_PATH"] == "/usr/local/nvidia/lib64"
+
+  # Check volume mounts
+  volume_mounts = {
+      vm["name"]: vm["mountPath"] for vm in gpu_container["volumeMounts"]
+  }
+  assert volume_mounts["tcpx-socket"] == "/tmp"
+  assert volume_mounts["libraries"] == "/usr/local/nvidia/lib64"
+  assert volume_mounts["dshm"] == "/dev/shm"
+
+  # Check non-GPU container is not updated
+  sidecar_container = pod_template_spec["containers"][1]
+  assert "env" not in sidecar_container
+  assert "volumeMounts" not in sidecar_container

xpk/core/workload_decorators/tcpxo_decorator.py

@@ -74,7 +74,8 @@ def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
   for job in manifest['spec']['replicatedJobs']:
     job_manifest = job['template']
     job_manifest = decorate_job(job_manifest, sub_networks)
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_result: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_result
 
 
 def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:

xpk/core/workload_test.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.
+"""
+
+from xpk.core.workload import get_jobsets_list_gcp_link
+
+
+def test_get_jobsets_list_gcp_link():
+  result = get_jobsets_list_gcp_link(
+      project='test-project',
+  )
+
+  assert (
+      result
+      == 'https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=test-project'
+  )

xpk/main.py

@@ -56,18 +56,17 @@ if (
       f' User currently is running {user_major_version}.{user_minor_version}'
   )
 
-# Create top level parser for xpk command.
-parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
-set_parser(parser=parser)
-
-xpk_print('Starting xpk', flush=True)
-validate_dependencies()
-main_args = parser.parse_args()
-main_args.enable_ray_cluster = False
-main_args.func(main_args)
-
 
 def main() -> None:
+  # Create top level parser for xpk command.
+  parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
+  set_parser(parser=parser)
+
+  xpk_print('Starting xpk', flush=True)
+  validate_dependencies()
+  main_args = parser.parse_args()
+  main_args.enable_ray_cluster = False
+  main_args.func(main_args)
   xpk_print('XPK Done.', flush=True)
 
 

xpk/parser/cluster.py

@@ -29,7 +29,7 @@ from ..commands.cluster import (
 from ..commands.config import xpk_cfg
 from ..core.config import CFG_BUCKET_KEY
 from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
-from .common import add_shared_arguments
+from .common import add_shared_arguments, ParserOrArgumentGroup
 from .validators import name_type
 
 
@@ -208,6 +208,14 @@ def set_cluster_create_pathways_parser(
       cluster_create_pathways_optional_arguments
   )
 
+  autoprovisioning_arguments = (
+      cluster_create_pathways_parser.add_argument_group(
+          'Autoprovisioning Arguments',
+          'Optional arguments for enabling autoprovisioning.',
+      )
+  )
+  add_autoprovisioning_arguments(autoprovisioning_arguments)
+
   ### Capacity arguments specific to "cluster create-pathways"
   cluster_create_pathways_capacity_arguments = (
       cluster_create_pathways_parser.add_argument_group(
@@ -529,15 +537,15 @@ def set_cluster_adapt_parser(cluster_adapt_parser: ArgumentParser):
   cluster_adapt_parser.set_defaults(func=cluster_adapt)
 
 
-def add_autoprovisioning_arguments(parser: ArgumentParser):
-  parser.add_argument(
+def add_autoprovisioning_arguments(parser_or_group: ParserOrArgumentGroup):
+  parser_or_group.add_argument(
       '--enable-autoprovisioning',
       action='store_true',
       help=(
           'Enable GKE features for autoprovisioning node pools in GKE clusters.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--autoprovisioning-min-chips',
       type=int,
       help=(
@@ -546,7 +554,7 @@ def add_autoprovisioning_arguments(parser: ArgumentParser):
           ' resources in the cluster as the minimum, and maximum.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--autoprovisioning-max-chips',
       type=int,
       help=(
@@ -557,13 +565,15 @@ def add_autoprovisioning_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_required_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_required_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared required arguments in cluster create and Pathways cluster create.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--cluster',
       type=name_type,
       default=None,
@@ -575,21 +585,23 @@ def add_shared_cluster_create_required_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_optional_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared optional arguments in cluster create and Pathways cluster create.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  add_shared_arguments(parser)
-  parser.add_argument(
+  add_shared_arguments(parser_or_group)
+  parser_or_group.add_argument(
       '--host-maintenance-interval',
       type=str,
       choices=['AS_NEEDED', 'PERIODIC'],
       default='AS_NEEDED',
       help='The maintenance policy of the cluster and respective clusters.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--gke-version',
       type=str,
       help=(
@@ -598,20 +610,20 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' recommended version.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--num-slices',
       type=int,
       default=1,
       help='The number of slices to run the job on, defaults to 1.',
       required=False,
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--pathways-gce-machine-type',
       type=str,
       default='n2-standard-64',
       help='The CPU type for Pathways CPU nodepools',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--default-pool-cpu-machine-type',
       type=str,
       default='e2-standard-16',
@@ -620,7 +632,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' regional clusters, all zones must support the machine type.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--cluster-cpu-machine-type',
       type=str,
       default='',
@@ -631,7 +643,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' cpu nodepools using --device-type.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--default-pool-cpu-num-nodes',
       type=int,
       default=6,
@@ -641,7 +653,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' over time.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--custom-cluster-arguments',
       type=str,
       default='',
@@ -652,7 +664,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           " --custom-cluster-arguments='--network=mtu9k --subnetwork=mtu9k'"
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--custom-nodepool-arguments',
       type=str,
       default='',
@@ -663,7 +675,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' --custom-nodepool-arguments="--disk-size=300"'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--force',
       action='store_true',
       help=(
@@ -671,7 +683,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' additional approval.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--custom-tpu-nodepool-arguments',
       type=str,
       default='',
@@ -682,7 +694,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' --custom-tpu-nodepool-arguments="--enable-ip-alias"'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--private',
       action='store_true',
       help=(
@@ -695,7 +707,7 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' clusters.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--authorized-networks',
       action='extend',
       nargs='+',
@@ -710,16 +722,16 @@ def add_shared_cluster_create_optional_arguments(parser: ArgumentParser):
           ' Example usage: --authorized-networks 1.2.3.0/24 1.2.4.5/32'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-workload-identity',
       action='store_true',
       help='Enable Workload Identity Federation on the cluster and node-pools.',
   )
-  add_driver_arguments(parser)
+  add_driver_arguments(parser_or_group)
 
 
-def add_driver_arguments(parser: ArgumentParser):
-  parser.add_argument(
+def add_driver_arguments(parser_or_group: ParserOrArgumentGroup):
+  parser_or_group.add_argument(
       '--enable-gcsfuse-csi-driver',
       action='store_true',
       help=(
@@ -728,42 +740,44 @@ def add_driver_arguments(parser: ArgumentParser):
           ' Identity is enabled by default.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-gcpfilestore-csi-driver',
       action='store_true',
       help='Enable GCPFilestore driver on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-parallelstore-csi-driver',
       action='store_true',
       help='Enable Parallelstore CSI driver on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-pd-csi-driver',
       action='store_true',
       help='Enable PersistentDisk CSI driver on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-lustre-csi-driver',
       action='store_true',
       help='Enable Lustre CSI driver on the cluster.',
   )
 
 
-def add_shared_cluster_create_tensorboard_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_tensorboard_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared tensorboard arguments in cluster create and Pathways cluster create.
   Note that this feature enables non-Pathways workloads to use tensorboard arguments
   on a Pathways cluster.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--create-vertex-tensorboard',
       action='store_true',
       help='Set this flag to create a Tensorboard instance in Vertex AI.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--tensorboard-region',
       type=str,
       default='us-central1',
@@ -774,7 +788,7 @@ def add_shared_cluster_create_tensorboard_arguments(parser: ArgumentParser):
           ' instance will be created in us-central1.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--tensorboard-name',
       type=str,
       required=False,
@@ -787,13 +801,15 @@ def add_shared_cluster_create_tensorboard_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_capacity_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared capacity arguments in cluster create and Pathways cluster create.
 
   Args:
-    parser: cluster create argument parser or argument group
+    parser_or_group: cluster create argument parser or argument group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--on-demand',
       action='store_true',
       help=(
@@ -802,7 +818,7 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' types.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--reservation',
       type=str,
       help=(
@@ -811,7 +827,7 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' `--flex` or `--on-demand` for other capacity types.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--spot',
       action='store_true',
       help=(
@@ -820,7 +836,7 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' capacity types.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--flex',
       action='store_true',
       help=(
@@ -831,18 +847,20 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
   )
 
 
-def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
+def add_shared_cluster_create_mtc_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
   """Add shared Multi-tier Checkpointing arguments in cluster create and Pathways cluster create.
 
   Args:
-      List of cluster create MTC arguments parsers
+      List of cluster create MTC arguments parsers or group
   """
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--enable-mtc',
       action='store_true',
       help='Enable MTC on the cluster.',
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--mtc-ramdisk-size',
       type=str,
       default=None,
@@ -851,7 +869,7 @@ def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
           ' used for multi-tier checkpointing. e.g. "64Mi" '
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--mtc-gcs-bucket',
       type=str,
       default=None,
@@ -860,7 +878,7 @@ def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
           ' multi-tier checkpointing.'
       ),
   )
-  parser.add_argument(
+  parser_or_group.add_argument(
       '--mtc-toleration-key',
       type=str,
       default=None,