Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support multiple clusters in sync_tron_from_k8s #991

Merged
merged 3 commits into from
Jul 19, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions tools/sync_tron_state_from_k8s.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import base64
import hashlib
import logging
import os
import subprocess
import sys
from typing import Any
Expand Down Expand Up @@ -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",
Expand All @@ -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.")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i'm assuming this will also do a sys.exit() or something similar?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup that's been my experience, it's using https://docs.python.org/3/library/argparse.html#argparse.ArgumentParser.error :

$ ./tools/sync_tron_state_from_k8s.py --kubeconfig-path /etc/kubernetes/paasta.conf kubeconfig_eks_infrastage.conf --kubecontext pnw-devc infrastage
usage: sync_tron_state_from_k8s.py [-h] [--kubeconfig-path KUBECONFIG_PATH [KUBECONFIG_PATH ...]] [--kubecontext [KUBECONTEXT [KUBECONTEXT ...]]] [--do-work] [--tron-url TRON_URL]
                                   [--tronctl-wrapper TRONCTL_WRAPPER] [-n NUM_RUNS] [-v]
sync_tron_state_from_k8s.py: error: You can only specify a single --kubeconfig-path if specifying --kubecontext arguments.
$ echo $?
2


# tron's base level is critical, not info, adjust accoringly
if args.verbose:
level = logging.DEBUG
Expand All @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

whoops, sorry - i should have also allowed for passing this through as an arg

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
Expand Down Expand Up @@ -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)
Loading