From 165ba1b5e8f1b065cf39ee56f0a497505adcb110 Mon Sep 17 00:00:00 2001
From: Jen Patague <jfong@yelp.com>
Date: Fri, 12 Jul 2024 17:19:15 -0700
Subject: [PATCH 1/3] Support multiple clusters in sync_tron_from_k8s

---
 tools/sync_tron_state_from_k8s.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tools/sync_tron_state_from_k8s.py b/tools/sync_tron_state_from_k8s.py
index 524496686..8eee8c40e 100644
--- a/tools/sync_tron_state_from_k8s.py
+++ b/tools/sync_tron_state_from_k8s.py
@@ -58,7 +58,12 @@ 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(
         "--do-work",
         dest="do_work",
@@ -214,7 +219,9 @@ 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 = {}
+    for kubeconfig in args.kubeconfig_path:
+        pods.update(fetch_pods(kubeconfig))
     log.debug(f"Found {len(pods.keys())} pods.")
 
     update_tron_from_pods(jobs, pods, args.tronctl_wrapper, args.do_work)

From 2d145c222dd6708a23e1d3994b69383fd19c6bce Mon Sep 17 00:00:00 2001
From: Jen Patague <jfong@yelp.com>
Date: Tue, 16 Jul 2024 10:37:54 -0700
Subject: [PATCH 2/3] Support running on devbox too

---
 tools/sync_tron_state_from_k8s.py | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

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)

From 172ad0ca933f34eb71e6341e730db6be0d8edb82 Mon Sep 17 00:00:00 2001
From: Jen Patague <jfong@yelp.com>
Date: Tue, 16 Jul 2024 16:44:05 -0700
Subject: [PATCH 3/3] Minor fix to help text for kubecontext param

---
 tools/sync_tron_state_from_k8s.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tools/sync_tron_state_from_k8s.py b/tools/sync_tron_state_from_k8s.py
index 29702c48d..60f4624d0 100644
--- a/tools/sync_tron_state_from_k8s.py
+++ b/tools/sync_tron_state_from_k8s.py
@@ -91,7 +91,7 @@ def 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.")
+        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: