dsv41-offload / offload /dspark /d2 /decode_copy.py
ProprietaryLegal's picture
PLI Labs public snapshot 2026-10-28 (history squashed)
0693474
Raw History Blame Contribute Delete
17.6 kB
"""Coordinator-only frozen decode copy. Run on x86host; standard library only.
A receipt certifies an unchanged nonce-bound snapshot, not sequence usability.
Errored sequences and newest-writer row validation remain dataset assembly gates.
"""
from __future__ import annotations
import argparse
import hashlib
import inspect
import json
import math
import os
import re
import shlex
import subprocess
import threading
import time
import urllib.request
from collections.abc import Callable
from pathlib import Path
from typing import TypeVar, cast
from offload.dspark import copy_to_x86host as transfer
Record = dict[str, object]
Inventory = dict[str, Record]
def require(condition: bool, message: str) -> None:
if not condition:
raise RuntimeError(message)
def scan(root_name: str, capture: bool, nonce: str) -> Inventory:
"""Runs locally in fake hosts or as self-contained stdlib code on source host.
Spark payloads are NEVER read through the page cache, including metadata.
Stat signatures detect same-size/new tagged writes as well as size changes.
"""
root = Path(root_name)
require(root.is_dir() and not root.is_symlink(), 'source is not a real directory')
result: Inventory = {}
metadata: dict[str, bytes] = {}
paths = sorted(root.rglob('*'))
for path in paths:
require(not path.is_symlink(), 'source symlink refused: '+str(path))
if path.is_dir():
continue
require(path.is_file(), 'nonregular source refused')
relative = str(path.relative_to(root))
st = path.stat()
entry: Record = {'size': st.st_size, 'mtime_ns': st.st_mtime_ns,
'ctime_ns': st.st_ctime_ns, 'inode': st.st_ino}
# Payload stream hashes are acquired by the capped transfer, not an
# uncapped second large-file read on a serving Spark.
payload = capture and path.name in ('taps.bin', 'ids.bin')
if not payload:
if capture:
data = subprocess.check_output(['dd', 'if='+str(path), 'iflag=direct',
'bs=4096', 'status=none'])
entry['sha256'] = hashlib.sha256(data).hexdigest()
metadata[relative] = data
else:
with path.open('rb') as stream:
entry['sha256'] = hashlib.file_digest(stream, 'sha256').hexdigest()
result[relative] = entry
if capture:
flushed = json.loads(metadata['FLUSHED.json'])
request = json.loads(metadata['FLUSH.json'])
require(flushed['nonce'] == nonce and request['nonce'] == nonce, 'flush nonce mismatch')
errors = [json.loads(row) for row in metadata.get('ERRORS.jsonl', b'').splitlines()]
count = json.loads(metadata['ERROR_COUNT.json'])
require(type(flushed['errors']) is int and flushed['errors'] == len(errors)
and count['errors'] == len(errors), 'flush error-count mismatch')
require(flushed.get('schema') == 'd2.decode-capture.v1', 'wrong flush schema')
sequences = flushed['sequences']
require(isinstance(sequences, dict) and bool(sequences), 'empty or malformed flush sequences')
expected = {'FLUSH.json', 'FLUSHED.json', 'ERROR_COUNT.json'}
if 'ERRORS.jsonl' in metadata:
expected.add('ERRORS.jsonl')
for seq, row in sequences.items():
require(re.fullmatch(r'[A-Za-z0-9_-]+', seq) is not None, 'unsafe sequence')
for name in ('CAPTURE.json', 'ledger.jsonl', 'taps.bin', 'ids.bin'):
expected.add(seq+'/'+name)
require(result[seq+'/ledger.jsonl']['sha256'] == row['ledger_sha256'],
'ledger changed after flush: '+seq)
require(result[seq+'/taps.bin']['size'] == row['taps_bytes']
and result[seq+'/ids.bin']['size'] == row['ids_bytes'], 'flush size drift: '+seq)
header = json.loads(metadata[seq+'/CAPTURE.json'])
require(header.get('schema') == 'd2.decode-capture.v1' and header['seq'] == seq
and header['taps_source'] == 'decode_capture', 'capture identity mismatch')
require({str(p.relative_to(root)) for p in paths if p.is_dir()} == set(sequences),
'new or unexpected capture directory after flush')
require(set(result) == expected, 'unexpected files/new tagged writes after flush')
require(bool(result), 'empty source tree')
return result
# Keep the remote scan implementation identical to the tested local function.
SCAN_CODE = ('from __future__ import annotations\nimport hashlib,json,re,subprocess,sys\n'
'from pathlib import Path\nRecord = dict[str, object]\nInventory = dict[str, Record]\n'
+ inspect.getsource(require)+'\n'+inspect.getsource(scan)
+ '\nprint(json.dumps(scan(sys.argv[1], sys.argv[2]=="capture", sys.argv[3])))\n')
class ProductionBusy(transfer.CopyInterrupted):
"""Valid metrics prove demand; only this condition is resumable."""
def production_idle() -> None:
with urllib.request.urlopen(transfer.METRICS, timeout=3) as response:
body = response.read().decode()
busy = False
for name in ('running', 'waiting'):
values = re.findall(r'^vllm:num_requests_'+name+r'(?:\{[^\n]*\})?\s+(\S+)\s*$', body, re.M)
require(bool(values), 'missing production gauges; copy aborted')
numbers = [float(value) for value in values]
require(all(math.isfinite(value) and value >= 0 for value in numbers),
'invalid production gauges; copy aborted')
busy = busy or any(value > 0 for value in numbers)
if busy:
raise ProductionBusy('production busy; transfer paused')
def destination_memory() -> None:
available = int(next(line.split()[1] for line in Path('/proc/meminfo').read_text().splitlines()
if line.startswith('MemAvailable:'))) * 1024
require(available >= 16*1024**3, 'x86host has less than 16 GiB available; copy aborted')
T = TypeVar('T')
class PauseBudget:
"""Cumulative wait budget shared by every file/inventory in this invocation."""
def __init__(self, cap_seconds: float = 4*3600, idle_seconds: float = 30,
check: Callable[[], None] = production_idle,
memory: Callable[[], None] = destination_memory,
clock: Callable[[], float] = time.monotonic,
sleep: Callable[[float], None] = time.sleep) -> None:
require(math.isfinite(cap_seconds) and cap_seconds > 0, 'positive finite wait cap required')
require(math.isfinite(idle_seconds) and idle_seconds > 0, 'positive finite idle interval required')
self.cap_seconds = cap_seconds
self.idle_seconds = idle_seconds
self.check = check
self.memory = memory
self.clock = clock
self.sleep = sleep
self.waited = 0.0
self.events: list[Record] = []
def wait(self) -> None:
start = self.clock()
stable: float | None = None
event: Record = {'started_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}
self.events.append(event)
print('COPY PAUSED: waiting for continuous production idle', flush=True)
try:
while True:
elapsed = self.clock()-start
require(self.waited+elapsed < self.cap_seconds, 'production pause total wait cap exhausted')
self.memory() # memory failures remain fatal even while paused
try:
self.check()
except ProductionBusy:
stable = None
else:
if stable is None:
stable = self.clock()
if self.clock()-stable >= self.idle_seconds:
print('COPY RESUMING: continuous idle verified; rechecking snapshot', flush=True)
return
self.sleep(min(0.5, self.cap_seconds-self.waited-elapsed))
finally:
duration = self.clock()-start
self.waited += duration
event['wait_seconds'] = duration
def run(self, action: Callable[[], T], validate: Callable[[], None] | None = None) -> T:
while True:
try:
return action()
except transfer.CopyInterrupted as error:
# The transport preserves the monitor's cause. Missing metrics,
# HTTP errors, memory and all non-demand failures are hard aborts.
if not isinstance(error, ProductionBusy) and not isinstance(error.__cause__, ProductionBusy):
raise
while True:
self.wait()
try:
if validate is not None:
validate()
break
except ProductionBusy:
continue
class Remote:
def __init__(self, ssh: list[str], capture: bool) -> None:
self.ssh = ssh
self.capture = capture
self.idle_check: Callable[[], None] = transfer.idle
self.restart_partial = False
def inventory(self, root: str, nonce: str) -> Inventory:
self.idle_check() # before ANY Spark NIC traffic
command = shlex.join(['python3', '-c', SCAN_CODE, root,
'capture' if self.capture else 'state', nonce])
if not self.restart_partial:
return cast(Inventory, json.loads(subprocess.check_output(self.ssh+[command], text=True)))
# Inventories may scan many ledgers: monitor their traffic too. Only the
# SSH child started here is terminated when a demand/error guard trips.
process = subprocess.Popen(self.ssh+[command], stdout=subprocess.PIPE, text=True)
stop = threading.Event()
errors: list[Exception] = []
def monitor() -> None:
while not stop.wait(0.5):
try:
self.idle_check()
except Exception as error:
errors.append(error)
process.terminate()
return
watcher = threading.Thread(target=monitor, daemon=True)
watcher.start()
try:
output, _ = process.communicate()
if errors:
raise errors[0]
require(process.returncode == 0, 'inventory read failed; copy aborted')
return cast(Inventory, json.loads(output))
finally:
stop.set()
watcher.join()
if process.poll() is None:
process.terminate()
process.wait()
if process.stdout is not None:
process.stdout.close()
def copy(self, source: str, destination: Path, entry: Record) -> Record:
# Reuse the reviewed rate cap, continuous production-idle and source
# memory monitors, stream hash + independent destination re-read.
previous = transfer.SSH
transfer.SSH = self.ssh
try:
# Only this invocation's demand-interrupted file may be restarted.
try:
return transfer.copy(source, destination, cast(str | None, entry.get('sha256')),
cast(int, entry['size']), idle_check=self.idle_check)
except transfer.CopyInterrupted as error:
if self.restart_partial and (isinstance(error, ProductionBusy)
or isinstance(error.__cause__, ProductionBusy)):
destination.with_name(destination.name+'.s2-partial').unlink(missing_ok=True)
raise
finally:
transfer.SSH = previous
def frozen_copy(capture: Remote, state: Remote, source: str, generator: str,
destination: Path, nonce: str, top20: str | None = None,
reread: Callable[[Path], str] = transfer.sha_file,
pause: PauseBudget | None = None) -> Record:
started = time.monotonic()
require(not destination.exists(), 'destination exists; NEW tree required')
require(bool(re.fullmatch(r'[A-Za-z0-9_-]{1,120}', nonce)), 'unsafe nonce')
groups = [('capture', capture, source), ('generator', state, generator)]
if top20 is not None:
groups.append(('top20', state, top20))
budget = pause or PauseBudget()
for _, remote, _ in groups:
remote.idle_check = budget.check
remote.restart_partial = True
before: dict[str, Inventory] = {}
def validate_snapshot() -> None:
for name, remote, root in groups:
if name in before:
require(remote.inventory(root, nonce) == before[name],
'snapshot changed during pause: '+name+'; copied tree INVALID')
for name, remote, root in groups:
before[name] = budget.run(lambda: remote.inventory(root, nonce), validate_snapshot)
os.umask(0o077)
destination.mkdir(mode=0o700, parents=True, exist_ok=False)
receipts: list[Record] = []
for name, remote, root in groups:
for relative, entry in before[name].items():
require(not Path(relative).is_absolute() and '..' not in Path(relative).parts,
'unsafe inventory path')
target = destination/name/relative
file_receipt = budget.run(lambda: remote.copy(root+'/'+relative, target, entry), validate_snapshot)
require(target.stat().st_size == entry['size'], 'destination size mismatch')
require(reread(target) == file_receipt['sha256'], 'destination re-read hash mismatch')
receipts.append({'tree': name, 'relative': relative, **file_receipt})
for name, remote, root in groups:
require(budget.run(lambda: remote.inventory(root, nonce), validate_snapshot) == before[name],
'snapshot changed during copy: '+name+'; copied tree INVALID')
receipt: Record = {'schema': 'd2.decode-copy.v1', 'nonce': nonce, 'frozen': True,
'source': source, 'generator': generator, 'top20': top20,
'inventories': before, 'files': receipts,
'elapsed_seconds_including_snapshots_and_verification': time.monotonic()-started,
'rate_cap_bytes_per_second': transfer.RATE,
'pause_policy': {'continuous_idle_seconds': budget.idle_seconds,
'total_wait_cap_seconds': budget.cap_seconds,
'total_wait_seconds': budget.waited, 'events': budget.events,
'partial_file_policy': 'restart interrupted file from byte zero'},
'assembly_required': 'newest-writer row hashes, committed IDs, error attribution'}
# Marker written LAST; failed partial trees can never pass as a completed copy.
with (destination/'COPY_RECEIPT.json').open('x') as output:
json.dump(receipt, output, indent=2)
output.write('\n')
output.flush()
os.fsync(output.fileno())
return receipt
FLUSH_CODE = r'''
import json,os,pathlib,sys,time
p=pathlib.Path(sys.argv[1]); nonce=sys.argv[2]
for control in ('FLUSH.json', 'FLUSHED.json'):
if (p/control).exists() and json.loads((p/control).read_text()).get('nonce') == nonce:
raise SystemExit('nonce already requested/acknowledged; choose a NEW nonce')
with (p/'FLUSH.json.tmp').open('x') as f:
json.dump({'nonce':nonce}, f); f.flush(); os.fsync(f.fileno())
(p/'FLUSH.json.tmp').replace(p/'FLUSH.json')
deadline=time.monotonic()+30
while time.monotonic()<deadline:
if (p/'FLUSHED.json').exists() and json.loads((p/'FLUSHED.json').read_text()).get('nonce')==nonce:
break
time.sleep(.2)
else: raise SystemExit('capture flush timed out')
'''
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--source-root', required=True)
parser.add_argument('--nonce', required=True)
parser.add_argument('--flush', action='store_true')
parser.add_argument('--run')
parser.add_argument('--generator-root')
parser.add_argument('--top20-root')
parser.add_argument('--pause-wait-cap-seconds', type=float, default=4*3600)
args = parser.parse_args()
require(os.environ.get('D2_COORDINATOR') == 'Deepseek/coordinator', 'coordinator execution only')
require(bool(re.fullmatch(r'[A-Za-z0-9_-]{1,120}', args.nonce)), 'unsafe nonce')
capture = Remote(transfer.SSH, True)
if args.flush:
transfer.idle()
subprocess.run(transfer.SSH+[shlex.join(['python3', '-c', FLUSH_CODE,
args.source_root, args.nonce])], check=True)
capture.inventory(args.source_root, args.nonce)
print('FLUSH VERIFIED '+args.nonce)
return
require(bool(args.run and re.fullmatch(r'[A-Za-z0-9_-]{1,80}', args.run)), 'safe run required')
require(bool(args.generator_root), 'frozen generator path required')
state = Remote(['ssh', '-o', 'BatchMode=yes', '-o', 'StrictHostKeyChecking=yes', 'controller'], False)
destination = Path('/opt/dsv41/x86host-data/d2')/('decode-'+args.run)
frozen_copy(capture, state, args.source_root, args.generator_root, destination, args.nonce, args.top20_root,
pause=PauseBudget(cap_seconds=args.pause_wait_cap_seconds))
print('COPY VERIFIED '+str(destination/'COPY_RECEIPT.json'))
if __name__ == '__main__':
main()