File size: 13,733 Bytes
e1acc15 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 | """Run on x86host: capped direct-link O_DIRECT copy, exact extraction, receipts.
Only standard-library Python is needed; no CUDA or tensor library is loaded.
"""
from __future__ import annotations
import hashlib
import json
import os
import re
import shlex
import struct
import subprocess
import sys
import threading
import time
import urllib.request
from collections.abc import Callable
from pathlib import Path
from typing import NotRequired, TypedDict
class TensorRow(TypedDict):
name: str
shape: list[int]
dtype: str
resident_packed_bytes: int
source_byte_range: list[int]
source_file: NotRequired[str]
representation: NotRequired[str]
exl3_k: NotRequired[int | None]
source_file_sha256: NotRequired[str]
baseline_identical: NotRequired[bool]
identity_evidence: NotRequired[dict[str, object]]
provenance: NotRequired[dict[str, object]]
source_byte_range_sha256: NotRequired[str]
destination_file: NotRequired[str]
destination_byte_range: NotRequired[list[int]]
ROOT = Path("/opt/dsv41/x86host-data")
RATE = 1_000_000_000 # decimal bytes/s; stricter than pv -L 1g
METRICS = "http://HEAD_IP:8000/metrics"
SSH = ["ssh", "-b", "10.202.28.1", "-c", "aes128-gcm@openssh.com",
"-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes",
"-o", "HostKeyAlias=HEAD_IP", "rank0user@10.202.28.2"]
REMOTE = r'''
import json,subprocess,sys,time
path=sys.argv[1];offset=int(sys.argv[2])
minimum=2**63
def available():
return int(next(x.split()[1] for x in open('/proc/meminfo') if x.startswith('MemAvailable:')))*1024
if available()<4*1024**3:
raise RuntimeError('MemAvailable below 4 GiB before dd')
p=subprocess.Popen(['dd','if='+path,'iflag=direct,skip_bytes','skip='+str(offset),'bs=64M','status=none'])
while True:
value=available()
if value<minimum:
minimum=value
print('S2_SOURCE_MEMORY_SAMPLE '+json.dumps({'path':path,'min_mem_available_bytes':minimum}),file=sys.stderr,flush=True)
if value<4*1024**3:
p.terminate();p.wait()
raise RuntimeError('MemAvailable below 4 GiB; aborted dd: '+str(value))
code=p.poll()
if code is not None:break
time.sleep(0.1)
print('S2_SOURCE_MEMORY '+json.dumps({'path':path,'min_mem_available_bytes':minimum}),file=sys.stderr,flush=True)
sys.exit(code)
'''
def idle() -> None:
with urllib.request.urlopen(METRICS, timeout=3) as response:
body = response.read().decode()
for name in ("running", "waiting"):
values = re.findall(r"^vllm:num_requests_"+name+r"(?:\{[^\n]*\})?\s+([\d.eE+-]+)$", body, re.M)
if not values or any(float(v) != 0 for v in values):
raise CopyInterrupted(f"serve is not idle: {name}={values}")
def sha_file(path: Path) -> str:
with path.open("rb") as stream:
return hashlib.file_digest(stream, "sha256").hexdigest()
class CopyInterrupted(RuntimeError):
"""A production-traffic guard interrupted this lane's transfer."""
def copy(source: str, destination: Path, expected: str | None, size: int | None,
*, resume: bool = False, idle_check: Callable[[], None] | None = None) -> dict[str, object]:
check_idle = idle if idle_check is None else idle_check
if destination.exists():
actual = sha_file(destination)
if expected is None or actual != expected or (size is not None and destination.stat().st_size != size):
raise RuntimeError(f"refusing unverified existing destination: {destination}")
return {"file": str(destination), "size": destination.stat().st_size, "sha256": actual, "reused_verified": True}
check_idle()
if int(next(line.split()[1] for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith('MemAvailable:'))) < 16*1024**2:
raise RuntimeError("x86host has less than 16 GiB available")
destination.parent.mkdir(parents=True, exist_ok=True)
partial = destination.with_name(destination.name+".s2-partial")
if partial.exists() and not resume:
raise RuntimeError(f"partial transfer requires inspection before retry: {partial}")
start = time.monotonic()
started_utc = time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
resume_from = partial.stat().st_size if partial.exists() else 0
if resume_from and (expected is None or size is None or resume_from > size):
raise RuntimeError("resume requires a known whole-file hash and size, with no oversized partial")
if resume_from and resume_from == size:
actual = sha_file(partial)
if actual != expected:
raise RuntimeError(f"verification failed for complete checkpoint: {partial}")
partial.rename(destination)
return {"file": str(destination), "source": source, "size": size, "sha256": actual,
"resumed_from_bytes": resume_from, "bytes_transferred": 0, "complete_checkpoint_verified": True}
digest = hashlib.sha256()
# The new read starts on a 4 KiB boundary; discard any already-saved overlap.
aligned_offset = resume_from // 4096 * 4096
command = "python3 -c "+shlex.quote(REMOTE)+" "+shlex.quote(source)+" "+str(aligned_offset)
check_idle() # Recheck immediately before NIC traffic.
transfer_start = time.monotonic()
process = subprocess.Popen(SSH+[command], stdout=subprocess.PIPE)
stop = threading.Event()
errors: list[Exception] = []
def monitor() -> None:
while not stop.wait(0.5):
try:
check_idle()
except Exception as error:
errors.append(error)
process.terminate()
return
watcher = threading.Thread(target=monitor, daemon=True)
watcher.start()
count = resume_from
try:
assert process.stdout is not None
overlap = resume_from-aligned_offset
if len(process.stdout.read(overlap)) != overlap:
raise EOFError("short source read while discarding checkpoint overlap")
with partial.open("ab" if resume_from else "xb") as output:
while chunk := process.stdout.read(1024*1024):
# Pace each receive before consuming the next chunk. SSH/pipe buffers
# bound any read-ahead; the source uses a single 64 MiB dd block.
delay = (count-resume_from+len(chunk))/RATE - (time.monotonic()-transfer_start)
if delay > 0:
time.sleep(delay)
output.write(chunk)
if not resume_from:
digest.update(chunk)
count += len(chunk)
output.flush()
os.fsync(output.fileno())
code = process.wait()
if code or errors:
message = f"copy aborted: ssh={code}, guards={errors}, bytes={count}"
if errors:
raise CopyInterrupted(message) from errors[0]
raise RuntimeError(message)
# A resumed prefix is untrusted until the complete on-disk file matches
# the authoritative manifest. Never infer identity from its size alone.
actual = sha_file(partial) if resume_from else digest.hexdigest()
if (expected is not None and actual != expected) or (size is not None and count != size):
raise RuntimeError(f"verification failed: {source}, size={count}/{size}, sha256={actual}/{expected}")
# Independently re-read destination, rather than only trusting the stream hash.
if not resume_from and sha_file(partial) != actual:
raise RuntimeError(f"destination disk hash mismatch: {partial}")
partial.rename(destination)
elapsed = time.monotonic()-start
receipt = {"file": str(destination), "source": source, "size": count, "sha256": actual,
"started_utc": started_utc, "finished_utc": time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
"elapsed_seconds_including_verification": elapsed, "bytes_per_second": (count-resume_from)/elapsed,
"resumed_from_bytes": resume_from, "bytes_transferred": count-resume_from,
"rate_cap_bytes_per_second": RATE, "transport": "10.202.28.2 -> 10.202.28.1 SSH aes128-gcm", "read": "dd iflag=direct bs=64M"}
print(json.dumps(receipt), flush=True)
return receipt
finally:
stop.set()
watcher.join()
if process.poll() is None:
process.terminate()
process.wait()
if process.stdout is not None:
process.stdout.close()
def extract(source: Path, destination: Path, rows: list[TensorRow]) -> None:
header: dict[str, object] = {}
offset = 0
for row in rows:
length = row["resident_packed_bytes"]
start, end = row["source_byte_range"]
if start < 0 or end-start != length or length < 0 or row["name"] in header:
raise ValueError(f"invalid source range, length or duplicate name: {row['name']}")
header[str(row["name"])] = {"shape": row["shape"], "dtype": row["dtype"], "data_offsets": [offset, offset+length]}
offset += length
encoded = json.dumps(header, separators=(",", ":")).encode()
encoded += b" " * (-len(encoded) % 8)
partial = destination.with_suffix(".s2-partial")
with source.open("rb") as incoming, partial.open("xb") as outgoing:
outgoing.write(struct.pack("<Q", len(encoded))+encoded)
for row in rows:
start, end = row["source_byte_range"]
incoming.seek(start)
remaining = end-start
digest = hashlib.sha256()
destination_start = outgoing.tell()
while remaining:
data = incoming.read(min(8*1024*1024, remaining))
if not data:
raise EOFError(str(source))
digest.update(data)
outgoing.write(data)
remaining -= len(data)
row["source_byte_range_sha256"] = digest.hexdigest()
row["destination_file"] = destination.name
row["destination_byte_range"] = [destination_start, outgoing.tell()]
outgoing.flush()
os.fsync(outgoing.fileno())
# Verify each output range independently after serialization.
with partial.open("rb") as stream:
for row in rows:
start, end = row["destination_byte_range"]
stream.seek(start)
digest = hashlib.sha256()
remaining = end-start
while remaining:
chunk = stream.read(min(8*1024*1024, remaining))
if not chunk:
raise EOFError(str(partial))
digest.update(chunk)
remaining -= len(chunk)
if digest.hexdigest() != row["source_byte_range_sha256"]:
raise RuntimeError(f"extraction hash mismatch: {row['name']}")
partial.rename(destination)
def main() -> None:
inventory = json.loads(Path(sys.argv[1]).read_text())
source = json.loads(Path(sys.argv[2]).read_text())
drafter = ROOT / "dspark"
receipts: list[dict[str, object]] = []
files = sorted({row["source_file"] for row in inventory["tensors"]})
for filename in files:
entry = source["manifest"][filename]
if not source["files"][filename]["baseline_same_inode"]:
raise RuntimeError(f"baseline provenance needs legal-tree hash: {filename}")
incoming = drafter / "source" / filename
receipts.append(copy(source["source_root"]+"/"+filename, incoming, entry["sha256"], entry["size"]))
rows = [row for row in inventory["tensors"] if row["source_file"] == filename]
extract(incoming, drafter / filename, rows)
for filename in ("config.json", "exl3_moe_plan.json", "tokenizer.json", "tokenizer_config.json", "CANDIDATE_MANIFEST.json"):
receipts.append(copy(source["source_root"]+"/"+filename, drafter / filename, source["metadata_sha256"][filename], None))
cache_a = Path(sys.argv[3])
cache_b = Path(sys.argv[4])
receipts.append(copy("/home/rank0user/.cache/dsv41-serve/markov_top_tokens.json", drafter / "markov_top_tokens.json", sha_file(cache_a), cache_a.stat().st_size))
if cache_a.read_bytes() != cache_b.read_bytes() or (drafter / "markov_top_tokens.json").read_bytes() != cache_a.read_bytes():
raise RuntimeError("Markov token caches differ")
(drafter / "TENSORS.json").write_text(json.dumps(inventory, indent=2)+"\n")
index = {"metadata": {"total_size": inventory["resident_packed_bytes"]},
"weight_map": {row["name"]: row["destination_file"] for row in inventory["tensors"]}}
(drafter / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)+"\n")
done = {"files": receipts, "extracted_files": [{"file": f, "sha256": sha_file(drafter/f), "size": (drafter/f).stat().st_size} for f in files],
"tensors": len(inventory["tensors"]), "markov_caches_byte_identical": True}
(drafter / "DSPARK_COPY_DONE").write_text(json.dumps(done, indent=2)+"\n")
print("DSPARK_COPY_DONE", flush=True)
engram_receipts = []
for filename in ("model-00047-of-00048.safetensors", "model-00048-of-00048.safetensors"):
entry = source["manifest"][filename]
engram_receipts.append(copy(source["source_root"]+"/"+filename, ROOT/"engram"/filename, entry["sha256"], entry["size"]))
engram_receipts.append(copy(source["source_root"]+"/model.safetensors.index.json", ROOT/"engram"/"model.safetensors.index.json", source["metadata_sha256"]["model.safetensors.index.json"], None))
(ROOT/"engram"/"ENGRAM_COPY_DONE").write_text(json.dumps({"files": engram_receipts}, indent=2)+"\n")
print("ENGRAM_COPY_DONE", flush=True)
if __name__ == "__main__":
main()
|