gen_sites.py

#!/usr/bin/env python
"""Generate everything that depends on the site list from sites.yaml:

  docker-compose.yml   one TRE container + one FLARE client container per site,
                       each pair on its own internal network; only the FLARE
                       client also joins the federation network
  flare/project.yml    NVIDIA FLARE provisioning project (one client per site)

Run:  python scripts/gen_sites.py
"""
from __future__ import annotations

import sys
from pathlib import Path

import yaml

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from scripts.sites import ROOT, load_sites  # noqa: E402

HEADER = "# GENERATED by scripts/gen_sites.py from sites.yaml -- do not edit by hand.\n"


def compose(cfg: dict) -> dict:
    srv = cfg["server"]
    services: dict = {
        srv["name"]: {
            "build": {"context": ".", "dockerfile": "flare/server/Dockerfile"},
            "container_name": srv["name"],
            "hostname": srv["host"],
            "networks": ["federation"],
            # published so remote clients (Gefion, NextCloud) and the admin console can reach it
            "ports": [f"{srv['fed_learn_port']}:{srv['fed_learn_port']}", f"{srv['admin_port']}:{srv['admin_port']}"],
            "environment": {"SERVER_OUT": "/server_out"},
            "volumes": [
                "./flare/workspace:/workspace",
                "./server/out:/server_out",
            ],
            "profiles": ["flare"],
        }
    }
    networks: dict = {"federation": {"driver": "bridge"}}
    for s in cfg["sites"]:
        tid, net = s["tre_id"], f"tre_{s['tre_id']}_internal"
        # The TRE itself: read-only data, internal-only network => no route to the internet.
        services[f"tre-{tid}"] = {
            "build": {"context": ".", "dockerfile": f"tres/{s['adapter']}/Dockerfile"},
            "container_name": f"tre-{tid}",
            "environment": {"TRE_ID": tid, "DATA_PATH": f"/data/{tid}.csv", "TRE_REGION": s["region"]},
            "volumes": [f"./data/sites/{tid}.csv:/data/{tid}.csv:ro"],
            "networks": [net],
            "read_only": True,
            "tmpfs": ["/tmp"],
            "healthcheck": {
                "test": ["CMD", "curl", "-fsS", "http://localhost:8000/health"],
                "interval": "5s",
                "timeout": "3s",
                "retries": 10,
            },
        }
        # The FLARE client: the only process inside the TRE boundary that may dial out,
        # and only to the federation network (server). Audit log persisted per site.
        services[f"flare-{tid}"] = {
            "build": {"context": ".", "dockerfile": "flare/client/Dockerfile"},
            "container_name": f"flare-{tid}",
            "environment": {
                "TRE_ID": tid,
                "TRE_API_URL": s["api_url"],
                "TRE_REGION": s["region"],
                "FLARE_SERVER": f"{srv['host']}:{srv['fed_learn_port']}",
            },
            "volumes": [
                f"./flare/workspace/{cfg['project']['name']}/prod_00/{tid}:/client_kit:ro",
                f"./audit/{tid}:/audit",
            ],
            "networks": [net, "federation"],
            "depends_on": {f"tre-{tid}": {"condition": "service_healthy"}},
        }
        networks[net] = {"internal": True}
    return {"services": services, "networks": networks}


def flare_project(cfg: dict) -> dict:
    srv = cfg["server"]
    participants = [
        {
            "name": srv["host"],
            "type": "server",
            "org": srv["org"],
            "fed_learn_port": srv["fed_learn_port"],
            "admin_port": srv["admin_port"],
        }
    ]
    for s in cfg["sites"]:
        participants.append({"name": s["tre_id"], "type": "client", "org": s["org"]})
    participants.append(
        {"name": f"admin@{srv['org']}.org", "type": "admin", "org": srv["org"], "role": "project_admin"}
    )
    return {
        "api_version": 3,
        "name": cfg["project"]["name"],
        "description": "Cross-TRE federated analysis (generated from sites.yaml)",
        "participants": participants,
        "builders": [
            {
                "path": "nvflare.lighter.impl.workspace.WorkspaceBuilder",
                "args": {"template_file": ["master_template.yml"]},
            },
            {"path": "nvflare.lighter.impl.static_file.StaticFileBuilder", "args": {"config_folder": "config"}},
            {"path": "nvflare.lighter.impl.cert.CertBuilder"},
            {"path": "nvflare.lighter.impl.signature.SignatureBuilder"},
        ],
    }


def main() -> None:
    cfg = load_sites()
    cfg.setdefault("project", {"name": "federated_apis"})
    ids = [s["tre_id"] for s in cfg["sites"]]

    out = ROOT / "docker-compose.yml"
    out.write_text(HEADER + yaml.safe_dump(compose(cfg), sort_keys=False))
    print(f"wrote {out.relative_to(ROOT)}  ({len(ids)} sites: {', '.join(ids)})")

    out = ROOT / "flare" / "project.yml"
    out.write_text(HEADER + yaml.safe_dump(flare_project(cfg), sort_keys=False))
    print(f"wrote {out.relative_to(ROOT)}")

    for tid in ids:
        (ROOT / "audit" / tid).mkdir(parents=True, exist_ok=True)
    (ROOT / "data" / "sites").mkdir(parents=True, exist_ok=True)


if __name__ == "__main__":
    main()