serve: no SIGPIPE on disconnect, tokenizer EOS, producer cache flags (#214, #229, #241)

opt_prime/examples/pd/serve_batched.py

@@ -30,9 +30,14 @@
 def prefill_worker(a, slots, config):
     """the PD prefill side from the flags, or None (aggregated serving)"""
     if a.pd_prefill_device:
+        chunk = getattr(a, "pd_prefill_chunk", 2048)
+        if chunk <= 0:
+            raise ValueError("PD prefill chunk must be positive")
         from opt_prime.accel.pd.pd_batching import PrefillWorker, TorchPrefillEngine
         m = AutoModelForCausalLM.from_pretrained(a.model, dtype=getattr(torch, a.dtype)).to(a.pd_prefill_device).eval()
-        return PrefillWorker(TorchPrefillEngine(m, max_context=a.max_context, out_device=a.device
+        return PrefillWorker(TorchPrefillEngine(m, max_context=a.max_context, chunk=chunk,
+                                                prefix_cache=not getattr(a, "pd_prefill_no_prefix_cache", False),
+                                                out_device=a.device
                                                 if a.backend == "torch" else "cpu"))
     if a.pd_prefill_python:
         import os, shlex
@@ -95,10 +100,14 @@ def build(a):
         else:
             b = ContinuousBatcher(model, prefill_chunk=a.chunk, **kw)
     eos = model.generation_config.eos_token_id
+    eos = tuple(eos) if isinstance(eos, (list, tuple)) else ((eos,) if eos is not None else ())
+    # Some chat checkpoints retain a base-model EOS in generation_config.
+    if tok.eos_token_id is not None:
+        eos = tuple(dict.fromkeys((*eos, tok.eos_token_id)))
     svc = BatchedService(b, model=a.model, vocab_size=model.config.get_text_config().vocab_size,
                          max_output=a.max_output, render=lambda t: tok.decode(t, skip_special_tokens=True),
                          tokenizer=tok,
-                         eos_token_ids=tuple(eos) if isinstance(eos, (list, tuple)) else ((eos,) if eos is not None else ()))
+                         eos_token_ids=eos)
     return BatchedOwner(svc)
 
 
@@ -114,6 +123,10 @@ def build(a):
     ap.add_argument("--fxb", default=None, help="RNGD: furiosa-llm compiled artifact")
     ap.add_argument("--rngd-devices", default="npu:0"); ap.add_argument("--no-prefix-cache", action="store_true")
     ap.add_argument("--pd-prefill-device", default=None, help="PD: torch prefill on this device (e.g. cuda:0)")
+    ap.add_argument("--pd-prefill-chunk", type=int, default=2048,
+                    help="Torch --pd-prefill-device producer chunk size (default: 2048)")
+    ap.add_argument("--pd-prefill-no-prefix-cache", action="store_true",
+                    help="disable prefix reuse in the Torch --pd-prefill-device producer (decoder policy unchanged)")
     ap.add_argument("--pd-prefill-python", default=None, help="PD: prefill_server under this interpreter")
     ap.add_argument("--pd-prefill-args", default="--engine torch --device cpu")
     ap.add_argument("--pd-prefill-pythonpath", default=None)

opt_prime/opt_prime/accel/serve/llmd.py

@@ -8,6 +8,7 @@
 from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 import json
+from io import BufferedIOBase
 import queue
 import select
 import socket
@@ -485,11 +486,32 @@ def run(self):
 POLL = 0.5          # seconds between client-disconnect checks of a waiting handler (results wake it at once)
 _DONE = object()    # end-of-stream marker in a handler's token queue
 
+class _NoSignalSocketWriter(BufferedIOBase):
+    """Unbuffered HTTP writes that retain EPIPE without invoking native SIGPIPE handlers."""
+    def __init__(self, sock):
+        self._sock = sock
+
+    def writable(self):
+        return True
+
+    def write(self, data):
+        self._sock.sendall(data, getattr(socket, "MSG_NOSIGNAL", 0))
+        with memoryview(data) as view:
+            return view.nbytes
+
+    def fileno(self):
+        return self._sock.fileno()
+
+
 def serve_stage(owner, stage, port):
     if stage not in ("prefill", "decode"):
         raise ValueError("invalid stage")
 
     class Handler(BaseHTTPRequestHandler):
+        def setup(self):
+            super().setup()
+            self.wfile = _NoSignalSocketWriter(self.connection)
+
         def log_message(self, *args):
             pass
 

opt_prime/tests/test_http_sigpipe.py

@@ -0,0 +1,62 @@
+"""Preserve BrokenPipeError without invoking a vendor-installed SIGPIPE handler."""
+
+import signal
+import socket
+
+import pytest
+from opt_prime.accel.serve.llmd import _NoSignalSocketWriter
+
+
+def test_writer_preserves_bytes_count_and_socket_descriptor():
+    left, right = socket.socketpair()
+    try:
+        writer = _NoSignalSocketWriter(left)
+        data = memoryview(bytearray(b"abcdef")).cast("B", shape=[2, 3])
+        assert writer.write(data) == 6
+        assert right.recv(6) == b"abcdef"
+        assert writer.writable() and writer.fileno() == left.fileno()
+    finally:
+        left.close()
+        right.close()
+
+
+@pytest.mark.skipif(
+    not hasattr(socket, "MSG_NOSIGNAL"), reason="Linux per-write signal suppression"
+)
+def test_closed_peer_keeps_exception_but_does_not_signal():
+    received = []
+    before = signal.signal(signal.SIGPIPE, lambda *_: received.append(True))
+    left, right = socket.socketpair()
+    right.close()
+    try:
+        with pytest.raises(BrokenPipeError):
+            _NoSignalSocketWriter(left).write(b"x")
+        assert received == []
+    finally:
+        left.close()
+        signal.signal(signal.SIGPIPE, before)
+
+
+def test_serve_stage_handler_writes_through_no_signal_writer(monkeypatch):
+    import types
+    import urllib.request
+    from opt_prime.accel.serve import llmd
+
+    writes = []
+    original = llmd._NoSignalSocketWriter.write
+
+    def record(self, data):
+        writes.append(bytes(data))
+        return original(self, data)
+
+    monkeypatch.setattr(llmd._NoSignalSocketWriter, "write", record)
+    owner = types.SimpleNamespace(started=True, service=types.SimpleNamespace(state="ready"))
+    server, thread = llmd.serve_stage(owner, "decode", 0)
+    try:
+        with urllib.request.urlopen(f"http://127.0.0.1:{server.server_port}/health", timeout=5) as response:
+            assert response.status == 200
+        assert any(b'"ready": true' in w or b'"ready":true' in w for w in writes)
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join(5)

opt_prime/tests/test_serving_eos_factory.py

@@ -0,0 +1,75 @@
+"""The public factory must honor model EOS and tokenizer chat turn end."""
+
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+
+@pytest.mark.parametrize(
+    "configured,tokenizer_eos,expected",
+    [
+        (8, 9, (8, 9)),
+        ([8, 9], 9, (8, 9)),
+        ((8, 10), 9, (8, 10, 9)),
+        (None, 9, (9,)),
+        (8, None, (8,)),
+        (None, None, ()),
+    ],
+)
+def test_factory_preserves_all_eos_and_ignore_override(
+    monkeypatch, configured, tokenizer_eos, expected
+):
+    path = Path(__file__).parents[1] / "examples/pd/serve_batched.py"
+    spec = importlib.util.spec_from_file_location("serve_eos_test", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    tokenizer = SimpleNamespace(
+        eos_token_id=tokenizer_eos, decode=lambda ids, **kw: str(ids)
+    )
+    model = SimpleNamespace(
+        generation_config=SimpleNamespace(eos_token_id=configured),
+        config=SimpleNamespace(get_text_config=lambda: SimpleNamespace(vocab_size=32)),
+    )
+    model.to = lambda device: model
+    model.eval = lambda: model
+    monkeypatch.setattr(
+        module.AutoTokenizer, "from_pretrained", lambda *a, **kw: tokenizer
+    )
+    monkeypatch.setattr(
+        module.AutoModelForCausalLM, "from_pretrained", lambda *a, **kw: model
+    )
+    monkeypatch.setattr(
+        module,
+        "ContinuousBatcher",
+        lambda *a, **kw: SimpleNamespace(max_context=64, paged=False),
+    )
+    monkeypatch.setattr(module, "prefill_worker", lambda *a: None)
+    monkeypatch.setattr(module, "remote_tier", lambda *a: None)
+    args = SimpleNamespace(
+        model="test",
+        kv_events=None,
+        backend="torch",
+        dtype="float32",
+        device="cpu",
+        slots=1,
+        max_context=64,
+        kv="slot",
+        page_size=16,
+        prefix_cache=False,
+        kv_host_memory_gb=0,
+        reload_bandwidth=1,
+        prefill_rate=1,
+        chunk=16,
+        max_output=16,
+    )
+    service = module.build(args).service
+    request, _ = service._request(
+        "r", {"model": "test", "prompt": [1], "max_tokens": 4}
+    )
+    assert request.sampling.stop_token_ids == expected
+    request, _ = service._request(
+        "r2", {"model": "test", "prompt": [1], "max_tokens": 4, "ignore_eos": True}
+    )
+    assert request.sampling.stop_token_ids == ()

opt_prime/tests/test_serving_prefill_cache.py

@@ -0,0 +1,74 @@
+"""The public factory must make remote producer cache policy controllable."""
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from opt_prime.accel.pd import PDRequest
+from transformers import LlamaConfig, LlamaForCausalLM
+
+
+@pytest.mark.parametrize("disabled,expected_hit", [(None, 64), (False, 64), (True, 0)])
+def test_producer_cache_policy_controls_real_repeated_prefill(tmp_path, disabled, expected_hit):
+    torch.set_num_threads(2)
+    torch.manual_seed(241)
+    model = LlamaForCausalLM(LlamaConfig(
+        vocab_size=97, hidden_size=32, intermediate_size=48, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=256,
+    )).eval()
+    model.save_pretrained(tmp_path)
+    path = Path(__file__).parents[1] / "examples/pd/serve_batched.py"
+    spec = importlib.util.spec_from_file_location("serve_cache_test", path)
+    factory = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(factory)
+    args = SimpleNamespace(model=str(tmp_path), pd_prefill_device="cpu", dtype="float32",
+                           backend="torch", device="cpu", max_context=256)
+    if disabled is not None:
+        args.pd_prefill_no_prefix_cache = disabled
+    worker = factory.prefill_worker(args, 1, model.config)
+    try:
+        prompt = tuple((i % 90) + 1 for i in range(128))
+        first = worker.submit(PDRequest("first", prompt, 1)).result(timeout=10)
+        repeated = worker.submit(PDRequest("repeated", prompt, 1)).result(timeout=10)
+        assert first.prefill_cached == 0
+        assert repeated.prefill_cached == expected_hit
+        assert first.token == repeated.token
+        assert first.keys.shape == repeated.keys.shape == (1, 2, 128, 8)
+    finally:
+        worker.close()
+
+
+def test_producer_chunk_controls_real_prefill_boundaries(tmp_path, monkeypatch):
+    import opt_prime.accel.pd.pd_batching as pd
+
+    model = LlamaForCausalLM(LlamaConfig(
+        vocab_size=97, hidden_size=32, intermediate_size=48, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=256,
+    )).eval()
+    model.save_pretrained(tmp_path)
+    path = Path(__file__).parents[1] / "examples/pd/serve_batched.py"
+    spec = importlib.util.spec_from_file_location("serve_chunk_test", path)
+    factory = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(factory)
+    spans = []
+    original = pd.prefill_chunk
+
+    def observe(model, kv, slot, ids, start, end, device):
+        spans.append((start, end))
+        return original(model, kv, slot, ids, start, end, device)
+
+    monkeypatch.setattr(pd, "prefill_chunk", observe)
+    args = SimpleNamespace(model=str(tmp_path), pd_prefill_device="cpu", dtype="float32",
+                           backend="torch", device="cpu", max_context=256, pd_prefill_chunk=32)
+    worker = factory.prefill_worker(args, 1, model.config)
+    try:
+        out = worker.submit(PDRequest("chunk", tuple(range(1, 65)), 1)).result(timeout=10)
+        assert spans == [(0, 32), (32, 64)]
+        assert out.keys.shape == (1, 2, 64, 8)
+    finally:
+        worker.close()
+    for chunk in (0, -1):
+        args.pd_prefill_chunk = chunk
+        with pytest.raises(ValueError, match="positive"):
+            factory.prefill_worker(args, 1, model.config)