| |
| """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() |
|
|
