File size: 3,384 Bytes
3fd1a35 | 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 | #include "ling3/model_package.h"
#include <algorithm>
#include <cstring>
#include <fcntl.h>
#include <iostream>
#include <stdexcept>
#include <unistd.h>
void Check(bool ok, const char * message) { if (!ok) throw std::runtime_error(message); }
int main() {
using namespace ling3;
const auto page = static_cast<std::size_t>(sysconf(_SC_PAGESIZE));
std::vector<std::byte> data(page * 8, std::byte{0x5a});
PackageHeader h {};
std::copy(kPackageMagic.begin(), kPackageMagic.end(), h.magic);
h.version = kPackageVersion; h.header_bytes = sizeof(h); h.endian_tag = kEndianTag;
h.tensor_count = 2; h.tensor_entry_bytes = sizeof(TensorEntry);
h.tensor_table_offset = sizeof(h); h.string_table_offset = page;
h.string_table_bytes = 15; h.payload_offset = page * 2;
h.payload_bytes = data.size() - h.payload_offset; h.file_bytes = data.size();
h.header_crc32 = Crc32(reinterpret_cast<const std::byte *>(&h), sizeof(h));
std::memcpy(data.data(), &h, sizeof(h));
const char names[] = "weightembedding";
std::memcpy(data.data() + page, names, 15);
TensorEntry entries[2] {};
entries[0].name_bytes = 6; entries[0].role = static_cast<unsigned>(TensorRole::kLinearWeight);
entries[0].dtype = static_cast<unsigned>(DataType::kInt4Low);
entries[0].data_offset = 3 * page + 123;
entries[0].data_bytes = 2 * page + 500;
entries[1].name_offset = 6; entries[1].name_bytes = 9;
entries[1].role = static_cast<unsigned>(TensorRole::kEmbedding);
entries[1].data_offset = 6 * page; entries[1].data_bytes = page;
std::memcpy(data.data() + h.tensor_table_offset, entries, sizeof(entries));
char name[] = "/tmp/ling3-reclaim-test-XXXXXX";
const int fd = mkstemp(name); Check(fd >= 0, "mkstemp");
struct Guard { int fd; char * name; ~Guard() { close(fd); unlink(name); } } guard {fd, name};
Check(write(fd, data.data(), data.size()) == static_cast<ssize_t>(data.size()), "write fixture");
Check(fsync(fd) == 0, "sync fixture");
ModelPackage package(name);
const auto & weight = package.tensor("weight");
std::vector<std::byte> native(weight.data, weight.data + weight.entry->data_bytes);
package.DiscardCopiedLinearWeight(weight);
Check(package.discarded_weight_bytes() == page, "discard crossed a partial-page boundary");
Check(package.header().file_bytes == data.size(), "header damaged");
Check(package.tensor("embedding").data[0] == std::byte{0x5a}, "CPU embedding damaged");
Check(native.front() == std::byte{0x5a} && native.back() == std::byte{0x5a}, "consumer copy damaged");
// Advisory eviction never changes the backing bytes. A deliberate future
// consumer can still fault them back in; normal inference must not do so.
Check(std::equal(native.begin(), native.end(), weight.data), "source bytes changed");
bool rejected = false;
try { package.DiscardCopiedLinearWeight(package.tensor("embedding")); }
catch (const std::invalid_argument &) { rejected = true; }
Check(rejected, "CPU-only tensor could be discarded through weight API");
std::vector<std::byte> disk(data.size());
Check(pread(fd, disk.data(), disk.size(), 0) == static_cast<ssize_t>(disk.size()), "read fixture");
Check(disk == data, "model file was modified");
std::cout << "page boundaries, CPU assets, copied weights and unchanged model passed\n";
}
|