76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
# GPU Roster Loader - reads gpu_roster.yaml via PyYAML
|
|
# Hot-reloadable via /admin/roster/reload endpoint
|
|
|
|
import os, json, threading, time
|
|
|
|
try:
|
|
import yaml
|
|
except ImportError:
|
|
yaml = None
|
|
|
|
ROSTER_PATH = os.environ.get('ROSTER_PATH', '/app/gpu_roster.yaml')
|
|
_last_mtime = 0
|
|
_lock = threading.Lock()
|
|
|
|
GPU_URLS = {}
|
|
GPU_SIDECARS = {}
|
|
GPU_LABELS = {}
|
|
GPU_MAX_CONCURRENT = {}
|
|
GPU_CONTEXT = {}
|
|
TIER_MODELS = {}
|
|
HOSTS = {}
|
|
|
|
def load_roster(path=None):
|
|
global GPU_URLS, GPU_SIDECARS, GPU_LABELS, GPU_MAX_CONCURRENT
|
|
global GPU_CONTEXT, TIER_MODELS, HOSTS, _last_mtime
|
|
if yaml is None:
|
|
return False, 'PyYAML not installed - run: pip3 install pyyaml'
|
|
path = path or ROSTER_PATH
|
|
try:
|
|
with open(path) as f:
|
|
data = yaml.safe_load(f)
|
|
with _lock:
|
|
GPU_URLS.clear()
|
|
GPU_SIDECARS.clear()
|
|
GPU_LABELS.clear()
|
|
GPU_MAX_CONCURRENT.clear()
|
|
GPU_CONTEXT.clear()
|
|
TIER_MODELS.clear()
|
|
models = data.get('models', {})
|
|
for name, cfg in models.items():
|
|
GPU_URLS[name] = cfg.get('gpu_url', '')
|
|
GPU_SIDECARS[name] = cfg.get('sidecar_url', '')
|
|
GPU_LABELS[name] = cfg.get('label', name)
|
|
GPU_MAX_CONCURRENT[name] = cfg.get('max_concurrent', 1)
|
|
GPU_CONTEXT[name] = cfg.get('context', 65536)
|
|
tiers = cfg.get('tiers', ['enterprise'])
|
|
for t in tiers:
|
|
if t not in TIER_MODELS:
|
|
TIER_MODELS[t] = []
|
|
if name not in TIER_MODELS[t]:
|
|
TIER_MODELS[t].append(name)
|
|
HOSTS.clear()
|
|
HOSTS.update(data.get('hosts', {}))
|
|
_last_mtime = os.path.getmtime(path)
|
|
return True, 'Loaded {} models, {} hosts'.format(len(GPU_URLS), len(HOSTS))
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
def check_reload():
|
|
global _last_mtime
|
|
try:
|
|
mtime = os.path.getmtime(ROSTER_PATH)
|
|
if mtime > _last_mtime:
|
|
success, msg = load_roster()
|
|
if success:
|
|
print('[ROSTER] Auto-reloaded:', msg)
|
|
except:
|
|
pass
|
|
|
|
def reload_thread(interval=30):
|
|
while True:
|
|
time.sleep(interval)
|
|
check_reload()
|
|
|
|
threading.Thread(target=reload_thread, daemon=True).start()
|