165 lines
6.3 KiB
Python
Executable File
165 lines
6.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
|
||
import yaml
|
||
import sys
|
||
import os
|
||
import shutil
|
||
from jinja2 import Environment, FileSystemLoader
|
||
|
||
def load_cluster_config(config_path):
|
||
"""Load cluster configuration from YAML file"""
|
||
with open(config_path, 'r') as f:
|
||
return yaml.safe_load(f)
|
||
|
||
def generate_inventory(cluster_config, inventory_path):
|
||
"""Generate Ansible inventory file from cluster configuration"""
|
||
with open(inventory_path, 'w') as f:
|
||
f.write("# This file is auto-generated by tools/generate-inventory.py\n")
|
||
f.write("# Do not edit – your changes will be overwritten.\n")
|
||
f.write("# Edit cluster.yml and re-run the generator instead.\n\n")
|
||
|
||
# Master group
|
||
f.write("[airflow_master]\n")
|
||
for hostname, config in cluster_config['master'].items():
|
||
line = f"{hostname} ansible_host={config['ip']}"
|
||
if 'port' in config:
|
||
line += f" ansible_port={config['port']}"
|
||
f.write(line + "\n")
|
||
|
||
f.write("\n")
|
||
|
||
# Workers group
|
||
f.write("[airflow_workers]\n")
|
||
for hostname, config in cluster_config['workers'].items():
|
||
line = f"{hostname} ansible_host={config['ip']}"
|
||
if 'port' in config:
|
||
line += f" ansible_port={config['port']}"
|
||
f.write(line + "\n")
|
||
|
||
def generate_host_vars(cluster_config, host_vars_dir):
|
||
"""Generate host-specific variables"""
|
||
# Create host_vars directory if it doesn't exist
|
||
os.makedirs(host_vars_dir, exist_ok=True)
|
||
|
||
# Clear existing host_vars files to avoid stale configurations
|
||
for filename in os.listdir(host_vars_dir):
|
||
file_path = os.path.join(host_vars_dir, filename)
|
||
try:
|
||
if os.path.isfile(file_path) or os.path.islink(file_path):
|
||
os.unlink(file_path)
|
||
elif os.path.isdir(file_path):
|
||
shutil.rmtree(file_path)
|
||
except Exception as e:
|
||
print(f'Failed to delete {file_path}. Reason: {e}')
|
||
|
||
# Get master IP for Redis configuration from the new structure
|
||
master_ip = list(cluster_config['master'].values())[0]['ip']
|
||
|
||
# Get global proxy definitions
|
||
shadowsocks_proxies = cluster_config.get('shadowsocks_proxies', {})
|
||
|
||
# Combine master and worker nodes for processing
|
||
all_nodes = {**cluster_config['master'], **cluster_config['workers']}
|
||
|
||
for hostname, config in all_nodes.items():
|
||
host_vars_file = os.path.join(host_vars_dir, f"{hostname}.yml")
|
||
|
||
# Per-node list of proxies to USE
|
||
worker_proxies = config.get('proxies', [])
|
||
|
||
with open(host_vars_file, 'w') as f:
|
||
f.write("---\n")
|
||
f.write(f"# Variables for {hostname}\n")
|
||
f.write(f"master_host_ip: {master_ip}\n")
|
||
f.write("redis_port: 52909\n")
|
||
|
||
# Write the global proxy definitions for deployment
|
||
if shadowsocks_proxies:
|
||
f.write("shadowsocks_proxies:\n")
|
||
for name, proxy_config in shadowsocks_proxies.items():
|
||
f.write(f" {name}:\n")
|
||
f.write(f" server: \"{proxy_config['server']}\"\n")
|
||
f.write(f" server_port: {proxy_config['server_port']}\n")
|
||
f.write(f" local_port: {proxy_config['local_port']}\n")
|
||
f.write(f" vault_password_key: \"{proxy_config['vault_password_key']}\"\n")
|
||
|
||
# Write the per-node list of proxies to USE
|
||
if worker_proxies:
|
||
f.write("worker_proxies:\n")
|
||
for proxy in worker_proxies:
|
||
f.write(f" - \"{proxy}\"\n")
|
||
|
||
def generate_group_vars(cluster_config, group_vars_dir):
|
||
"""Generate group-level variables"""
|
||
# Create group_vars directory if it doesn't exist
|
||
os.makedirs(group_vars_dir, exist_ok=True)
|
||
|
||
# Create group_vars/all directory if it doesn't exist
|
||
all_vars_dir = os.path.join(group_vars_dir, "all")
|
||
os.makedirs(all_vars_dir, exist_ok=True)
|
||
|
||
# Define path for the generated file and remove it if it exists to avoid stale data.
|
||
# This is safer than removing the whole directory, which would delete vault.yml.
|
||
all_vars_file = os.path.join(all_vars_dir, "generated_vars.yml")
|
||
if os.path.exists(all_vars_file):
|
||
os.remove(all_vars_file)
|
||
|
||
global_vars = cluster_config.get('global_vars', {})
|
||
external_ips = cluster_config.get('external_access_ips', [])
|
||
|
||
# Get master IP for Redis configuration
|
||
master_ip = list(cluster_config['master'].values())[0]['ip']
|
||
|
||
# Combine master and worker nodes to create a hostvars-like structure
|
||
all_nodes = {**cluster_config.get('master', {}), **cluster_config.get('workers', {})}
|
||
|
||
# Prepare data for YAML dump
|
||
generated_data = {
|
||
'master_host_ip': master_ip,
|
||
'redis_port': 52909,
|
||
'external_access_ips': external_ips if external_ips else [],
|
||
'hostvars': all_nodes
|
||
}
|
||
generated_data.update(global_vars)
|
||
|
||
with open(all_vars_file, 'w') as f:
|
||
f.write("---\n")
|
||
f.write("# This file is auto-generated by tools/generate-inventory.py\n")
|
||
f.write("# Do not edit – your changes will be overwritten.\n")
|
||
yaml.dump(generated_data, f, default_flow_style=False)
|
||
|
||
def main():
|
||
if len(sys.argv) != 2:
|
||
print("Usage: python3 generate-inventory.py <cluster-config-file>")
|
||
sys.exit(1)
|
||
|
||
config_path = sys.argv[1]
|
||
|
||
# Check if config file exists
|
||
if not os.path.exists(config_path):
|
||
print(f"Error: Configuration file {config_path} not found")
|
||
sys.exit(1)
|
||
|
||
# Load cluster configuration
|
||
cluster_config = load_cluster_config(config_path)
|
||
|
||
# Generate inventory file
|
||
inventory_path = "ansible/inventory.ini"
|
||
generate_inventory(cluster_config, inventory_path)
|
||
print(f"Generated {inventory_path}")
|
||
|
||
# Generate host variables
|
||
host_vars_dir = "ansible/host_vars"
|
||
generate_host_vars(cluster_config, host_vars_dir)
|
||
print(f"Generated host variables in {host_vars_dir}")
|
||
|
||
# Generate group variables
|
||
group_vars_dir = "ansible/group_vars"
|
||
generate_group_vars(cluster_config, group_vars_dir)
|
||
print(f"Generated group variables in {group_vars_dir}")
|
||
|
||
print("Inventory generation complete!")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|