diff --git a/aiopslab/config.py b/aiopslab/config.py index 5fc61736..aa0adf2b 100644 --- a/aiopslab/config.py +++ b/aiopslab/config.py @@ -19,6 +19,41 @@ def get(self, key, default=None): return self.config.get(key, default) +def get_kube_context(): + """Get the kubernetes context from config.yml with consistent priority logic + + Priority (highest to lowest): + 1. Explicit kube_context setting + 2. If k8s_host is 'kind', construct from kind_cluster_name + 3. No context (return None to use system default) + + Returns: + str or None: Context name if should be specified, None if should use default + """ + try: + # Import BASE_DIR inside function to avoid circular import + from aiopslab.paths import BASE_DIR + config_yaml = Config(BASE_DIR / "config.yml") + + # Priority 1: Explicit kube_context setting + kube_context = config_yaml.get("kube_context") + if kube_context: + return kube_context + + # Priority 2: If k8s_host is kind, construct from kind_cluster_name + k8s_host = config_yaml.get("k8s_host") + if k8s_host == "kind": + cluster_name = config_yaml.get("kind_cluster_name", "kind") + return f"kind-{cluster_name}" + + # Priority 3: No context specified, use system default + return None + + except Exception: + # If config reading fails, use system default + return None + + # Usage example # config = Config(Path("config.yml")) # data_dir = config.get("data_dir") diff --git a/aiopslab/config.yml.example b/aiopslab/config.yml.example index 3da81468..648e593b 100644 --- a/aiopslab/config.yml.example +++ b/aiopslab/config.yml.example @@ -1,5 +1,17 @@ # Kubernetes control node k8s_host: control_node_hostname # Put `localhost` if running on cluster; put `kind` if using kind cluster; put a hostname if managing a remote cluster, e.g., apt035.apt.emulab.net + +# Kubernetes context configuration (optional, for advanced users) +# Priority 1: If specified, this context will be used for all Helm operations +# kube_context: my-custom-context + +# kind cluster name (only used when k8s_host is "kind" and kube_context is not set) +# Priority 2: When k8s_host is "kind", this will construct context as "kind-{kind_cluster_name}" +kind_cluster_name: kind # Name of your kind cluster (default: "kind") + +# Priority 3: If neither kube_context nor k8s_host="kind" is set, +# Helm will use the default kubectl context (no --kube-context flag) + k8s_user: your_username # ssh key path diff --git a/aiopslab/generators/workload/wrk.py b/aiopslab/generators/workload/wrk.py index ffb0a43d..6497d11b 100644 --- a/aiopslab/generators/workload/wrk.py +++ b/aiopslab/generators/workload/wrk.py @@ -5,6 +5,7 @@ from kubernetes import client, config from aiopslab.paths import BASE_DIR +from aiopslab.config import get_kube_context import yaml import time @@ -18,7 +19,9 @@ def __init__(self, rate, dist="norm", connections=2, duration=6, threads=2, late self.threads = threads self.latency = latency - config.load_kube_config() + kube_context = get_kube_context() + config.load_kube_config(context=kube_context) + def create_configmap(self, name, namespace, payload_script_path): with open(payload_script_path, "r") as script_file: diff --git a/aiopslab/observer/__init__.py b/aiopslab/observer/__init__.py index 44a63f2e..73b25471 100644 --- a/aiopslab/observer/__init__.py +++ b/aiopslab/observer/__init__.py @@ -8,6 +8,7 @@ from kubernetes import config, client from yaml import full_load +from aiopslab.config import get_kube_context root_path = pathlib.Path(__file__).parent sys.path.append(root_path) @@ -37,7 +38,7 @@ def get_services_list(v1, namespace="default"): return services_names -config.kube_config.load_kube_config(config_file=monitor_config["kubernetes_path"]) +config.kube_config.load_kube_config(config_file=monitor_config["kubernetes_path"], context=get_kube_context()) v1 = client.CoreV1Api() # pod_list = [ diff --git a/aiopslab/observer/log_api.py b/aiopslab/observer/log_api.py index 6a60b15f..b1923ea9 100644 --- a/aiopslab/observer/log_api.py +++ b/aiopslab/observer/log_api.py @@ -8,13 +8,14 @@ from ssl import create_default_context from enum import Enum from typing import Union -from kubernetes import client +from kubernetes import client, config import pandas as pd from elasticsearch import Elasticsearch from elasticsearch.exceptions import ConnectionTimeout from . import monitor_config, root_path, get_pod_list, get_services_list +from aiopslab.config import get_kube_context from .utils.extract import merge_csv @@ -43,6 +44,8 @@ def __init__(self, url: str, username: str, password: str): def initialize_pod_and_service_lists(self, custom_namespace=None): namespace = custom_namespace or monitor_config["namespace"] + kube_context = get_kube_context() + config.load_kube_config(context=kube_context) v1 = client.CoreV1Api() pod_list = [ pod diff --git a/aiopslab/observer/metric_api.py b/aiopslab/observer/metric_api.py index 587ef55a..57912118 100644 --- a/aiopslab/observer/metric_api.py +++ b/aiopslab/observer/metric_api.py @@ -10,13 +10,14 @@ from datetime import datetime from typing import Union from datetime import datetime, timedelta -from kubernetes import client +from kubernetes import client, config import pandas as pd import pytz from prometheus_api_client import PrometheusConnect from aiopslab.observer import monitor_config, root_path, get_pod_list, get_services_list +from aiopslab.config import get_kube_context normal_metrics = [ # cpu @@ -247,6 +248,8 @@ def cleanup(self): def initialize_pod_and_service_lists(self, custom_namespace=None): namespace = custom_namespace or monitor_config["namespace"] + kube_context = get_kube_context() + config.load_kube_config(context=kube_context) v1 = client.CoreV1Api() pod_list = [ pod diff --git a/aiopslab/observer/trace_api.py b/aiopslab/observer/trace_api.py index e123f83e..16be5fa7 100644 --- a/aiopslab/observer/trace_api.py +++ b/aiopslab/observer/trace_api.py @@ -83,14 +83,16 @@ def is_port_in_use(self, port): def get_jaeger_pod_name(self): try: - result = subprocess.check_output( - ["kubectl", "get", "pods", "-n", self.namespace, - "-l", "app.kubernetes.io/name=jaeger", - "-o", "jsonpath={.items[0].metadata.name}"], - text=True - ) + from aiopslab.service.kubectl import KubeCtl + kubectl = KubeCtl() + + # Use kubectl service to get pod name with automatic context support + command = (f"kubectl get pods -n {self.namespace} " + f"-l app.kubernetes.io/name=jaeger " + f"-o jsonpath={{.items[0].metadata.name}}") + result = kubectl.exec_command(command) return result.strip() - except subprocess.CalledProcessError as e: + except Exception as e: print("Error getting Jaeger pod name:", e) raise diff --git a/aiopslab/service/helm.py b/aiopslab/service/helm.py index 2cc2c8f6..6eee8f94 100644 --- a/aiopslab/service/helm.py +++ b/aiopslab/service/helm.py @@ -7,6 +7,10 @@ import time from aiopslab.service.kubectl import KubeCtl +from aiopslab.config import Config, get_kube_context +from aiopslab.paths import BASE_DIR + +config = Config(BASE_DIR / "config.yml") class Helm: @@ -30,9 +34,13 @@ def install(**args): extra_args = args.get("extra_args") remote_chart = args.get("remote_chart", False) + kube_context = get_kube_context() + if not remote_chart: # Install dependencies for chart before installation dependency_command = f"helm dependency update {chart_path}" + if kube_context: + dependency_command += f" --kube-context {kube_context}" dependency_process = subprocess.Popen( dependency_command, shell=True, @@ -42,6 +50,8 @@ def install(**args): dependency_output, dependency_error = dependency_process.communicate() command = f"helm install {release_name} {chart_path} -n {namespace} --create-namespace" + if kube_context: + command += f" --kube-context {kube_context}" if version: command += f" --version {version}" @@ -68,12 +78,16 @@ def uninstall(**args): print("== Helm Uninstall ==") release_name = args.get("release_name") namespace = args.get("namespace") + + kube_context = get_kube_context() if not Helm.exists_release(release_name, namespace): print(f"Release {release_name} does not exist. Skipping uninstall.") return command = f"helm uninstall {release_name} -n {namespace}" + if kube_context: + command += f" --kube-context {kube_context}" process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE) output, error = process.communicate() @@ -93,7 +107,10 @@ def exists_release(release_name: str, namespace: str) -> bool: Returns: bool: True if release exists """ + kube_context = get_kube_context() command = f"helm list -n {namespace}" + if kube_context: + command += f" --kube-context {kube_context}" process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE) output, error = process.communicate() @@ -141,6 +158,8 @@ def upgrade(**args): namespace = args.get("namespace") values_file = args.get("values_file") set_values = args.get("set_values", {}) + + kube_context = get_kube_context() command = [ "helm", @@ -152,6 +171,9 @@ def upgrade(**args): "-f", values_file, ] + + if kube_context: + command.extend(["--kube-context", kube_context]) # Add --set options if provided for key, value in set_values.items(): @@ -179,7 +201,12 @@ def add_repo(name: str, url: str): url (str): URL of the repository """ print(f"== Helm Repo Add: {name} ==") + kube_context = get_kube_context() command = f"helm repo add {name} {url}" + # Note: helm repo add doesn't typically need --kube-context + # as it operates on local helm configuration, but keeping for consistency + if kube_context: + command += f" --kube-context {kube_context}" process = subprocess.Popen( command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) diff --git a/aiopslab/service/kubectl.py b/aiopslab/service/kubectl.py index d06f73b2..c5e54e3c 100644 --- a/aiopslab/service/kubectl.py +++ b/aiopslab/service/kubectl.py @@ -9,14 +9,20 @@ from rich.console import Console from kubernetes import client, config from kubernetes.client.rest import ApiException +from aiopslab.config import Config, get_kube_context +from aiopslab.paths import BASE_DIR + +config_yaml = Config(BASE_DIR / "config.yml") class KubeCtl: def __init__(self): """Initialize the KubeCtl object and load the Kubernetes configuration.""" - config.load_kube_config() + kube_context = get_kube_context() + config.load_kube_config(context=kube_context) self.core_v1_api = client.CoreV1Api() self.apps_v1_api = client.AppsV1Api() + def list_namespaces(self): """Return a list of all namespaces in the cluster.""" @@ -58,7 +64,10 @@ def get_pod_logs(self, pod_name, namespace): def get_service_json(self, service_name, namespace, deserialize=True): """Retrieve the JSON description of a specified service within a namespace.""" + kube_context = get_kube_context() command = f"kubectl get service {service_name} -n {namespace} -o json" + if kube_context: + command += f" --context {kube_context}" result = self.exec_command(command) return json.loads(result) if deserialize else result @@ -205,20 +214,27 @@ def update_configmap(self, name, namespace, data): def apply_configs(self, namespace: str, config_path: str): """Apply Kubernetes configurations from a specified path to a namespace.""" + kube_context = get_kube_context() command = f"kubectl apply -Rf {config_path} -n {namespace}" + if kube_context: + command += f" --context {kube_context}" self.exec_command(command) def delete_configs(self, namespace: str, config_path: str): """Delete Kubernetes configurations from a specified path in a namespace.""" try: - exists_resource = self.exec_command( - f"kubectl get all -n {namespace} -o name" - ) + kube_context = get_kube_context() + + get_command = f"kubectl get all -n {namespace} -o name" + if kube_context: + get_command += f" --context {kube_context}" + exists_resource = self.exec_command(get_command) + if exists_resource: print(f"Deleting K8S configs in namespace: {namespace}") - command = ( - f"kubectl delete -Rf {config_path} -n {namespace} --timeout=10s" - ) + command = f"kubectl delete -Rf {config_path} -n {namespace} --timeout=10s" + if kube_context: + command += f" --context {kube_context}" self.exec_command(command) else: print(f"No resources found in: {namespace}. Skipping deletion.") @@ -253,7 +269,14 @@ def create_namespace_if_not_exist(self, namespace: str): print(f"Error checking/creating namespace '{namespace}': {e}") def exec_command(self, command: str, input_data=None): - """Execute an arbitrary kubectl command.""" + """Execute an arbitrary kubectl command with automatic context support.""" + # If the command contains kubectl and doesn't already have --context, add it + if "kubectl" in command and "--context" not in command: + kube_context = get_kube_context() + if kube_context: + # Insert --context after kubectl command + command = command.replace("kubectl", f"kubectl --context {kube_context}", 1) + if input_data is not None: input_data = input_data.encode("utf-8") try: diff --git a/aiopslab/service/shell.py b/aiopslab/service/shell.py index af482bac..48e60344 100644 --- a/aiopslab/service/shell.py +++ b/aiopslab/service/shell.py @@ -22,7 +22,9 @@ def exec(command: str, input_data=None, cwd=None): k8s_host = config.get("k8s_host", "localhost") # Default to localhost if k8s_host == "kind": - return Shell.docker_exec("kind-control-plane", command) + kind_cluster_name = config.get("kind_cluster_name", "kind") + container_name = f"{kind_cluster_name}-control-plane" + return Shell.docker_exec(container_name, command) elif k8s_host == "localhost": # print(