@modular-prompt/driver 0.15.0 → 0.17.0

package/src/pytorch/templates/cpu-minimal/tests/test_server.py

@@ -0,0 +1,141 @@
+import json
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+
+class _Tokenizer:
+    bos_token = "<bos>"
+    chat_template = "template"
+
+    def apply_chat_template(self, messages, **kwargs):
+        del messages, kwargs
+        return "prefix"
+
+
+class _Backend:
+    model_kind = "lm"
+
+    def __init__(self):
+        self.tokenizer = _Tokenizer()
+        self.cache = None
+        self.generate_prompt = None
+        self.pending_cache_write_tokens = 0
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def cache_prefill(self, cache_path, prompt, **kwargs):
+        del prompt, kwargs
+        self.cache = object()
+        self.pending_cache_write_tokens = 2
+        return {
+            "cache_path": cache_path,
+            "token_count": 2,
+            "cache_write_tokens": 2,
+        }
+
+    def load_cache_from_file(self, cache_path, **kwargs):
+        del kwargs
+        return self.cache if cache_path == "memory://prefix" else None
+
+    def get_cache_offset(self, prompt_cache):
+        return 2 if prompt_cache is self.cache else 0
+
+    def consume_cache_write_tokens(self, cache_path):
+        del cache_path
+        result = self.pending_cache_write_tokens
+        self.pending_cache_write_tokens = 0
+        return result
+
+    def tokenize_prompt(self, prompt):
+        assert prompt == "prefix suffix"
+        return [1, 2, 3]
+
+    def stream_generate(self, prompt, options, images=None, prompt_cache=None):
+        del options, images
+        self.generate_prompt = prompt
+        assert prompt_cache is self.cache
+        yield SimpleNamespace(
+            text="ok",
+            prompt_tokens=3,
+            generation_tokens=1,
+            cache_read_tokens=2,
+        )
+
+
+def _make_server():
+    from server import Server
+
+    backend = MagicMock()
+    return Server(backend, {"methods": ["capabilities"], "model_kind": "lm"})
+
+
+def test_cache_prefill_dispatches_to_handler(capsys):
+    server = _make_server()
+
+    with patch("server.handle_cache_prefill") as mock_prefill:
+        server._dispatch(
+            {
+                "method": "cache_prefill",
+                "cache_path": "memory://prefix",
+                "messages": [{"role": "user", "content": "hello"}],
+            }
+        )
+
+    mock_prefill.assert_called_once()
+    assert mock_prefill.call_args.args[2:] == (
+        "memory://prefix",
+        [{"role": "user", "content": "hello"}],
+    )
+    assert capsys.readouterr().out == ""
+
+
+def test_cache_prefill_requires_path_and_messages(capsys):
+    server = _make_server()
+
+    server._dispatch({"method": "cache_prefill"})
+
+    assert capsys.readouterr().out.endswith("\0")
+
+
+def test_cache_prefill_then_generate_with_cache_via_server(capsys):
+    from server import Server
+
+    backend = _Backend()
+    server = Server(backend, {"methods": ["cache_prefill"], "model_kind": "lm"})
+
+    server._dispatch(
+        {
+            "method": "cache_prefill",
+            "cache_path": "memory://prefix",
+            "messages": [{"role": "user", "content": "hello"}],
+        }
+    )
+    prefill_output = capsys.readouterr().out
+    assert json.loads(prefill_output.split("\0", 1)[0]) == {
+        "cache_path": "memory://prefix",
+        "token_count": 2,
+        "cache_write_tokens": 2,
+    }
+
+    server._dispatch(
+        {
+            "method": "generate",
+            "prompt": "prefix suffix",
+            "options": {"max_tokens": 1},
+            "cache_path": "memory://prefix",
+        }
+    )
+    generate_output = capsys.readouterr().out
+    assert generate_output.startswith("ok")
+    meta = json.loads(
+        generate_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,
+    }
+    assert backend.generate_prompt == [3]

package/src/pytorch/templates/cpu-minimal/tests/test_transformers_errors.py

@@ -0,0 +1,89 @@
+from types import SimpleNamespace
+
+import pytest
+
+from utils.transformers_errors import (
+    MIN_TRANSFORMERS_VERSION,
+    extract_unsupported_model_type,
+    unsupported_model_type_error,
+)
+
+
+def test_extracts_model_type_from_transformers_value_error():
+    error = ValueError(
+        "The checkpoint has model type `qwen3_5` but Transformers does not "
+        "recognize this architecture."
+    )
+
+    assert extract_unsupported_model_type(error) == "qwen3_5"
+
+
+def test_extracts_model_type_from_older_registry_key_error():
+    assert extract_unsupported_model_type(KeyError("qwen3_5")) == "qwen3_5"
+
+
+def test_leaves_unrelated_errors_unchanged():
+    assert extract_unsupported_model_type(ValueError("invalid weights")) is None
+    assert extract_unsupported_model_type(KeyError("model_type")) is None
+
+
+def test_error_includes_runtime_requirement_and_setup_guidance():
+    error = unsupported_model_type_error("Qwen/Qwen3.5-0.8B", "qwen3_5", "4.57.6")
+    message = str(error)
+
+    assert "qwen3_5" in message
+    assert "transformers 4.57.6" in message
+    assert f"transformers>={MIN_TRANSFORMERS_VERSION}" in message
+    assert "setup-pytorch" in message
+    assert "modular-prompt-runtime sync pytorch" in message
+
+
+def test_backend_wraps_unknown_architecture_error(monkeypatch):
+    pytest.importorskip("torch")
+    import backends.transformers_lm as backend_module
+
+    tokenizer = SimpleNamespace(pad_token=None, eos_token="<eos>")
+    backend = backend_module.TransformersLmBackend(device="cpu")
+
+    def load_tokenizer(*args, **kwargs):
+        return tokenizer
+
+    def load_model(*args, **kwargs):
+        raise ValueError(
+            "The checkpoint has model type `qwen3_5` but Transformers does not "
+            "recognize this architecture."
+        )
+
+    monkeypatch.setattr(backend_module.AutoTokenizer, "from_pretrained", load_tokenizer)
+    monkeypatch.setattr(
+        backend_module.AutoModelForCausalLM,
+        "from_pretrained",
+        load_model,
+    )
+
+    with pytest.raises(RuntimeError, match="qwen3_5") as raised:
+        backend.load("Qwen/Qwen3.5-0.8B")
+
+    assert f"transformers {backend_module.transformers.__version__}" in str(raised.value)
+    assert "transformers>=5.14.0" in str(raised.value)
+
+
+def test_backend_wraps_unknown_architecture_error_from_tokenizer(monkeypatch):
+    pytest.importorskip("torch")
+    import backends.transformers_lm as backend_module
+
+    backend = backend_module.TransformersLmBackend(device="cpu")
+
+    def load_tokenizer(*args, **kwargs):
+        raise ValueError(
+            "The checkpoint has model type `qwen3_5` but Transformers does not "
+            "recognize this architecture."
+        )
+
+    monkeypatch.setattr(backend_module.AutoTokenizer, "from_pretrained", load_tokenizer)
+
+    with pytest.raises(RuntimeError, match="qwen3_5") as raised:
+        backend.load("Qwen/Qwen3.5-0.8B")
+
+    assert f"transformers {backend_module.transformers.__version__}" in str(raised.value)
+    assert "transformers>=5.14.0" in str(raised.value)