bert_base/scripts/generate_supervisor_conf.py

104 lines
2.8 KiB
Python

import os
import sys
from dotenv import load_dotenv
load_dotenv()
def discover_models():
"""从环境变量中发现所有模型"""
models = {}
for key, value in os.environ.items():
if key.endswith('_MODEL_ID'):
prefix = key[:-9]
models[prefix.lower()] = {
'prefix': prefix,
'model_id': value,
'port': os.getenv(f'{prefix}_PORT', '5002'),
'service_name': os.getenv(f'{prefix}_SERVICE_NAME', f'{prefix.lower()}_classifier'),
'model_dir': os.getenv(f'{prefix}_MODEL_DIR', f'/app/services/{prefix.lower()}/model')
}
return models
def generate_supervisor_conf(models):
"""生成supervisor配置文件"""
conf = []
# supervisor全局配置
conf.append("""[supervisord]
nodaemon=true
logfile=/var/log/supervisor/supervisord.log
pidfile=/tmp/supervisord.pid
childlogdir=/var/log/supervisor
""")
# 为每个模型生成program配置
programs = []
for name, config in models.items():
program_name = config['service_name']
programs.append(program_name)
conf.append(f"""[program:{program_name}]
command=gunicorn --bind 0.0.0.0:{config['port']} --workers %(ENV_MAX_WORKERS)s --threads 2 --timeout 120 services.{name}.app:app
directory=/app
autostart=true
autorestart=true
startretries=3
stdout_logfile=/app/logs/{program_name}.log
stdout_logfile_maxbytes=50MB
stderr_logfile=/app/logs/{program_name}_error.log
stderr_logfile_maxbytes=50MB
environment=SERVICE_NAME="{program_name}",SERVICE_PORT="{config['port']}",MODEL_ID="{config['model_id']}"
user=appuser
""")
# 生成group配置
conf.append(f"""[group:bert_services]
programs={','.join(programs)}
""")
return ''.join(conf)
def main():
supervisor_conf_path = './supervisor.conf'
# 检查文件是否已存在
if os.path.exists(supervisor_conf_path):
print("=" * 50)
print("supervisor.conf already exists, skipping generation...")
print(f"File: {os.path.abspath(supervisor_conf_path)}")
print("=" * 50)
sys.exit(0)
print("=" * 50)
print("Generating supervisor.conf dynamically...")
print("=" * 50)
models = discover_models()
if not models:
print("⚠️ No models found, generating empty config")
sys.exit(1)
print(f"Found {len(models)} model(s):")
for name, config in models.items():
print(f" - {name}: {config['model_id']} (port {config['port']})")
# 生成配置
conf_content = generate_supervisor_conf(models)
# 写入文件
with open(supervisor_conf_path, 'w') as f:
f.write(conf_content)
print(f"\n✓ supervisor.conf generated successfully!")
print(f" File: {os.path.abspath(supervisor_conf_path)}")
print("=" * 50)
if __name__ == "__main__":
main()