pyink

Author: DannyLiCom

src/xpk/commands/cluster.py

@@ -1322,6 +1322,7 @@ def prepare_gpus(system: SystemCharacteristics):
     if err_code > 0:
       xpk_exit(err_code)
 
+
 def install_cert_manager(version: str = 'v1.13.0'):
   """
   Apply the cert-manager manifest.
@@ -1331,7 +1332,8 @@ def install_cert_manager(version: str = 'v1.13.0'):
   """
 
   command = (
-      f'kubectl apply -f https://github.com/cert-manager/cert-manager/releases/download/'
+      'kubectl apply -f'
+      ' https://github.com/cert-manager/cert-manager/releases/download/'
       f'{version}/cert-manager.yaml'
   )
 
@@ -1344,38 +1346,45 @@ def install_cert_manager(version: str = 'v1.13.0'):
 
   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'
+      'gcloud artifacts generic download'
+      ' --repository=mldiagnostics-webhook-and-operator-yaml --location=us'
+      f' --package={package_name} --version={version} --destination=./'
+      ' --project=ai-on-gke'
   )
 
   return_code, return_output = run_command_for_value(
-      command, f'Starting gcloud artifacts download for {package_name} {version}...'
+      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.')
+      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('Artifact download completed successfully.')
   return return_code
 
+
 def create_mldiagnostics_namespace():
   """
   Creates the 'gke-diagon' namespace.
-    
+
   Returns:
     0 if successful and 1 otherwise.
   """
@@ -1396,6 +1405,7 @@ def create_mldiagnostics_namespace():
   xpk_print('gke-diagon Namespace created or already exists.')
   return return_code
 
+
 def install_mldiagnostics_yaml(artifact_filename: str):
   """
   Applies the mldiagnostics injection webhook YAML manifest.
@@ -1422,25 +1432,32 @@ def install_mldiagnostics_yaml(artifact_filename: str):
       xpk_print(f'Successfully deleted local file: {artifact_filename}')
 
     except PermissionError:
-      xpk_print(f'Failed to delete file {artifact_filename} due to Permission Error.')
+      xpk_print(
+          f'Failed to delete file {artifact_filename} due to Permission Error.'
+      )
 
   else:
-    xpk_print(f'File {artifact_filename} does not exist locally. Skipping deletion (Cleanup assumed).')
+    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 = 'kubectl label namespace default diagon-enabled=true'
 
   return_code = run_command_with_updates(
-      command, 'Starting kubectl label namespace default with diagon-enabled=true...'
+      command,
+      'Starting kubectl label namespace default with diagon-enabled=true...',
   )
 
   if return_code != 0:
@@ -1450,6 +1467,7 @@ def label_default_namespace_mldiagnostics():
   xpk_print('default Namespace successfully labeled.')
   return return_code
 
+
 def install_diagon_prerequisites():
   """
   Diagon installation requirements.
@@ -1468,22 +1486,28 @@ def install_diagon_prerequisites():
     if return_code != 0:
       return return_code
 
-    cert_webhook_ready = wait_for_deployment_ready(cert_webhook_deployment_name, cert_webhook_namespace_name)
+    cert_webhook_ready = wait_for_deployment_ready(
+        cert_webhook_deployment_name, cert_webhook_namespace_name
+    )
     if cert_webhook_ready:
 
       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)
+      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)
+      return_code = install_mldiagnostics_yaml(
+          artifact_filename=webhook_filename
+      )
       if return_code != 0:
         return return_code
 
@@ -1496,26 +1520,39 @@ def install_diagon_prerequisites():
       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)
+      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)
+      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!')
+      xpk_print(
+          'All diagon installation and setup steps have been successfully'
+          ' completed!'
+      )
       return return_code
     else:
       xpk_print('The cert-manager-webhook installation failed.')
       xpk_exit(1)
   else:
-    xpk_print(f'Application {deployment_name} failed to become ready within the timeout.')
+    xpk_print(
+        f'Application {deployment_name} failed to become ready within the'
+        ' timeout.'
+    )
     xpk_exit(1)
 
-def wait_for_deployment_ready(deployment_name: str, namespace: str, timeout_seconds: int = 300) -> bool:
+
+def wait_for_deployment_ready(
+    deployment_name: str, namespace: str, timeout_seconds: int = 300
+) -> bool:
   """
-  Polls the Kubernetes Deployment status using kubectl rollout status 
+  Polls the Kubernetes Deployment status using kubectl rollout status
   until it successfully rolls out (all replicas are ready) or times out.
 
   Args:
@@ -1532,7 +1569,10 @@ def wait_for_deployment_ready(deployment_name: str, namespace: str, timeout_seco
       f' --timeout={timeout_seconds}s'
   )
 
-  print(f'Waiting for deployment {deployment_name} in namespace {namespace} to successfully roll out...')
+  print(
+      f'Waiting for deployment {deployment_name} in namespace {namespace} to'
+      ' successfully roll out...'
+  )
 
   return_code, return_output = run_command_for_value(
       command, f'Checking status of deployment {deployment_name}...'

src/xpk/commands/cluster_test.py

@@ -257,4 +257,3 @@ def test_cluster_create_args_for_install_diagon_prerequisites(
   install_diagon_prerequisites()
 
   assert mocks.common_print_mock.call_count == 0
-  

src/xpk/parser/cluster.py

@@ -156,9 +156,9 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
       action='store_true',
       default=False,
       help=(
-          '[Optional] Enables the installation of required ML Diagnostics components: '
-          'cert-manager, injection-webhook, and connection-operator. '
-          'This feature is OFF by default.'
+          '[Optional] Enables the installation of required ML Diagnostics'
+          ' components: cert-manager, injection-webhook, and'
+          ' connection-operator. This feature is OFF by default.'
       ),
   )
 
@@ -239,9 +239,9 @@ def set_cluster_create_pathways_parser(
       action='store_true',
       default=False,
       help=(
-          '[Optional] Enables the installation of required ML Diagnostics components: '
-          'cert-manager, injection-webhook, and connection-operator. '
-          'This feature is OFF by default.'
+          '[Optional] Enables the installation of required ML Diagnostics'
+          ' components: cert-manager, injection-webhook, and'
+          ' connection-operator. This feature is OFF by default.'
       ),
   )
 

src/xpk/parser/cluster_test.py

@@ -70,8 +70,12 @@ def test_cluster_create_managed_mldiagnostics():
   parser = argparse.ArgumentParser()
 
   set_cluster_create_parser(parser)
-  args = parser.parse_args(
-      ["--cluster", "test-cluster", "--tpu-type", "v5p-8", "--managed-mldiagnostics"]
-  )
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--tpu-type",
+      "v5p-8",
+      "--managed-mldiagnostics",
+  ])
 
   assert args.managed_mldiagnostics is True