make-deployments-table.py

#!/usr/bin/env python3
from datetime import datetime
from textwrap import dedent

import yaml


MD_FILENAME = "DEPLOYMENTS.md"

def stringify_nodeselector(selector: dict[str, str]):
    # separate with markdown linebreak
    return " * ".join([f"{k}={v}" for k, v in selector.items()])


def get_deployments(filename):
    # filename: yaml filename
    with open(filename) as f:
        values = yaml.safe_load(f)
    deployments = values["deployments"]
    return {depl["name"]: stringify_nodeselector(depl["nodeSelector"]) for depl in deployments}


def main():
    gcp_deployments = get_deployments(filename="chart/values.yaml")
    az_deployments = get_deployments(filename="chart/values-azure.yaml")

    md_table = [
        ["Model ID", "GCP", "Azure", "Node Selector (GCP)", "Node Selector (Azure)"],
    ]

    deployment_info = {}
    for model_id in gcp_deployments | az_deployments:
        deployment_info.setdefault(
            model_id,
            {"GCP": "❌", "Azure": "❌", "Node Selector (GCP)": "", "Node Selector (Azure)": ""},
        )
        if info := az_deployments.get(model_id):
            deployment_info[model_id]["Azure"] = "✅"
            deployment_info[model_id]["Node Selector (Azure)"] = info
        if info := gcp_deployments.get(model_id):
            deployment_info[model_id]["GCP"] = "✅"
            deployment_info[model_id]["Node Selector (GCP)"] = info

    for model_id, info in deployment_info.items():
        md_table.append([model_id, info["GCP"], info["Azure"], info["Node Selector (GCP)"], info["Node Selector (Azure)"]])

    with open(MD_FILENAME, "w") as f:
        f.write(dedent(f"""\
        # Deployments Table

        Generated by {__file__.rsplit("/", maxsplit=1)[-1]} on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.

        """))

        header = "| " + " | ".join(md_table[0]) + " |\n"
        f.write(header)
        sep = "| " + " | ".join(["---"] * len(md_table[0])) + " |\n"
        f.write(sep)

        for row in md_table[1:]:
            f.write("| " + " | ".join(row) + " |\n")

    print("Wrote", MD_FILENAME)


if __name__ == "__main__":
    main()