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

xpk/core/workload.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import re
 from ..utils.console import xpk_exit, xpk_print
 from .commands import run_command_for_value
 from .gcloud_context import zone_to_region
@@ -240,3 +241,13 @@ def wait_for_job_completion(args) -> int:
     xpk_print('Your workload did not complete successfully')
     return 125
   return 0
+
+
+GCP_NAME_FILTER_VALUE_REGEX = re.compile(r'[a-z0-9\-]+')
+"""Defines correct name prefix value (contains only letters, numbers and dashes) that can be used in GCP filter chips."""
+
+
+def get_jobsets_list_gcp_link(project: str) -> str:
+  """Returns a link to Cloud Console JobSets list"""
+
+  return f'https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project={project}'

xpk/core/workload_decorators/rdma_decorator.py

@@ -18,7 +18,7 @@ import yaml
 from ...utils.yaml import literal_string
 
 
-def decorate_kjob_template(job_manifest) -> str:
+def decorate_kjob_template(job_manifest: dict) -> dict:
   spec = (
       job_manifest.setdefault('spec', {})
       .setdefault('template', {})
@@ -64,7 +64,8 @@ def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
     add_tolerations(job_manifest)
     update_gpu_containers(job_manifest)
 
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_str: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_str
 
 
 def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:

xpk/core/workload_decorators/storage_decorator.py

@@ -36,7 +36,8 @@ def decorate_jobset(jobset_manifest_str, storages) -> str:
     job_manifest = job['template']
     add_annotations(job_manifest, storages)
     add_volumes(job_manifest, storage_volumes)
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_result: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_result
 
 
 def add_annotations(job_manifest, storages):

xpk/core/workload_decorators/tcpx_decorator.py

@@ -55,7 +55,8 @@ def decorate_jobset(jobset_manifest_str: str) -> str:
   for job in manifest['spec']['replicatedJobs']:
     job_manifest = job['template']
     job_manifest = decorate_job(job_manifest)
-  return yaml.dump(manifest, sort_keys=False)
+  yaml_str: str = yaml.dump(manifest, sort_keys=False)
+  return yaml_str
 
 
 def get_interfaces_annotation() -> dict:
@@ -131,6 +132,7 @@ def add_volumes(job_manifest: dict):
   })
   volumes.append({'name': 'sys', 'hostPath': {'path': '/sys'}})
   volumes.append({'name': 'proc-sys', 'hostPath': {'path': '/proc/sys'}})
+  volumes.append({'name': 'tcpx-socket', 'hostPath': {'path': '/run/tcpx'}})
   volumes.append(
       {'name': 'dshm', 'emptyDir': {'medium': 'Memory', 'sizeLimit': '128Gi'}}
   )
@@ -168,7 +170,7 @@ def add_tcpx_daemon_container(job_manifest):
   spec['initContainers'].append(tcpxo_daemon_container)
 
 
-def update_gpu_containers(job_manifest):
+def update_gpu_containers(job_manifest) -> None:
   for container in job_manifest['spec']['template']['spec']['containers']:
     if 'nvidia.com/gpu' in container.get('resources', {}).get('limits', {}):
       env: list = container.setdefault('env', [])

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)