146 lines
5.6 KiB
Python
146 lines
5.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""``--no-context-shift`` launch-flag contract.
|
|
|
|
When llama-server runs with its default context-shift behavior, the UI
|
|
has no way to tell the user that the KV cache has been rotated --
|
|
earlier turns silently vanish from the conversation. The Studio
|
|
backend always passes ``--no-context-shift`` so the server returns a
|
|
clean error instead, and the chat adapter can point the user at the
|
|
``Context Length`` input in the settings panel.
|
|
|
|
This file is a static read of the launch command: we ask
|
|
``LlamaCppBackend`` to assemble its ``cmd`` list and assert the flag
|
|
is always present. Testing via the real subprocess would require an
|
|
actual GGUF on disk, which is out of scope for the fast test suite.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import sys
|
|
import types as _types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Same external-dep stubs as the other llama_cpp tests.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
_loggers_stub = _types.ModuleType("loggers")
|
|
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
|
sys.modules.setdefault("loggers", _loggers_stub)
|
|
|
|
_structlog_stub = _types.ModuleType("structlog")
|
|
sys.modules.setdefault("structlog", _structlog_stub)
|
|
|
|
_httpx_stub = _types.ModuleType("httpx")
|
|
for _exc in (
|
|
"ConnectError",
|
|
"TimeoutException",
|
|
"ReadTimeout",
|
|
"ReadError",
|
|
"RemoteProtocolError",
|
|
"CloseError",
|
|
):
|
|
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
|
|
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
|
|
_httpx_stub.Client = type(
|
|
"C",
|
|
(),
|
|
{
|
|
"__init__": lambda s, **kw: None,
|
|
"__enter__": lambda s: s,
|
|
"__exit__": lambda s, *a: None,
|
|
},
|
|
)
|
|
sys.modules.setdefault("httpx", _httpx_stub)
|
|
|
|
from core.inference import llama_cpp as llama_cpp_module
|
|
|
|
|
|
def _load_model_source() -> str:
|
|
"""Return the source of ``LlamaCppBackend.load_model`` PLUS the
|
|
internal ``_load_model_impl_locked`` body it delegates to.
|
|
|
|
Studio's diffusion PR split ``load_model`` into a thin wrapper
|
|
that publishes ``_loading_model_identifier`` under
|
|
``_serial_load_lock`` and an inner ``_load_model_impl_locked``
|
|
body that actually spawns llama-server. The launch flags and the
|
|
``_wait_for_vram_settle`` call now live in the inner method, so
|
|
inspecting only ``load_model`` would miss them. Concatenating the
|
|
two sources keeps these source-inspection regression tests
|
|
working without weakening the scope (we still only look at the
|
|
two load entry points, not the entire module).
|
|
"""
|
|
parts = [inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)]
|
|
impl = getattr(llama_cpp_module.LlamaCppBackend, "_load_model_impl_locked", None)
|
|
if impl is not None:
|
|
parts.append(inspect.getsource(impl))
|
|
return "\n".join(parts)
|
|
|
|
|
|
def test_no_context_shift_is_in_load_model():
|
|
"""The flag is part of the static launch-command template.
|
|
|
|
We check the source of ``load_model`` rather than mocking the whole
|
|
call chain (GPU probing, GGUF stat, etc.): the flag is written as
|
|
a literal in one place and any regression has to delete it, which
|
|
a text search will catch.
|
|
"""
|
|
assert '"--no-context-shift"' in _load_model_source(), (
|
|
"llama-server must be launched with --no-context-shift so the "
|
|
"UI can surface a clean 'context full' error instead of silently "
|
|
"losing old turns to a KV-cache rotation."
|
|
)
|
|
|
|
|
|
def test_flag_sits_inside_the_base_cmd_list():
|
|
"""Pin the flag's location so a future refactor can't accidentally
|
|
move it into a branch that only fires on some code paths.
|
|
|
|
We slice from ``cmd = [`` to the first ``]`` at the same indent.
|
|
Using ``inspect.getsource`` means the function lives in its own
|
|
string and there are no siblings to worry about, so a plain
|
|
bracket search would also work -- anchoring on the trailing indent
|
|
just keeps the slice from wandering into a later expression if the
|
|
opening literal ever grows an in-line comment trailing it.
|
|
"""
|
|
source = _load_model_source()
|
|
start = source.find("cmd = [")
|
|
assert start >= 0, "could not find the base cmd = [...] block"
|
|
# Find the first line containing only ``]`` (possibly indented).
|
|
# Works for any indentation style the formatter picks.
|
|
rest = source[start:]
|
|
end_rel = -1
|
|
for line_start, line in _iter_lines_with_offset(rest):
|
|
if line_start == 0:
|
|
# Skip the opening ``cmd = [`` line itself.
|
|
continue
|
|
if line.strip() == "]":
|
|
end_rel = line_start
|
|
break
|
|
assert end_rel > 0, "could not find end of cmd = [...] block"
|
|
block = rest[:end_rel]
|
|
assert '"--no-context-shift"' in block, (
|
|
"--no-context-shift must be in the base cmd list, not in a "
|
|
"conditional branch -- otherwise some code paths would still "
|
|
"run with silent context shift enabled."
|
|
)
|
|
# Also pin that it is next to -c / --ctx so the grouping makes sense.
|
|
assert '"-c"' in block
|
|
assert '"--flash-attn"' in block
|
|
|
|
|
|
def _iter_lines_with_offset(text: str):
|
|
"""Yield (offset, line) pairs over ``text`` without losing offsets."""
|
|
offset = 0
|
|
for line in text.splitlines(keepends = True):
|
|
yield offset, line
|
|
offset += len(line)
|