131 lines
5.7 KiB
Python
131 lines
5.7 KiB
Python
"""One launch, fresh resource guards, only owned scratch/process-group cleanup."""
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
import signal
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from datetime import datetime, timezone
|
|
|
|
ROOT = Path('/home/frank14f/CelerisLab')
|
|
OUT = Path(__file__).resolve().parent
|
|
PYTHON = '/home/frank14f/anaconda3/envs/pycuda_3_10/bin/python'
|
|
UUID = 'GPU-e8853ae7-a302-6a17-bcfa-aa767e9ed20c'
|
|
|
|
|
|
def command(args):
|
|
return subprocess.check_output(args, text=True, timeout=10)
|
|
|
|
|
|
def gpu():
|
|
return command(['nvidia-smi', '--query-gpu=index,uuid,name,memory.used,memory.free,utilization.gpu', '--format=csv,noheader,nounits'])
|
|
|
|
|
|
def apps():
|
|
return command(['nvidia-smi', '--query-compute-apps=gpu_uuid,pid,process_name,used_memory', '--format=csv,noheader,nounits'])
|
|
|
|
|
|
def ram():
|
|
for line in Path('/proc/meminfo').read_text().splitlines():
|
|
if line.startswith('MemAvailable:'):
|
|
return int(line.split()[1])*1024
|
|
raise RuntimeError('MemAvailable unavailable')
|
|
|
|
|
|
def disk(paths):
|
|
return sum(p.stat().st_size for root in paths for p in root.rglob('*') if p.is_file())
|
|
|
|
|
|
def group_rss(pgid):
|
|
total = 0
|
|
for proc in Path('/proc').iterdir():
|
|
if not proc.name.isdigit():
|
|
continue
|
|
try:
|
|
fields = (proc/'stat').read_text().rsplit(') ', 1)[1].split()
|
|
if int(fields[2]) == pgid:
|
|
for line in (proc/'status').read_text().splitlines():
|
|
if line.startswith('VmRSS:'):
|
|
total += int(line.split()[1])*1024
|
|
except (OSError, ValueError, IndexError):
|
|
pass
|
|
return total
|
|
|
|
|
|
def main():
|
|
assert not (OUT/'gpu_run.json').exists(), 'Do not repeat an existing launch'
|
|
data = dict(start_utc=datetime.now(timezone.utc).isoformat(), command=[PYTHON, str(OUT/'validate.py')],
|
|
cwd=str(ROOT), physical_gpu_uuid=UUID, note='Sampled guards, not an OS hard isolation; no fixed RSS/AS cap.')
|
|
scratch = None
|
|
process = None
|
|
started = time.monotonic()
|
|
code = 1
|
|
try:
|
|
frozen = json.loads((OUT/'frozen.json').read_text())
|
|
for path, digest in frozen['hashes'].items():
|
|
assert hashlib.sha256((ROOT/path).read_bytes()).hexdigest() == digest, path
|
|
data['preflight_gpu'], data['preflight_apps'], data['preflight_ram_bytes'] = gpu(), apps(), ram()
|
|
fields = next(line.split(', ') for line in data['preflight_gpu'].splitlines() if UUID in line)
|
|
used, free, util = map(int, fields[-3:])
|
|
assert fields[0] == '0' and used <= 16 and free >= 2048 and util == 0
|
|
assert UUID not in data['preflight_apps'] and data['preflight_ram_bytes'] >= 4*1024**3
|
|
limits = dict(wall_seconds=180, temp_bytes=2_000_000_000,
|
|
ram_reserve_bytes=max(4*1024**3, data['preflight_ram_bytes']//10),
|
|
vram_reserve_MiB=max(512, free//10))
|
|
data['limits'] = limits
|
|
scratch = Path(tempfile.mkdtemp(prefix='celeris-joint24-'))
|
|
env = dict(PYTHONPATH=str(ROOT/'src'), PYTHONNOUSERSITE='1', PYTHONDONTWRITEBYTECODE='1',
|
|
OPENBLAS_NUM_THREADS='1', CUDA_VISIBLE_DEVICES=UUID,
|
|
CUDA_CACHE_DISABLE='1', TMPDIR=str(scratch))
|
|
data['env'] = env
|
|
data['frozen_sha256'] = hashlib.sha256((OUT/'frozen.json').read_bytes()).hexdigest()
|
|
data['peak_group_rss_bytes'] = 0
|
|
data['peak_temp_bytes'] = disk([OUT,scratch])
|
|
data['min_ram_bytes'] = data['preflight_ram_bytes']
|
|
data['min_vram_free_MiB'] = free
|
|
data['termination_reason'] = None
|
|
with (OUT/'stdout.txt').open('w') as output:
|
|
process = subprocess.Popen(data['command'], cwd=ROOT, env={**os.environ, **env},
|
|
stdout=output, stderr=subprocess.STDOUT, start_new_session=True)
|
|
while process.poll() is None:
|
|
available, occupied = ram(), disk([OUT,scratch])
|
|
free = int(next(line.split(', ')[-2] for line in gpu().splitlines() if UUID in line))
|
|
data['peak_group_rss_bytes'] = max(data['peak_group_rss_bytes'], group_rss(process.pid))
|
|
data['peak_temp_bytes'] = max(data['peak_temp_bytes'], occupied)
|
|
data['min_ram_bytes'] = min(data['min_ram_bytes'], available)
|
|
data['min_vram_free_MiB'] = min(data['min_vram_free_MiB'], free)
|
|
if (time.monotonic()-started > limits['wall_seconds'] or occupied > limits['temp_bytes']
|
|
or available < limits['ram_reserve_bytes'] or free < limits['vram_reserve_MiB']):
|
|
data['termination_reason'] = 'Time/disk/RAM/VRAM guard triggered'
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
try:
|
|
process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
break
|
|
time.sleep(.5)
|
|
code = process.wait(timeout=10)
|
|
data['postflight_gpu'], data['postflight_apps'] = gpu(), apps()
|
|
except BaseException as error:
|
|
data['runner_error'] = repr(error)
|
|
if process is not None and process.poll() is None:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
process.wait(timeout=10)
|
|
raise
|
|
finally:
|
|
data['exit_code'] = code
|
|
data['elapsed_seconds'] = time.monotonic()-started
|
|
if scratch is not None:
|
|
shutil.rmtree(scratch)
|
|
data['scratch_removed'] = scratch is None or not scratch.exists()
|
|
(OUT/'gpu_run.json').write_text(json.dumps(data, indent=2)+'\n')
|
|
print(json.dumps(data, indent=2), flush=True)
|
|
raise SystemExit(code)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|