[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

ChatQnA/benchmark/performance/kubernetes/intel/gaudi/deploy.py

@@ -1,13 +1,18 @@
-import os
-import subprocess
-import json
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
 import argparse
-import yaml
 import glob
+import json
+import os
 import shutil
+import subprocess
 import sys
+
+import yaml
 from generate_helm_values import generate_helm_values
 
+
 def run_kubectl_command(command):
     """Run a kubectl command and return the output."""
     try:
@@ -17,12 +22,14 @@ def run_kubectl_command(command):
         print(f"Error running command: {command}\n{e.stderr}")
         exit(1)
 
+
 def get_all_nodes():
     """Get the list of all nodes in the Kubernetes cluster."""
     command = ["kubectl", "get", "nodes", "-o", "json"]
     output = run_kubectl_command(command)
     nodes = json.loads(output)
-    return [node['metadata']['name'] for node in nodes['items']]
+    return [node["metadata"]["name"] for node in nodes["items"]]
+
 
 def add_label_to_node(node_name, label):
     """Add a label to the specified node."""
@@ -31,6 +38,7 @@ def add_label_to_node(node_name, label):
     run_kubectl_command(command)
     print(f"Label {label} added to node {node_name} successfully.")
 
+
 def add_labels_to_nodes(node_count=None, label=None, node_names=None):
     """Add a label to the specified number of nodes or to specified nodes."""
 
@@ -49,9 +57,10 @@ def add_labels_to_nodes(node_count=None, label=None, node_names=None):
         for node_name in selected_nodes:
             add_label_to_node(node_name, label)
 
+
 def clear_labels_from_nodes(label, node_names=None):
     """Clear the specified label from specific nodes if provided, otherwise from all nodes."""
-    label_key = label.split('=')[0]  # Extract key from 'key=value' format
+    label_key = label.split("=")[0]  # Extract key from 'key=value' format
 
     # If specific nodes are provided, use them; otherwise, get all nodes
     nodes_to_clear = node_names if node_names else get_all_nodes()
@@ -63,7 +72,7 @@ def clear_labels_from_nodes(label, node_names=None):
         node_metadata = json.loads(node_info)
 
         # Check if the label exists on this node
-        labels = node_metadata['metadata'].get('labels', {})
+        labels = node_metadata["metadata"].get("labels", {})
         if label_key in labels:
             # Remove the label from the node
             command = ["kubectl", "label", "node", node_name, f"{label_key}-"]
@@ -73,6 +82,7 @@ def clear_labels_from_nodes(label, node_names=None):
         else:
             print(f"Label {label_key} not found on node {node_name}, skipping.")
 
+
 def add_helm_repo(repo_name, repo_url):
     # Add the repo if it does not exist
     add_command = ["helm", "repo", "add", repo_name, repo_url]
@@ -82,6 +92,7 @@ def add_helm_repo(repo_name, repo_url):
     except subprocess.CalledProcessError as e:
         print(f"Failed to add Helm repo {repo_name}: {e}")
 
+
 def delete_helm_repo(repo_name):
     """Delete Helm repo if it exists."""
     command = ["helm", "repo", "remove", repo_name]
@@ -91,29 +102,29 @@ def delete_helm_repo(repo_name):
     except subprocess.CalledProcessError:
         print(f"Failed to delete Helm repo {repo_name}. It may not exist.")
 
+
 def configmap_exists(name, namespace):
     """Check if a ConfigMap exists in the specified namespace."""
     check_command = ["kubectl", "get", "configmap", name, "-n", namespace]
-    result = subprocess.run(
-        check_command,
-        check=False,
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL
-    )
+    result = subprocess.run(check_command, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
     return result.returncode == 0
 
+
 def create_configmap(name, namespace, data):
     """Create a ConfigMap if it does not already exist."""
     if configmap_exists(name, namespace):
         print(f"ConfigMap '{name}' already exists in namespace '{namespace}', skipping creation.")
     else:
-        create_command = ["kubectl", "create", "configmap", name] + [
-            f"--from-literal={k}={v}" for k, v in data.items()
-        ] + ["-n", namespace]
+        create_command = (
+            ["kubectl", "create", "configmap", name]
+            + [f"--from-literal={k}={v}" for k, v in data.items()]
+            + ["-n", namespace]
+        )
         print(f"Creating ConfigMap '{name}' in namespace '{namespace}'...")
         subprocess.run(create_command, check=True)
         print(f"ConfigMap '{name}' created successfully.")
 
+
 def delete_configmap(name, namespace):
     """Delete a ConfigMap if it exists."""
     if configmap_exists(name, namespace):
@@ -124,9 +135,9 @@ def delete_configmap(name, namespace):
     else:
         print(f"ConfigMap '{name}' does not exist in namespace '{namespace}', skipping deletion.")
 
+
 def install_helm_release(release_name, chart_name, namespace, values_file, device_type):
-    """
-    Deploy a Helm release with a specified name and chart.
+    """Deploy a Helm release with a specified name and chart.
 
     Parameters:
     - release_name: The name of the Helm release.
@@ -162,7 +173,7 @@ def install_helm_release(release_name, chart_name, namespace, values_file, devic
         subprocess.run(["helm", "pull", chart_name, "--untar"], check=True)
 
         # Determine the directory name (get the actual chart_name if chart_name is in the format 'repo_name/chart_name', else use chart_name directly)
-        chart_dir_name = chart_name.split('/')[-1] if '/' in chart_name else chart_name
+        chart_dir_name = chart_name.split("/")[-1] if "/" in chart_name else chart_name
 
         # Find the untarred directory (assumes only one directory matches chart_dir_name)
         untar_dirs = glob.glob(f"{chart_dir_name}*")
@@ -175,10 +186,7 @@ def install_helm_release(release_name, chart_name, namespace, values_file, devic
             return
 
     # Prepare the Helm install command
-    command = [
-        "helm", "install", release_name, chart_name,
-        "--namespace", namespace
-    ]
+    command = ["helm", "install", release_name, chart_name, "--namespace", namespace]
 
     # Append additional values file for gaudi if it exists
     if hw_values_file:
@@ -201,6 +209,7 @@ def install_helm_release(release_name, chart_name, namespace, values_file, devic
         shutil.rmtree(untar_dir)
         print("Temporary directory removed successfully.")
 
+
 def uninstall_helm_release(release_name, namespace=None):
     """Uninstall a Helm release and clean up resources, optionally delete the namespace if not 'default'."""
     # Default to 'default' namespace if none is specified
@@ -229,48 +238,61 @@ def uninstall_helm_release(release_name, namespace=None):
     except subprocess.CalledProcessError as e:
         print(f"Error occurred while uninstalling Helm release or deleting namespace: {e}")
 
+
 def main():
     parser = argparse.ArgumentParser(description="Manage Helm Deployment.")
-    parser.add_argument("--release-name", type=str, default="chatqna",
-                        help="The Helm release name created during deployment (default: chatqna).")
-    parser.add_argument("--chart-name", type=str, default="opea/chatqna",
-                        help="The chart name to deploy, composed of repo name and chart name (default: opea/chatqna).")
-    parser.add_argument("--namespace", default="default",
-                        help="Kubernetes namespace (default: default).")
-    parser.add_argument("--hf-token",
-                        help="Hugging Face API token.")
-    parser.add_argument("--model-dir",
-                        help="Model directory, mounted as volumes for service access to pre-downloaded models")
-    parser.add_argument("--repo-name", default="opea",
-                        help="Helm repo name to add/delete (default: opea).")
-    parser.add_argument("--repo-url", default="https://opea-project.github.io/GenAIInfra",
-                        help="Helm repository URL (default: https://opea-project.github.io/GenAIInfra).")
-    parser.add_argument("--user-values",
-                        help="Path to a user-specified values.yaml file.")
-    parser.add_argument("--create-values-only", action="store_true",
-                        help="Only create the values.yaml file without deploying.")
-    parser.add_argument("--uninstall", action="store_true",
-                        help="Uninstall the Helm release.")
-    parser.add_argument("--num-nodes", type=int, default=1,
-                        help="Number of nodes to use (default: 1).")
-    parser.add_argument("--node-names", nargs='*',
-                        help="Optional specific node names to label.")
-    parser.add_argument("--add-label", action="store_true",
-                        help="Add label to specified nodes if this flag is set.")
-    parser.add_argument("--delete-label", action="store_true",
-                        help="Delete label from specified nodes if this flag is set.")
-    parser.add_argument("--label", default="node-type=opea-benchmark",
-                        help="Label to add/delete (default: node-type=opea-benchmark).")
-    parser.add_argument("--with-rerank", action="store_true",
-                        help="Include rerank service in the deployment.")
-    parser.add_argument("--tuned", action="store_true",
-                        help="Modify resources for services and change extraCmdArgs when creating values.yaml.")
-    parser.add_argument("--add-repo", action="store_true",
-                        help="Add the Helm repo specified by --repo-url.")
-    parser.add_argument("--delete-repo", action="store_true",
-                        help="Delete the Helm repo specified by --repo-name.")
-    parser.add_argument("--device-type", type=str, choices=["cpu", "gaudi"], default="gaudi",
-                        help="Specify the device type for deployment (choices: 'cpu', 'gaudi'; default: gaudi).")
+    parser.add_argument(
+        "--release-name",
+        type=str,
+        default="chatqna",
+        help="The Helm release name created during deployment (default: chatqna).",
+    )
+    parser.add_argument(
+        "--chart-name",
+        type=str,
+        default="opea/chatqna",
+        help="The chart name to deploy, composed of repo name and chart name (default: opea/chatqna).",
+    )
+    parser.add_argument("--namespace", default="default", help="Kubernetes namespace (default: default).")
+    parser.add_argument("--hf-token", help="Hugging Face API token.")
+    parser.add_argument(
+        "--model-dir", help="Model directory, mounted as volumes for service access to pre-downloaded models"
+    )
+    parser.add_argument("--repo-name", default="opea", help="Helm repo name to add/delete (default: opea).")
+    parser.add_argument(
+        "--repo-url",
+        default="https://opea-project.github.io/GenAIInfra",
+        help="Helm repository URL (default: https://opea-project.github.io/GenAIInfra).",
+    )
+    parser.add_argument("--user-values", help="Path to a user-specified values.yaml file.")
+    parser.add_argument(
+        "--create-values-only", action="store_true", help="Only create the values.yaml file without deploying."
+    )
+    parser.add_argument("--uninstall", action="store_true", help="Uninstall the Helm release.")
+    parser.add_argument("--num-nodes", type=int, default=1, help="Number of nodes to use (default: 1).")
+    parser.add_argument("--node-names", nargs="*", help="Optional specific node names to label.")
+    parser.add_argument("--add-label", action="store_true", help="Add label to specified nodes if this flag is set.")
+    parser.add_argument(
+        "--delete-label", action="store_true", help="Delete label from specified nodes if this flag is set."
+    )
+    parser.add_argument(
+        "--label", default="node-type=opea-benchmark", help="Label to add/delete (default: node-type=opea-benchmark)."
+    )
+    parser.add_argument("--with-rerank", action="store_true", help="Include rerank service in the deployment.")
+    parser.add_argument(
+        "--tuned",
+        action="store_true",
+        help="Modify resources for services and change extraCmdArgs when creating values.yaml.",
+    )
+    parser.add_argument("--add-repo", action="store_true", help="Add the Helm repo specified by --repo-url.")
+    parser.add_argument("--delete-repo", action="store_true", help="Delete the Helm repo specified by --repo-name.")
+    parser.add_argument(
+        "--device-type",
+        type=str,
+        choices=["cpu", "gaudi"],
+        default="gaudi",
+        help="Specify the device type for deployment (choices: 'cpu', 'gaudi'; default: gaudi).",
+    )
 
     args = parser.parse_args()
 
@@ -309,25 +331,25 @@ def main():
     else:
         if not args.hf_token:
             parser.error("--hf-token are required")
-        node_selector = {args.label.split('=')[0]: args.label.split('=')[1]}
+        node_selector = {args.label.split("=")[0]: args.label.split("=")[1]}
         values_file_path = generate_helm_values(
             with_rerank=args.with_rerank,
             num_nodes=args.num_nodes,
             hf_token=args.hf_token,
             model_dir=args.model_dir,
             node_selector=node_selector,
-            tune=args.tuned
+            tune=args.tuned,
         )
 
     # Read back the generated YAML file for verification
-    with open(values_file_path, 'r') as file:
+    with open(values_file_path, "r") as file:
         print("Generated YAML contents:")
         print(file.read())
 
-
     # Deploy unless --create-values-only is specified
     if not args.create_values_only:
         install_helm_release(args.release_name, args.chart_name, args.namespace, values_file_path, args.device_type)
 
+
 if __name__ == "__main__":
     main()