File size: 3,904 Bytes
0693474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Supervise locally spawned fake SSH, never a GPU or remote host."""
from __future__ import annotations
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import unittest
HERE = Path(__file__).resolve().parents[1]
FAKE_SSH = '''#!/usr/bin/env python3
import os, sys, time
mode = os.environ.get('FAKE_MODE', 'good')
if 'nvidia-smi' in sys.argv[-1]:
uuids = ['GPU-c4e7a371-d8a8-ebf0-918c-da3e38c87a93', 'GPU-7744b44b-8a13-b06d-a67e-f958fd74b8fc']
while True:
for uuid in uuids:
print('now, ' + uuid + ', 4000, 100, 40', flush=True)
time.sleep(.05)
else:
for index in range(1000):
if mode == 'closed': sys.exit(0)
mem = 5242879 if mode == 'low' else 6000000
swap = 1000 - (index if mode == 'swap' else 0)
print(f'MemAvailable: {mem} kB;SwapFree: {swap} kB;', flush=True)
time.sleep(.05)
'''
class Supervisor(unittest.TestCase):
def invoke(self, root: Path, mode: str, command: list[str], timeout: float = 5) -> subprocess.CompletedProcess[str]:
fake = root / 'ssh'; fake.write_text(FAKE_SSH); fake.chmod(0o755)
env = dict(os.environ, PATH=str(root) + ':' + os.environ['PATH'], FAKE_MODE=mode)
return subprocess.run([sys.executable, str(HERE / 'checks.py'), 'supervise', '--out', str(root),
'--name', 'trial', '--hosts', 'fake-rank0', 'fake-rank1', '--x86host', 'fake-x86host',
'--timeout', str(timeout), '--', *command], env=env, capture_output=True, text=True, timeout=15)
def test_success_records_both_cards_and_memory(self) -> None:
with tempfile.TemporaryDirectory() as td:
root = Path(td)
p = self.invoke(root, 'good', [sys.executable, '-c', 'print("owned-command")'])
self.assertEqual(p.returncode, 0, p.stderr)
self.assertIn('owned-command', p.stdout)
self.assertFalse((root / 'STOP').exists())
self.assertIn('GPU-c4e7a371', (root / 'trial.gpu-fake-x86host.log').read_text())
self.assertIn('GPU-7744b44b', (root / 'trial.gpu-fake-x86host.log').read_text())
self.assertEqual(len(json.loads((root / 'swap-baseline.json').read_text())), 2)
def test_low_or_closed_prevents_command_start(self) -> None:
for mode in ('low', 'closed'):
with self.subTest(mode=mode), tempfile.TemporaryDirectory() as td:
root = Path(td)
p = self.invoke(root, mode, [sys.executable, '-c', f'open({str(root / "must-not-exist")!r}, "w").close()'])
self.assertNotEqual(p.returncode, 0)
self.assertTrue((root / 'STOP').exists())
self.assertFalse((root / 'must-not-exist').exists())
def test_command_failure_and_timeout(self) -> None:
for code, timeout in [('raise SystemExit(29)', 5), ('import time; time.sleep(20)', .1)]:
with self.subTest(code=code), tempfile.TemporaryDirectory() as td:
root = Path(td)
p = self.invoke(root, 'good', [sys.executable, '-c', code], timeout)
self.assertNotEqual(p.returncode, 0)
self.assertTrue((root / 'STOP').exists())
def test_swap_abort_preserves_unrelated_process(self) -> None:
unrelated = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(20)'])
try:
with tempfile.TemporaryDirectory() as td:
root = Path(td)
p = self.invoke(root, 'swap', [sys.executable, '-c', 'import time; time.sleep(20)'])
self.assertNotEqual(p.returncode, 0)
self.assertIn('swap growth', (root / 'STOP').read_text())
self.assertIsNone(unrelated.poll())
finally:
unrelated.terminate(); unrelated.wait()
if __name__ == '__main__':
unittest.main()
|