Add kube context and custom cluster support

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")

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

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:

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 = [

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

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

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
 

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
         )

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:

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(