File size: 2,305 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 | """Coordinator experiment-only stats flush; emits no model data and imports no engine.
Uses the existing proxy housekeeping FLUSH_REQ/ACK control words. No tensor bytes
are read or written. Never use this against production or an unreviewed layout.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import struct
import time
PAGE = 65536
MAGIC = 0x31565344
FLUSH_REQ = 448
FLUSH_ACK = 512
def flush(fd: int, *, timeout: float = 2) -> None:
def word(offset: int) -> int:
return int(struct.unpack('<I',os.pread(fd,4,PAGE+offset))[0])
if word(0) != MAGIC or word(4) != 1 or word(192) != 1 or word(256) != 0:
raise ValueError('proxy control layout/state mismatch')
request = (word(FLUSH_REQ)+1) & 0xffffffff
# The engine is idle and this runner is the sole coordinator/flush requester.
os.pwrite(fd,struct.pack('<I',request),PAGE+FLUSH_REQ)
end = time.monotonic()+timeout
while word(FLUSH_ACK) != request:
if time.monotonic() >= end:
raise TimeoutError('proxy stats flush ACK not received')
time.sleep(.01)
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument('--pid',type=int,required=True)
args = ap.parse_args()
process = Path('/proc')/str(args.pid)
maps = (process/'maps').read_text()
if any(name in maps for name in ('libtorch','libc10','libcuda','libcudart')):
raise ValueError('proxy loaded torch/CUDA library')
candidates = [p for p in (process/'fd').iterdir() if os.readlink(p).startswith('/memfd:dsv41-remote-dspark')]
if not candidates or len({os.stat(p).st_ino for p in candidates}) != 1:
raise ValueError('expected one distinct proxy memfd (duplicate handles allowed)')
fd = os.open(candidates[0],os.O_RDWR)
try:
if struct.unpack('<I',os.pread(fd,4,PAGE+320))[0] != args.pid:
raise ValueError('proxy PID/control mismatch')
flush(fd)
finally:
os.close(fd)
data = json.loads(Path('/tmp/dsv41-remote-dspark.json').read_text())
if data.get('async_proxy_process',{}).get('pid') != args.pid:
raise ValueError('stats belong to another proxy')
print(json.dumps(data,allow_nan=False))
if __name__ == '__main__':
main()
|