File size: 7,140 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 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 | """Focused checks for exact extraction and fail-closed traffic protection."""
from __future__ import annotations
import hashlib
import io
import json
import struct
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
import copy_to_x86host as transfer
from inventory_source import direct_read
class TransferTests(unittest.TestCase):
def test_extract_preserves_names_dtypes_shapes_and_packed_bytes(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
root = Path(directory)
source = root / "source"
source.write_bytes(b"ignored-header"+bytes(range(128)))
rows: list[transfer.TensorRow] = [
{"name": "mtp.0.attn.wo_a.weight", "shape": [2, 4], "dtype": "F8_E4M3",
"resident_packed_bytes": 8, "source_byte_range": [14, 22]},
{"name": "mtp.0.attn.wo_a.scale", "shape": [2], "dtype": "U8",
"resident_packed_bytes": 2, "source_byte_range": [100, 102]},
]
destination = root / "extracted.safetensors"
transfer.extract(source, destination, rows)
with destination.open("rb") as stream:
size = struct.unpack("<Q", stream.read(8))[0]
header = json.loads(stream.read(size))
payload = stream.read()
self.assertEqual(set(header), {row["name"] for row in rows})
self.assertEqual(header[rows[0]["name"]]["dtype"], "F8_E4M3")
self.assertEqual(header[rows[0]["name"]]["shape"], [2, 4])
self.assertEqual(payload, bytes(range(8))+bytes(range(86, 88)))
for row in rows:
start, end = row["source_byte_range"]
self.assertEqual(row["source_byte_range_sha256"], hashlib.sha256(source.read_bytes()[start:end]).hexdigest())
def test_short_source_never_publishes_destination(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
root = Path(directory)
source = root / "source"
source.write_bytes(b"short")
destination = root / "extracted.safetensors"
with self.assertRaises(EOFError):
transfer.extract(source, destination, [{"name": "mtp.0.test", "shape": [12], "dtype": "U8",
"resident_packed_bytes": 12, "source_byte_range": [0, 12]}])
self.assertFalse(destination.exists())
def test_idle_guard_rejects_busy_missing_and_bad_metrics(self) -> None:
idle = 'vllm:num_requests_running{engine="0"} 0.0\nvllm:num_requests_waiting{engine="0"} 0.0\n'
with patch("copy_to_x86host.urllib.request.urlopen", return_value=io.BytesIO(idle.encode())):
transfer.idle()
for body in (idle.replace('0.0', '1.0', 1), idle.replace('waiting', 'missing'), 'invalid'):
with self.subTest(body=body), patch("copy_to_x86host.urllib.request.urlopen", return_value=io.BytesIO(body.encode())):
with self.assertRaises(RuntimeError):
transfer.idle()
def test_existing_bad_file_is_never_reused(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
path = Path(directory) / "wrong"
path.write_bytes(b"wrong")
with self.assertRaisesRegex(RuntimeError, "refusing unverified"):
transfer.copy("unused", path, "0"*64, 5)
def test_bad_range_rejected_before_creating_output(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
root = Path(directory)
destination = root / "invalid.safetensors"
with self.assertRaisesRegex(ValueError, "invalid source range"):
transfer.extract(root/"unused", destination, [{"name": "mtp.bad", "shape": [4],
"dtype": "U8", "resident_packed_bytes": 4, "source_byte_range": [12, 17]}])
self.assertFalse(destination.exists())
def test_resume_uses_direct_aligned_read_and_exact_overlap(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
root = Path(directory)
source = (root/"original").resolve()
data = bytes(range(256))*32+b"tail!"
source.write_bytes(data)
destination = root/"resumed"
partial = destination.with_name(destination.name+".s2-partial")
partial.write_bytes(data[:4123])
with patch.object(transfer, "SSH", ["bash", "-c"]), patch.object(transfer, "idle"):
receipt = transfer.copy(str(source), destination, hashlib.sha256(data).hexdigest(), len(data), resume=True)
self.assertEqual(destination.read_bytes(), data)
self.assertEqual(receipt["resumed_from_bytes"], 4123)
self.assertEqual(receipt["bytes_transferred"], len(data)-4123)
def test_bad_checkpoint_prefix_fails_complete_hash(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
root = Path(directory)
source = (root/"original").resolve()
data = bytes(range(256))*32
source.write_bytes(data)
destination = root/"bad-resume"
destination.with_name(destination.name+".s2-partial").write_bytes(b"corrupt")
with patch.object(transfer, "SSH", ["bash", "-c"]), patch.object(transfer, "idle"):
with self.assertRaisesRegex(RuntimeError, "verification failed"):
transfer.copy(str(source), destination, hashlib.sha256(data).hexdigest(), len(data), resume=True)
self.assertFalse(destination.exists())
def test_complete_checkpoint_verified_without_source_traffic(self) -> None:
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
destination = Path(directory)/"finished"
data = b"completed before disconnect"
destination.with_name(destination.name+".s2-partial").write_bytes(data)
with patch.object(transfer, "idle"), patch("copy_to_x86host.subprocess.Popen") as process:
result = transfer.copy("unused", destination, hashlib.sha256(data).hexdigest(), len(data), resume=True)
process.assert_not_called()
self.assertTrue(result["complete_checkpoint_verified"])
self.assertEqual(destination.read_bytes(), data)
def test_direct_read_alignment_and_short_read(self) -> None:
import os
with tempfile.TemporaryDirectory(dir=Path(__file__).parent) as directory:
path = Path(directory) / "direct"
data = bytes(range(256))*32
path.write_bytes(data)
fd = os.open(path, os.O_RDONLY | os.O_DIRECT)
try:
self.assertEqual(direct_read(fd, 4091, 50), data[4091:4141])
with self.assertRaises(EOFError):
direct_read(fd, 8000, 500)
finally:
os.close(fd)
if __name__ == "__main__":
unittest.main()
|