Arena under a hard memory limit also needs its headroom: a limit just above the arena no longer kills the engine

On l620 a 49 GiB cgroup limit and the 46.84 GiB arena: the engine's check (arena + blob <= limit) passed, the arena is pinned and never given back, the
engine's own memory (about 3 GiB) came on top, and the cgroup OOM killer ended the engine at the start (journal: "A process of this unit has been
killed by the OOM killer", twice). With the fallback allowed the arena now needs arena + blob + max(RAM to keep, the fallback headroom) under the
hard limit (cgroup, container, job object, --memory-limit-mib), else the RAM copy runs. Without a fallback (hand-set --ple-io ram, shared arena file,
STRATA_ARENA_RAM_FALLBACK=0) the comparison is unchanged. arena_hard_limit_need in ple_residency.hpp, unit cases for a limit just above the arena.

include/strata/core/expert_source.hpp

@@ -1316,6 +1316,9 @@ public:
     /// --ram-keep-gib W, in bytes (0: none): the arena is chosen only if this much RAM stays available after it (ple_residency.hpp,
     /// arena_need_with_keep).  Set before `open`.  With the fallback not allowed the shortfall is only warned about, as for any other.
     uint64_t ram_keep_bytes = 0;
+    /// The headroom the arena must find under a HARD memory limit (cgroup, container, job object) beyond its own size when the fallback is
+    /// allowed: the RAM copy's headroom (ple_residency.hpp, arena_hard_limit_need).  0: none.  Set before `open`.
+    uint64_t hard_limit_headroom_bytes = 0;
     bool ram_short() const { return ram_short_; }
     /// #633: set by `open` when the RAM available is less than the arena needs (a recommendation; it still loads)
     const std::string& ram_warning() const { return ram_warning_; }

include/strata/core/ple_residency.hpp

@@ -95,6 +95,16 @@ inline uint64_t arena_need_with_keep(uint64_t need, uint64_t keep) {
     return keep > UINT64_MAX - need ? UINT64_MAX : need + keep;
 }
 
+/// What the arena needs under a HARD memory limit (a cgroup, a container, a job object, --memory-limit-mib).  The arena is pinned: it is
+/// never given back to make room, and the engine's own memory (about 3 GiB: buffers, the draft head, the PLE and PLE row cache, the CUDA
+/// context) comes on top, so a limit just above the arena kills the engine at the start (cgroup OOM, found on l620: a 49 GiB limit and a
+/// 46.84 GiB arena).  When the fallback is allowed the arena needs its headroom under the limit as well, max(keep, the fallback headroom);
+/// when it is not allowed (a hand-set --ple-io ram, a shared arena file, STRATA_ARENA_RAM_FALLBACK=0) there is nothing to fall back to, so
+/// the comparison stays as it was (the arena and the RAM to keep).
+inline uint64_t arena_hard_limit_need(uint64_t base_need, uint64_t keep, uint64_t headroom, bool fallback_allowed) {
+    return arena_need_with_keep(base_need, fallback_allowed ? (keep > headroom ? keep : headroom) : keep);
+}
+
 /// The headroom the RAM-copy fallback leaves free.  `env_headroom`: STRATA_RESIDENT_HEADROOM_GIB when it is set (wins, even below the safety headroom).
 inline uint64_t fallback_headroom(std::optional<uint64_t> env_headroom, uint64_t keep, uint64_t safety) {
     if (env_headroom) return *env_headroom;

src/core/expert_source.cpp

@@ -5874,15 +5874,17 @@ bool ArenaExpertSource::open(const std::string& pack_dir, int64_t n_layers, int6
     if (detail::HostMemory hm; shared_arena_file.empty() && detail::host_available_memory(hm)) {
         // --ram-keep-gib W: the arena is chosen only if W GiB stay free after it, so its need here is the arena plus W
         const uint64_t need = arena_need_with_keep(want + (uint64_t) blob, ram_keep_bytes);
+        // under a hard limit the arena also needs the headroom (the engine's own memory comes on top of it, and the arena is never given back)
+        const uint64_t limit_need = arena_hard_limit_need(want + (uint64_t) blob, ram_keep_bytes, hard_limit_headroom_bytes, fall_back_when_ram_short);
         const double gib = 1073741824.0;
         char buf[512];
         char keep_note[96] = "";
         if (ram_keep_bytes > 0)
             std::snprintf(keep_note, sizeof keep_note, ", with the %.2f GiB --ram-keep-gib asks to leave free", (double) ram_keep_bytes / gib);
-        if (arena_ram_verdict(need, ~uint64_t{0}, hm.cgroup_limit, fall_back_when_ram_short) != ArenaRamVerdict::Load) {
+        if (arena_ram_verdict(limit_need, ~uint64_t{0}, hm.cgroup_limit, fall_back_when_ram_short) != ArenaRamVerdict::Load) {
             if (fall_back_when_ram_short) {   // #4846: the caller reads the experts from the file instead of refusing
-                std::snprintf(buf, sizeof buf, "it needs %.2f GiB%s and this process's memory limit is %.2f GiB",
-                              (double) need / gib, keep_note, (double) hm.cgroup_limit / gib);
+                std::snprintf(buf, sizeof buf, "it needs %.2f GiB (with the %.2f GiB the engine's own memory and the RAM to keep take on top of it) and this process's memory limit is %.2f GiB",
+                              (double) limit_need / gib, (double) (limit_need - (want + (uint64_t) blob)) / gib, (double) hm.cgroup_limit / gib);
                 err = buf;
                 ram_short_ = true;
                 return false;

src/core/ple_residency_test.cpp

@@ -128,6 +128,30 @@ int main() {
         check(!ram_keep_plan(need, f446, 0, std::nullopt, win).arena && ram_keep_plan(need, f446, 0, std::nullopt, win).budget > 36 * GiB,
               "44.6 GiB free, W=0: the fallback with a 36.6 GiB copy");
     }
+    // A HARD memory limit just above the arena (R5070, found on l620: a 49 GiB cgroup limit, the 46.84 GiB arena chosen, the engine killed by the
+    // cgroup OOM killer at the start).  With the fallback allowed the arena needs its headroom under the limit too.
+    {
+        const uint64_t need = 46ull * GiB + 860ull * (1ull << 20) + 2660000ull;   // IQ3_S: 46.84 GiB of experts and the spare blob
+        const uint64_t head4 = 4 * GiB, head8 = 8 * GiB;
+        auto verdict = [&](uint64_t limit, uint64_t keep, uint64_t headroom, bool fallback) {
+            return arena_ram_verdict(arena_hard_limit_need(need, keep, headroom, fallback), ~uint64_t{0}, limit, fallback);
+        };
+        check(arena_hard_limit_need(need, 0, head4, true) == need + head4, "fallback allowed: the arena needs its headroom under the limit");
+        check(arena_hard_limit_need(need, 8 * GiB, head4, true) == need + 8 * GiB, "the RAM to keep counts when it is larger than the headroom");
+        check(arena_hard_limit_need(need, 0, head4, false) == need, "fallback not allowed: the comparison stays as it was");
+        check(arena_hard_limit_need(need, 8 * GiB, head4, false) == need + 8 * GiB, "fallback not allowed: the RAM to keep is still asked");
+        check(verdict(49ull * GiB, 0, head4, true) == ArenaRamVerdict::FallBack, "49 GiB limit, 4 GiB headroom: the copy, not an arena that the OOM killer ends");
+        check(verdict(49ull * GiB, 0, head8, true) == ArenaRamVerdict::FallBack, "49 GiB limit, 8 GiB headroom (Windows): the copy");
+        check(verdict(51ull * GiB, 0, head4, true) == ArenaRamVerdict::Load, "51 GiB limit, 4 GiB headroom: the arena");
+        check(verdict(need + head4, 0, head4, true) == ArenaRamVerdict::Load, "a limit of exactly arena + headroom: the arena");
+        check(verdict(need + head4 - 1, 0, head4, true) == ArenaRamVerdict::FallBack, "one byte under: the copy");
+        check(verdict(need - 1, 0, head4, true) == ArenaRamVerdict::FallBack, "a limit under the arena: the copy, as before");
+        check(verdict(49ull * GiB, 0, head4, false) == ArenaRamVerdict::Load, "fallback not allowed, 49 GiB limit: as before (no new refusal)");
+        check(verdict(45ull * GiB, 0, head4, false) == ArenaRamVerdict::Refuse, "fallback not allowed, limit under the arena: refused as before");
+        check(verdict(60ull * GiB, 8 * GiB, head4, true) == ArenaRamVerdict::Load, "60 GiB limit, 8 GiB to keep: the arena");
+        check(verdict(54ull * GiB, 8 * GiB, head4, true) == ArenaRamVerdict::FallBack, "54 GiB limit, 8 GiB to keep: the copy");
+        check(verdict(std::numeric_limits<uint64_t>::max(), 0, head4, true) == ArenaRamVerdict::Load, "no limit: the arena");
+    }
     std::printf("ple_residency_test: %d checks passed\n", checks);
     return 0;
 }

src/program/generate.cpp

@@ -7082,6 +7082,7 @@ int main(int argc, char** argv) {
         strata::core::pin_reserve_take();
         arena_src.fall_back_when_ram_short = ram_fallback_ok && !ram_fallback_failed;
         arena_src.ram_keep_bytes = o.ram_keep;
+        arena_src.hard_limit_headroom_bytes = ram_fallback_headroom;
         if (o.ram_keep > 0)
             std::fprintf(stderr, "strata generate: --ram-keep-gib %.2f: the expert arena is used only if %.2f GiB of RAM stay available after it; "
                                  "otherwise the RAM copy of the hottest experts leaves %.2f GiB free (the larger of that and the %.0f GiB safety headroom, "