Resolve conflicts

Author: DannyLiCom

src/xpk/commands/cluster.py

@@ -84,7 +84,7 @@
 from ..utils.templates import get_templates_absolute_path
 import shutil
 import os
-from . import managed_ml_diagnostics
+from .managed_ml_diagnostics import install_mldiagnostics_prerequisites
 
 CLUSTER_PREHEAT_JINJA_FILE = 'cluster_preheat.yaml.j2'
 
@@ -424,8 +424,8 @@ def cluster_create(args) -> None:
       f' https://console.cloud.google.com/kubernetes/clusters/details/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/details?project={args.project}'
   )
 
-  if args.managed_ml_diagnostics:
-    return_code = managed_ml_diagnostics.install_mldiagnostics_prerequisites()
+  if args.managed_mldiagnostics:
+    return_code = install_mldiagnostics_prerequisites()
     if return_code != 0:
       xpk_print('Installation of MLDiagnostics failed.')
       xpk_exit(return_code)

src/xpk/commands/cluster_test.py

@@ -124,7 +124,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       cluster_cpu_machine_type='',
       create_vertex_tensorboard=False,
       enable_autoprovisioning=False,
-      managed_ml_diagnostics=False,
+      managed_mldiagnostics=False,
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)

src/xpk/commands/managed_ml_diagnostics.py

@@ -24,9 +24,16 @@
 _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 = Version('v1.13.0')) -> int:
+def _install_cert_manager(version: Version = _CERT_MANAGER_VERSION) -> int:
   """
   Apply the cert-manager manifest.
 
@@ -111,13 +118,11 @@ def _install_mldiagnostics_yaml(artifact_filename: str) -> int:
 
   command = f'kubectl apply -f {full_artifact_path} -n gke-mldiagnostics'
 
-  return_code = run_command_with_updates(
+  return run_command_with_updates(
       command,
       f'Install {full_artifact_path}...',
   )
 
-  return return_code
-
 
 def _label_default_namespace_mldiagnostics() -> int:
   """
@@ -129,13 +134,11 @@ def _label_default_namespace_mldiagnostics() -> int:
 
   command = 'kubectl label namespace default managed-mldiagnostics-gke=true'
 
-  return_code = run_command_with_updates(
+  return run_command_with_updates(
       command,
       'Label default namespace with managed-mldiagnostics-gke=true',
   )
 
-  return return_code
-
 
 def install_mldiagnostics_prerequisites() -> int:
   """
@@ -166,12 +169,8 @@ def install_mldiagnostics_prerequisites() -> int:
     xpk_print('The cert-manager-webhook installation failed.')
     return 1
 
-  webhook_package = 'mldiagnostics-injection-webhook'
-  webhook_version = Version('v0.5.0')
-  webhook_filename = f'{webhook_package}-v{webhook_version}.yaml'
-
   return_code = _download_mldiagnostics_yaml(
-      package_name=webhook_package, version=webhook_version
+      package_name=_WEBHOOK_PACKAGE, version=_WEBHOOK_VERSION
   )
   if return_code != 0:
     return return_code
@@ -180,25 +179,23 @@ def install_mldiagnostics_prerequisites() -> int:
   if return_code != 0:
     return return_code
 
-  return_code = _install_mldiagnostics_yaml(artifact_filename=webhook_filename)
+  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
 
-  operator_package = 'mldiagnostics-connection-operator'
-  operator_version = Version('v0.5.0')
-  operator_filename = f'{operator_package}-v{operator_version}.yaml'
-
   return_code = _download_mldiagnostics_yaml(
-      package_name=operator_package, version=operator_version
+      package_name=_OPERATOR_PACKAGE, version=_OPERATOR_VERSION
   )
   if return_code != 0:
     return return_code
 
-  return_code = _install_mldiagnostics_yaml(artifact_filename=operator_filename)
+  return_code = _install_mldiagnostics_yaml(
+      artifact_filename=_OPERATOR_FILENAME
+  )
   if return_code != 0:
     return return_code
 

src/xpk/commands/managed_ml_diagnostics_test.py

@@ -51,64 +51,64 @@ def mocks(mocker) -> _Mocks:
       commands_tester=CommandsTester(
           mocker,
           run_command_with_updates_path=(
-              'xpk.commands.managed_ml_diagnostics.run_command_with_updates'
+              'xpk.commands.cluster.run_command_with_updates'
           ),
           run_command_for_value_path=(
-              'xpk.commands.managed_ml_diagnostics.run_command_for_value'
+              'xpk.commands.cluster.run_command_for_value'
           ),
       ),
   )
 
 
-def construct_args(**kwargs: Any) -> Namespace:
-  args_dict = dict(
-      managed_mldiagnostics=False,
-  )
-  args_dict.update(kwargs)
-  return Namespace(**args_dict)
-
-
 def test_install_mldiagnostics_prerequisites_commands_executed(
     mocks: _Mocks,
     mocker,
 ):
-
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'kubectl',
       'rollout',
       'status',
       'deployment/kueue-controller-manager',
+      '-n',
+      'kueue-system',
+      '--timeout=300s',
   )
 
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'kubectl',
       'rollout',
       'status',
       'deployment/cert-manager-webhook',
   )
 
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'kubectl',
       'apply',
       '-f',
       'https://github.com/cert-manager/cert-manager/releases/',
   )
 
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'gcloud',
       'artifacts',
       'generic',
       'download',
   )
 
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'kubectl',
       'create',
       'namespace',
       'gke-mldiagnostics',
   )
 
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'kubectl',
       'apply',
       '-f',
@@ -117,6 +117,7 @@ def test_install_mldiagnostics_prerequisites_commands_executed(
   )
 
   mocks.commands_tester.set_result_for_command(
+      (0, ''),
       'kubectl',
       'label',
       'namespace',

src/xpk/parser/cluster.py

@@ -692,6 +692,11 @@ def add_shared_cluster_create_optional_arguments(
           ' regional clusters, all zones must support the machine type.'
       ),
   )
+  parser_or_group.add_argument(
+      '--managed-mldiagnostics',
+      action='store_true',
+      help='Enables the installation of required ML Diagnostics components.',
+  )
   parser_or_group.add_argument(
       '--cluster-cpu-machine-type',
       type=str,