diff --git a/opt_prime/examples/pd/serve_batched.py b/opt_prime/examples/pd/serve_batched.py index b8debf46..5f3f38cf 100644 --- a/opt_prime/examples/pd/serve_batched.py +++ b/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) diff --git a/opt_prime/opt_prime/accel/serve/llmd.py b/opt_prime/opt_prime/accel/serve/llmd.py index fad81b23..52713c6d 100644 --- a/opt_prime/opt_prime/accel/serve/llmd.py +++ b/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 diff --git a/opt_prime/tests/test_http_sigpipe.py b/opt_prime/tests/test_http_sigpipe.py new file mode 100644 index 00000000..e088c735 --- /dev/null +++ b/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) diff --git a/opt_prime/tests/test_serving_eos_factory.py b/opt_prime/tests/test_serving_eos_factory.py new file mode 100644 index 00000000..4ab801aa --- /dev/null +++ b/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 == () diff --git a/opt_prime/tests/test_serving_prefill_cache.py b/opt_prime/tests/test_serving_prefill_cache.py new file mode 100644 index 00000000..d2d4f8af --- /dev/null +++ b/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)