diff --git a/tools/sync_tron_state_from_k8s.py b/tools/sync_tron_state_from_k8s.py index 8eee8c40e..29702c48d 100644 --- a/tools/sync_tron_state_from_k8s.py +++ b/tools/sync_tron_state_from_k8s.py @@ -10,6 +10,7 @@ import base64 import hashlib import logging +import os import subprocess import sys from typing import Any @@ -64,6 +65,12 @@ def parse_args(): help="KUBECONFIG path; multiple can be specified to find pods in multiple clusters", nargs="+", ) + parser.add_argument( + "--kubecontext", + dest="kubecontext", + help="kubecontext to use from specified kubeconfig. multiple can be specified to find pods in multiple clusters, ONLY if a single kubeconfig-path is provided", + nargs="*", + ) parser.add_argument( "--do-work", dest="do_work", @@ -82,6 +89,10 @@ def parse_args(): parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", default=False, help="Verbose logging") args = parser.parse_args() + # We can only have multiple kubeconfigs, or multiple contexts with a single config + if len(args.kubeconfig_path) > 1 and args.kubecontext: + parser.error("You can only specify a single --kubeconfig-path if specifying multiple --kubecontext arguments.") + # tron's base level is critical, not info, adjust accoringly if args.verbose: level = logging.DEBUG @@ -102,7 +113,10 @@ def parse_args(): return args -def fetch_pods(kubeconfig_path: str) -> Dict[str, V1Pod]: +def fetch_pods(kubeconfig_path: str, kubecontext: Optional[str]) -> Dict[str, V1Pod]: + if kubecontext: + # KubeClient only uses the environment variable + os.environ["KUBECONTEXT"] = kubecontext kube_client = KubeClient(kubeconfig_path=kubeconfig_path, user_agent="sync_tron_state_from_k8s") # Bit of a hack, no helper to fetch pods so reach into core api @@ -220,8 +234,15 @@ def update_tron_from_pods( log.debug(f"Found {len(jobs)} jobs.") pods = {} - for kubeconfig in args.kubeconfig_path: - pods.update(fetch_pods(kubeconfig)) + kube_client_args = ( + [(args.kubeconfig_path[0], kubecontext) for kubecontext in args.kubecontext] + if args.kubecontext + else [(kubeconfig_path, None) for kubeconfig_path in args.kubeconfig_path] + ) + + for kubeconfig_path, kubecontext in kube_client_args: + pods.update(fetch_pods(kubeconfig_path, kubecontext)) + log.debug(f"Found {len(pods.keys())} pods.") update_tron_from_pods(jobs, pods, args.tronctl_wrapper, args.do_work)