@modular-prompt/driver 0.16.0 → 0.17.0

package/src/pytorch/templates/cuda/handlers/generate.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from handlers.cancel import poll_cancel
+
+
+def _stream_to_stdout(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict,
+    images: list | None = None,
+    primer: str | None = None,
+    prompt_cache=None,
+    cache_loaded: bool | None = None,
+    cache_read_tokens: int = 0,
+    cache_write_tokens: int = 0,
+) -> None:
+    if images:
+        raise ValueError("PyTorch LIP backend does not support images in Phase 1")
+
+    if primer is not None:
+        print(primer, end="", flush=True)
+
+    first_prompt_tokens = None
+    first_cache_read_tokens = None
+    first_cache_write_tokens = None
+    reported_generation_tokens = None
+    for response in backend.stream_generate(
+        prompt,
+        options,
+        images,
+        prompt_cache=prompt_cache,
+    ):
+        if poll_cancel():
+            break
+        response_prompt_tokens = getattr(response, "prompt_tokens", None)
+        if first_prompt_tokens is None and response_prompt_tokens is not None:
+            first_prompt_tokens = response_prompt_tokens
+        if (
+            first_cache_read_tokens is None
+            and getattr(response, "cache_read_tokens", None) is not None
+        ):
+            first_cache_read_tokens = response.cache_read_tokens
+        if (
+            first_cache_write_tokens is None
+            and getattr(response, "cache_write_tokens", None) is not None
+        ):
+            first_cache_write_tokens = response.cache_write_tokens
+        response_generation_tokens = getattr(response, "generation_tokens", None)
+        if response_generation_tokens is not None:
+            reported_generation_tokens = max(
+                reported_generation_tokens or 0,
+                int(response_generation_tokens),
+            )
+        print(response.text.replace("\0", "").replace("\x1e", ""), end="", flush=True)
+
+    meta: dict = {}
+    if first_prompt_tokens is not None:
+        meta["prompt_tokens"] = first_prompt_tokens
+    if reported_generation_tokens is not None:
+        meta["generation_tokens"] = max(0, reported_generation_tokens)
+    if first_cache_read_tokens is not None:
+        meta["cache_read_tokens"] = first_cache_read_tokens
+    if first_cache_write_tokens is not None:
+        meta["cache_write_tokens"] = first_cache_write_tokens
+    if cache_read_tokens > 0 and "cache_read_tokens" not in meta:
+        meta["cache_read_tokens"] = cache_read_tokens
+    if cache_write_tokens > 0 and "cache_write_tokens" not in meta:
+        meta["cache_write_tokens"] = cache_write_tokens
+    if cache_loaded is not None:
+        meta["cache_loaded"] = cache_loaded
+
+    if meta:
+        print(f"\x1e__META__:{json.dumps(meta)}", end="\0", flush=True)
+    else:
+        print("", end="\0", flush=True)
+
+
+def handle_generate(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
+    primer: str | None = None,
+    cache_path: str | None = None,
+    cache_trim_tokens: int | None = None,
+) -> None:
+    """LIP generate: 整形済み prompt のストリーム推論"""
+    if cache_trim_tokens is not None:
+        raise ValueError(
+            "PyTorch LIP backend does not support cache trimming in Phase 1"
+        )
+
+    if options is None:
+        options = {}
+
+    final_options = dict(options)
+    final_options.pop("trust_remote_code", None)
+
+    prompt_cache = None
+    cache_loaded = None
+    cache_read_tokens = 0
+    cache_write_tokens = 0
+    if cache_path:
+        if images:
+            cache_loaded = False
+        else:
+            prompt_cache = backend.load_cache_from_file(
+                cache_path,
+                images=images,
+                max_image_size=max_image_size,
+                prompt=prompt,
+            )
+            cache_loaded = prompt_cache is not None
+
+        if prompt_cache is not None:
+            cache_read_tokens = backend.get_cache_offset(prompt_cache)
+            if cache_read_tokens <= 0:
+                prompt_cache = None
+                cache_loaded = False
+            elif isinstance(prompt, (str, list)):
+                full_tokens = (
+                    backend.tokenize_prompt(prompt)
+                    if isinstance(prompt, str)
+                    else [int(token_id) for token_id in prompt]
+                )
+                if cache_read_tokens < len(full_tokens):
+                    prompt = full_tokens[cache_read_tokens:]
+                else:
+                    # A cache covering the complete prompt cannot be passed
+                    # with an empty input_ids tensor.  The safe fallback is a
+                    # cold generation for this Phase 1 handler.
+                    prompt_cache = None
+                    cache_loaded = False
+                    cache_read_tokens = 0
+                if prompt_cache is not None:
+                    cache_write_tokens = backend.consume_cache_write_tokens(cache_path)
+
+    _stream_to_stdout(
+        backend,
+        prompt,
+        final_options,
+        images=images,
+        primer=primer,
+        prompt_cache=prompt_cache,
+        cache_loaded=cache_loaded,
+        cache_read_tokens=cache_read_tokens,
+        cache_write_tokens=cache_write_tokens,
+    )

package/src/pytorch/templates/cuda/handlers/render.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from utils.template_render import apply_chat_template_prompt
+
+
+def handle_render(
+    backend: ModelBackend,
+    messages: list,
+    options: dict | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """LIP render: apply_chat_template のみ（推論しない）"""
+    if options is None:
+        options = {}
+
+    result: dict = {
+        "formatted_prompt": None,
+        "error": None,
+    }
+
+    try:
+        trust_remote_code = options.get("trust_remote_code")
+        primer = options.get("primer")
+        prompt = apply_chat_template_prompt(
+            backend,
+            messages,
+            primer=primer,
+            tools=tools,
+            reasoning_effort=reasoning_effort,
+            trust_remote_code=trust_remote_code,
+        )
+        result["formatted_prompt"] = prompt
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

package/src/pytorch/templates/cuda/handlers/tokenize.py

@@ -0,0 +1,63 @@
+import json
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_tokenize(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """メッセージをchat template適用後にトークン化して返す"""
+    tokenizer = backend.get_tokenizer()
+
+    result = {
+        "token_ids": None,
+        "token_count": 0,
+        "error": None,
+    }
+
+    try:
+        # apply_chat_templateのfallbackパターン (chat.py L165-188 と同じ)
+        # add_generation_prompt=False で、アシスタントの開始トークンは含めない
+        extra_kwargs = {}
+        if tools is not None:
+            extra_kwargs["tools"] = tools
+        if reasoning_effort is not None:
+            extra_kwargs["reasoning_effort"] = reasoning_effort
+
+        if supports_chat_template(tokenizer):
+            # chat.py と同じfallbackチェーン
+            prompt = None
+            for kwargs in [extra_kwargs, {k: v for k, v in extra_kwargs.items() if k == "tools"}, {}]:
+                try:
+                    prompt = tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=False,
+                        tokenize=False,
+                        **kwargs,
+                    )
+                    break
+                except TypeError:
+                    continue
+
+            if prompt is None:
+                prompt = str(messages)
+        else:
+            prompt = generate_merged_prompt(messages, capabilities)
+
+        # トークン化
+        add_special = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token or ""
+        )
+        token_ids = tokenizer.encode(prompt, add_special_tokens=add_special)
+
+        result["token_ids"] = token_ids
+        result["token_count"] = len(token_ids)
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

package/src/pytorch/templates/cuda/pyproject.toml

@@ -0,0 +1,37 @@
+[project]
+name = "pytorch_driver"
+version = "0.1.0"
+description = "PyTorch (Transformers) driver for modular-prompt — CUDA runtime"
+requires-python = ">=3.10,<3.14"
+dependencies = [
+    "safetensors==0.8.0",
+    "tokenizers==0.22.2",
+    "transformers>=5.14.0",
+]
+
+[dependency-groups]
+dev = ["pytest>=9.0"]
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.setuptools]
+py-modules = ["__main__", "server"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["backends*", "handlers*", "utils*"]
+
+# setup-pytorch --variant cuda は CUDA wheel を明示インストールする。
+# デフォルトは cu124。別の index を使う場合は --cuda <version> を指定する。
+[[tool.uv.index]]
+name = "pytorch-cuda"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true
+
+[tool.uv.sources]
+torch = { index = "pytorch-cuda" }

package/src/pytorch/templates/cuda/server.py

@@ -0,0 +1,158 @@
+"""JSON-RPC風サーバー: stdin/stdoutベースのリクエストディスパッチ"""
+import json
+import sys
+
+from backends.base import ModelBackend
+from handlers import handle_cache_prefill, handle_capabilities, handle_completion, handle_format_test, handle_generate, handle_render, handle_tokenize
+from handlers.cancel import request_cancel, reset_cancel
+
+
+MAX_READ_LINES = 10000
+
+
+def read():
+    lines = []
+    while True:
+        line = sys.stdin.readline()
+        if not line:
+            return None
+        lines.append(line)
+        if len(lines) > MAX_READ_LINES:
+            sys.stderr.write(f"Error: read buffer exceeded {MAX_READ_LINES} lines, discarding\n")
+            lines.clear()
+            continue
+        try:
+            return json.loads(''.join(lines))
+        except json.JSONDecodeError:
+            continue
+
+
+def _is_valid_generate_prompt(prompt) -> bool:
+    """prompt は非空文字列または非空のトークン ID リスト"""
+    if isinstance(prompt, str):
+        return bool(prompt)
+    if isinstance(prompt, list):
+        return len(prompt) > 0
+    return False
+
+
+class Server:
+    def __init__(self, backend: ModelBackend, capabilities: dict):
+        self.backend = backend
+        self.capabilities = capabilities
+
+    def run(self):
+        while True:
+            req = read()
+            if req is None:
+                break
+            self._dispatch(req)
+
+    def _error_response(self, message: str) -> None:
+        sys.stderr.write(f"Error: {message}\n")
+        print(json.dumps({"error": message}), end='\0', flush=True)
+
+    def _dispatch(self, req: dict):
+        method = req.get('method')
+        if not method:
+            self._error_response("'method' field is required")
+            return
+
+        if method == 'cancel':
+            request_cancel()
+            return
+
+        reset_cancel()
+
+        try:
+            if method == 'capabilities':
+                handle_capabilities(self.capabilities)
+
+            elif method == 'format_test':
+                messages = req.get('messages')
+                if not messages:
+                    self._error_response("'messages' field is required for format_test method")
+                    return
+                handle_format_test(self.backend, self.capabilities, messages, req.get('options', {}), req.get('tools'))
+
+            elif method == 'tokenize':
+                messages = req.get('messages')
+                if messages is None:
+                    self._error_response("'messages' field is required for tokenize method")
+                    return
+                handle_tokenize(
+                    self.backend, self.capabilities, messages,
+                    tools=req.get('tools'),
+                    reasoning_effort=req.get('reasoning_effort'),
+                )
+
+            elif method == 'cache_prefill':
+                cache_path = req.get('cache_path')
+                messages = req.get('messages')
+                if not cache_path or not messages:
+                    self._error_response("'cache_path' and 'messages' fields are required for cache_prefill")
+                    return
+                handle_cache_prefill(
+                    self.backend,
+                    self.capabilities,
+                    cache_path,
+                    messages,
+                    base_cache_path=req.get('base_cache_path'),
+                    trim_to_tokens=req.get('trim_to_tokens'),
+                    prefix_offsets=req.get('prefix_offsets'),
+                    prefix_hashes=req.get('prefix_hashes'),
+                    tools=req.get('tools'),
+                    reasoning_effort=req.get('reasoning_effort'),
+                    images=req.get('images'),
+                    max_image_size=req.get('maxImageSize', 768),
+                )
+
+            elif method == 'render':
+                messages = req.get('messages')
+                if not messages:
+                    self._error_response("'messages' field is required for render method")
+                    return
+                handle_render(
+                    self.backend,
+                    messages,
+                    options=req.get('options', {}),
+                    tools=req.get('tools'),
+                    reasoning_effort=req.get('reasoning_effort'),
+                )
+
+            elif method == 'generate':
+                prompt = req.get('prompt')
+                if not _is_valid_generate_prompt(prompt):
+                    self._error_response("'prompt' field is required for generate method")
+                    return
+                images = req.get('images', [])
+                handle_generate(
+                    self.backend,
+                    prompt,
+                    options=req.get('options', {}),
+                    images=images if images else None,
+                    max_image_size=req.get('maxImageSize', 768),
+                    primer=req.get('primer'),
+                    cache_path=req.get('cache_path'),
+                    cache_trim_tokens=req.get('cache_trim_tokens'),
+                )
+
+            elif method == 'completion':
+                prompt = req.get('prompt')
+                if not prompt:
+                    self._error_response("'prompt' field is required for completion method")
+                    return
+                images = req.get('images', [])
+                handle_completion(
+                    self.backend,
+                    prompt,
+                    options=req.get('options', {}),
+                    images=images if images else None,
+                    max_image_size=req.get('maxImageSize', 768),
+                )
+
+            else:
+                self._error_response(f"Unknown method '{method}'")
+
+        except Exception as e:
+            self._error_response(f"Error processing request: {e}")

package/src/pytorch/templates/cuda/tests/test_cache_handler.py

@@ -0,0 +1,207 @@
+import pytest
+
+pytest.importorskip("torch")
+
+import json
+from types import SimpleNamespace
+
+from handlers.cache import handle_cache_prefill
+from handlers.generate import handle_generate
+
+
+class _Tokenizer:
+    bos_token = "<bos>"
+    chat_template = "template"
+
+    def __init__(self):
+        self.calls = []
+
+    def apply_chat_template(self, messages, **kwargs):
+        self.calls.append((messages, kwargs))
+        return "prefix"
+
+
+class _Cache:
+    pass
+
+
+class _Backend:
+    model_kind = "lm"
+
+    def __init__(self, cache=None, multi_chunk=False, cache_write_tokens=0):
+        self.tokenizer = _Tokenizer()
+        self.cache = cache
+        self.multi_chunk = multi_chunk
+        self.pending_cache_write_tokens = cache_write_tokens
+        self.calls = []
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def cache_prefill(self, cache_path, prompt, **kwargs):
+        self.calls.append(("prefill", cache_path, prompt, kwargs))
+        self.pending_cache_write_tokens = 2
+        return {
+            "cache_path": cache_path,
+            "token_count": 2,
+            "cache_write_tokens": 2,
+        }
+
+    def consume_cache_write_tokens(self, cache_path):
+        self.calls.append(("consume-write", cache_path))
+        result = self.pending_cache_write_tokens
+        self.pending_cache_write_tokens = 0
+        return result
+
+    def load_cache_from_file(self, cache_path, **kwargs):
+        self.calls.append(("load", cache_path, kwargs))
+        return self.cache
+
+    def get_cache_offset(self, prompt_cache):
+        return 2 if prompt_cache is self.cache else 0
+
+    def tokenize_prompt(self, prompt):
+        assert prompt == "prefix suffix"
+        return [1, 2, 3]
+
+    def stream_generate(self, prompt, options, images=None, prompt_cache=None):
+        self.calls.append(("generate", prompt, options, images, prompt_cache))
+        cache_read_tokens = (
+            2 if self.cache is not None and prompt_cache is self.cache else None
+        )
+        if self.multi_chunk:
+            for index, text in enumerate(("a", "b", "c"), start=1):
+                yield SimpleNamespace(
+                    text=text,
+                    prompt_tokens=3 if index == 1 else None,
+                    generation_tokens=index,
+                    cache_read_tokens=cache_read_tokens if index == 1 else None,
+                )
+            return
+        yield SimpleNamespace(
+            text="ok",
+            prompt_tokens=3,
+            generation_tokens=1,
+            cache_read_tokens=cache_read_tokens,
+        )
+
+
+def _json_response(output):
+    return json.loads(output.split("\0", 1)[0])
+
+
+def test_cache_prefill_renders_without_generation_prompt(capsys):
+    backend = _Backend()
+
+    handle_cache_prefill(
+        backend,
+        {"special_tokens": {}},
+        "memory://prefix",
+        [{"role": "user", "content": "hello"}],
+    )
+
+    assert _json_response(capsys.readouterr().out) == {
+        "cache_path": "memory://prefix",
+        "token_count": 2,
+        "cache_write_tokens": 2,
+    }
+    assert backend.calls == [
+        ("prefill", "memory://prefix", "prefix", {
+            "base_cache_path": None,
+            "trim_to_tokens": None,
+            "prefix_offsets": None,
+            "prefix_hashes": None,
+            "images": None,
+            "max_image_size": 768,
+        })
+    ]
+    assert backend.tokenizer.calls[0][1]["add_generation_prompt"] is False
+
+
+def test_generate_loads_cache_and_only_generates_suffix(capsys):
+    cache = _Cache()
+    backend = _Backend(cache, cache_write_tokens=2)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 1},
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert output.startswith("ok")
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta == {
+        "prompt_tokens": 3,
+        "generation_tokens": 1,
+        "cache_read_tokens": 2,
+        "cache_write_tokens": 2,
+        "cache_loaded": True,
+    }
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == [3]
+    assert generate_call[4] is cache
+    load_call = next(call for call in backend.calls if call[0] == "load")
+    assert load_call[2]["prompt"] == "prefix suffix"
+    assert load_call[2]["images"] is None
+    assert load_call[2]["max_image_size"] == 768
+    assert ("consume-write", "memory://prefix") in backend.calls
+
+
+def test_generate_preserves_usage_meta_across_multiple_chunks(capsys):
+    cache = _Cache()
+    backend = _Backend(cache, multi_chunk=True, cache_write_tokens=2)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 3},
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert output.startswith("abc")
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta == {
+        "prompt_tokens": 3,
+        "generation_tokens": 3,
+        "cache_read_tokens": 2,
+        "cache_write_tokens": 2,
+        "cache_loaded": True,
+    }
+
+
+def test_generate_uses_cold_path_for_missing_cache(capsys):
+    backend = _Backend(cache=None)
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        cache_path="memory://missing",
+    )
+
+    output = capsys.readouterr().out
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta["cache_loaded"] is False
+    assert "cache_read_tokens" not in meta
+    assert "cache_write_tokens" not in meta
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == "prefix suffix"
+    assert generate_call[4] is None
+
+
+def test_generate_removes_cached_prefix_from_token_ids(capsys):
+    cache = _Cache()
+    backend = _Backend(cache)
+
+    handle_generate(
+        backend,
+        [1, 2, 3],
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert '"cache_loaded": true' in output
+    generate_call = next(call for call in backend.calls if call[0] == "generate")
+    assert generate_call[1] == [3]

package/src/pytorch/templates/cuda/tests/test_capabilities.py

@@ -0,0 +1,14 @@
+from utils.token_utils import get_capabilities
+
+
+class _Tokenizer:
+    apply_chat_template = None
+    chat_template = None
+    special_tokens_map = {}
+    added_tokens_encoder = {}
+
+
+def test_capabilities_advertise_cache_prefill():
+    capabilities = get_capabilities(_Tokenizer())
+
+    assert "cache_prefill" in capabilities["methods"]