| """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() |
|
|
