xpk 0.17.3__py3-none-any.whl → 1.1.0__py3-none-any.whl

xpk/core/system_characteristics.py

@@ -131,6 +131,8 @@ class SystemCharacteristics:
     supports_super_slicing: Whether the Super-slicing feature is supported.
     requires_workload_policy: A boolean indicating if a GCE resource
       workload policy is required. This is automatically set to True for GPUs.
+    parallel_containers: The number of containers running on a single VM.
+
   """
 
   topology: str
@@ -142,9 +144,11 @@ 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
+  parallel_containers: int = 1
 
   def __post_init__(self):
     if self.accelerator_type == AcceleratorType.GPU:
@@ -233,10 +237,12 @@ 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,
     super_slicing_topologies: set[str] | None = None,
+    parallel_containers: int = 1,
 ) -> dict[str, SystemCharacteristics]:
   system_characteristics_map = {}
   default_topologies = default_topologies or set()
@@ -259,7 +265,9 @@ 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,
+        parallel_containers=parallel_containers,
     )
     system_characteristics_map[f'{prefix}-{topology}'] = system
     if (
@@ -303,6 +311,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,
     ),
@@ -316,6 +325,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,
     ),
@@ -329,6 +339,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,
     ),
@@ -342,6 +353,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,
     ),
@@ -356,6 +368,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,
     ),
@@ -369,6 +382,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,
     ),
@@ -382,6 +396,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,
     ),
@@ -395,6 +410,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,
     ),
@@ -408,6 +424,7 @@ 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,
@@ -426,6 +443,7 @@ 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,
@@ -444,6 +462,7 @@ 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,
@@ -462,6 +481,7 @@ 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,
@@ -481,6 +501,7 @@ 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,
@@ -500,6 +521,7 @@ 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,
@@ -516,6 +538,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(
@@ -524,7 +547,9 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='tpu7x',
         machine_type='tpu7x-standard-4t',
         tpu_type_requires_workload_policy=True,
+        supports_accelerator_network_profile=False,
         docker_platform=AMD_PLATFORM,
+        parallel_containers=2,
         supported_topologies=generate_tpu_topologies(max_cubes=144),
         super_slicing_topologies=set(['4x4x4']),
         default_topologies=set([
@@ -635,6 +660,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',
@@ -644,6 +670,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',
@@ -652,6 +679,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',
@@ -758,6 +786,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',
@@ -768,6 +797,7 @@ UserFacingNameToSystemCharacteristics = {
         supported_topologies=generate_tpu_topologies(
             max_cubes=64, enforce_nondecreasing=False
         ),
+        supports_accelerator_network_profile=False,
         default_topologies=set([
             '2x2x1',
             '2x2x2',
@@ -796,6 +826,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
@@ -809,6 +840,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(
@@ -821,6 +853,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(
@@ -833,6 +866,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(
@@ -845,6 +879,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(
@@ -857,6 +892,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(
@@ -869,6 +905,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(
@@ -881,6 +918,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(
@@ -893,6 +931,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(
@@ -905,6 +944,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(
@@ -917,6 +957,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(
@@ -929,6 +970,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(
@@ -941,6 +983,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(
@@ -953,6 +996,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/telemetry.py

@@ -30,7 +30,7 @@ from dataclasses import dataclass
 from .config import get_config, CLIENT_ID_KEY, SEND_TELEMETRY_KEY, __version__ as xpk_version
 from ..utils.execution_context import is_dry_run
 from ..utils.user_agent import get_user_agent
-from ..utils.feature_flags import FeatureFlags
+from ..utils.feature_flags import FeatureFlags, is_tester
 
 
 def should_send_telemetry():
@@ -114,6 +114,8 @@ def _clearcut_flush(file_path: str) -> None:
 
 
 class MetricsEventMetadataKey(Enum):
+  """Represents available metadata keys."""
+
   SESSION_ID = "XPK_SESSION_ID"
   DRY_RUN = "XPK_DRY_RUN"
   PYTHON_VERSION = "XPK_PYTHON_VERSION"
@@ -125,6 +127,7 @@ class MetricsEventMetadataKey(Enum):
   RUNNING_AS_PIP = "XPK_RUNNING_AS_PIP"
   RUNNING_FROM_SOURCE = "XPK_RUNNING_FROM_SOURCE"
   LATENCY_SECONDS = "XPK_LATENCY_SECONDS"
+  TESTER = "XPK_TESTER"
 
 
 @dataclass
@@ -230,6 +233,9 @@ def _get_base_event_metadata() -> dict[MetricsEventMetadataKey, str]:
       MetricsEventMetadataKey.RUNNING_FROM_SOURCE: str(
           _is_running_from_source()
       ).lower(),
+      MetricsEventMetadataKey.TESTER: str(
+          is_tester() or _is_trash_execution()
+      ).lower(),
   }
 
 
@@ -241,6 +247,10 @@ def _get_base_concord_event() -> dict[str, str]:
   }
 
 
+def _is_trash_execution() -> bool:
+  return os.getenv("TELEMETRY_TRASH_EXECUTION") == "true"
+
+
 def _is_running_as_pip() -> bool:
   return os.path.basename(sys.argv[0]) == "xpk"
 

xpk/core/telemetry_test.py

@@ -30,7 +30,9 @@ def setup_mocks(mocker: MockerFixture):
   mocker.patch('time.time', side_effect=itertools.count())
   mocker.patch('platform.python_version', return_value='99.99.99')
   mocker.patch('os.path.basename', return_value='xpk.py')
+  mocker.patch('os.getenv', return_value='false')
   mocker.patch('os.path.abspath', return_value='/home/xpk_user')
+  mocker.patch('xpk.core.telemetry.is_tester', return_value=False)
   set_dry_run(False)
   get_config().set(CLIENT_ID_KEY, 'client_id')
   yield
@@ -76,6 +78,7 @@ def test_metrics_collector_logs_start_event_correctly():
           {'key': 'XPK_PYTHON_VERSION', 'value': '99.99.99'},
           {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
           {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
+          {'key': 'XPK_TESTER', 'value': 'false'},
           {'key': 'XPK_COMMAND', 'value': 'test'},
           {'key': 'XPK_LATENCY_SECONDS', 'value': '0'},
       ],
@@ -107,6 +110,7 @@ def test_metrics_collector_logs_complete_event_correctly():
           {'key': 'XPK_PYTHON_VERSION', 'value': '99.99.99'},
           {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
           {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
+          {'key': 'XPK_TESTER', 'value': 'false'},
           {'key': 'XPK_EXIT_CODE', 'value': '2'},
           {'key': 'XPK_LATENCY_SECONDS', 'value': '0'},
       ],
@@ -131,6 +135,7 @@ def test_metrics_collector_logs_custom_event_correctly():
           {'key': 'XPK_PYTHON_VERSION', 'value': '99.99.99'},
           {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
           {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
+          {'key': 'XPK_TESTER', 'value': 'false'},
           {'key': 'XPK_PROVISIONING_MODE', 'value': 'flex'},
           {'key': 'XPK_LATENCY_SECONDS', 'value': '0'},
       ],
@@ -219,6 +224,40 @@ def test_metrics_collectors_logs_correct_running_from_source_value(
   assert _get_metadata_value(payload, 'XPK_RUNNING_FROM_SOURCE') == expected
 
 
+@pytest.mark.parametrize(
+    argnames='tester,expected',
+    argvalues=[
+        (True, 'true'),
+        (False, 'false'),
+    ],
+)
+def test_metrics_collectors_logs_correct_tester_value_for_is_tester_variable(
+    tester: bool, expected: str, mocker: MockerFixture
+):
+  mocker.patch('xpk.core.telemetry.is_tester', return_value=tester)
+  MetricsCollector.log_start(command='test')
+  payload = MetricsCollector.flush()
+  assert _get_metadata_value(payload, 'XPK_TESTER') == expected
+
+
+@pytest.mark.parametrize(
+    argnames='trash_execution,expected',
+    argvalues=[
+        ('true', 'true'),
+        ('false', 'false'),
+        ('', 'false'),
+        (None, 'false'),
+    ],
+)
+def test_metrics_collectors_logs_correct_tester_value_for_trash_variable(
+    trash_execution: str, expected: str, mocker: MockerFixture
+):
+  mocker.patch('os.getenv', return_value=trash_execution)
+  MetricsCollector.log_start(command='test')
+  payload = MetricsCollector.flush()
+  assert _get_metadata_value(payload, 'XPK_TESTER') == expected
+
+
 def _get_metadata_value(payload_str: str, key: str) -> str | None:
   payload = json.loads(payload_str)
   metadata = json.loads(payload['log_event'][0]['source_extension_json'])[

xpk/core/testing/commands_tester.py

@@ -17,6 +17,8 @@ limitations under the License.
 import re
 from pytest_mock import MockerFixture
 
+from ..commands import FailedCommand
+
 
 class CommandsTester:
   """Tester class useful for mocking and asserting command runs."""
@@ -27,6 +29,7 @@ class CommandsTester:
       run_command_for_value_path: str | None = None,
       run_command_with_updates_path: str | None = None,
       run_command_with_updates_retry_path: str | None = None,
+      run_command_batch_path: str | None = None,
   ):
     self.__results: dict[re.Pattern, tuple[int, str]] = {}
     self.commands_history: list[str] = []
@@ -45,6 +48,11 @@ class CommandsTester:
           run_command_with_updates_retry_path,
           wraps=self.__fake_run_command_with_updates_retry,
       )
+    if run_command_batch_path:
+      mocker.patch(
+          run_command_batch_path,
+          wraps=self.__fake_run_command_batch,
+      )
 
   def set_result_for_command(
       self, result: tuple[int, str], *command_parts: str
@@ -111,6 +119,24 @@ class CommandsTester:
   ) -> tuple[int, str]:
     return self.__common_fake_run_command(command, (0, dry_run_return_val))
 
+  def __fake_run_command_batch(
+      self,
+      commands: list[str],
+      jobname: str,
+      per_command_name: list[str],
+      output_logs: list[str],
+  ) -> FailedCommand | None:
+    for i, command in enumerate(commands):
+      result = self.__common_fake_run_command(command, (0, ""))[0]
+      if result != 0:
+        return FailedCommand(
+            return_code=result,
+            name=per_command_name[i],
+            command=command,
+            logfile=output_logs[i],
+        )
+    return None
+
   # pylint: enable=unused-argument
 
   def __common_fake_run_command(

xpk/core/testing/commands_tester_test.py

@@ -17,7 +17,7 @@ limitations under the License.
 import pytest
 from pytest_mock import MockerFixture
 
-from xpk.core.commands import run_command_for_value, run_command_with_updates_retry
+from xpk.core.commands import run_command_for_value, run_command_with_updates_retry, run_command_batch
 from xpk.core.testing.commands_tester import CommandsTester
 
 
@@ -31,6 +31,9 @@ def mock_commands(mocker: MockerFixture) -> CommandsTester:
       run_command_with_updates_retry_path=(
           "xpk.core.testing.commands_tester_test.run_command_with_updates_retry"
       ),
+      run_command_batch_path=(
+          "xpk.core.testing.commands_tester_test.run_command_batch"
+      ),
   )
 
 
@@ -54,6 +57,22 @@ def test_run_command_with_updates_retry_default_result(
   mock_commands.assert_command_run("cmd", "bar")
 
 
+def test_run_command_batch_default_result(
+    mock_commands: CommandsTester,
+):
+  result = run_command_batch(
+      commands=["cmd1 foo bar", "cmd2 foo bar"],
+      jobname="Test command",
+      per_command_name=["cmd1", "cmd2"],
+      output_logs=["log1", "log2"],
+  )
+
+  assert result is None
+  mock_commands.assert_command_run("foo bar", times=2)
+  mock_commands.assert_command_run("cmd1")
+  mock_commands.assert_command_run("cmd2")
+
+
 def test_set_result_for_command(mock_commands: CommandsTester):
   mock_commands.set_result_for_command((17, "Error!"), "cmd", "--err")
 

xpk/core/workload_decorators/rdma_decorator.py

@@ -84,6 +84,12 @@ def add_volumes(job_manifest):
   volumes.append(
       {'name': 'gib', 'hostPath': {'path': '/home/kubernetes/bin/gib'}}
   )
+  volumes.append({
+      'name': 'dshm',
+      'emptyDir': {
+          'medium': 'Memory',
+      },
+  })
 
 
 def add_tolerations(job_manifest):
@@ -111,3 +117,6 @@ def update_gpu_containers(job_manifest):
       container['volumeMounts'].append(
           {'name': 'gib', 'mountPath': '/usr/local/gib'}
       )
+      container['volumeMounts'].append(
+          {'name': 'dshm', 'mountPath': '/dev/shm'}
+      )

xpk/parser/cluster.py

@@ -338,7 +338,10 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
   add_resource_limits(cluster_create_resource_limits)
 
   cluster_create_ray_parser.set_defaults(
-      func=cluster_create_ray_cluster, sub_slicing=False, super_slicing=False
+      func=cluster_create_ray_cluster,
+      sub_slicing=False,
+      super_slicing=False,
+      num_cubes=None,
   )
 
 
@@ -503,6 +506,13 @@ def set_cluster_adapt_parser(cluster_adapt_parser: ArgumentParser):
   )
   add_driver_arguments(cluster_adapt_optional_arguments)
   add_shared_arguments(cluster_adapt_optional_arguments)
+  add_resource_limits(cluster_adapt_optional_arguments)
+
+  if FeatureFlags.SUB_SLICING_ENABLED:
+    add_cluster_create_sub_slicing_arguments(cluster_adapt_optional_arguments)
+
+  if FeatureFlags.SUPER_SLICING_ENABLED:
+    add_cluster_create_super_slicing_arguments(cluster_adapt_optional_arguments)
 
   cluster_adapt_capacity_arguments = cluster_adapt_parser.add_argument_group(
       'Capacity Arguments', 'Arguments related to capacity for cluster create.'