Merge pull request #991 from Yelp/u/jfong/sync_tron_from_k8s_multiclu…

…ster

Support multiple clusters in sync_tron_from_k8s

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
@@ -58,7 +59,18 @@ def limit_size_with_hash(name: str, limit: int = 63, suffix: int = 4) -> str:
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--kubeconfig-path", dest="kubeconfig_path", help="KUBECONFIG path")
+    parser.add_argument(
+        "--kubeconfig-path",
+        dest="kubeconfig_path",
+        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",
@@ -77,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 --kubecontext arguments.")
+
     # tron's base level is critical, not info, adjust accoringly
     if args.verbose:
         level = logging.DEBUG
@@ -97,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
@@ -214,7 +233,16 @@ def update_tron_from_pods(
     jobs = get_tron_state_from_api(args.tron_url, args.num_runs)
     log.debug(f"Found {len(jobs)} jobs.")
 
-    pods = fetch_pods(args.kubeconfig_path)
+    pods = {}
+    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)