mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-10 03:41:36 -10:00
Niko1221
16a6d77ed4
The pinned expert arena (the default tier) was allocated without looking at the RAM: on Linux the anonymous mapping succeeds whatever the host has, and a container whose memory limit is below the arena was killed by the OOM killer part-way through the load, without a message. - detail::host_available_memory: the file tier's probe (MemAvailable lowered to the tightest cgroup v2 limit's room), now with its files as parameters (tested on fake trees), cgroup v1's memory controller (memory.limit_in_bytes - memory.usage_in_bytes), and no cgroup line at all meaning MemAvailable alone (it was "cannot determine", which refused a --resident-budget-gib on a v1 host). It also returns the tightest hard limit itself. - ArenaExpertSource::open, after the size check and before the allocation: a hard cgroup limit below the arena is refused with both numbers (a certain OOM kill); less RAM available than the arena is a warning with the numbers and what runs with less (--mmap-experts with --resident-budget-gib), and the load goes on. Not for a shared arena. - file_expert_source_test: v2 limit / max / malformed / missing, v1 limit / unlimited / not mounted, no cgroup.
361 lines
18 KiB
C++
361 lines
18 KiB
C++
#include "strata/core/expert_source.hpp"
|
|
#include "strata/kernels/cpu/expert.hpp"
|
|
#include "strata/kernels/cpu/expert_layout.hpp"
|
|
|
|
#if defined(STRATA_NATIVE_EXPERTS)
|
|
#include "strata/kernels/cpu/native_expert.hpp"
|
|
#include "ggml.h"
|
|
#endif
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <limits>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <system_error>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
namespace {
|
|
|
|
void require(bool ok, const std::string& message) {
|
|
if (!ok) throw std::runtime_error(message);
|
|
}
|
|
|
|
struct TempDirectory {
|
|
fs::path path;
|
|
|
|
TempDirectory() {
|
|
const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count();
|
|
path = fs::temp_directory_path() / ("strata-file-source-test-" + std::to_string(stamp));
|
|
fs::create_directories(path);
|
|
}
|
|
|
|
~TempDirectory() {
|
|
std::error_code ignored;
|
|
fs::remove_all(path, ignored);
|
|
}
|
|
};
|
|
|
|
void create_pack(const fs::path& dir, uint64_t bytes, const std::vector<std::pair<uint64_t, char>>& markers = {}) {
|
|
const fs::path file = dir / "experts.bin";
|
|
{
|
|
std::ofstream out(file, std::ios::binary | std::ios::trunc);
|
|
require((bool) out, "could not create synthetic experts.bin");
|
|
if (bytes > 0) {
|
|
out.seekp((std::streamoff) (bytes - 1));
|
|
out.put('\0');
|
|
}
|
|
require((bool) out, "could not size synthetic experts.bin");
|
|
}
|
|
for (const auto& [offset, marker] : markers) {
|
|
require(offset < bytes, "synthetic marker lies beyond experts.bin");
|
|
std::fstream file_out(file, std::ios::binary | std::ios::in | std::ios::out);
|
|
require((bool) file_out, "could not open synthetic experts.bin for markers");
|
|
file_out.seekp((std::streamoff) offset);
|
|
file_out.put(marker);
|
|
require((bool) file_out, "could not write synthetic expert marker");
|
|
}
|
|
}
|
|
|
|
void check_file_size_rejection(strata::core::FileExpertSource& source, const fs::path& dir, int64_t layers,
|
|
int64_t experts, uint64_t expected_bytes) {
|
|
std::string err;
|
|
source.close();
|
|
fs::resize_file(dir / "experts.bin", expected_bytes - 1);
|
|
require(!source.open(dir.string(), layers, experts, err), "a truncated expert file was accepted");
|
|
require(!source.mapped() && !err.empty(), "truncated-file rejection left the source mapped or unreported");
|
|
|
|
fs::resize_file(dir / "experts.bin", expected_bytes + 1);
|
|
err.clear();
|
|
require(!source.open(dir.string(), layers, experts, err), "an oversized expert file was accepted");
|
|
require(!source.mapped() && !err.empty(), "oversized-file rejection left the source mapped or unreported");
|
|
}
|
|
|
|
void test_canonical_layout() {
|
|
using namespace strata::core;
|
|
using namespace strata::kernels::cpu;
|
|
constexpr int64_t layers = 2;
|
|
constexpr int64_t experts = 3;
|
|
const uint64_t layer_bytes = (uint64_t) experts * BLOB;
|
|
const uint64_t total = (uint64_t) layers * layer_bytes;
|
|
TempDirectory dir;
|
|
|
|
std::string err;
|
|
const bool layout_ok = expert_layout_load(dir.path.string(), layers, experts, err);
|
|
require(layout_ok, "could not load canonical layout: " + err);
|
|
create_pack(dir.path, total, {{0, 'a'}, {(uint64_t) BLOB, 'b'}, {layer_bytes, 'c'},
|
|
{layer_bytes + (uint64_t) BLOB, 'd'}});
|
|
|
|
FileExpertSource source;
|
|
bool opened = source.open(dir.path.string(), layers, experts, err);
|
|
require(opened, "could not map canonical pack: " + err);
|
|
require(source.blobs() == layers * experts, "canonical blob count is wrong");
|
|
const uint8_t* first = source.blob(0, 0);
|
|
const uint8_t* second = source.blob(0, 1);
|
|
const uint8_t* next_layer = source.blob(1, 0);
|
|
require(first && second && next_layer, "valid canonical blob lookup failed");
|
|
require(second - first == (ptrdiff_t) BLOB && next_layer - first == (ptrdiff_t) layer_bytes,
|
|
"canonical expert or layer stride is wrong");
|
|
require(first[0] == 'a' && second[0] == 'b' && next_layer[0] == 'c',
|
|
"canonical blob lookup returned bytes from the wrong expert");
|
|
require(source.blob(-1, 0) == nullptr && source.blob(0, -1) == nullptr &&
|
|
source.blob(layers, 0) == nullptr && source.blob(0, experts) == nullptr,
|
|
"canonical bounds check accepted an invalid layer or expert");
|
|
require(source.reads() == 3, "invalid canonical lookups changed the read count");
|
|
|
|
check_file_size_rejection(source, dir.path, layers, experts, total);
|
|
create_pack(dir.path, total, {{layer_bytes, 'c'}});
|
|
err.clear();
|
|
opened = source.open(dir.path.string(), layers, experts, err);
|
|
require(opened, "source failed to reopen after size errors: " + err);
|
|
const uint8_t* reopened = source.blob(1, 0);
|
|
require(reopened && source.reads() == 1 && reopened[0] == 'c', "reopened canonical source kept stale state");
|
|
source.close();
|
|
}
|
|
|
|
#if defined(STRATA_NATIVE_EXPERTS)
|
|
void test_native_variable_layout() {
|
|
using namespace strata::core;
|
|
using namespace strata::kernels::cpu;
|
|
constexpr int64_t layers = 2;
|
|
constexpr int64_t experts = 3;
|
|
constexpr int64_t embedding = 2560;
|
|
constexpr int64_t feed_forward = 640;
|
|
TempDirectory dir;
|
|
|
|
NativeFmt first_fmt, second_fmt;
|
|
std::string err;
|
|
const bool first_ok = native_fmt(GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ4_NL, embedding, feed_forward, first_fmt, err);
|
|
require(first_ok, "IQ3_XXS/IQ4_NL synthetic format is unavailable: " + err);
|
|
err.clear();
|
|
const bool second_ok = native_fmt(GGML_TYPE_IQ2_XS, GGML_TYPE_IQ4_NL, embedding, feed_forward, second_fmt, err);
|
|
require(second_ok, "IQ2_XS/IQ4_NL synthetic format is unavailable: " + err);
|
|
require(first_fmt.bytes != second_fmt.bytes, "chosen native formats do not exercise variable layer sizes");
|
|
|
|
const uint64_t layer0_bytes = (uint64_t) first_fmt.bytes * experts;
|
|
const uint64_t total = layer0_bytes + (uint64_t) second_fmt.bytes * experts;
|
|
{
|
|
std::ofstream metadata(dir.path / "native_experts.txt");
|
|
require((bool) metadata, "could not create synthetic native_experts.txt");
|
|
metadata << "0 " << GGML_TYPE_IQ3_XXS << ' ' << GGML_TYPE_IQ4_NL << " 0 " << first_fmt.bytes << '\n';
|
|
metadata << "1 " << GGML_TYPE_IQ2_XS << ' ' << GGML_TYPE_IQ4_NL << ' ' << layer0_bytes << ' '
|
|
<< second_fmt.bytes << '\n';
|
|
require((bool) metadata, "could not write synthetic native_experts.txt");
|
|
}
|
|
create_pack(dir.path, total, {{0, 'a'}, {(uint64_t) first_fmt.bytes, 'b'},
|
|
{layer0_bytes, 'c'}, {layer0_bytes + (uint64_t) second_fmt.bytes, 'd'}});
|
|
|
|
const bool layout_ok = expert_layout_load(dir.path.string(), layers, experts, err);
|
|
require(layout_ok, "could not load synthetic native layout: " + err);
|
|
FileExpertSource source;
|
|
bool opened = source.open(dir.path.string(), layers, experts, err);
|
|
require(opened, "could not map native pack: " + err);
|
|
require(source.blobs() == layers * experts, "native blob count is wrong");
|
|
const uint8_t* first = source.blob(0, 0);
|
|
const uint8_t* second = source.blob(0, 1);
|
|
const uint8_t* next_layer = source.blob(1, 0);
|
|
const uint8_t* next_layer_second = source.blob(1, 1);
|
|
require(first && second && next_layer && next_layer_second, "valid native blob lookup failed");
|
|
require(second - first == (ptrdiff_t) first_fmt.bytes && next_layer - first == (ptrdiff_t) layer0_bytes &&
|
|
next_layer_second - next_layer == (ptrdiff_t) second_fmt.bytes,
|
|
"native per-layer expert stride is wrong");
|
|
require(first[0] == 'a' && second[0] == 'b' && next_layer[0] == 'c' && next_layer_second[0] == 'd',
|
|
"native blob lookup returned bytes from the wrong expert");
|
|
require(source.blob(-1, 0) == nullptr && source.blob(0, -1) == nullptr &&
|
|
source.blob(layers, 0) == nullptr && source.blob(0, experts) == nullptr,
|
|
"native bounds check accepted an invalid layer or expert");
|
|
require(source.reads() == 4, "invalid native lookups changed the read count");
|
|
|
|
check_file_size_rejection(source, dir.path, layers, experts, total);
|
|
create_pack(dir.path, total, {{layer0_bytes, 'c'}});
|
|
err.clear();
|
|
opened = source.open(dir.path.string(), layers, experts, err);
|
|
require(opened, "source failed to reopen after size errors: " + err);
|
|
const uint8_t* reopened = source.blob(1, 0);
|
|
require(reopened && source.reads() == 1 && reopened[0] == 'c', "reopened native source kept stale state");
|
|
source.close();
|
|
}
|
|
#endif
|
|
|
|
void test_complement_plan() {
|
|
using namespace strata::core::detail;
|
|
std::vector<uint64_t> offsets;
|
|
uint64_t bytes = 0;
|
|
std::string error;
|
|
require(make_cache_complement_plan(2, 3, {3, 5}, {{0, 1}, {1, 2}}, {}, offsets, bytes, error), error);
|
|
require(bytes == 16 && offsets == std::vector<uint64_t>{0, kNoCacheComplement, 3, 6, 11, kNoCacheComplement},
|
|
"wrong compact offsets for variable native layer sizes");
|
|
const uint8_t resident[16] = {}, mapped[5] = {};
|
|
require(cache_complement_blob_or_fallback(1, offsets, resident, mapped) == mapped,
|
|
"GPU-resident expert lost mmap fallback");
|
|
require(cache_complement_blob_or_fallback(4, offsets, resident, mapped) == resident + 11,
|
|
"CPU miss did not use resident complement");
|
|
require(!make_cache_complement_plan(2, 3, {3, 5}, {{0, 1}, {0, 1}}, {}, offsets, bytes, error)
|
|
&& offsets.empty() && bytes == 0, "duplicate pair accepted");
|
|
require(!make_cache_complement_plan(2, 3, {3, 5}, {{0, 1}}, {{0, 1}}, offsets, bytes, error),
|
|
"overlapping tiers accepted");
|
|
require(!make_cache_complement_plan(2, 3, {3, 5}, {{2, 0}}, {}, offsets, bytes, error),
|
|
"out-of-range pair accepted");
|
|
require(!make_cache_complement_plan(2, 3, {0, 5}, {}, {}, offsets, bytes, error), "zero-size layer accepted");
|
|
}
|
|
|
|
void test_resident_lend_region() {
|
|
using namespace strata::core::detail;
|
|
// four slots (sizes 5, 3, 3, 5); the experts no slot holds take 10 bytes
|
|
const std::vector<uint64_t> slots{5, 3, 3, 5};
|
|
require(choose_resident_keep_from(slots, 10, 9, 1) == -1, "a base larger than the budget was accepted");
|
|
require(choose_resident_keep_from(slots, 10, 10, 1) == 4, "no room: the lend region must stay on the file");
|
|
require(choose_resident_keep_from(slots, 10, 15, 1) == 3, "the last slot fits and must be kept first");
|
|
require(choose_resident_keep_from(slots, 10, 17, 1) == 3, "a slot that does not fit ends the lend region's walk");
|
|
require(choose_resident_keep_from(slots, 10, 18, 1) == 2, "two slots fit");
|
|
require(choose_resident_keep_from(slots, 10, 100, 1) == 1, "the walk must stop at the lend region's first slot");
|
|
require(choose_resident_keep_from(slots, 10, 100, -1) == 4, "no lend region kept slots in RAM");
|
|
require(choose_resident_keep_from(slots, 10, 100, 9) == 4, "an out-of-range lend region kept slots in RAM");
|
|
require(choose_resident_keep_from({}, 0, 0, 0) == 0, "an empty cache");
|
|
}
|
|
|
|
void test_resident_exchange() {
|
|
using namespace strata::core;
|
|
using namespace strata::core::detail;
|
|
// 2 layers x 3 experts; the GPU holds (0,1) and (1,2): the compact copy holds the other four
|
|
std::vector<uint64_t> offsets;
|
|
uint64_t bytes = 0;
|
|
std::string error;
|
|
require(make_cache_complement_plan(2, 3, {3, 5}, {{0, 1}, {1, 2}}, {}, offsets, bytes, error), error);
|
|
const std::vector<uint64_t> before = offsets;
|
|
// (1,0) moves into the GPU, (1,2) leaves it: (1,2) takes (1,0)'s bytes' place
|
|
require(exchange_cache_complement(offsets, 3, 5), "a valid exchange was refused");
|
|
require(offsets[5] == before[3] && offsets[3] == kNoCacheComplement, "the exchange did not move the place");
|
|
for (size_t i : {0u, 1u, 2u, 4u}) require(offsets[i] == before[i], "an exchange touched another expert");
|
|
// the copy's size and its set of places are unchanged
|
|
std::vector<uint64_t> a, b;
|
|
for (uint64_t o : before) if (o != kNoCacheComplement) a.push_back(o);
|
|
for (uint64_t o : offsets) if (o != kNoCacheComplement) b.push_back(o);
|
|
std::sort(a.begin(), a.end());
|
|
std::sort(b.begin(), b.end());
|
|
require(a == b, "an exchange changed the compact copy's places");
|
|
const std::vector<uint64_t> after = offsets;
|
|
require(!exchange_cache_complement(offsets, 3, 0), "an `in` the copy does not hold was accepted");
|
|
require(!exchange_cache_complement(offsets, 0, 2), "an `out` the copy holds already was accepted");
|
|
require(!exchange_cache_complement(offsets, 0, 0), "a self-exchange was accepted");
|
|
require(!exchange_cache_complement(offsets, 0, 6), "an out-of-range `out` was accepted");
|
|
require(offsets == after, "a refused exchange changed the offsets");
|
|
// and back: the copy is again what the plan makes for the original placement
|
|
require(exchange_cache_complement(offsets, 5, 3), "the reverse exchange was refused");
|
|
require(offsets == before, "exchanging back did not restore the plan");
|
|
}
|
|
|
|
void test_cgroup_memory_budget() {
|
|
using namespace strata::core::detail;
|
|
constexpr uint64_t GiB = 1ull << 30;
|
|
uint64_t bytes = 0;
|
|
|
|
CgroupMemoryStat clean_cache{40 * GiB, 32 * GiB, 0, 0, true};
|
|
require(cgroup_available_bytes(56 * GiB, clean_cache, bytes), "valid cgroup memory.stat was rejected");
|
|
require(bytes == 48 * GiB, "clean inactive file cache was not credited against the cgroup cap");
|
|
|
|
CgroupMemoryStat dirty_cache{40 * GiB, 32 * GiB, 4 * GiB, 2 * GiB, true};
|
|
require(cgroup_available_bytes(56 * GiB, dirty_cache, bytes), "valid dirty-cache counters were rejected");
|
|
require(bytes == 42 * GiB, "dirty/writeback pages were incorrectly counted as reclaimable");
|
|
|
|
CgroupMemoryStat oversized_inactive{10 * GiB, 20 * GiB, 0, 0, true};
|
|
require(cgroup_available_bytes(12 * GiB, oversized_inactive, bytes),
|
|
"valid oversized inactive-file counter was rejected");
|
|
require(bytes == 12 * GiB, "inactive-file accounting exceeded charged current usage");
|
|
|
|
const uint64_t max = std::numeric_limits<uint64_t>::max();
|
|
CgroupMemoryStat large_counters{max, max, max - 1, max, true};
|
|
require(cgroup_available_bytes(max, large_counters, bytes),
|
|
"large memory.stat counters were rejected");
|
|
require(bytes == 0, "dirty/writeback subtraction overflowed or escaped the usage cap");
|
|
|
|
CgroupMemoryStat usage_over_limit{60 * GiB, 0, 0, 0, true};
|
|
require(cgroup_available_bytes(56 * GiB, usage_over_limit, bytes),
|
|
"valid over-limit memory counters were rejected");
|
|
require(bytes == 0, "over-limit cgroup usage produced a positive budget");
|
|
|
|
CgroupMemoryStat missing_stat{};
|
|
bytes = 123;
|
|
require(!cgroup_available_bytes(56 * GiB, missing_stat, bytes) && bytes == 0,
|
|
"missing memory.stat counters did not fail closed");
|
|
}
|
|
|
|
// #633: the host RAM probe with fake /proc and cgroup trees: v2 (a limit, "max", a missing or malformed limit), v1
|
|
// (a limit, unlimited), and no cgroup line at all. Elsewhere than Linux it reads the machine's RAM.
|
|
void test_host_memory() {
|
|
using namespace strata::core::detail;
|
|
constexpr uint64_t GiB = 1ull << 30;
|
|
HostMemory m;
|
|
#if defined(__linux__)
|
|
TempDirectory t;
|
|
auto put = [&](const fs::path& rel, const std::string& text) {
|
|
fs::create_directories((t.path / rel).parent_path());
|
|
std::ofstream(t.path / rel) << text;
|
|
};
|
|
put("meminfo", "MemTotal: 134217728 kB\nMemAvailable: 104857600 kB\n"); // 100 GiB available
|
|
const std::string mi = (t.path / "meminfo").string(), cg = (t.path / "cgroup").string(),
|
|
root = (t.path / "fs").string();
|
|
auto probe = [&](const std::string& self) {
|
|
put("cgroup", self);
|
|
return host_available_memory(m, mi, cg, root);
|
|
};
|
|
// no cgroup line at all: MemAvailable alone (before #633: "cannot determine")
|
|
require(probe("") && m.available == 100 * GiB && m.cgroup_limit == ~0ull, "no cgroup: MemAvailable alone");
|
|
// v2, a 48 GiB limit with 16 GiB charged, 4 GiB of it clean cache
|
|
put("fs/cgroup.controllers", "memory\n");
|
|
put("fs/box/memory.max", std::to_string(48 * GiB) + "\n");
|
|
put("fs/box/memory.current", std::to_string(16 * GiB) + "\n");
|
|
put("fs/box/memory.stat", "inactive_file " + std::to_string(4 * GiB) + "\nfile_dirty 0\nfile_writeback 0\n");
|
|
require(probe("0::/box\n") && m.available == 36 * GiB && m.cgroup_limit == 48 * GiB, "v2 limit");
|
|
put("fs/box/memory.max", "max\n");
|
|
require(probe("0::/box\n") && m.available == 100 * GiB && m.cgroup_limit == ~0ull, "v2 max");
|
|
put("fs/box/memory.max", "48G\n");
|
|
require(!probe("0::/box\n"), "v2 malformed limit accepted");
|
|
fs::remove(t.path / "fs/box/memory.max");
|
|
require(!probe("0::/box\n"), "v2 group without memory.max accepted");
|
|
// v1: the memory controller's group (a hybrid line list), a 32 GiB limit with 10 GiB used
|
|
put("fs/memory/docker/abc/memory.limit_in_bytes", std::to_string(32 * GiB) + "\n");
|
|
put("fs/memory/docker/abc/memory.usage_in_bytes", std::to_string(10 * GiB) + "\n");
|
|
put("fs/memory/memory.limit_in_bytes", "9223372036854771712\n"); // the root: unlimited
|
|
put("fs/memory/memory.usage_in_bytes", std::to_string(50 * GiB) + "\n");
|
|
require(probe("12:cpu,cpuacct:/docker/abc\n4:memory:/docker/abc\n1:name=systemd:/docker/abc\n") &&
|
|
m.available == 22 * GiB && m.cgroup_limit == 32 * GiB, "v1 limit");
|
|
put("fs/memory/docker/abc/memory.limit_in_bytes", "9223372036854771712\n");
|
|
require(probe("4:memory:/docker/abc\n") && m.available == 100 * GiB && m.cgroup_limit == ~0ull, "v1 unlimited");
|
|
// a v1 group not visible here (another namespace): MemAvailable alone
|
|
require(probe("4:memory:/elsewhere\n") && m.available == 100 * GiB, "v1 group not mounted");
|
|
#else
|
|
require(host_available_memory(m) && m.available > 0 && m.cgroup_limit == ~0ull, "this PC's RAM");
|
|
(void) GiB;
|
|
#endif
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main() {
|
|
try {
|
|
test_complement_plan();
|
|
test_resident_lend_region();
|
|
test_resident_exchange();
|
|
test_cgroup_memory_budget();
|
|
test_host_memory();
|
|
test_canonical_layout();
|
|
#if defined(STRATA_NATIVE_EXPERTS)
|
|
test_native_variable_layout();
|
|
#endif
|
|
std::cout << "file_expert_source_test: PASS\n";
|
|
return 0;
|
|
} catch (const std::exception& error) {
|
|
std::cerr << "file_expert_source_test: " << error.what() << '\n';
|
|
return 1;
|
|
}
|
|
}
|