xpk 0.14.0__py3-none-any.whl → 0.14.1__py3-none-any.whl

integration/gcluster_test.py

@@ -0,0 +1,107 @@
+"""
+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.
+"""
+
+from xpk.core.docker_manager import DockerManager
+from xpk.core.gcluster_manager import GclusterManager
+from xpk.core.blueprint.blueprint_generator import BlueprintGenerator
+import os
+import pytest
+import shutil
+
+ctk_gcloud_cfg = os.getenv("GCLOUD_CFG_PATH")
+project_id = os.getenv("PROJECT_ID")
+region = os.getenv("REGION")
+zone = os.getenv("ZONE")
+auth_cidr = os.getenv("AUTH_CIDR")
+cluster_name = os.getenv("GKE_ML_TEST_CLUSTER_NAME")
+
+uploads_dir = "uploads"
+
+
+def prepare_test(docker_path: str, bp_path: str) -> None:
+  if not os.path.exists(docker_path):
+    os.makedirs(docker_path)
+  if not os.path.exists(bp_path):
+    os.makedirs(bp_path)
+
+
+@pytest.mark.skip(reason="Credentails not working. Skipping for now")
+def test_create_deployment():
+  assert project_id is not None
+  assert region is not None
+  assert zone is not None
+  assert auth_cidr is not None
+  assert ctk_gcloud_cfg is not None
+  assert cluster_name is not None
+
+  pwd = os.getcwd()
+  test_docker_working_dir = os.path.join(
+      pwd, "xpkclusters/tests/xpk_test_docker_dir"
+  )
+  test_bp_dir = os.path.join(pwd, "xpkclusters/tests/xpk_test_bp_dir")
+  prepare_test(test_docker_working_dir, test_bp_dir)
+  blueprint_name = "my-test-blueprint"
+  prefix = "prefix"
+
+  docker_manager = DockerManager(
+      gcloud_cfg_path=ctk_gcloud_cfg, working_dir=test_docker_working_dir
+  )
+  docker_manager.initialize()
+
+  bpm = BlueprintGenerator(storage_path=test_bp_dir)
+  ml_gke_blueprint = bpm.generate_gke_ml_blueprint(
+      cluster_name=cluster_name,
+      blueprint_name=blueprint_name,
+      prefix=prefix,
+      region=region,
+      project_id=project_id,
+      auth_cidr=auth_cidr,
+  )
+  blueprint_test_path = os.path.join(
+      test_bp_dir, prefix, f"{blueprint_name}.yaml"
+  )
+  # there are no files in ghcp stage for this blueprint
+  blueprint_deps_test_path = ""
+
+  assert ml_gke_blueprint.blueprint_file == blueprint_test_path
+  assert ml_gke_blueprint.blueprint_dependencies == blueprint_deps_test_path
+
+  assert os.path.exists(blueprint_test_path)
+
+  gcluster_manager = GclusterManager(
+      gcluster_command_runner=docker_manager, remote_state_client=None
+  )
+
+  staged_bp_path = gcluster_manager.stage_files(
+      blueprint_file=ml_gke_blueprint.blueprint_file,
+      blueprint_dependencies=ml_gke_blueprint.blueprint_dependencies,
+      prefix=prefix,
+  )
+
+  assert staged_bp_path == os.path.join(
+      "/out", uploads_dir, prefix, f"{blueprint_name}.yaml"
+  )
+
+  gcluster_manager.deploy(
+      blueprint_path=staged_bp_path,
+      deployment_name=blueprint_name,
+      prefix=prefix,
+  )
+  gcluster_manager.destroy_deployment(
+      deployment_name=blueprint_name, prefix=prefix
+  )
+  shutil.rmtree(test_docker_working_dir)
+  shutil.rmtree(test_bp_dir)

xpk/commands/cluster.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 from tabulate import tabulate
 
+from ..utils.feature_flags import FeatureFlags
 from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE
 from ..core.cluster import (
     get_all_clusters_programmatic,
@@ -75,9 +76,9 @@ from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from . import cluster_gcluster
-from .common import set_cluster_command
+from .common import set_cluster_command, validate_sub_slicing_system
 from jinja2 import Environment, FileSystemLoader
-from ..utils.templates import TEMPLATE_PATH
+from ..utils.templates import get_templates_absolute_path
 import shutil
 import os
 
@@ -200,6 +201,11 @@ def cluster_adapt(args) -> None:
   xpk_exit(0)
 
 
+def _validate_cluster_create_args(args, system: SystemCharacteristics):
+  if FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing:
+    validate_sub_slicing_system(system)
+
+
 def cluster_create(args) -> None:
   """Function around cluster creation.
 
@@ -212,12 +218,14 @@ def cluster_create(args) -> None:
         SystemDependency.KJOB,
         SystemDependency.GCLOUD,
     ])
-  system, return_code = get_system_characteristics(args)
 
+  system, return_code = get_system_characteristics(args)
   if return_code > 0 or system is None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
+  _validate_cluster_create_args(args, system)
+
   xpk_print(f'Starting cluster create for cluster {args.cluster}:', flush=True)
   add_zone_and_project(args)
 
@@ -426,7 +434,9 @@ def cluster_cacheimage(args) -> None:
       system.accelerator_type
   ].accelerator_label
 
-  template_env = Environment(loader=FileSystemLoader(TEMPLATE_PATH))
+  template_env = Environment(
+      loader=FileSystemLoader(searchpath=get_templates_absolute_path())
+  )
   cluster_preheat_yaml = template_env.get_template(CLUSTER_PREHEAT_JINJA_FILE)
   rendered_yaml = cluster_preheat_yaml.render(
       cachekey=args.cache_key,
@@ -1251,6 +1261,9 @@ def install_kueue(args, system: SystemCharacteristics, autoprovisioning_config):
           memory_limit=args.memory_limit,
           cpu_limit=args.cpu_limit,
           is_pathways_cluster=args.enable_pathways,
+          configure_sub_slicing=(
+              FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing
+          ),
       ),
   )
 

xpk/commands/cluster_gcluster.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 import os
 
+from ..utils.feature_flags import FeatureFlags
 from ..utils.execution_context import is_dry_run
 from ..core.kueue_manager import KueueConfig, KueueManager
 from ..core.nap import enable_autoprovisioning_on_cluster
@@ -159,6 +160,9 @@ def __install_kueue(args) -> int:
           cpu_limit=args.cpu_limit,
           is_pathways_cluster=args.enable_pathways,
           flex=args.flex,
+          configure_sub_slicing=(
+              FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing
+          ),
       ),
       tolerations=tolerations,
   )

xpk/commands/cluster_test.py

@@ -0,0 +1,92 @@
+"""
+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 argparse import Namespace
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+import pytest
+
+from xpk.commands.cluster import _validate_cluster_create_args
+from xpk.core.system_characteristics import SystemCharacteristics, UserFacingNameToSystemCharacteristics
+from xpk.utils.feature_flags import FeatureFlags
+
+
+@dataclass
+class _Mocks:
+  common_print_mock: MagicMock
+  common_exit_mock: MagicMock
+
+
+@pytest.fixture
+def mock_common_print_and_exit(mocker):
+  common_print_mock = mocker.patch(
+      'xpk.commands.common.xpk_print',
+      return_value=None,
+  )
+  common_exit_mock = mocker.patch(
+      'xpk.commands.common.xpk_exit',
+      return_value=None,
+  )
+  return _Mocks(
+      common_print_mock=common_print_mock, common_exit_mock=common_exit_mock
+  )
+
+
+DEFAULT_TEST_SYSTEM: SystemCharacteristics = (
+    UserFacingNameToSystemCharacteristics['l4-1']
+)
+SUB_SLICING_SYSTEM: SystemCharacteristics = (
+    UserFacingNameToSystemCharacteristics['v6e-4x4']
+)
+
+
+def test_validate_cluster_create_args_for_correct_args_pass(
+    mock_common_print_and_exit: _Mocks,
+):
+  args = Namespace()
+
+  _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
+
+  assert mock_common_print_and_exit.common_print_mock.call_count == 0
+  assert mock_common_print_and_exit.common_exit_mock.call_count == 0
+
+
+def test_validate_cluster_create_args_for_correct_sub_slicing_args_pass(
+    mock_common_print_and_exit: _Mocks,
+):
+  FeatureFlags.SUB_SLICING_ENABLED = True
+  args = Namespace(sub_slicing=True)
+
+  _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
+
+  assert mock_common_print_and_exit.common_print_mock.call_count == 0
+  assert mock_common_print_and_exit.common_exit_mock.call_count == 0
+
+
+def test_validate_cluster_create_args_for_not_supported_system_throws(
+    mock_common_print_and_exit: _Mocks,
+):
+  FeatureFlags.SUB_SLICING_ENABLED = True
+  args = Namespace(sub_slicing=True)
+
+  _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
+
+  assert mock_common_print_and_exit.common_print_mock.call_count == 1
+  assert (
+      mock_common_print_and_exit.common_print_mock.call_args[0][0]
+      == 'Error: l4-1 does not support Sub-slicing.'
+  )
+  assert mock_common_print_and_exit.common_exit_mock.call_count == 1

xpk/commands/common.py

@@ -67,3 +67,9 @@ def is_TAS_possible(
       system_characteristics.device_type != H100_MEGA_DEVICE_TYPE
       or capacity_type == CapacityType.RESERVATION
   )
+
+
+def validate_sub_slicing_system(system: SystemCharacteristics):
+  if not system.supports_sub_slicing:
+    xpk_print(f'Error: {system.device_type} does not support Sub-slicing.')
+    xpk_exit(1)

xpk/commands/kind.py

@@ -110,6 +110,7 @@ def cluster_create(args) -> None:
           cpu_limit=0,
           is_pathways_cluster=False,
           flex=False,
+          configure_sub_slicing=False,
       ),
   )
 

xpk/commands/workload.py

@@ -52,7 +52,7 @@ from ..core.pathways import (
     get_user_workload_for_pathways,
     try_to_delete_pathwaysjob_first,
 )
-from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
+from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics, SystemCharacteristics
 from ..core.resources import CLUSTER_METADATA_CONFIGMAP, get_cluster_configmap
 from ..core.scheduling import (
     check_if_workload_can_schedule,
@@ -78,6 +78,7 @@ from ..core.storage import (
 from ..core.system_characteristics import (
     AcceleratorType,
     get_system_characteristics,
+    compute_vms_per_slice,
 )
 from ..core.vertex import create_vertex_experiment
 from ..core.workload import (
@@ -98,7 +99,8 @@ from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from . import cluster_gcluster
-from .common import is_TAS_possible
+from .common import is_TAS_possible, validate_sub_slicing_system
+from ..utils.topology import is_topology_contained
 from ..utils.feature_flags import FeatureFlags
 
 WORKLOAD_CREATE_YAML = """apiVersion: jobset.x-k8s.io/v1alpha2
@@ -120,8 +122,8 @@ spec:
       replicas: {args.num_slices}
       template:
         spec:
-          parallelism: {system.vms_per_slice}    # Equal to the number of VMs per slice
-          completions: {system.vms_per_slice}    # Same as the above.
+          parallelism: {vms_per_slice}    # Equal to the number of VMs per slice (or sub-slice).
+          completions: {vms_per_slice}    # Same as the above.
           backoffLimit: 0   # When any pod fails, the job is failed
           {pod_failure_policy}
           template:
@@ -280,6 +282,8 @@ PW_WORKLOAD_CREATE_YAML = """
       {user_workload}
 """
 
+SUB_SLICING_TOPOLOGIES = ['2x2', '2x4', '4x4', '4x8', '8x8', '8x16', '16x16']
+
 
 def workload_create_pathways(args) -> None:
   """Run jobset apply command for a file, specifically for Pathways.
@@ -330,13 +334,14 @@ def workload_create(args) -> None:
     )
     xpk_exit(1)
 
-  xpk_print('Starting workload create', flush=True)
   system, return_code = get_system_characteristics(args)
-
   if return_code > 0 or system is None:
     xpk_print('Fetching system characteristics failed!')
     xpk_exit(return_code)
 
+  if FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing_topology is not None:
+    _validate_sub_slicing_topology(system, args.sub_slicing_topology)
+
   if not check_if_workload_can_schedule(args, system):
     xpk_exit(1)
 
@@ -558,8 +563,14 @@ def workload_create(args) -> None:
     )
     yml_string = WORKLOAD_CREATE_YAML.format(
         args=args,
-        system=system,
         container=container,
+        vms_per_slice=(
+            compute_vms_per_slice(args.sub_slicing_topology)
+            if system.accelerator_type == AcceleratorType['TPU']
+            and FeatureFlags.SUB_SLICING_ENABLED
+            and args.sub_slicing_topology is not None
+            else system.vms_per_slice
+        ),
         affinity=get_cpu_affinity(system.accelerator_type),
         accelerator_label=create_accelerator_label(
             system.accelerator_type, system
@@ -667,6 +678,29 @@ def workload_create(args) -> None:
   xpk_exit(0)
 
 
+def _validate_sub_slicing_topology(
+    system_characteristics: SystemCharacteristics, sub_slicing_topology: str
+) -> None:
+  if sub_slicing_topology not in SUB_SLICING_TOPOLOGIES:
+    xpk_print(
+        f'Error: --sub-slicing-topology={sub_slicing_topology} shape is'
+        f' invalid. It has to be one of: {", ".join(SUB_SLICING_TOPOLOGIES)}.'
+    )
+    xpk_exit(1)
+
+  if not is_topology_contained(
+      contained=sub_slicing_topology, container=system_characteristics.topology
+  ):
+    xpk_print(
+        f'Error: --sub-slicing-topology={sub_slicing_topology} shape is too'
+        ' large. The shape cannot be bigger than'
+        f' {system_characteristics.topology}.'
+    )
+    xpk_exit(1)
+
+  validate_sub_slicing_system(system_characteristics)
+
+
 def get_restart_exit_codes(args) -> list:
   exit_codes = [42]
   exit_codes.extend(range(127, 256, 1))

xpk/commands/workload_test.py

@@ -0,0 +1,81 @@
+"""
+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 dataclasses
+from unittest.mock import MagicMock, patch
+import pytest
+from ..core.system_characteristics import SystemCharacteristics
+from .workload import _validate_sub_slicing_topology
+
+
+SYSTEM_CHARACTERISTICS = SystemCharacteristics(
+    topology='8x8',
+    vms_per_slice=1,
+    gke_accelerator='nvidia-l4',
+    gce_machine_type='g2-standard-12',
+    chips_per_vm=1,
+    accelerator_type=1,
+    device_type='l4-1',
+    supports_sub_slicing=True,
+    requires_workload_policy=False,
+)
+
+
+@pytest.fixture(autouse=True)
+def xpk_print(mocker):
+  return mocker.patch('xpk.commands.workload.xpk_print')
+
+
+def test_validate_sub_slicing_topology_exits_for_unsupported_topology(
+    xpk_print,
+):
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '2x1')
+
+  assert (
+      'shape is invalid. It has to be one of' in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_validate_sub_slicing_topology_exits_for_too_large_topology(xpk_print):
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '16x16')
+
+  assert (
+      'shape is too large. The shape cannot be'
+      in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_validate_sub_slicing_topology_does_nothing_for_supported_topology():
+  _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '4x4')
+
+
+@patch('xpk.commands.common.xpk_print')
+def test_validate_sub_slicing_topology_fails_for_unsupported_system(
+    common_xpk_print: MagicMock,
+):
+  unsupported_system = dataclasses.replace(
+      SYSTEM_CHARACTERISTICS, supports_sub_slicing=False
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_topology(unsupported_system, '4x4')
+
+  assert (
+      'l4-1 does not support Sub-slicing.'
+      in common_xpk_print.mock_calls[0].args[0]
+  )

xpk/core/blueprint/testing/__init__.py

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

xpk/core/config.py

@@ -22,7 +22,7 @@ from ..utils import file
 from ..utils.console import xpk_print
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.14.0'
+__version__ = 'v0.14.1'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 

xpk/core/kueue_manager.py

@@ -39,14 +39,16 @@ from ..core.commands import (
 )
 from ..utils.file import write_tmp_file
 from ..utils.console import xpk_print, xpk_exit
-from ..utils.templates import TEMPLATE_PATH
+from ..utils.templates import TEMPLATE_PATH, get_templates_absolute_path
 
 WAIT_FOR_KUEUE_TIMEOUT = "10m"
 CLUSTER_QUEUE_NAME = "cluster-queue"
 LOCAL_QUEUE_NAME = "multislice-queue"
+SUB_SLICE_TOPOLOGY_NAME = "sub-slice-topology"
 KUEUE_CONFIG_JINJA_FILE = "kueue_config.yaml.j2"
-KUEUE_TOPOLOGY_JINJA_FILE = "kueue_topology.yaml.j2"
+KUEUE_GKE_DEFAULT_TOPOLOGY_JINJA_FILE = "kueue_gke_default_topology.yaml.j2"
 KUEUE_CONTROLLER_MANAGER_JINJA_FILE = "kueue_controller_manager.yaml.j2"
+KUEUE_SUB_SLICING_TOPOLOGY_JINJA_FILE = "kueue_sub_slicing_topology.yaml.j2"
 MEMORY_SIZE_PER_VM = 1.2
 MIN_MEMORY_LIMIT_SIZE = 4096
 KUEUE_VERSION = "v0.14.1"
@@ -58,12 +60,19 @@ class KueueConfig:
   total_chips: int
   cpu_limit: int
   memory_limit: str
+  configure_sub_slicing: bool
   is_pathways_cluster: bool = False
   autoprovisioning_enabled: bool = False
   flex: bool = False
   num_slices: int = 1
 
 
+@dataclass
+class _NameAndYaml:
+  name: str
+  yaml: str
+
+
 class KueueManager:
   """Manages the installation and configuration of Kueue on an XPK cluster."""
 
@@ -73,7 +82,12 @@ class KueueManager:
       template_path=TEMPLATE_PATH,
   ):
     self.kueue_version = kueue_version
-    self.template_env = Environment(loader=FileSystemLoader(template_path))
+
+    self.template_env = Environment(
+        loader=FileSystemLoader(
+            searchpath=get_templates_absolute_path(template_path)
+        )
+    )
 
   def install_or_upgrade(
       self,
@@ -87,7 +101,7 @@ class KueueManager:
     Args:
         tolerations: An optional list of tolerations to apply to the kueue-controller-manager.
     """
-    return_code, installed_version = self.__get_installed_kueue_version()
+    return_code, installed_version = self.get_installed_kueue_version()
 
     if return_code == 0:
       if installed_version and installed_version > self.kueue_version:
@@ -107,7 +121,7 @@ class KueueManager:
 
     return self.__configure(kueue_config)
 
-  def __get_installed_kueue_version(self) -> tuple[int, str | None]:
+  def get_installed_kueue_version(self) -> tuple[int, str | None]:
     command = (
         "kubectl get deployment kueue-controller-manager -n kueue-system -o"
         " jsonpath='{.spec.template.spec.containers[0].image}'"
@@ -208,6 +222,13 @@ class KueueManager:
     """
     template = self.template_env.get_template(KUEUE_CONFIG_JINJA_FILE)
 
+    topology_name_and_yaml = self.__get_topology_name_and_yaml(
+        kueue_config.system, kueue_config.configure_sub_slicing
+    )
+    topology_name = (
+        topology_name_and_yaml.name if topology_name_and_yaml else None
+    )
+
     # The manager builds the context internally based on its opinionated logic
     context = self.__build_template_context(
         system=kueue_config.system,
@@ -218,18 +239,16 @@ class KueueManager:
         num_slices=kueue_config.num_slices,
         cpu_limit=kueue_config.cpu_limit,
         memory_limit=kueue_config.memory_limit,
+        topology_name=topology_name,
     )
 
-    rendered_manifest = template.render(context)
+    config_yaml = template.render(context)
+    yamls = [config_yaml]
 
-    if kueue_config.system.device_type in [
-        H100_MEGA_DEVICE_TYPE,
-        H200_DEVICE_TYPE,
-        B200_DEVICE_TYPE,
-    ]:
-      topology_yaml = self.template_env.get_template(KUEUE_TOPOLOGY_JINJA_FILE)
-      rendered_manifest = topology_yaml.render() + rendered_manifest
+    if topology_name_and_yaml:
+      yamls.append(topology_name_and_yaml.yaml)
 
+    rendered_manifest = "\n---\n".join(yamls)
     return_code = self.__apply_manifest(rendered_manifest)
     if return_code != 0:
       return return_code
@@ -246,6 +265,7 @@ class KueueManager:
       num_slices: int,
       cpu_limit: int,
       memory_limit: str,
+      topology_name: str | None,
   ) -> Dict[str, Any]:
     """Prepares the context for the Jinja2 template."""
     # Main accelerator flavor
@@ -267,13 +287,7 @@ class KueueManager:
       key, value = machine_label.split(":", 1)
       node_labels_dict[key] = value.strip()
 
-    topology_label = ""
-    if system.device_type in [
-        H100_MEGA_DEVICE_TYPE,
-        H200_DEVICE_TYPE,
-        B200_DEVICE_TYPE,
-    ]:
-      topology_label = 'topologyName: "gke-default"'
+    topology_label = f"topologyName: {topology_name}" if topology_name else ""
 
     flavors = [{
         "name": main_flavor_name,
@@ -335,6 +349,32 @@ class KueueManager:
         "admission_checks": admission_checks,
     }
 
+  def __get_topology_name_and_yaml(
+      self, system: SystemCharacteristics, configure_sub_slicing: bool
+  ) -> _NameAndYaml | None:
+    if system.device_type in [
+        H100_MEGA_DEVICE_TYPE,
+        H200_DEVICE_TYPE,
+        B200_DEVICE_TYPE,
+    ]:
+      return _NameAndYaml(
+          name="gke-default",
+          yaml=self.template_env.get_template(
+              KUEUE_GKE_DEFAULT_TOPOLOGY_JINJA_FILE
+          ).render(),
+      )
+    elif configure_sub_slicing:
+      return _NameAndYaml(
+          name=SUB_SLICE_TOPOLOGY_NAME,
+          yaml=self.template_env.get_template(
+              KUEUE_SUB_SLICING_TOPOLOGY_JINJA_FILE
+          ).render({
+              "sub_slice_topology_name": SUB_SLICE_TOPOLOGY_NAME,
+          }),
+      )
+    else:
+      return None
+
   def __apply_manifest(self, manifest: str) -> int:
     task = "Applying Kueue Custom Resources"
     if is_dry_run():