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