Add Diagon installation during cluster creation and modify the workload.py

Add wait_for_deployment_ready()

Added unit test

update goldens.yaml

update goldens.yaml

update goldens.yaml

Fixed parser/cluster.py

update goldens.yaml

fixed linter

fixed linter

pyink

Test unit test

Author: DannyLiCom

src/xpk/commands/cluster.py

@@ -81,6 +81,7 @@
 from ..utils.templates import get_templates_absolute_path
 import shutil
 import os
+from . import managed_ml_diagnostics
 
 CLUSTER_PREHEAT_JINJA_FILE = 'cluster_preheat.yaml.j2'
 
@@ -407,6 +408,13 @@ def cluster_create(args) -> None:
       # pylint: disable=line-too-long
       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 return_code != 0:
+      xpk_print('Installation of MLDiagnostics failed.')
+      xpk_exit(return_code)
+
   xpk_exit(0)
 
 

src/xpk/commands/cluster_test.py

@@ -21,6 +21,7 @@
 import pytest
 
 from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command
+from xpk.commands.managed_ml_diagnostics import install_mldiagnostics_prerequisites
 from xpk.core.system_characteristics import SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
 from xpk.utils.feature_flags import FeatureFlags
@@ -56,6 +57,15 @@ def mocks(mocker) -> _Mocks:
           run_command_with_updates_path=(
               'xpk.commands.cluster.run_command_with_updates'
           ),
+          managed_ml_diagnostics_run_command_with_updates_path=(
+              'xpk.commands.managed_ml_diagnostics.run_command_with_updates'
+          ),
+          run_command_for_value_path=(
+              'xpk.commands.cluster.run_command_for_value'
+          ),
+          managed_ml_diagnostics_run_command_for_value_path=(
+              'xpk.commands.managed_ml_diagnostics.run_command_for_value'
+          ),
       ),
   )
 
@@ -87,6 +97,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       memory_limit='100Gi',
       cpu_limit=100,
       cluster_cpu_machine_type='',
+      managed_mldiagnostics=False,
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)
@@ -247,3 +258,148 @@ def test_run_gke_cluster_create_command_with_gke_version_has_no_autoupgrade_flag
   mocks.commands_tester.assert_command_run(
       'clusters create', ' --no-enable-autoupgrade'
   )
+
+
+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',
+  )
+
+  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',
+      '-n',
+      'gke-mldiagnostics',
+  )
+
+  mocks.commands_tester.set_result_for_command(
+      (0, ''),
+      'kubectl',
+      'label',
+      'namespace',
+      'default',
+      'managed-mldiagnostics-gke=true',
+  )
+
+  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',
+      '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',
+      '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
+  )