# SPDX-License-Identifier: Apache-2.0
"""Stop only the named rank after three consecutive seconds below 8 GiB available."""
import json
from pathlib import Path
import re
import subprocess
import sys
import time


def main():
    name = sys.argv[1]
    if not re.fullmatch(r'glm53-qualified-r[01]', name):
        raise SystemExit('Unexpected rank name')
    log_dir = Path(__file__).resolve().parents[1] / 'runtime'
    log_dir.mkdir(exist_ok=True)
    seen, low, minimum = False, 0, None
    start = time.monotonic()
    with (log_dir / (name + '-memory.jsonl')).open('a', buffering=1) as log:
        while True:
            state = subprocess.run(['docker', 'inspect', '--format', '{{.State.Running}}', name],
                capture_output=True, text=True, timeout=10)
            running = state.returncode == 0 and state.stdout.strip() == 'true'
            if not running:
                if seen or time.monotonic() - start > 120:
                    return
                time.sleep(1)
                continue
            seen = True
            available = next(int(line.split()[1]) * 1024 for line in Path('/proc/meminfo').read_text().splitlines()
                             if line.startswith('MemAvailable:'))
            minimum = available if minimum is None else min(available, minimum)
            low = low + 1 if available < 8 * 2**30 else 0
            vmstat = dict(line.split() for line in Path('/proc/vmstat').read_text().splitlines())
            log.write(json.dumps({'unix': time.time(), 'available_gib': available / 2**30,
                'minimum_gib': minimum / 2**30, 'consecutive_low': low,
                'oom_kill': int(vmstat.get('oom_kill', 0))}) + '\n')
            if low >= 3:
                subprocess.run(['docker', 'kill', name], check=True, timeout=30)
                return
            time.sleep(1)


if __name__ == '__main__':
    main()