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

xpk/core/system_characteristics.py

@@ -80,15 +80,6 @@ class GpuConfig:
 
   requires_topology: bool
   gpu_direct_name: Literal['fastrak', 'rdma', 'tcpx', 'tcpxo'] = 'fastrak'
-  kjob_decorator_fn: Optional[Callable[[dict], dict]] = None
-  """A function to decorate the kjob template for GPU-specific configurations.
-
-  Args:
-    job_manifest (dict): The kjob manifest as a dictionary.
-
-  Returns:
-    dict: The modified kjob manifest as a dictionary.
-  """
   nccl_installer: Optional[str] = None
   jobset_decorator_fn: Optional[Callable[[str, list[str]], str]] = None
   """A function to decorate the jobset for GPU-specific configurations.
@@ -106,7 +97,7 @@ class GpuConfig:
     parts = []
     for f in dataclasses.fields(self):
       value = getattr(self, f.name)
-      if f.name in ('kjob_decorator_fn', 'jobset_decorator_fn') and value:
+      if f.name in ('jobset_decorator_fn') and value:
         parts.append(f'{f.name}=<function {value.__name__}>')
       else:
         parts.append(f'{f.name}={repr(value)}')
@@ -151,6 +142,7 @@ class SystemCharacteristics:
   device_type: str
   supports_sub_slicing: bool
   supports_super_slicing: bool
+  supports_accelerator_network_profile: bool
   docker_platform: DockerPlatform
   requires_workload_policy: bool = False
   gpu_config: Optional[GpuConfig] = None
@@ -242,6 +234,7 @@ def get_tpu_system_characteristics_map(
     machine_type: str,
     supported_topologies: list[str],
     docker_platform: DockerPlatform,
+    supports_accelerator_network_profile: bool,
     tpu_type_requires_workload_policy: bool = False,
     default_topologies: set[str] | None = None,
     sub_slicing_topologies: set[str] | None = None,
@@ -268,6 +261,7 @@ def get_tpu_system_characteristics_map(
         and vms_per_slice > 1,
         supports_sub_slicing=topology in sub_slicing_topologies,
         supports_super_slicing=topology in super_slicing_topologies,
+        supports_accelerator_network_profile=supports_accelerator_network_profile,
         docker_platform=docker_platform,
     )
     system_characteristics_map[f'{prefix}-{topology}'] = system
@@ -312,6 +306,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='l4-1',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -325,6 +320,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='l4-2',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -338,6 +334,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='l4-4',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -351,6 +348,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='l4-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -365,6 +363,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='a100-40gb-1',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -378,6 +377,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='a100-40gb-2',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -391,6 +391,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='a100-40gb-4',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -404,6 +405,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='a100-40gb-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -417,10 +419,10 @@ UserFacingNameToSystemCharacteristics = {
         device_type='gb200-4',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=True,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA_A4X,
-            kjob_decorator_fn=rdma_decorator.decorate_kjob_template,
             jobset_decorator_fn=rdma_decorator.decorate_jobset,
             gpu_direct_name='rdma',
         ),
@@ -436,10 +438,10 @@ UserFacingNameToSystemCharacteristics = {
         device_type='gb200-4',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=True,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA_A4X,
-            kjob_decorator_fn=rdma_decorator.decorate_kjob_template,
             jobset_decorator_fn=rdma_decorator.decorate_jobset,
             gpu_direct_name='rdma',
         ),
@@ -455,10 +457,10 @@ UserFacingNameToSystemCharacteristics = {
         device_type='b200-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=True,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA,
-            kjob_decorator_fn=rdma_decorator.decorate_kjob_template,
             jobset_decorator_fn=rdma_decorator.decorate_jobset,
             gpu_direct_name='rdma',
         ),
@@ -474,10 +476,10 @@ UserFacingNameToSystemCharacteristics = {
         device_type='h200-141gb-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=True,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA,
-            kjob_decorator_fn=rdma_decorator.decorate_kjob_template,
             jobset_decorator_fn=rdma_decorator.decorate_jobset,
             gpu_direct_name='rdma',
         ),
@@ -494,10 +496,10 @@ UserFacingNameToSystemCharacteristics = {
         device_type='h100-80gb-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=True,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_TCPX,
-            kjob_decorator_fn=tcpx_decorator.decorate_kjob_template,
             jobset_decorator_fn=tcpx_decorator.decorate_jobset,
             gpu_direct_name='tcpx',
         ),
@@ -514,10 +516,10 @@ UserFacingNameToSystemCharacteristics = {
         device_type='h100-mega-80gb-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=True,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_TCPXO,
-            kjob_decorator_fn=tcpxo_decorator.decorate_kjob_template,
             jobset_decorator_fn=tcpxo_decorator.decorate_jobset,
             gpu_direct_name='tcpxo',
         ),
@@ -531,6 +533,7 @@ UserFacingNameToSystemCharacteristics = {
         machine_type='tpu7x-standard-1t',
         supported_topologies=['1x1x1'],
         tpu_type_requires_workload_policy=True,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     **get_tpu_system_characteristics_map(
@@ -539,6 +542,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='tpu7x',
         machine_type='tpu7x-standard-4t',
         tpu_type_requires_workload_policy=True,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
         supported_topologies=generate_tpu_topologies(max_cubes=144),
         super_slicing_topologies=set(['4x4x4']),
@@ -650,6 +654,7 @@ UserFacingNameToSystemCharacteristics = {
         machine_type='ct6e-standard-1t',
         supported_topologies=['1x1'],
         docker_platform=AMD_PLATFORM,
+        supports_accelerator_network_profile=True,
     ),
     **get_tpu_system_characteristics_map(
         prefix='v6e',
@@ -659,6 +664,7 @@ UserFacingNameToSystemCharacteristics = {
         supported_topologies=['2x2'] + SUB_SLICING_TOPOLOGIES,
         sub_slicing_topologies=set(SUB_SLICING_TOPOLOGIES),
         docker_platform=AMD_PLATFORM,
+        supports_accelerator_network_profile=True,
     ),
     **get_tpu_system_characteristics_map(
         prefix='v5p',
@@ -667,6 +673,7 @@ UserFacingNameToSystemCharacteristics = {
         machine_type='ct5p-hightpu-4t',
         docker_platform=AMD_PLATFORM,
         supported_topologies=generate_tpu_topologies(max_cubes=140),
+        supports_accelerator_network_profile=False,
         default_topologies=set([
             '2x2x1',
             '2x2x2',
@@ -773,6 +780,7 @@ UserFacingNameToSystemCharacteristics = {
         machine_type='ct5lp-hightpu-4t',
         docker_platform=AMD_PLATFORM,
         supported_topologies=['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'],
+        supports_accelerator_network_profile=False,
     ),
     **get_tpu_system_characteristics_map(
         prefix='v4',
@@ -783,6 +791,7 @@ UserFacingNameToSystemCharacteristics = {
         supported_topologies=generate_tpu_topologies(
             max_cubes=64, enforce_nondecreasing=False
         ),
+        supports_accelerator_network_profile=False,
         default_topologies=set([
             '2x2x1',
             '2x2x2',
@@ -811,6 +820,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='m1-megamem-96-1',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     # n2-standard-#vCPUs-#VMs
@@ -824,6 +834,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-64-1',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-1': SystemCharacteristics(
@@ -836,6 +847,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-1',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-2': SystemCharacteristics(
@@ -848,6 +860,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-2',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-4': SystemCharacteristics(
@@ -860,6 +873,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-4',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-8': SystemCharacteristics(
@@ -872,6 +886,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-8',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-16': SystemCharacteristics(
@@ -884,6 +899,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-16',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-32': SystemCharacteristics(
@@ -896,6 +912,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-32',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-64': SystemCharacteristics(
@@ -908,6 +925,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-64',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-128': SystemCharacteristics(
@@ -920,6 +938,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-128',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-256': SystemCharacteristics(
@@ -932,6 +951,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-256',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-512': SystemCharacteristics(
@@ -944,6 +964,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-512',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-1024': SystemCharacteristics(
@@ -956,6 +977,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-1024',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-2048': SystemCharacteristics(
@@ -968,6 +990,7 @@ UserFacingNameToSystemCharacteristics = {
         device_type='n2-standard-32-2048',
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
     ),
 }

xpk/core/system_characteristics_test.py

@@ -34,6 +34,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol
       supported_topologies=["1x1"],
       docker_platform=DockerPlatform.AMD,
       tpu_type_requires_workload_policy=False,
+      supports_accelerator_network_profile=False,
   )
 
   expected_system_characteristics = SystemCharacteristics(
@@ -48,6 +49,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol
       supports_super_slicing=False,
       docker_platform=DockerPlatform.AMD,
       requires_workload_policy=False,
+      supports_accelerator_network_profile=False,
   )
   assert result == {
       "test-1": expected_system_characteristics,
@@ -62,6 +64,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["2x2"],
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       tpu_type_requires_workload_policy=True,
   )
@@ -76,6 +79,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol
       device_type="test-8",
       supports_sub_slicing=False,
       supports_super_slicing=False,
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       requires_workload_policy=False,
   )
@@ -94,6 +98,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top
       supported_topologies=["2x2x2"],
       docker_platform=DockerPlatform.AMD,
       tpu_type_requires_workload_policy=True,
+      supports_accelerator_network_profile=False,
   )
 
   expected_system_characteristics = SystemCharacteristics(
@@ -106,6 +111,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top
       device_type="test-16",
       supports_sub_slicing=False,
       supports_super_slicing=False,
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       requires_workload_policy=True,
   )
@@ -122,6 +128,7 @@ def test_get_tpu_system_characteristics_map_sets_sub_slicing_support():
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["4x4x4", "4x4x8", "4x4x16"],
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       sub_slicing_topologies=set(["4x4x8", "4x4x16"]),
   )
@@ -138,6 +145,7 @@ def test_get_tpu_system_characteristics_map_sets_super_slicing_support():
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["4x4x4", "4x4x8", "4x4x16"],
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       super_slicing_topologies=set(["4x4x8", "4x4x16"]),
   )
@@ -154,6 +162,7 @@ def test_get_tpu_system_characteristics_map_prefers_default_topologies():
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["4x4x4", "4x4x32", "4x8x16", "8x8x8"],
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       default_topologies=set(["4x8x16"]),
   )
@@ -206,6 +215,7 @@ def test_system_characteristics_post_init_sets_workload_policy_for_gpu():
       device_type="l4-1",
       supports_sub_slicing=False,
       supports_super_slicing=False,
+      supports_accelerator_network_profile=False,
       docker_platform=DockerPlatform.AMD,
       gpu_config=GpuConfig(requires_topology=False),
   )
@@ -225,5 +235,6 @@ def test_system_characteristics_post_init_throws_for_gpu_without_config():
         device_type="l4-1",
         supports_sub_slicing=False,
         supports_super_slicing=False,
+        supports_accelerator_network_profile=False,
         docker_platform=DockerPlatform.AMD,
     )

xpk/core/workload_decorators/rdma_decorator.py

@@ -18,21 +18,6 @@ import yaml
 from ...utils.yaml import literal_string
 
 
-def decorate_kjob_template(job_manifest: dict) -> dict:
-  spec = (
-      job_manifest.setdefault('spec', {})
-      .setdefault('template', {})
-      .setdefault('spec', {})
-  )
-  spec.setdefault('tolerations', [])
-  spec.setdefault('volumes', [])
-
-  add_volumes(job_manifest)
-  add_tolerations(job_manifest)
-  update_gpu_containers(job_manifest)
-  return job_manifest
-
-
 def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
   """
   Decorates a JobSet manifest with the necessary components for rdma-daemon.

xpk/core/workload_decorators/tcpx_decorator.py

@@ -22,14 +22,6 @@ from ...utils.yaml import literal_string
 tcpx = 'v2.0.11'
 
 
-def decorate_kjob_template(job_manifest: dict) -> dict:
-  add_volumes(job_manifest)
-  add_tolerations(job_manifest)
-  add_tcpx_daemon_container(job_manifest)
-  update_gpu_containers(job_manifest)
-  return job_manifest
-
-
 def decorate_job(job_manifest: dict) -> dict:
   add_annotations(job_manifest)
   add_volumes(job_manifest)

xpk/core/workload_decorators/tcpx_decorator_test.py

@@ -47,24 +47,6 @@ spec:
                 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": {
@@ -205,63 +187,3 @@ def test_decorate_job():
   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

@@ -22,22 +22,6 @@ from ...utils.yaml import literal_string
 rxdm = 'v1.0.12'
 
 
-def decorate_kjob_template(job_manifest: dict) -> dict:
-  spec = (
-      job_manifest.setdefault('spec', {})
-      .setdefault('template', {})
-      .setdefault('spec', {})
-  )
-  spec.setdefault('tolerations', [])
-  spec.setdefault('volumes', [])
-
-  add_volumes(job_manifest)
-  add_tolerations(job_manifest)
-  add_tcpxo_daemon_container(job_manifest)
-  update_gpu_containers(job_manifest)
-  return job_manifest
-
-
 def decorate_job(job_manifest: dict, sub_networks: list[str]) -> dict:
   job_manifest.setdefault('spec', {}).setdefault('template', {}).setdefault(
       'metadata', {}

xpk/parser/common.py

@@ -144,23 +144,6 @@ def add_cluster_arguments(
   )
 
 
-def add_kind_cluster_arguments(
-    custom_parser_or_group: ParserOrArgumentGroup,
-) -> None:
-  """Add kind cluster arguments to the parser or argument group.
-
-  Args:
-    custom_parser_or_group: parser or argument group to add shared arguments to.
-  """
-  custom_parser_or_group.add_argument(
-      '--kind-cluster',
-      type=bool,
-      action=argparse.BooleanOptionalAction,
-      default=False,
-      help='Apply command to a local test cluster.',
-  )
-
-
 def add_global_arguments(custom_parser_or_group: ParserOrArgumentGroup):
   """Add global - no cloud dependent -  arguments to the parser.
 

xpk/parser/core.py

@@ -23,13 +23,8 @@ from .cluster import set_cluster_parser
 from .inspector import set_inspector_parser
 from .storage import set_storage_parser
 from .workload import set_workload_parsers
-from .batch import set_batch_parser
-from .job import set_job_parser
 from .info import set_info_parser
-from .kind import set_kind_parser
-from .shell import set_shell_parser
 from .version import set_version_parser
-from .run import set_run_parser
 
 
 def set_parser(parser: argparse.ArgumentParser):
@@ -54,20 +49,6 @@ def set_parser(parser: argparse.ArgumentParser):
       "info",
       help="Commands around listing kueue clusterqueues and localqueues.",
   )
-  batch_parser = xpk_subcommands.add_parser(
-      "batch",
-      help="commands around running batch job",
-  )
-  job_parser = xpk_subcommands.add_parser(
-      "job", help="commands around listing, cancelling and investigating jobs"
-  )
-  kind_parser = xpk_subcommands.add_parser(
-      "kind",
-      help="commands around Kind cluster management",
-  )
-  shell_parser = xpk_subcommands.add_parser(
-      "shell", help="Commands around configuring and using interactive shell."
-  )
   version_parser = xpk_subcommands.add_parser(
       "version", help="Command to get xpk version"
   )
@@ -76,11 +57,6 @@ def set_parser(parser: argparse.ArgumentParser):
       "config", help="Commands to set and retrieve values from xpk config."
   )
 
-  run_parser = xpk_subcommands.add_parser(
-      "run",
-      help="Command to run parallel jobs",
-  )
-
   def default_subcommand_function(
       _args,
   ) -> int:  # args is unused, so pylint: disable=invalid-name
@@ -96,14 +72,9 @@ def set_parser(parser: argparse.ArgumentParser):
     parser.print_help()
     cluster_parser.print_help()
     workload_parser.print_help()
-    batch_parser.print_help()
     info_parser.print_help()
-    job_parser.print_help()
-    shell_parser.print_help()
     version_parser.print_help()
-    kind_parser.print_help()
     config_parser.print_help()
-    run_parser.print_help()
 
     storage_parser.print_help()
     return 0
@@ -111,25 +82,15 @@ def set_parser(parser: argparse.ArgumentParser):
   parser.set_defaults(func=default_subcommand_function)
   workload_parser.set_defaults(func=default_subcommand_function)
   cluster_parser.set_defaults(func=default_subcommand_function)
-  batch_parser.set_defaults(func=default_subcommand_function)
   info_parser.set_defaults(func=default_subcommand_function)
-  job_parser.set_defaults(func=default_subcommand_function)
-  kind_parser.set_defaults(func=default_subcommand_function)
-  shell_parser.set_defaults(func=default_subcommand_function)
   storage_parser.set_defaults(func=default_subcommand_function)
   version_parser.set_defaults(func=default_subcommand_function)
   config_parser.set_defaults(func=default_subcommand_function)
-  run_parser.set_defaults(func=default_subcommand_function)
 
   set_workload_parsers(workload_parser=workload_parser)
   set_cluster_parser(cluster_parser=cluster_parser)
   set_inspector_parser(inspector_parser=inspector_parser)
-  set_batch_parser(batch_parser=batch_parser)
   set_info_parser(info_parser=info_parser)
-  set_job_parser(job_parser=job_parser)
-  set_kind_parser(kind_parser=kind_parser)
-  set_shell_parser(shell_parser=shell_parser)
   set_storage_parser(storage_parser=storage_parser)
   set_version_parser(version_parser=version_parser)
   set_config_parsers(config_parser=config_parser)
-  set_run_parser(run_parser=run_parser)