serve: return 503 when the elastic image encoder cannot start (#2435)

When the elastic image encoder failed to restart, its exception left request preparation and the client got a closed connection. It now goes
through the EngineDied path: HTTP 503 with the reason, a later picture tries again, text requests and a later good start are unaffected. An idle
encoder stop is logged even when the cache shrink was refused, and a layer split gets a line saying its encoder VRAM has to be kept free with
--vram-reserve-mib. Ported from PR #2435 by i-Vladimir-i with its tests (chat, messages and responses, streamed and not, plus a real encoder
that exits with code 71 on Linux) and the docs/VRAM_ELASTIC.md text.

docs/VRAM_ELASTIC.md

@@ -113,7 +113,9 @@ and starts the encoder with the model, as before) the GPU encoder is not started
 - **A picture's turn** (the same FIFO turn an encode always takes): if the encoder is not running, the engine's
   `VRAM` command first keeps `"elastic_mib"` (default 2048) free on top of the reserve - the last `POST /v1/vram`, or 0
   when none was made - and then the encoder starts. If the cache cannot shrink (the resident low-RAM mode, an engine
-  without `--vram-elastic`) it starts anyway, as a non-elastic encoder would, and the log says why.
+  without `--vram-elastic`, a layer split) it starts anyway, as a non-elastic encoder would, and the log says why; on a
+  layer split the server says at start that the encoder's VRAM has to be kept free with `--vram-reserve-mib`. An
+  encoder that fails to start (out of VRAM) ends the request with a `503`; the next picture tries again.
 - **Idle** for `"elastic_idle_s"` (default 120) without a picture: a thread stops the encoder between requests and
   sends `VRAM` with that reserve again, so the cache grows back to the size it had (0, not the start's reserve: that
   would leave it a segment or more smaller).

serve/server.py

@@ -4202,7 +4202,10 @@ class Service:
         if self.vision_held:
             time.sleep(0.2)          # the driver's free-VRAM counter shows the cache's release a moment later
         free0, t0 = self.free_vram_mib(), time.time()
-        v.restart()
+        try:
+            v.restart()
+        except Exception as e:      # a 503 as for a non-elastic encoder at a load, not a dropped connection
+            raise EngineDied(f"the image encoder failed to start: {e}") from e
         free1 = self.free_vram_mib()
         took = f", {free0 - free1} MiB of VRAM" if free0 is not None and free1 is not None else ""
         print(f"[strata] images: the image encoder started in {time.time() - t0:.1f} s{took}", flush=True)
@@ -4220,11 +4223,14 @@ class Service:
             return "busy"
         try:
             v.unload()
+            back = ""
             if self.vision_held and hasattr(self.engine, "vram") and self.loaded():
                 out = self.engine.vram(self._engine_reserve())
-                print(f"[strata] images: the image encoder stopped after {v.elastic_idle_s:g} s without a picture; "
-                      f"the expert cache took its VRAM back ({out.get('expert_cache_mib')} of "
-                      f"{out.get('expert_cache_full_mib')} MiB)", flush=True)
+                back = (f"; the expert cache took its VRAM back ({out.get('expert_cache_mib')} of "
+                        f"{out.get('expert_cache_full_mib')} MiB)")
+            # also when the cache never made room (a layer split, an engine without --vram-elastic): the stop happened
+            print(f"[strata] images: the image encoder stopped after {v.elastic_idle_s:g} s without a picture{back}",
+                  flush=True)
             self.vision_held = False
             return "down"
         finally:
@@ -8377,6 +8383,12 @@ def main(argv=None, *, lifecycle_event=None, interrupt_exit_code: int = 0) -> in
                 print('[strata] the vision section\'s "elastic" needs "vram_elastic": true (the expert cache gives the '
                       "encoder its VRAM): the encoder starts with the model instead", flush=True)
                 vcfg["elastic"] = False
+            elif vcfg.get("elastic") is True and layer_split_of(cfg):
+                # the engine turns --vram-elastic off on a layer split (a cache per GPU): the encoder still starts
+                # with the first picture and stops when idle, but its VRAM has to be free already
+                print('[strata] images: on a layer split the expert cache cannot make room for the image encoder (the '
+                      'engine turns --vram-elastic off); keep its VRAM free with "--vram-reserve-mib" on the '
+                      "encoder's GPU, or the first picture may not get it", flush=True)
             vision = Vision(vcfg, log=open(cfg["log"], "a", encoding="utf-8") if cfg.get("log") else None,
                             env=vision_env(cfg, env), lazy=lazy)
         print("model unloaded; the first request loads it ..." if lazy else

serve/test_vision_elastic.py

@@ -2,13 +2,15 @@
 expert cache keep the encoder's VRAM free and then starts it, later pictures use it as it is, an idle encoder stops and
 the cache takes its VRAM back, a POST /v1/vram reserve stacks on the encoder's share, and the engine's load never starts
 it on its own."""
+import contextlib
+import io
 import tempfile
 import time
 import unittest
 from pathlib import Path
 
 from serve.frontend import ChatTemplate
-from serve.server import ByteTokenizer, Service, Vision
+from serve.server import ByteTokenizer, EngineDied, Service, Vision
 from serve.test_vram import VramEngine
 
 ROOT = Path(__file__).resolve().parents[1]
@@ -99,6 +101,27 @@ class Elastic(unittest.TestCase):
         self.assertEqual(self.svc._vision_idle_down(), "down")
         self.assertEqual(self.engine.calls, [2048])         # the refused shrink only: nothing to give back
 
+    def test_an_encoder_that_fails_to_start_is_a_503_not_a_dropped_connection(self):
+        def fail():
+            raise RuntimeError("the vision encoder did not start: cudaMalloc failed: out of memory")
+        self.vision.restart = fail
+        with self.assertRaises(EngineDied) as e:            # the handlers answer EngineDied with a 503
+            self.picture()
+        self.assertIn("the image encoder failed to start", str(e.exception))
+        self.svc.last_image_at = 0
+        self.assertEqual(self.svc._vision_idle_down(), "down")
+        self.assertEqual(self.engine.calls, [2048, 0])      # the share made for it goes back
+
+    def test_the_stop_is_logged_when_the_cache_never_made_room(self):
+        self.engine.elastic = False                         # a layer split: the engine refuses --vram-elastic
+        self.picture()
+        self.svc.last_image_at = 0
+        out = io.StringIO()
+        with contextlib.redirect_stdout(out):
+            self.assertEqual(self.svc._vision_idle_down(), "down")
+        self.assertIn("the image encoder stopped after 60 s without a picture", out.getvalue())
+        self.assertNotIn("expert cache", out.getvalue())
+
     def test_the_load_never_starts_an_elastic_encoder(self):
         self.assertFalse(self.svc._vision_down())
         self.svc.vision = ElasticVision(tempfile.mkdtemp(), elastic=False)

serve/test_vision_elastic_http.py

@@ -0,0 +1,140 @@
+"""Elastic encoder failures reach clients, can be retried, and do not prevent recovery."""
+import json
+from pathlib import Path
+import sys
+import tempfile
+import unittest
+import urllib.error
+import urllib.request
+
+from serve.frontend import ChatTemplate
+from serve.server import ByteTokenizer, Service, Vision, serve
+from serve.test_vision_elastic import ElasticVision
+from serve.test_vram import VramEngine
+
+
+class EncoderFailureHTTP(unittest.TestCase):
+    def setUp(self):
+        tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(tmp.cleanup)
+        tok = ByteTokenizer()
+        self.engine = VramEngine(tok)
+        self.engine.loaded = True
+        self.engine.elastic = False  # as on a layer split: the cache cannot make room
+        self.vision = ElasticVision(tmp.name)
+        self.svc = Service(self.engine, tok, ChatTemplate(Path(__file__).parent / 'chat_template.jinja'),
+                           vision=self.vision)
+        httpd = serve(self.svc, port=0)
+        self.addCleanup(httpd.server_close)
+        self.addCleanup(httpd.shutdown)
+        self.addCleanup(self.drop_embeddings)
+        self.base = f'http://127.0.0.1:{httpd.server_address[1]}'
+        self.attempts = 0
+        self.original_restart = self.vision.restart
+
+    def drop_embeddings(self):
+        path = getattr(getattr(self.svc, 'embeddings', None), 'path', None)
+        if path is not None:
+            path.unlink(missing_ok=True)
+
+    def post(self, path, body):
+        req = urllib.request.Request(self.base + path, json.dumps(body).encode(),
+                                     {'Content-Type': 'application/json'})
+        return urllib.request.urlopen(req, timeout=10)
+
+    @staticmethod
+    def chat_picture():
+        return {'model': 'm', 'max_tokens': 8, 'messages': [{'role': 'user', 'content': [
+            {'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,eA=='}}]}]}
+
+    def fail_with(self, error_type):
+        def fail():
+            self.attempts += 1
+            raise error_type('synthetic encoder startup failure')
+        self.vision.restart = fail
+
+    def assert_503(self, path, body):
+        with self.assertRaises(urllib.error.HTTPError) as got:
+            self.post(path, body)
+        with got.exception as response:
+            self.assertEqual(response.code, 503)
+            payload = json.load(response)
+            self.assertIn('synthetic encoder startup failure', payload['error']['message'])
+
+    def test_start_failure_and_retry_on_all_three_apis(self):
+        cases = [
+            ('/v1/chat/completions', self.chat_picture()),
+            ('/v1/messages', {'model': 'm', 'max_tokens': 8, 'messages': [{'role': 'user', 'content': [
+                {'type': 'image', 'source': {'type': 'base64', 'media_type': 'image/png', 'data': 'eA=='}}]}]}),
+            ('/v1/responses', {'model': 'm', 'max_output_tokens': 8, 'input': [{'role': 'user', 'content': [
+                {'type': 'input_image', 'image_url': 'data:image/png;base64,eA=='}]}]}),
+        ]
+        for error_type in (RuntimeError, OSError):
+            self.fail_with(error_type)
+            for path, body in cases:
+                for stream in (False, True):
+                    for retry in range(2):
+                        with self.subTest(error=error_type.__name__, path=path, stream=stream, retry=retry):
+                            before = self.attempts
+                            self.assert_503(path, {**body, 'stream': stream})
+                            self.assertEqual(self.attempts, before + 1)
+
+    def test_text_requests_and_image_recovery_after_a_failed_start(self):
+        self.fail_with(RuntimeError)
+        text = {'model': 'm', 'max_tokens': 8, 'messages': [{'role': 'user', 'content': 'hi'}]}
+        with self.post('/v1/chat/completions', text) as response:
+            self.assertEqual(response.status, 200)
+            response.read()
+        self.assertEqual(self.attempts, 0)
+        self.assert_503('/v1/chat/completions', self.chat_picture())
+        self.assertEqual(self.attempts, 1)
+        self.vision.restart = self.original_restart
+        with self.post('/v1/chat/completions', self.chat_picture()) as response:
+            self.assertEqual(response.status, 200)
+            response.read()
+        self.assertEqual(self.vision.starts, 1)
+        with self.post('/v1/chat/completions', text) as response:
+            self.assertEqual(response.status, 200)
+            response.read()
+        self.assertEqual(self.vision.starts, 1)
+
+    @unittest.skipUnless(sys.platform.startswith('linux'), 'child reaping is checked through /proc')
+    def test_real_encoder_failure_retries_and_reaps_each_child(self):
+        with tempfile.TemporaryDirectory() as directory:
+            started = Path(directory) / 'started-pids'
+            executable = Path(directory) / 'encoder-exits'
+            executable.write_text(
+                '#!' + sys.executable + '\nimport os\n'
+                'with open(' + repr(str(started)) + ', "a") as output:\n'
+                '    output.write(str(os.getpid()) + "\\n")\n'
+                'raise SystemExit(71)\n')
+            executable.chmod(0o700)
+            vision = Vision({'exe': str(executable), 'mmproj': 'unused', 'model': 'unused',
+                             'elastic': True}, lazy=True)
+            self.addCleanup(vision.unload)
+            self.svc.vision = vision
+            cases = [
+                ('/v1/chat/completions', self.chat_picture()),
+                ('/v1/messages', {'model': 'm', 'max_tokens': 8, 'messages': [{'role': 'user', 'content': [
+                    {'type': 'image', 'source': {'type': 'base64', 'media_type': 'image/png', 'data': 'eA=='}}]}]}),
+                ('/v1/responses', {'model': 'm', 'max_output_tokens': 8, 'input': [{'role': 'user', 'content': [
+                    {'type': 'input_image', 'image_url': 'data:image/png;base64,eA=='}]}]}),
+            ]
+            for path, body in cases:
+                for stream in (False, True):
+                    for retry in range(2):
+                        with self.subTest(path=path, stream=stream, retry=retry):
+                            previous = started.read_text().splitlines() if started.exists() else []
+                            with self.assertRaises(urllib.error.HTTPError) as got:
+                                self.post(path, {**body, 'stream': stream})
+                            with got.exception as response:
+                                self.assertEqual(response.code, 503)
+                                self.assertIn('the image encoder failed to start', json.load(response)['error']['message'])
+                            children = started.read_text().splitlines()
+                            self.assertEqual(len(children), len(previous) + 1)
+                            self.assertFalse(vision.alive())
+                            self.assertFalse(Path('/proc', children[-1]).exists())
+
+
+if __name__ == '__main__':
+    unittest.main()