xpk 0.15.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

xpk/commands/cluster_test.py

@@ -35,6 +35,7 @@ class _Mocks:
   common_print_mock: MagicMock
   commands_print_mock: MagicMock
   commands_get_reservation_deployment_type: MagicMock
+  commands_get_pathways_machine_types: MagicMock
   commands_tester: CommandsTester
 
 
@@ -75,15 +76,23 @@ def mocks(mocker) -> _Mocks:
       'xpk.commands.cluster.get_reservation_deployment_type',
       return_value='DENSE',
   )
+  commands_get_pathways_machine_types = mocker.patch(
+      'xpk.commands.cluster.get_pathways_machine_types',
+      return_value=(0, []),
+  )
   return _Mocks(
       common_print_mock=common_print_mock,
       commands_get_reservation_deployment_type=commands_get_reservation_deployment_type,
       commands_print_mock=commands_print_mock,
+      commands_get_pathways_machine_types=commands_get_pathways_machine_types,
       commands_tester=CommandsTester(
           mocker,
           run_command_with_updates_path=(
               'xpk.commands.cluster.run_command_with_updates'
           ),
+          run_command_for_value_path=(
+              'xpk.commands.cluster.run_command_for_value'
+          ),
       ),
   )
 
@@ -104,6 +113,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       gke_version='',
       private=False,
       authorized_networks=None,
+      pathways_gce_machine_type='n2-standard-64',
       enable_pathways=False,
       enable_ray_cluster=False,
       enable_workload_identity=False,
@@ -121,6 +131,21 @@ def construct_args(**kwargs: Any) -> Namespace:
       cluster_cpu_machine_type='',
       create_vertex_tensorboard=False,
       enable_autoprovisioning=False,
+      sub_slicing_topology='2x2x2',
+      use_vertex_tensorboard=False,
+      env_file='',
+      env=None,
+      use_pathways=False,
+      debug_dump_gcs=False,
+      storage='',
+      restart_on_exit_codes=None,
+      ttl_seconds_after_finished=0,
+      max_restarts=1,
+      priority=0,
+      termination_grace_period_seconds=0,
+      docker_image_pull_secret='',
+      managed_mldiagnostics=False,
+      output_manifest_file='',
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)
@@ -288,6 +313,64 @@ def test_validate_cluster_create_args_for_invalid_reservation(
   )
 
 
+def test_validate_cluster_create_args_for_enable_pathways_set_to_false(
+    mocks: _Mocks,
+):
+  args = construct_args(enable_pathways=False)
+  mocks.commands_get_pathways_machine_types.return_value = (1, [])
+
+  _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 0
+
+
+def test_validate_cluster_create_args_for_errored_pathways_machine_types_retrieval(
+    mocks: _Mocks,
+):
+  args = construct_args(enable_pathways=True)
+  mocks.commands_get_pathways_machine_types.return_value = (1, [])
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 1
+  assert 'Unable to retrieve' in mocks.commands_print_mock.call_args[0][0]
+
+
+def test_validate_cluster_create_args_for_invalid_pathways_machine_type(
+    mocks: _Mocks,
+):
+  args = construct_args(
+      enable_pathways=True, pathways_gce_machine_type='n2-standard-32'
+  )
+  mocks.commands_get_pathways_machine_types.return_value = (
+      0,
+      ['n2-standard-64'],
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 2
+  assert 'Available machine types' in mocks.commands_print_mock.call_args[0][0]
+
+
+def test_validate_cluster_create_args_for_valid_pathways_machine_type(
+    mocks: _Mocks,
+):
+  args = construct_args(
+      enable_pathways=True, pathways_gce_machine_type='n2-standard-32'
+  )
+  mocks.commands_get_pathways_machine_types.return_value = (
+      0,
+      ['n2-standard-32'],
+  )
+
+  _validate_cluster_create_args(args, TPU_TEST_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 0
+
+
 @patch('xpk.commands.cluster.KueueManager.install_or_upgrade')
 def test_install_kueue_returns_kueue_installation_code(
     mock_kueue_manager_install: MagicMock,
@@ -358,6 +441,48 @@ def test_run_gke_cluster_create_command_with_gke_version_has_no_autoupgrade_flag
   )
 
 
+def test_run_gke_cluster_create_command_with_lustre_runs_correct_command(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(
+          enable_lustre_csi_driver=True, enable_legacy_lustre_port=False
+      ),
+      gke_control_plane_version='1.2.3',
+      system=TPU_TEST_SYSTEM,
+      release_channel=ReleaseChannel.REGULAR,
+  )
+
+  assert result == 0
+  commands = mocks.commands_tester.get_matching_commands('clusters create')
+  assert len(commands) == 1
+  command = commands[0]
+  assert (
+      '--addons=LustreCsiDriver' in command
+      and '--enable-legacy-lustre-port' not in command
+  )
+
+
+def test_run_gke_cluster_create_command_with_lustre_legacy_port_adds_correct_flag(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(
+          enable_lustre_csi_driver=True, enable_legacy_lustre_port=True
+      ),
+      gke_control_plane_version='1.2.3',
+      system=TPU_TEST_SYSTEM,
+      release_channel=ReleaseChannel.REGULAR,
+  )
+
+  assert result == 0
+  mocks.commands_tester.assert_command_run(
+      'clusters create',
+      '--enable-legacy-lustre-port',
+      '--addons=LustreCsiDriver',
+  )
+
+
 def test_log_cluster_create_telemetry_does_not_log_when_feature_flag_is_disabled():
   FeatureFlags.TELEMETRY_ENABLED = False
   _log_cluster_create_telemetry(construct_args())

xpk/commands/config.py

@@ -14,14 +14,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.config import xpk_config
+from ..core.config import get_config as get_xpk_config
 from ..utils.console import xpk_print
 
 
 def set_config(args):
-  xpk_config.set(args.set_config_args[0], args.set_config_args[1])
+  get_xpk_config().set(args.set_config_args[0], args.set_config_args[1])
 
 
 def get_config(args):
-  value = xpk_config.get(args.get_config_key[0])
+  value = get_xpk_config().get(args.get_config_key[0])
   xpk_print(value)

xpk/commands/inspector.py

@@ -18,7 +18,7 @@ from ..core.cluster import get_cluster_credentials
 from ..core.commands import run_command_for_value
 from ..core.gcloud_context import add_zone_and_project, get_cluster_location
 from ..core.kueue_manager import CLUSTER_QUEUE_NAME, LOCAL_QUEUE_NAME
-from ..core.resources import CLUSTER_METADATA_CONFIGMAP, CLUSTER_RESOURCES_CONFIGMAP
+from ..core.resources import ConfigMapType, get_config_map_name
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import append_tmp_file, write_tmp_file
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -162,14 +162,16 @@ def inspector(args) -> None:
       (
           (
               'kubectl get configmap'
-              f' {args.cluster}-{CLUSTER_METADATA_CONFIGMAP} -o yaml'
+              f' {get_config_map_name(args.cluster, ConfigMapType.METADATA)} -o'
+              ' yaml'
           ),
           'GKE: Cluster Metadata ConfigMap Details',
       ),
       (
           (
               'kubectl get configmap'
-              f' {args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP} -o yaml'
+              f' {get_config_map_name(args.cluster, ConfigMapType.RESOURCES)} -o'
+              ' yaml'
           ),
           'GKE: Cluster Resources ConfigMap Details',
       ),

xpk/commands/kind.py

@@ -30,6 +30,7 @@ from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
     SystemCharacteristics,
     AcceleratorType,
+    DockerPlatform,
 )
 from ..utils.console import (xpk_exit, xpk_print)
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -97,6 +98,7 @@ def cluster_create(args) -> None:
       AcceleratorType.CPU,
       'kind',
       supports_sub_slicing=False,
+      docker_platform=DockerPlatform.ARM,
   )
 
   kueue_manager = KueueManager(project='', zone='')

xpk/commands/managed_ml_diagnostics.py

@@ -0,0 +1,249 @@
+"""
+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 packaging.version import Version
+from ..core.commands import run_command_for_value, run_command_with_updates
+from ..utils.console import xpk_print
+import os
+import tempfile
+
+_KUEUE_DEPLOYMENT_NAME = 'kueue-controller-manager'
+_KUEUE_NAMESPACE_NAME = 'kueue-system'
+_CERT_WEBHOOK_DEPLOYMENT_NAME = 'cert-manager-webhook'
+_CERT_WEBHOOK_NAMESPACE_NAME = 'cert-manager'
+_WEBHOOK_PACKAGE = 'mldiagnostics-injection-webhook'
+_WEBHOOK_VERSION = Version('v0.5.0')
+_WEBHOOK_FILENAME = f'{_WEBHOOK_PACKAGE}-v{_WEBHOOK_VERSION}.yaml'
+_OPERATOR_PACKAGE = 'mldiagnostics-connection-operator'
+_OPERATOR_VERSION = Version('v0.5.0')
+_OPERATOR_FILENAME = f'{_OPERATOR_PACKAGE}-v{_OPERATOR_VERSION}.yaml'
+_CERT_MANAGER_VERSION = Version('v1.13.0')
+
+
+def _install_cert_manager(version: Version = _CERT_MANAGER_VERSION) -> int:
+  """
+  Apply the cert-manager manifest.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  command = (
+      'kubectl apply -f'
+      ' https://github.com/cert-manager/cert-manager/releases/download/'
+      f'v{version}/cert-manager.yaml'
+  )
+
+  return_code = run_command_with_updates(
+      command, f'Applying cert-manager {version} manifest...'
+  )
+
+  return return_code
+
+
+def _download_mldiagnostics_yaml(package_name: str, version: Version) -> int:
+  """
+  Downloads the mldiagnostics injection webhook YAML from Artifact Registry.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  command = (
+      'gcloud artifacts generic download'
+      ' --repository=mldiagnostics-webhook-and-operator-yaml --location=us'
+      f' --package={package_name} --version=v{version} --destination=/tmp/'
+      ' --project=ai-on-gke'
+  )
+
+  return_code, return_output = run_command_for_value(
+      command,
+      f'Download {package_name} {version}...',
+  )
+
+  if return_code != 0:
+    if 'already exists' in return_output:
+      xpk_print(
+          f'Artifact file for {package_name} {version} already exists locally.'
+          ' Skipping download.'
+      )
+      return 0
+
+  return return_code
+
+
+def _create_mldiagnostics_namespace() -> int:
+  """
+  Creates the 'gke-mldiagnostics' namespace.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  command = 'kubectl create namespace gke-mldiagnostics'
+
+  return_code, return_output = run_command_for_value(
+      command, 'Create gke-mldiagnostics namespace...'
+  )
+
+  if return_code != 0:
+    if 'already exists' in return_output:
+      xpk_print('Namespace already exists. Skipping creation.')
+      return 0
+
+  return return_code
+
+
+def _install_mldiagnostics_yaml(artifact_filename: str) -> int:
+  """
+  Applies the mldiagnostics injection webhook YAML manifest.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  full_artifact_path = os.path.join(tempfile.gettempdir(), artifact_filename)
+
+  command = f'kubectl apply -f {full_artifact_path} -n gke-mldiagnostics'
+
+  return run_command_with_updates(
+      command,
+      f'Install {full_artifact_path}...',
+  )
+
+
+def _label_default_namespace_mldiagnostics() -> int:
+  """
+  Labels the 'default' namespace with 'managed-mldiagnostics-gke=true'.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  command = 'kubectl label namespace default managed-mldiagnostics-gke=true'
+
+  return run_command_with_updates(
+      command,
+      'Label default namespace with managed-mldiagnostics-gke=true',
+  )
+
+
+def install_mldiagnostics_prerequisites() -> int:
+  """
+  Mldiagnostics installation requirements.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  if not _wait_for_deployment_ready(
+      deployment_name=_KUEUE_DEPLOYMENT_NAME, namespace=_KUEUE_NAMESPACE_NAME
+  ):
+    xpk_print(
+        f'Application {_KUEUE_DEPLOYMENT_NAME} failed to become ready within'
+        ' the timeout.'
+    )
+    return 1
+
+  return_code = _install_cert_manager()
+  if return_code != 0:
+    return return_code
+
+  cert_webhook_ready = _wait_for_deployment_ready(
+      deployment_name=_CERT_WEBHOOK_DEPLOYMENT_NAME,
+      namespace=_CERT_WEBHOOK_NAMESPACE_NAME,
+  )
+  if not cert_webhook_ready:
+    xpk_print('The cert-manager-webhook installation failed.')
+    return 1
+
+  return_code = _download_mldiagnostics_yaml(
+      package_name=_WEBHOOK_PACKAGE, version=_WEBHOOK_VERSION
+  )
+  if return_code != 0:
+    return return_code
+
+  return_code = _create_mldiagnostics_namespace()
+  if return_code != 0:
+    return return_code
+
+  return_code = _install_mldiagnostics_yaml(artifact_filename=_WEBHOOK_FILENAME)
+  if return_code != 0:
+    return return_code
+
+  return_code = _label_default_namespace_mldiagnostics()
+  if return_code != 0:
+    return return_code
+
+  return_code = _download_mldiagnostics_yaml(
+      package_name=_OPERATOR_PACKAGE, version=_OPERATOR_VERSION
+  )
+  if return_code != 0:
+    return return_code
+
+  return_code = _install_mldiagnostics_yaml(
+      artifact_filename=_OPERATOR_FILENAME
+  )
+  if return_code != 0:
+    return return_code
+
+  xpk_print(
+      'All mldiagnostics installation and setup steps have been'
+      ' successfully completed!'
+  )
+  return 0
+
+
+def _wait_for_deployment_ready(
+    deployment_name: str, namespace: str, timeout_seconds: int = 300
+) -> bool:
+  """
+  Polls the Kubernetes Deployment status using kubectl rollout status
+  until it successfully rolls out (all replicas are ready) or times out.
+
+  Args:
+      deployment_name: The name of the Kubernetes Deployment (e.g., 'kueue-controller-manager').
+      namespace: The namespace where the Deployment is located (e.g., 'kueue-system').
+      timeout_seconds: Timeout duration in seconds (default is 300s / 5 minutes).
+
+  Returns:
+      bool: True if the Deployment successfully rolled out, False otherwise (timeout or error).
+  """
+
+  command = (
+      f'kubectl rollout status deployment/{deployment_name} -n {namespace}'
+      f' --timeout={timeout_seconds}s'
+  )
+
+  return_code = run_command_with_updates(
+      command, f'Checking status of deployment {deployment_name}...'
+  )
+
+  if return_code != 0:
+    return False
+
+  # When the status changes to 'running,' it might need about 10 seconds to fully stabilize.
+  stabilization_seconds = 30
+  stabilization_command = f'sleep {stabilization_seconds}'
+  stabilization_code = run_command_with_updates(
+      stabilization_command,
+      f'Deployment {deployment_name} is ready. Waiting {stabilization_seconds}'
+      ' seconds for full stabilization',
+      verbose=True,
+  )
+  if stabilization_code != 0:
+    return False
+
+  return True

xpk/commands/managed_ml_diagnostics_test.py

@@ -0,0 +1,146 @@
+"""
+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 dataclasses import dataclass
+from unittest.mock import MagicMock
+import pytest
+from xpk.commands.managed_ml_diagnostics import install_mldiagnostics_prerequisites
+from xpk.core.testing.commands_tester import CommandsTester
+
+
+@dataclass
+class _Mocks:
+  common_print_mock: MagicMock
+  commands_print_mock: MagicMock
+  commands_get_reservation_deployment_type: MagicMock
+  commands_tester: CommandsTester
+
+
+@pytest.fixture
+def mocks(mocker) -> _Mocks:
+  common_print_mock = mocker.patch(
+      'xpk.commands.common.xpk_print',
+      return_value=None,
+  )
+  commands_print_mock = mocker.patch(
+      'xpk.commands.cluster.xpk_print', return_value=None
+  )
+  commands_get_reservation_deployment_type = mocker.patch(
+      'xpk.commands.cluster.get_reservation_deployment_type',
+      return_value='DENSE',
+  )
+  return _Mocks(
+      common_print_mock=common_print_mock,
+      commands_get_reservation_deployment_type=commands_get_reservation_deployment_type,
+      commands_print_mock=commands_print_mock,
+      commands_tester=CommandsTester(
+          mocker,
+          run_command_with_updates_path=(
+              'xpk.commands.managed_ml_diagnostics.run_command_with_updates'
+          ),
+          run_command_for_value_path=(
+              'xpk.commands.managed_ml_diagnostics.run_command_for_value'
+          ),
+      ),
+  )
+
+
+def test_install_mldiagnostics_prerequisites_commands_executed(
+    mocks: _Mocks,
+):
+
+  install_mldiagnostics_prerequisites()
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl',
+      'rollout',
+      'status',
+      'deployment/kueue-controller-manager',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl',
+      'apply',
+      '-f',
+      'https://github.com/cert-manager/cert-manager/',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl', 'rollout', 'status', 'deployment/cert-manager-webhook', times=1
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'gcloud',
+      'artifacts',
+      'generic',
+      'download',
+      '--package=mldiagnostics-injection-webhook',
+      '--version=v0.5.0',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl', 'create', 'namespace', 'gke-mldiagnostics', times=1
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl',
+      'apply',
+      '-f',
+      '/tmp/mldiagnostics-injection-webhook-v0.5.0.yaml',
+      '-n',
+      'gke-mldiagnostics',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl',
+      'label',
+      'namespace',
+      'default',
+      'managed-mldiagnostics-gke=true',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'gcloud',
+      'artifacts',
+      'generic',
+      'download',
+      '--package=mldiagnostics-connection-operator',
+      '--version=v0.5.0',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl',
+      'apply',
+      '-f',
+      '/tmp/mldiagnostics-connection-operator-v0.5.0.yaml',
+      '-n',
+      'gke-mldiagnostics',
+      times=1,
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'gcloud', 'artifacts', 'generic', 'download', times=2
+  )
+
+  mocks.commands_tester.assert_command_run(
+      'kubectl', 'apply', '-f', '-n', 'gke-mldiagnostics', times=2
+  )