gpu_capacity/state.py

import json
import logging
import os
import threading

from . import config

log = logging.getLogger(__name__)

# {instance_type: {region: {"status": ..., "azs": {...}, "checked_at": ...}}}
_state: dict[str, dict[str, dict]] = {}
_lock = threading.Lock()


def load():
    global _state
    try:
        with open(config.STATE_FILE) as f:
            _state = json.load(f)
        log.info("state restored from %s (%d types)", config.STATE_FILE, len(_state))
    except FileNotFoundError:
        _state = {}
    except json.JSONDecodeError:
        log.warning("corrupt state file %s, starting empty", config.STATE_FILE)
        _state = {}


def get() -> dict:
    with _lock:
        return json.loads(json.dumps(_state))


def newest_checked_at() -> float | None:
    with _lock:
        times = [r["checked_at"] for regions in _state.values() for r in regions.values()]
    return max(times) if times else None


def update(instance_type: str, region: str, result: dict):
    with _lock:
        _state.setdefault(instance_type, {})[region] = result
        data = json.dumps(_state, indent=1)
    tmp = config.STATE_FILE + ".tmp"
    try:
        with open(tmp, "w") as f:
            f.write(data)
        os.replace(tmp, config.STATE_FILE)
    except OSError:  # bucket-backed FUSE mounts may not support rename
        with open(config.STATE_FILE, "w") as f:
            f.write(data)