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

Author: DannyLiCom

src/xpk/commands/cluster.py

@@ -306,6 +306,8 @@ def cluster_create(args) -> None:
   if update_coredns_command_code != 0:
     xpk_exit(update_cluster_command_code)
 
+  install_diagon_prerequisites()
+
   if not is_dry_run():
     k8s_client = setup_k8s_env(args)
     install_storage_crd(k8s_client)
@@ -1317,3 +1319,183 @@ def prepare_gpus(system: SystemCharacteristics):
     err_code = disable_mglru_on_cluster()
     if err_code > 0:
       xpk_exit(err_code)
+
+def install_cert_manager(version: str = 'v1.13.0'):
+  """
+  Apply the cert-manager manifest.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  
+  command = (
+      f'kubectl apply -f https://github.com/cert-manager/cert-manager/releases/download/'
+      f'{version}/cert-manager.yaml'
+  )
+
+  return_code = run_command_with_updates(
+      command, f'Applying cert-manager {version} manifest...'
+  )
+
+  if return_code != 0:
+    xpk_print(f'Applying cert-manager returned with ERROR {return_code}.\n')
+
+  return return_code
+
+def download_mldiagnostics_yaml(package_name: str, version: str):
+  """
+  Downloads the mldiagnostics injection webhook YAML from Artifact Registry.
+  
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  command = (
+      f'gcloud artifacts generic download --repository=mldiagnostics-webhook-and-operator-yaml --location=us '
+      f'--package={package_name} --version={version} --destination=./ '
+      f'--project=ai-on-gke'
+  )
+  
+  return_code, return_output = run_command_for_value(
+      command, f'Starting gcloud artifacts download for {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
+    xpk_print(f'gcloud download returned with ERROR {return_code}.\n')
+    xpk_exit(return_code)
+
+  xpk_print(f'Artifact download completed successfully.')
+  return return_code
+
+def create_mldiagnostics_namespace():
+  """
+  Creates the 'gke-diagon' namespace.
+    
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  
+  command = f'kubectl create namespace gke-diagon'
+  
+  return_code, return_output = run_command_for_value(
+      command, f'Starting kubectl create namespace...'
+  )
+  
+  if return_code != 0:
+    if 'already exists' in return_output:
+      xpk_print('Namespace already exists locally. Skipping creation.')
+      return 0
+    xpk_print(f'Namespace creation returned with ERROR {return_code}.\n')
+    xpk_exit(return_code)
+
+  xpk_print(f'gke-diagon Namespace created or already exists.')
+  return return_code
+
+def install_mldiagnostics_yaml(artifact_filename: str):
+  """
+  Applies the mldiagnostics injection webhook YAML manifest.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  
+  command = f'kubectl apply -f {artifact_filename}'
+  
+  return_code = run_command_with_updates(
+      command, f'Starting kubectl apply -f {artifact_filename} ...'
+  )
+  
+  if return_code != 0:
+    xpk_print(f'kubectl apply returned with ERROR {return_code}.\n')
+    xpk_exit(return_code)
+  
+  xpk_print(f'{artifact_filename} applied successfully.')
+
+  if os.path.exists(artifact_filename):
+    try:
+      os.remove(artifact_filename)
+      xpk_print(f'Successfully deleted local file: {artifact_filename}')
+      
+    except PermissionError:
+      xpk_print(f'Failed to delete file {artifact_filename} due to Permission Error.')
+      
+    except Exception as e:
+      xpk_print(f'Failed to delete file {artifact_filename}. Unexpected error: {e}')
+
+  else:
+    xpk_print(f'File {artifact_filename} does not exist locally. Skipping deletion (Cleanup assumed).')
+  
+  return return_code
+
+def label_default_namespace_mldiagnostics():
+  """
+  Labels the 'default' namespace with 'diagon-enabled=true'.
+  
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  
+  command = f'kubectl label namespace default diagon-enabled=true'
+  
+  return_code = run_command_with_updates(
+      command, f"Starting kubectl label namespace default with diagon-enabled=true..."
+  )
+  
+  if return_code != 0:
+    xpk_print(f'Namespace labeling returned with ERROR {return_code}.\n')
+    xpk_exit(return_code)
+
+  xpk_print('default Namespace successfully labeled.')
+  return return_code
+
+def install_diagon_prerequisites():
+  """
+  Diagon installation requirements.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  return_code = install_cert_manager()
+  if return_code != 0:
+    return return_code
+  
+
+  webhook_package = "mldiagnostics-injection-webhook"
+  webhook_version = "v0.3.0"
+  webhook_filename = f"{webhook_package}-{webhook_version}.yaml"
+
+  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
+
+  # --- Install Operator ---
+  operator_package = "mldiagnostics-connection-operator"
+  operator_version = "v0.3.0"
+  operator_filename = f"{operator_package}-{operator_version}.yaml"
+  
+  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 diagon installation and setup steps have been successfully completed!")
+  return return_code

src/xpk/commands/workload.py

@@ -112,6 +112,7 @@
   labels:
     kueue.x-k8s.io/queue-name: {local_queue_name}  # Name of the LocalQueue
     xpk.google.com/workload: {args.workload}
+    {mldiagnostics_labels}
   annotations:
     alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
 spec:
@@ -400,6 +401,12 @@ def workload_create(args) -> None:
     if return_code != 0:
       xpk_exit(return_code)
 
+  mldiagnostics_labels = ''
+
+  if args.managed_mldiagnostics:
+    mldiagnostics_labels = """diagon-enabled: "true" """
+    xpk_print('Managed ML Diagnostics injection enabled. Adding Pod Label.')
+
   service_account = ''
   all_storages = []
   # Currently storage customization is not supported for Pathways workloads. b/408468941
@@ -601,6 +608,7 @@ def workload_create(args) -> None:
         """ if system.accelerator_type == AcceleratorType['TPU'] else '',
         failure_policy_rules=failure_policy_rules,
         pod_failure_policy=pod_failure_policy,
+        mldiagnostics_labels=mldiagnostics_labels,
     )
   tmp = write_tmp_file(yml_string)
   command = f'kubectl apply -f {str(tmp)}'

src/xpk/parser/workload.py

@@ -216,6 +216,15 @@ def set_workload_create_parser(workload_create_parser: ArgumentParser):
           ' conditions.'
       ),
   )
+  workload_create_parser_optional_arguments.add_argument(
+      '--managed-mldiagnostics',
+      action='store_true',
+      default=False,
+      help=(
+          '[Optional] Enables injection of GKE managed ML Diagnostics'
+          '  webhook metadata.'
+      ),
+  )
 
   add_shared_workload_create_required_arguments([
       workload_create_parser_required_arguments,