serve: wait for measured RAM before retrying reclaimed SAVE

docs/DETAILS.md

@@ -988,6 +988,13 @@ that floor or assume that freed allocations have reached the OS. Pinned prefixes
 stay intact, including their order and tie preference. Unknown RAM telemetry or an unknown allocation estimate
 does not evict anything. The usual admission check still runs before copying state.
 
+With this option the HTTP server also allows up to 30 seconds for physical RAM accounting to catch up after
+cache reclamation. Only an unpublished `memory` refusal from the SAVE RAM preflight is retried, at most three
+times, when measured available RAM reaches the reported estimate plus the original floor and 2 MiB for rounding.
+The same live engine process and service FIFO are held throughout; a process change, death, unknown telemetry
+or missing estimate ends the wait. This does not unload a model or trim another process's working set. Persistent
+pressure still returns the last refusal. The wait is outside the engine's reported `save_ms`.
+
 This can trade later prefix-cache hits for room to save. Released caches stay released even if admission or file
 I/O later fails; the current tokens, images, steering and live device state are unchanged. It cannot guarantee a
 SAVE under arbitrary memory pressure. With the option off, or enough RAM at the first probe, cache retention and

serve/server.py

@@ -604,6 +604,7 @@ class StrataEngine:
     def __init__(self, exe: str, args: list[str], cwd: str | None = None, log: str | None = None,
                  env: dict | None = None, lazy: bool = False):
         self.spawn = (exe, list(args), cwd, log, env)   # to start it again after it died (issue #27)
+        self.session_save_reclaim = '--session-save-reclaim' in args
         paths = {k: v for k, v in zip(args, args[1:]) if k in ("--native", "--pack")}
         self.model_path = paths.get("--native") or paths.get("--pack", "pack/full")
         self.log_path = log
@@ -2707,7 +2708,11 @@ class Service:
                                        else "restoring a session", started=time.time(), first_token=None,
                                        prompt_tokens=None, generated=None, max_tokens=None)
                 try:
-                    r = self.engine.session_file(action, path)
+                    if action == 'save' and getattr(self.engine, 'session_save_reclaim', False):
+                        from serve.session_save_retry import session_file
+                        r = session_file(self, action, path, SessionRefused)
+                    else:
+                        r = self.engine.session_file(action, path)
                 except SessionRefused as e:
                     body = error(e.status, str(e))
                     body[1]["error"]["kind"] = e.kind

serve/session_save_retry.py

@@ -0,0 +1,84 @@
+"""Bounded SAVE admission retries after opt-in engine cache reclamation.
+
+The caller holds the service FIFO. Only an unpublished, typed RAM preflight
+refusal is retried, on the same live engine process and with its floor intact.
+No model unload, working-set trim or unrelated process is involved.
+"""
+from __future__ import annotations
+
+import re
+import time
+
+SETTLE_SECONDS = 30.0
+MAX_RETRIES = 3
+POLL_SECONDS = 0.5
+
+
+def required_mib(error) -> int | None:
+    match = re.search(r'^not enough RAM to save the session '
+                      r'\((\d+) MiB plus a floor of (\d+) MiB needed,', str(error))
+    # Each diagnostic term rounds bytes down; two MiB preserve admission headroom.
+    return sum(map(int, match.groups())) + 2 if match else None
+
+
+def available_mib() -> int | None:
+    try:
+        import psutil
+        available = psutil.virtual_memory().available
+        if isinstance(available, bool) or not isinstance(available, int) or available < 0:
+            return None
+        return available // 2**20
+    except (ImportError, OSError, AttributeError, ValueError):
+        return None
+
+
+def retryable(action, error) -> bool:
+    return (action == 'save' and error.kind == 'memory' and not error.published
+            and required_mib(error) is not None)
+
+
+def session_file(service, action, path, refused_type):
+    """Called only when --session-save-reclaim is enabled; never changes RESTORE."""
+    engine = service.engine
+    process = getattr(engine, 'proc', None)
+    try:
+        return engine.session_file(action, path)
+    except refused_type as error:
+        if not retryable(action, error) or process is None:
+            raise
+        deadline = time.monotonic() + SETTLE_SECONDS
+        last_error = error
+
+        def same_process():
+            try:
+                return (service.engine is engine and engine.proc is process
+                        and not getattr(engine, 'ended', False) and process.poll() is None)
+            except OSError:
+                return False
+
+        print('[strata] SAVE waiting for RAM after cache reclaim: '
+              f'need={required_mib(error)} MiB, limit={SETTLE_SECONDS:g}s; floor unchanged', flush=True)
+        for attempt in range(MAX_RETRIES):
+            while True:
+                if not same_process() or time.monotonic() >= deadline:
+                    raise last_error
+                available = available_mib()
+                needed = required_mib(last_error)
+                if available is None or needed is None:
+                    raise last_error
+                if available >= needed:
+                    break
+                time.sleep(max(0.0, min(POLL_SECONDS, deadline - time.monotonic())))
+            # No replacement, restart, or blind retry when telemetry is unknown.
+            if not same_process():
+                raise last_error
+            try:
+                return engine.session_file(action, path)
+            except refused_type as failure:
+                if not retryable(action, failure):
+                    raise
+                last_error = failure
+                if attempt + 1 == MAX_RETRIES or time.monotonic() >= deadline:
+                    raise
+                time.sleep(max(0.0, min(POLL_SECONDS, deadline - time.monotonic())))
+        raise last_error

serve/test_session_save_retry.py

@@ -0,0 +1,192 @@
+"""SAVE memory settling on a fake clock and real HTTP/FIFO; no GPU or model."""
+from __future__ import annotations
+import io
+from pathlib import Path
+import tempfile
+import threading
+import time
+from types import SimpleNamespace
+import unittest
+from unittest import mock
+
+from serve import session_save_retry as retry
+from serve.frontend import ChatTemplate
+from serve.server import ByteTokenizer, EngineDied, Service, SessionRefused, serve
+from serve import test_slots as slot_tests
+from serve.test_slots import FakeEngine, FakeProc
+
+RESULT = {'tokens': 62993, 'bytes': 1234, 'ms': 401.5}
+
+def refused(need=355, floor=768):
+    return SessionRefused('memory', f'not enough RAM to save the session ({need} MiB plus a floor of {floor} MiB needed, 503 MiB available)')
+
+class Clock:
+    def __init__(self): self.now=0.0; self.on_sleep=lambda:None
+    def monotonic(self): return self.now
+    def sleep(self, seconds):
+        assert 0 <= seconds <= .5
+        self.now += seconds
+        self.on_sleep()
+
+class Retry(unittest.TestCase):
+    def setUp(self):
+        self.clock=Clock()
+        self.proc=SimpleNamespace(poll=lambda:None)
+        self.engine=SimpleNamespace(proc=self.proc,ended=False,session_file=mock.Mock())
+        self.service=SimpleNamespace(engine=self.engine)
+        self.clock_patch=mock.patch.object(retry,'time',self.clock)
+        self.clock_patch.start()
+        self.addCleanup(self.clock_patch.stop)
+        self.stdout=mock.patch('sys.stdout',new_callable=io.StringIO)
+        self.stdout.start(); self.addCleanup(self.stdout.stop)
+
+    def call(self, available=5000, action='save'):
+        with mock.patch.object(retry,'available_mib',side_effect=available if isinstance(available,list) else None,
+                               return_value=None if isinstance(available,list) else available):
+            return retry.session_file(self.service,action,'same.bin',SessionRefused)
+
+    def test_success_does_not_poll_or_wait(self):
+        self.engine.session_file.return_value=RESULT
+        with mock.patch.object(retry,'available_mib',side_effect=AssertionError('unexpected memory probe')):
+            self.assertEqual(retry.session_file(self.service,'save','same.bin',SessionRefused),RESULT)
+        self.assertEqual(self.clock.now,0)
+        self.engine.session_file.assert_called_once_with('save','same.bin')
+
+    def test_delayed_ram_reuses_same_process_path_and_floor(self):
+        self.engine.session_file.side_effect=[refused(),RESULT]
+        self.assertEqual(self.call([534,541,800,1400]),RESULT)
+        self.assertEqual(self.clock.now,1.5)
+        self.assertIs(self.engine.proc,self.proc)
+        self.assertEqual(self.engine.session_file.call_args_list,[mock.call('save','same.bin')]*2)
+        self.assertEqual(retry.required_mib(refused()),1125)
+
+    def test_persistent_pressure_expires_without_another_save(self):
+        error=refused(); self.engine.session_file.side_effect=error
+        with self.assertRaises(SessionRefused) as got: self.call(541)
+        self.assertIs(got.exception,error)
+        self.assertEqual(self.clock.now,30)
+        self.assertEqual(self.engine.session_file.call_count,1)
+
+    def test_admission_races_are_bounded_and_keep_latest_error(self):
+        errors=[refused(355+n) for n in range(4)]
+        self.engine.session_file.side_effect=errors
+        with self.assertRaises(SessionRefused) as got: self.call()
+        self.assertIs(got.exception,errors[-1])
+        self.assertEqual(self.engine.session_file.call_count,4)  # first SAVE plus three retries
+        self.assertEqual(self.clock.now,1)
+
+    def test_new_refusal_rechecks_its_larger_estimate(self):
+        self.engine.session_file.side_effect=[refused(),refused(2048),RESULT]
+        self.assertEqual(self.call([1400,1400,2900]),RESULT)
+        self.assertEqual(self.clock.now,1)
+
+    def test_unknown_telemetry_or_estimate_never_retries(self):
+        for error,available in [(refused(),None),
+              (SessionRefused('memory','not enough RAM to save the session'),5000),
+              (SessionRefused('memory','not enough RAM to save the session (unknown MiB plus a floor of 768 MiB needed,'),5000)]:
+            with self.subTest(error=error):
+                self.engine.session_file=mock.Mock(side_effect=error)
+                with self.assertRaises(SessionRefused) as got: self.call(available)
+                self.assertIs(got.exception,error)
+                self.assertEqual(self.engine.session_file.call_count,1)
+
+    def test_restore_other_categories_and_published_files_never_retry(self):
+        for action,error in [('restore',refused()),('save',SessionRefused('io',str(refused()))),
+               ('save',SessionRefused('storage',str(refused()))),('save',SessionRefused('invalid',str(refused()))),
+               ('save',SessionRefused('memory',str(refused()),published=True))]:
+            with self.subTest(action=action,error=error):
+                self.engine.session_file=mock.Mock(side_effect=error)
+                with self.assertRaises(SessionRefused) as got: self.call(action=action)
+                self.assertIs(got.exception,error)
+                self.assertEqual(self.engine.session_file.call_count,1)
+
+    def test_process_death_or_replacement_ends_wait(self):
+        for change in [lambda:setattr(self.proc,'poll',lambda:1),
+                       lambda:setattr(self.engine,'proc',SimpleNamespace(poll=lambda:None)),
+                       lambda:setattr(self.service,'engine',SimpleNamespace(proc=self.proc)),
+                       lambda:setattr(self.engine,'ended',True)]:
+            with self.subTest(change=change):
+                self.proc=SimpleNamespace(poll=lambda:None)
+                self.engine=SimpleNamespace(proc=self.proc,ended=False,session_file=mock.Mock(side_effect=refused()))
+                self.service.engine=self.engine
+                self.clock.now=0; self.clock.on_sleep=change
+                with self.assertRaises(SessionRefused): self.call(541)
+                self.assertEqual(self.engine.session_file.call_count,1)
+                self.assertEqual(self.clock.now,.5)
+
+    def test_non_admission_error_during_retry_is_returned(self):
+        for error in [EngineDied('gone'), SessionRefused('io','flushing directory',published=True)]:
+            self.engine.session_file=mock.Mock(side_effect=[refused(),error])
+            with self.assertRaises(type(error)) as got: self.call()
+            self.assertIs(got.exception,error)
+            self.assertEqual(self.engine.session_file.call_count,2)
+
+    def test_memory_probe_is_portable_and_unknown_fails_closed(self):
+        for value,expected in [(1400*2**20,1400),(0,0),(-1,None),(True,None),(None,None),(float('nan'),None)]:
+            with mock.patch.dict('sys.modules',{'psutil':SimpleNamespace(virtual_memory=lambda:SimpleNamespace(available=value))}):
+                self.assertEqual(retry.available_mib(),expected)
+
+class AdmissionProc(FakeProc):
+    def write(self,text):
+        if not self.sent and text.startswith('SAVE '):
+            self.sent.append(text)
+            self.engine.lines.put('SERR memory 0 '+str(refused())+'\n')
+            return
+        super().write(text)
+
+class HTTP(unittest.TestCase):
+    post=slot_tests.Slots.post
+    def setUp(self):
+        self.directory=tempfile.TemporaryDirectory()
+        self.addCleanup(self.directory.cleanup)
+        tok=ByteTokenizer(); self.engine=FakeEngine(); self.engine.session_save_reclaim=True
+        self.engine.proc=AdmissionProc(self.engine)
+        self.svc=Service(self.engine,tok,ChatTemplate(Path(__file__).parent/'chat_template.jinja'))
+        self.svc.slot_save_path=self.directory.name
+        self.httpd=serve(self.svc,port=0)
+        self.base=f'http://127.0.0.1:{self.httpd.server_address[1]}'
+        self.addCleanup(self.httpd.server_close); self.addCleanup(self.httpd.shutdown)
+
+    def test_opt_in_recovers_native_refusal_and_keeps_response_and_process(self):
+        proc=self.engine.proc
+        with mock.patch.object(retry,'available_mib',return_value=1400):
+            status,body=self.post('/slots/0?action=save',{'filename':'a.bin'})
+        self.assertEqual(status,200,body)
+        self.assertEqual(body,{'id_slot':0,'filename':'a.bin','n_saved':62993,'n_written':1234,'timings':{'save_ms':401.5}})
+        self.assertEqual(proc.sent,[f'SAVE {Path(self.directory.name,"a.bin")}\n']*2)
+        self.assertIs(self.engine.proc,proc); self.assertFalse(self.engine.ended)
+        self.assertEqual(Path(self.directory.name,'a.bin').read_bytes(),b'x'*1234)
+
+    def test_default_refusal_is_immediate_and_byte_identical(self):
+        self.engine.session_save_reclaim=False
+        with mock.patch.object(retry,'available_mib',side_effect=AssertionError('default must not probe')):
+            status,body=self.post('/slots/0?action=save',{'filename':'off.bin'})
+        self.assertEqual(status,503)
+        self.assertEqual(body,{'error':{'code':503,'message':str(refused()),'type':'server_error','kind':'memory'}})
+        self.assertEqual(len(self.engine.proc.sent),1)
+        self.assertFalse(Path(self.directory.name,'off.bin').exists())
+
+    def test_fifo_blocks_restore_while_save_waits_and_releases_afterwards(self):
+        waiting=threading.Event(); enough=threading.Event()
+        def available(): waiting.set(); return 1400 if enough.is_set() else 541
+        result={}
+        path=Path(self.directory.name,'fifo.bin'); path.write_bytes(b'old file')
+        a=threading.Thread(target=lambda:result.update(save=self.post('/slots/0?action=save',{'filename':'fifo.bin'})))
+        b=threading.Thread(target=lambda:result.update(restore=self.post('/slots/0?action=restore',{'filename':'fifo.bin'})))
+        with mock.patch.object(retry,'available_mib',side_effect=available), mock.patch.object(retry,'POLL_SECONDS',.01):
+            try:
+                a.start(); self.assertTrue(waiting.wait(2)); b.start()
+                deadline=time.monotonic()+2
+                while self.svc.status['queued'] != 1 and time.monotonic()<deadline: time.sleep(.01)
+                self.assertTrue(self.svc.status['busy']); self.assertEqual(self.svc.status['queued'],1)
+                self.assertEqual(len(self.engine.proc.sent),1)
+            finally:
+                enough.set(); a.join(5)
+                if b.ident is not None: b.join(5)
+        self.assertFalse(a.is_alive() or b.is_alive())
+        self.assertEqual(result['save'][0],200); self.assertEqual(result['restore'][0],200)
+        self.assertEqual(self.engine.proc.sent,[f'SAVE {path}\n',f'SAVE {path}\n',f'RESTORE {path}\n'])
+        self.assertFalse(self.svc.status['busy']); self.assertEqual(self.svc.status['queued'],0)
+        self.assertTrue(self.svc.fifo.acquire(blocking=False)); self.svc.fifo.release()
+
+if __name__=='__main__': unittest.main()