Download offload/dspark/d2/decode_copy.py from ProprietaryLegal/dsv41-offload: direct link, hf CLI and curl.
- Browser
- Download file 17.6 kB
-
https://huggingface.co/ProprietaryLegal/dsv41-offload/resolve/main/offload/dspark/d2/decode_copy.py
- Command line
-
hf download hf://ProprietaryLegal/dsv41-offload/offload/dspark/d2/decode_copy.py
-
curl -L -o decode_copy.py https://huggingface.co/ProprietaryLegal/dsv41-offload/resolve/main/offload/dspark/d2/decode_copy.py
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() | |