diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py index 870e9817d..5f59919fd 100644 --- a/mellea/backends/huggingface.py +++ b/mellea/backends/huggingface.py @@ -34,7 +34,11 @@ StoppingCriteriaList, ) from transformers.generation.streamers import AsyncTextIteratorStreamer - from transformers.generation.utils import GenerateDecoderOnlyOutput, GenerationMixin + from transformers.generation.utils import ( + GenerateBeamDecoderOnlyOutput, + GenerateDecoderOnlyOutput, + GenerationMixin, + ) from transformers.modeling_utils import PreTrainedModel from transformers.tokenization_utils_base import PreTrainedTokenizerBase from transformers.trainer_utils import set_seed @@ -804,6 +808,7 @@ async def granite_formatters_processing( if want_scores: if raw_hf_output_cell[0] is not None: mot.raw.response = raw_hf_output_cell[0] + raw_hf_output_cell[0] = None else: warn_key = "intrinsic_no_hf_output" if warn_key not in self._warned_about: @@ -1519,12 +1524,25 @@ class used during generation, if any. self.cache_put(cache_key, cache_info) # Clear KV cache and scores from HF output; retained via LRU cache above. - hf_output.past_key_values = None - hf_output.scores = None + # `ModelOutput` (`OrderedDict` subclass) does not sync `None` writes back + # to the mapping, so plain attribute assignment leaves the dict entry — and + # its tensor — alive. + for _field in ("past_key_values", "scores"): + if _field in hf_output: + dict.__delitem__(hf_output, _field) + try: + object.__delattr__(hf_output, _field) + except AttributeError: + pass # Clear the raw logits tensor (scores already cleared above if cached). if isinstance(hf_output, GenerateDecoderOnlyOutput): - hf_output.logits = None + if "logits" in hf_output: + dict.__delitem__(hf_output, "logits") + try: + object.__delattr__(hf_output, "logits") + except AttributeError: + pass # Only scan for tools if we are not doing structured output and tool calls were provided to the model. if _format is None and tool_calls: @@ -1602,20 +1620,26 @@ class used during generation, if any. if not self._use_caches and isinstance( mot.raw.response, GenerateDecoderOnlyOutput ): - import gc - hf_out = mot.raw.response - if hasattr(hf_out, "sequences") and hf_out.sequences is not None: - del hf_out.sequences - if hasattr(hf_out, "scores") and hf_out.scores is not None: - del hf_out.scores - if hasattr(hf_out, "logits") and hf_out.logits is not None: - del hf_out.logits + # GenerateDecoderOnlyOutput is a ModelOutput (OrderedDict subclass). + # ModelOutput.__setattr__ skips the dict write when value is None, and + # ModelOutput defines no __delattr__, so both `out.f = None` and + # `del out.f` leave the mapping entry — and its tensor — alive. + # Clear both the dict entry and the instance attribute to release them. + for field in ("sequences", "scores", "logits", "past_key_values"): + if field in hf_out: + dict.__delitem__(hf_out, field) + try: + object.__delattr__(hf_out, field) + except AttributeError: + pass mot.raw.response = None - # Force Python GC and return CUDA memory to device - gc.collect() - torch.cuda.empty_cache() + if torch.cuda.is_available(): + import gc + + gc.collect() + torch.cuda.empty_cache() # Generate the log for this ModelOutputThunk. generate_log = GenerateLog() @@ -1735,58 +1759,167 @@ async def _generate_from_raw( for i, sequence in enumerate(outputs.sequences) ] - decoded_results = self._tokenizer.batch_decode( - sequences_to_decode, skip_special_tokens=True - ) + try: + decoded_results = self._tokenizer.batch_decode( + sequences_to_decode, skip_special_tokens=True + ) - results = [] - agg_prompt = 0 - agg_completion = 0 - want_logits = bool(model_opts and model_opts.get(ModelOption.LOGITS)) - want_raw_logits = bool(model_opts and model_opts.get(ModelOption.RAW_LOGITS)) - for i, decoded_result in enumerate(decoded_results): - n_prompt_tokens = int(inputs["input_ids"][i].size(0)) - n_completion_tokens = len(sequences_to_decode[i]) - agg_prompt += n_prompt_tokens - agg_completion += n_completion_tokens - per_mot_usage: dict[str, Any] = { - "prompt_tokens": n_prompt_tokens, - "completion_tokens": n_completion_tokens, - "total_tokens": n_prompt_tokens + n_completion_tokens, - } - result = ModelOutputThunk(value=decoded_result) - result.generation.usage = per_mot_usage - result.generation.model = self._model_id - result.generation.provider = self._provider - result.raw.provider = self._provider - - if want_logits and outputs.scores is not None: - # Clone each slice so this MOT does not hold a view into the shared batch allocation. - result.generation.logits = tuple( - step_scores[i].detach().clone() for step_scores in outputs.scores - ) + results = [] + agg_prompt = 0 + agg_completion = 0 + want_logits = bool(model_opts and model_opts.get(ModelOption.LOGITS)) + want_raw_logits = bool( + model_opts and model_opts.get(ModelOption.RAW_LOGITS) + ) + for i, decoded_result in enumerate(decoded_results): + n_prompt_tokens = int(inputs["input_ids"][i].size(0)) + n_completion_tokens = len(sequences_to_decode[i]) + agg_prompt += n_prompt_tokens + agg_completion += n_completion_tokens + per_mot_usage: dict[str, Any] = { + "prompt_tokens": n_prompt_tokens, + "completion_tokens": n_completion_tokens, + "total_tokens": n_prompt_tokens + n_completion_tokens, + } + result = ModelOutputThunk(value=decoded_result) + result.generation.usage = per_mot_usage + result.generation.model = self._model_id + result.generation.provider = self._provider + result.raw.provider = self._provider + + # Extract per-MOT tensors once. These are shared by generation + # and, where supported, raw.response. + mot_scores: tuple[torch.Tensor, ...] | None = None + if outputs.scores is not None: + mot_scores = tuple( + score[i].detach().clone() for score in outputs.scores + ) + + mot_logits_raw: tuple[torch.Tensor, ...] | None = None + if outputs.logits is not None: + mot_logits_raw = tuple( + logits[i].detach().clone() for logits in outputs.logits + ) + + # Construct a per-MOT GenerateDecoderOnlyOutput slice holding clones of row i. + # past_key_values, attentions, and hidden_states are shared across the batch + # and cannot be sliced per MOT, so they are omitted. + # Note: beam-search outputs (GenerateBeamDecoderOnlyOutput) are not supported here; + # mot.raw.response will be None. + if ( + self._use_caches + and isinstance(outputs, GenerateDecoderOnlyOutput) + and isinstance(outputs.sequences, torch.Tensor) + ): + response_scores = ( + tuple(score.unsqueeze(0) for score in mot_scores) + if mot_scores is not None + else None + ) + response_logits = ( + tuple(logits.unsqueeze(0) for logits in mot_logits_raw) + if mot_logits_raw is not None + else None + ) + + if "raw_batch_response_fields_omitted" not in self._warned_about: + self._warned_about.add("raw_batch_response_fields_omitted") + MelleaLogger.get_logger().debug( + "mot.raw.response.past_key_values, .attentions, and " + ".hidden_states are not available on the raw batch path " + "and will always be None." + ) - if want_raw_logits and outputs.logits is not None: - result.generation.raw_logits = tuple( - step_logits[i].detach().clone() for step_logits in outputs.logits + result.raw.response = GenerateDecoderOnlyOutput( + sequences=cast( + "torch.LongTensor", + outputs.sequences[i : i + 1, :].detach().clone(), + ), + scores=cast("tuple[torch.FloatTensor] | None", response_scores), + logits=cast("tuple[torch.FloatTensor] | None", response_logits), + attentions=None, + hidden_states=None, + past_key_values=None, + ) + elif self._use_caches: + if ( + isinstance(outputs, GenerateBeamDecoderOnlyOutput) + and "raw_batch_beam_search_unsupported" + not in self._warned_about + ): + self._warned_about.add("raw_batch_beam_search_unsupported") + MelleaLogger.get_logger().debug( + "mot.raw.response is not available for beam-search outputs on the " + "raw batch path and will be None." + ) + elif ( + not isinstance(outputs.sequences, torch.Tensor) + and "raw_batch_non_tensor_sequences_unsupported" + not in self._warned_about + ): + self._warned_about.add( + "raw_batch_non_tensor_sequences_unsupported" + ) + MelleaLogger.get_logger().debug( + "mot.raw.response is not available because raw batch output " + "sequences are not a torch.Tensor and will be None." + ) + else: + # defensive fallback + if "raw_batch_response_unsupported" not in self._warned_about: + self._warned_about.add("raw_batch_response_unsupported") + MelleaLogger.get_logger().debug( + "mot.raw.response is not available for this raw batch output " + "and will be None." + ) + + # Reuse the cloned tensors. + if want_logits and mot_scores is not None: + result.generation.logits = mot_scores + + if want_raw_logits and mot_logits_raw is not None: + result.generation.raw_logits = mot_logits_raw + + action = actions[i] + result.parsed_repr = ( + action.parse(result) + if isinstance(action, Component) + else result.value ) - action = actions[i] - result.parsed_repr = ( - action.parse(result) if isinstance(action, Component) else result.value - ) + generate_log = GenerateLog() + generate_log.prompt = self.formatter.print(actions[i]) + generate_log.backend = f"hf::{self.model_id!s}" + generate_log.model_options = model_opts + generate_log.date = datetime.datetime.now() + generate_log.model_output = decoded_result + generate_log.extra = {"format": format, "seed": seed} + generate_log.action = action + + result._generate_log = generate_log + results.append(result) + finally: + # Drop all references that might pin the shared batch tensors. + # `sequences_to_decode` holds slice views of `outputs.sequences`, and + # `outputs` is a `GenerateDecoderOnlyOutput` — a `ModelOutput`, which + # subclasses `OrderedDict`. `del obj.attr` only removes the `__dict__` + # slot; the `OrderedDict` entry keeps a strong reference to the tensor. + # Setting `outputs = None` drops the whole container at once + del sequences_to_decode + outputs = None + if torch.cuda.is_available(): + import gc - generate_log = GenerateLog() - generate_log.prompt = self.formatter.print(actions[i]) - generate_log.backend = f"hf::{self.model_id!s}" - generate_log.model_options = model_opts - generate_log.date = datetime.datetime.now() - generate_log.model_output = decoded_result - generate_log.extra = {"format": format, "seed": seed} - generate_log.action = action - - result._generate_log = generate_log - results.append(result) + MelleaLogger.get_logger().debug( + "GPU memory before raw batch cleanup: %d bytes reserved", + torch.cuda.memory_reserved(), + ) + gc.collect() + torch.cuda.empty_cache() + MelleaLogger.get_logger().debug( + "GPU memory after raw batch cleanup: %d bytes reserved", + torch.cuda.memory_reserved(), + ) usage: dict[str, Any] | None = ( { diff --git a/test/backends/test_huggingface_raw_response_copy.py b/test/backends/test_huggingface_raw_response_copy.py new file mode 100644 index 000000000..d05112bab --- /dev/null +++ b/test/backends/test_huggingface_raw_response_copy.py @@ -0,0 +1,89 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for copy/deepcopy semantics of ModelOutputThunks that hold an HF +GenerateDecoderOnlyOutput in raw.response — the raw batch path specific case.""" + +import copy +from typing import Any + +import pytest + +torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]") +pytest.importorskip( + "transformers", reason="transformers not installed — install mellea[hf]" +) + +from transformers.generation.utils import GenerateDecoderOnlyOutput + +from mellea.core import ModelOutputThunk + + +def _make_mot_with_hf_raw_response() -> tuple[ModelOutputThunk, Any]: + """Build a MOT whose raw.response mirrors what the raw batch path stores. + + The batch path slices one row from the full-batch sequences tensor and + immediately calls .detach().clone(), so raw.response.sequences is always + an owning, contiguous tensor — never a view. + """ + full_batch = torch.arange(6, dtype=torch.long).reshape(2, 3) + sequences = full_batch[0:1, :].detach().clone() + hf_out = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + mot = ModelOutputThunk(value="hello") + mot.raw.response = hf_out + return mot, sequences + + +def test_shallow_copy_raw_response_is_same_object(): + """copy.copy(mot).raw.response is the same object as mot.raw.response.""" + mot, _ = _make_mot_with_hf_raw_response() + copied = copy.copy(mot) + assert copied.raw.response is mot.raw.response, ( + "shallow copy must keep raw.response as the same object" + ) + + +def test_shallow_copy_raw_response_sequences_shares_storage(): + """Shallow-copied MOT shares the same sequences tensor as the original.""" + mot, original_sequences = _make_mot_with_hf_raw_response() + copied = copy.copy(mot) + assert ( + copied.raw.response.sequences.untyped_storage().data_ptr() + == original_sequences.untyped_storage().data_ptr() + ), "shallow copy: raw.response.sequences must share storage with the original MOT" + + +def test_deepcopy_raw_response_is_distinct_object(): + """copy.deepcopy(mot).raw.response is a distinct object from mot.raw.response.""" + mot, _ = _make_mot_with_hf_raw_response() + deep = copy.deepcopy(mot) + assert deep.raw.response is not mot.raw.response, ( + "deepcopy must produce a new raw.response object" + ) + + +def test_deepcopy_raw_response_sequences_does_not_share_storage(): + """Deepcopy breaks tensor storage sharing for raw.response.sequences.""" + mot, original_sequences = _make_mot_with_hf_raw_response() + deep = copy.deepcopy(mot) + assert ( + deep.raw.response.sequences.untyped_storage().data_ptr() + != original_sequences.untyped_storage().data_ptr() + ), "deepcopy: raw.response.sequences must NOT share storage with the original" + + +def test_deepcopy_raw_response_sequences_preserves_values(): + """Deepcopy preserves tensor values in raw.response.sequences despite storage isolation.""" + mot, _ = _make_mot_with_hf_raw_response() + original_values = mot.raw.response.sequences.clone() + deep = copy.deepcopy(mot) + assert torch.equal(deep.raw.response.sequences, original_values), ( + "deepcopy must preserve tensor values in raw.response.sequences" + ) diff --git a/test/backends/test_huggingface_raw_response_memory.py b/test/backends/test_huggingface_raw_response_memory.py new file mode 100644 index 000000000..e779a9b1b --- /dev/null +++ b/test/backends/test_huggingface_raw_response_memory.py @@ -0,0 +1,579 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Retention and fidelity tests for `mot.raw.response` on the HF backend. + +Two code paths are covered: + +**Raw batch path** (`LocalHFBackend._generate_from_raw`): +`_generate_from_raw` populates `mot.raw.response` with a per-row slice of the +batch output. Tests pin the invariants that make that safe: + +- a MOT retains only its own row, never the whole batch allocation; +- once the caller holds nothing but MOTs, the batch tensors are actually freed; +- a CPU/MPS run pays neither a full GC nor a CUDA allocator flush; +- `raw.response` never misreports the type or the contents of what `generate()` + actually returned. + +**Chat / post-processing path** (`LocalHFBackend.post_processing`): + +- when caching is enabled, cleared fields (scores, logits) release their tensors + even though the `GenerateDecoderOnlyOutput` container is kept alive; +- when caching is disabled, `raw.response` is set to `None` and every tensor + reachable through the HF output graph (sequences, scores, logits, + past_key_values) is released. + +Sizes here are tiny; the assertions are on *which* storage is retained, not on +how much memory a real model would use. +""" + +import copy +import gc +import logging +import weakref +from typing import TYPE_CHECKING, Any +from unittest.mock import AsyncMock, MagicMock, patch + +from mellea.core.base import Component + +if TYPE_CHECKING: + import torch as _torch + +import pytest + +torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]") +pytest.importorskip( + "transformers", reason="transformers not installed — install mellea[hf]" +) +pytest.importorskip( + "llguidance", reason="llguidance not installed — install mellea[hf]" +) + +from transformers.cache_utils import CacheLayerMixin, DynamicCache +from transformers.generation.utils import ( + GenerateBeamDecoderOnlyOutput, + GenerateDecoderOnlyOutput, +) + +from mellea.backends import ModelOption +from mellea.backends.huggingface import LocalHFBackend +from mellea.core import ModelOutputThunk +from mellea.stdlib.components import Message + +_VOCAB = 64 +_PROMPT_LEN = 1 +_SEQ_LEN = 4 + + +def _make_backend(batch_size: int) -> LocalHFBackend: + """A CPU-only LocalHFBackend whose tokenizer returns a fixed batch encoding.""" + mock_tok = MagicMock(eos_token_id=0, vocab_size=_VOCAB) + mock_tok._tokenizer = MagicMock() + mock_tok._tokenizer.get_vocab_size.return_value = _VOCAB + mock_tok.__len__ = MagicMock(return_value=_VOCAB) + mock_model = MagicMock(vocab_size=_VOCAB) + with ( + patch("mellea.backends.huggingface.llguidance") as mock_llg, + patch("mellea.backends.huggingface.set_seed"), + ): + mock_llg.hf.from_tokenizer.return_value = MagicMock(vocab_size=_VOCAB) + backend = LocalHFBackend( + model_id="ibm-granite/granite-3.3-8b-instruct", + custom_config=(mock_tok, mock_model, torch.device("cpu")), + ) + + input_ids = torch.zeros(batch_size, _PROMPT_LEN, dtype=torch.long) + encoding = MagicMock() + encoding.__getitem__ = lambda _self, k: ( + input_ids + if k == "input_ids" + else torch.ones(batch_size, _PROMPT_LEN, dtype=torch.long) + ) + encoding.to = MagicMock(return_value=encoding) + backend._tokenizer = MagicMock(eos_token_id=0, vocab_size=_VOCAB) + backend._tokenizer.__len__ = MagicMock(return_value=_VOCAB) + backend._tokenizer.return_value = encoding + # A plain function, not a MagicMock: `batch_decode` is handed slices of the batch + # and a mock would record them in `call_args`, keeping the batch alive and + # defeating the retention tests below. + decoded = [f"result_{i}" for i in range(batch_size)] + + def _batch_decode(_sequences: Any, **_kwargs: Any) -> list[str]: + return decoded + + backend._tokenizer.batch_decode = _batch_decode + return backend + + +def _batch_output( + batch_size: int, *, n_steps: int = 0, raw_logits: bool = False +) -> GenerateDecoderOnlyOutput: + """A batch `generate()` result: `(batch_size, _SEQ_LEN)` sequences plus optional scores.""" + step_shape = (batch_size, _VOCAB) + return GenerateDecoderOnlyOutput( + sequences=torch.zeros(batch_size, _SEQ_LEN, dtype=torch.long), + scores=tuple(torch.zeros(*step_shape) for _ in range(n_steps)) or None, + logits=(tuple(torch.ones(*step_shape) for _ in range(n_steps)) or None) + if raw_logits + else None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + + +async def _drive( + backend: LocalHFBackend, + outputs: Any, + batch_size: int, + model_options: dict | None = None, +) -> list[ModelOutputThunk]: + """Run `generate_from_raw` with `model.generate` stubbed to return `outputs`. + + The stub is a plain async function rather than an `AsyncMock` so that no mock + call record keeps `outputs` alive after the call returns — the retention tests + below depend on this frame being the only strong reference. + """ + + async def _fake_to_thread(*_args: Any, **_kwargs: Any) -> Any: + return outputs + + actions: list[Component[Any]] = [ + Message("user", f"prompt {i}") for i in range(batch_size) + ] + with ( + patch("mellea.backends.huggingface.asyncio.to_thread", new=_fake_to_thread), + patch.object(backend, "do_generate_walks", new=AsyncMock()), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + return await backend.generate_from_raw( + actions, MagicMock(), model_options=model_options or {} + ) + + +def _row_bytes(dtype: "_torch.dtype", length: int) -> int: + return length * torch.empty(0, dtype=dtype).element_size() + + +# --- retention: a MOT must keep its own row, not the batch -------------------- + + +async def test_raw_response_sequences_retain_only_their_own_row(): + """`raw.response.sequences` must not keep the whole batch allocation alive. + + A `t[i:i+1]` view shares `t`'s storage, so every MOT would pin the full + `(batch, seq_len)` buffer for its entire lifetime. Slices handed to a MOT must + own compact storage — the same reason `generation.logits` is cloned. + """ + batch_size = 8 + backend = _make_backend(batch_size) + outputs = _batch_output(batch_size) + batch_bytes = outputs.sequences.untyped_storage().nbytes() + row_bytes = _row_bytes(torch.long, _SEQ_LEN) + assert batch_bytes > row_bytes, "test setup: batch must be larger than one row" + + results = await _drive(backend, outputs, batch_size) + + for i, result in enumerate(results): + held = result.raw.response.sequences.untyped_storage().nbytes() + assert held == row_bytes, ( + f"item {i}: raw.response.sequences retains {held} bytes of storage but its " + f"own row is only {row_bytes} bytes — it is a view pinning the whole " + f"{batch_bytes}-byte batch" + ) + + +async def test_raw_response_scores_retain_only_their_own_row(): + """`raw.response.scores` must not keep the whole batch scores tensors alive. + + Scores are the expensive field: one `(batch, vocab)` tensor per generated + token. A view into each pins every row of every step. + """ + batch_size, n_steps = 8, 3 + backend = _make_backend(batch_size) + outputs = _batch_output(batch_size, n_steps=n_steps) + batch_bytes = outputs.scores[0].untyped_storage().nbytes() + row_bytes = _row_bytes(torch.float32, _VOCAB) + + results = await _drive(backend, outputs, batch_size) + + for i, result in enumerate(results): + assert result.raw.response.scores is not None + for step, tensor in enumerate(result.raw.response.scores): + held = tensor.untyped_storage().nbytes() + assert held == row_bytes, ( + f"item {i} step {step}: raw.response.scores retains {held} bytes but its " + f"own row is {row_bytes} bytes — it is a view pinning the whole " + f"{batch_bytes}-byte step tensor" + ) + + +async def _run_holding_only_weakrefs( + backend: LocalHFBackend, batch_size: int, n_steps: int +) -> tuple[list[ModelOutputThunk], dict[str, Any]]: + """Run `generate_from_raw` and return the MOTs plus weakrefs to the batch tensors. + + Every strong reference to the batch (the tensors and the `GenerateDecoderOnlyOutput` + holding them) lives in this frame, so all of them are gone once it returns. + Anything still keeping the batch alive after that is held by the returned MOTs. + """ + outputs = _batch_output(batch_size, n_steps=n_steps) + assert outputs.scores is not None + refs = { + "sequences": weakref.ref(outputs.sequences), + "scores": [weakref.ref(s) for s in outputs.scores], + } + results = await _drive(backend, outputs, batch_size) + del outputs + return results, refs + + +async def test_batch_tensors_are_freed_once_only_mots_are_held(): + """Holding the returned MOTs must not keep the batch tensors alive. + + This is the invariant that matters in practice: a caller that keeps one MOT out + of a 32-prompt batch should not be pinning all 32 rows of sequences and of every + step's scores. The MOTs must still carry their own usable data afterwards. + """ + batch_size, n_steps = 4, 2 + backend = _make_backend(batch_size) + + results, refs = await _run_holding_only_weakrefs(backend, batch_size, n_steps) + gc.collect() + gc.collect() # second pass: release anything freed by breaking a cycle + + assert refs["sequences"]() is None, ( + "batch sequences tensor is still alive while only MOTs are held — " + "raw.response.sequences is a view into it" + ) + for step, ref in enumerate(refs["scores"]): + assert ref() is None, ( + f"batch scores tensor for step {step} is still alive while only MOTs are " + "held — raw.response.scores holds views into it" + ) + + # The data must survive the batch being freed. + for i, result in enumerate(results): + assert result.raw.response.sequences.shape == (1, _SEQ_LEN), ( + f"item {i}: sequences must remain usable after the batch is released" + ) + assert len(result.raw.response.scores) == n_steps + + +async def test_deepcopy_of_result_does_not_duplicate_the_batch(): + """`copy.deepcopy(mot)` must not allocate a whole extra batch. + + `torch.Tensor.__deepcopy__` deep-copies the *storage* and rebuilds the tensor + with the original size/stride/offset, so deep-copying a row view reallocates the + entire batch buffer. Asserting only that the data pointers differ hides this. + """ + batch_size = 8 + backend = _make_backend(batch_size) + outputs = _batch_output(batch_size) + batch_bytes = outputs.sequences.untyped_storage().nbytes() + row_bytes = _row_bytes(torch.long, _SEQ_LEN) + + results = await _drive(backend, outputs, batch_size) + deep = copy.deepcopy(results[0]) + + held = deep.raw.response.sequences.untyped_storage().nbytes() + assert held == row_bytes, ( + f"deepcopy allocated {held} bytes for a {row_bytes}-byte row — it duplicated " + f"the whole {batch_bytes}-byte batch because the slice was a view" + ) + + +# --- cleanup cost ------------------------------------------------------------- + + +async def test_generate_from_raw_does_not_force_gc_or_cuda_flush_without_cuda(): + """No full GC and no CUDA allocator flush on a CPU/MPS run. + + `gc.collect()` is a full generational collection and `torch.cuda.empty_cache()` + returns pooled blocks to the driver, making the next `generate()` pay fresh + allocations. Neither belongs on the per-call path, least of all when there is no + CUDA device at all. + """ + batch_size = 2 + backend = _make_backend(batch_size) + outputs = _batch_output(batch_size, n_steps=1) + + with ( + patch("torch.cuda.is_available", return_value=False), + patch("torch.cuda.empty_cache") as mock_empty_cache, + patch("gc.collect") as mock_collect, + ): + await _drive(backend, outputs, batch_size) + + assert mock_collect.call_count == 0, ( + f"generate_from_raw forced {mock_collect.call_count} full GC pass(es) on a " + "non-CUDA device" + ) + assert mock_empty_cache.call_count == 0, ( + f"generate_from_raw called torch.cuda.empty_cache() {mock_empty_cache.call_count} " + "time(s) with no CUDA device available" + ) + + +# --- fidelity: raw.response must describe what generate() returned ------------ + + +async def test_raw_response_preserves_beam_search_output_type_and_fields(): + """A beam-search batch must not be reported as a plain decoder-only output. + + With `num_beams > 1`, `generate()` returns `GenerateBeamDecoderOnlyOutput` with + `sequences_scores`/`beam_indices`, and `scores` has a leading dimension of + `batch * num_beams` — so per-item row indexing is wrong as well. Either mirror + the real type (carrying the beam fields) or leave `raw.response` unset; silently + relabelling it as `GenerateDecoderOnlyOutput` is the one unacceptable option. + """ + batch_size, num_beams, n_steps = 2, 3, 2 + backend = _make_backend(batch_size) + outputs = GenerateBeamDecoderOnlyOutput( + sequences=torch.zeros(batch_size, _SEQ_LEN, dtype=torch.long), + sequences_scores=torch.zeros(batch_size), + scores=tuple( + torch.zeros(batch_size * num_beams, _VOCAB) for _ in range(n_steps) + ), + logits=None, + beam_indices=torch.zeros(batch_size, _SEQ_LEN, dtype=torch.long), + attentions=None, + hidden_states=None, + past_key_values=None, + ) + + results = await _drive( + backend, outputs, batch_size, model_options={"num_beams": num_beams} + ) + + for i, result in enumerate(results): + raw = result.raw.response + if raw is None: + continue # opting out is acceptable; misreporting is not + assert isinstance(raw, GenerateBeamDecoderOnlyOutput), ( + f"item {i}: generate() returned GenerateBeamDecoderOnlyOutput but " + f"raw.response is {type(raw).__name__}" + ) + assert raw.sequences_scores is not None, ( + f"item {i}: beam sequences_scores was dropped" + ) + assert raw.beam_indices is not None, f"item {i}: beam_indices was dropped" + + +async def test_raw_response_is_never_emitted_with_null_sequences(): + """`raw.response` must never be a sequence-bearing output whose sequences are None. + + The raw path guards its slice with `isinstance(outputs.sequences, torch.Tensor)` + and falls back to `None`, casting it to `torch.LongTensor`. Decoding only needs + `sequences` to be indexable, so a non-tensor batch still produces MOTs with + values — and a `raw.response` that lies about having sequences. Either drop the + unreachable branch or leave `raw.response` unset. + """ + batch_size = 2 + backend = _make_backend(batch_size) + outputs = GenerateDecoderOnlyOutput( + sequences=[torch.zeros(_SEQ_LEN, dtype=torch.long) for _ in range(batch_size)], + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + + results = await _drive(backend, outputs, batch_size) + + for i, result in enumerate(results): + if result.raw.response is None: + continue + assert result.raw.response.sequences is not None, ( + f"item {i}: raw.response was populated with sequences=None" + ) + + +async def test_raw_response_scores_follow_generate_output_not_the_logits_option(): + """`raw.response.scores` mirrors `generate()`, independent of `ModelOption.LOGITS`. + + The implementation gates on `outputs.scores is not None`, never on the option, so + scores requested through a passthrough option (or any other route) still land in + `raw.response`. This is the correct behaviour — the point of the test is that + "scores is None when LOGITS is not requested" is not the invariant. + """ + batch_size, n_steps = 2, 2 + backend = _make_backend(batch_size) + outputs = _batch_output(batch_size, n_steps=n_steps) + + results = await _drive(backend, outputs, batch_size, model_options={}) + + for i, result in enumerate(results): + assert result.raw.response.scores is not None, ( + f"item {i}: generate() returned scores, so raw.response.scores must be set " + "even though ModelOption.LOGITS was not requested" + ) + assert len(result.raw.response.scores) == n_steps + assert result.generation.logits is None, ( + f"item {i}: generation.logits is the option-gated field and must stay unset" + ) + + +async def test_raw_response_raw_logits_retain_only_their_own_row(): + """`raw.response.logits` gets the same treatment as `scores`. + + Covers the `outputs.logits` branch, which the PR's tests never exercise. + """ + batch_size, n_steps = 4, 2 + backend = _make_backend(batch_size) + outputs = _batch_output(batch_size, n_steps=n_steps, raw_logits=True) + row_bytes = _row_bytes(torch.float32, _VOCAB) + + results = await _drive( + backend, outputs, batch_size, model_options={ModelOption.RAW_LOGITS: True} + ) + + for i, result in enumerate(results): + assert result.raw.response.logits is not None + for step, tensor in enumerate(result.raw.response.logits): + assert tensor.shape == (1, _VOCAB) + held = tensor.untyped_storage().nbytes() + assert held == row_bytes, ( + f"item {i} step {step}: raw.response.logits retains {held} bytes for a " + f"{row_bytes}-byte row" + ) + + +# --- one-time notice ---------------------------------------------------------- + + +async def test_omitted_fields_notice_is_logged_once_per_backend(caplog): + """The "fields omitted" notice fires once per backend, not once per item or call.""" + batch_size = 3 + backend = _make_backend(batch_size) + needle = "are not available on the raw batch path" + + with caplog.at_level(logging.DEBUG, logger="mellea"): + await _drive(backend, _batch_output(batch_size), batch_size) + await _drive(backend, _batch_output(batch_size), batch_size) + + hits = [r for r in caplog.records if needle in r.getMessage()] + assert len(hits) == 1, ( + f"expected the omitted-fields notice once per backend, saw {len(hits)} across " + f"two calls of {batch_size} items each" + ) + + +# --- chat path: the same "null it out" idiom, same defect --------------------- + + +def _make_kv_cache() -> DynamicCache: + """Return a small `DynamicCache` with real tensors for weakref testing.""" + cache = DynamicCache() + cache.update( + key_states=torch.zeros(1, 1, 1, 4), + value_states=torch.zeros(1, 1, 1, 4), + layer_idx=0, + ) + return cache + + +async def _post_process_holding_only_weakrefs( + backend: LocalHFBackend, n_steps: int +) -> tuple[ModelOutputThunk, dict[str, Any]]: + """Run `post_processing` and return the MOT plus weakrefs to the HF output tensors.""" + kv_cache = _make_kv_cache() + sequences = torch.zeros(1, _SEQ_LEN, dtype=torch.long) + hf_out = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=tuple(torch.zeros(1, _VOCAB) for _ in range(n_steps)), + logits=tuple(torch.ones(1, _VOCAB) for _ in range(n_steps)), + attentions=None, + hidden_states=None, + past_key_values=kv_cache, + ) + assert hf_out.logits is not None + assert hf_out.scores is not None + assert kv_cache.layers + assert all(isinstance(layer, CacheLayerMixin) for layer in kv_cache.layers), ( + "test setup: expected standard KV cache layers" + ) + + refs = { + "container": weakref.ref(hf_out), + "sequences": [weakref.ref(sequences)], + "logits": [weakref.ref(t) for t in hf_out.logits], + "scores": [weakref.ref(t) for t in hf_out.scores], + "past_key_values": [ + weakref.ref(t) + for layer in kv_cache.layers + if isinstance(layer, CacheLayerMixin) + for t in (layer.keys, layer.values) + if t is not None + ], + } + del kv_cache, sequences + + mot = ModelOutputThunk(value="hi") + mot._call.action = Message("user", "noop") + mot._call.model_options = {} + mot.raw.response = hf_out + del hf_out + + await backend.post_processing( + mot, [], None, False, {}, None, torch.zeros(1, _PROMPT_LEN, dtype=torch.long) + ) + return mot, refs + + +async def test_post_processing_clearing_raw_logits_actually_releases_them(): + """Clearing raw scores and logits must release their tensors. + + `GenerateDecoderOnlyOutput` is a `ModelOutput`, i.e. an `OrderedDict` + subclass that mirrors fields into the mapping. Clearing an attribute alone + can leave the mapping entry—and therefore its tensors—alive. When + `raw.response` is retained, post-processing must remove all references to + the cleared scores and logits. + """ + backend = _make_backend(1) + backend._use_caches = True # keeps raw.response, so the container survives + + mot, refs = await _post_process_holding_only_weakrefs(backend, n_steps=2) + gc.collect() + gc.collect() + + assert mot.raw.response is not None, "test setup: raw.response should be retained" + + for field in ("scores", "logits"): + assert getattr(mot.raw.response, field) is None, ( + f"test setup: {field} attribute was not cleared" + ) + + for step, ref in enumerate(refs[field]): + assert ref() is None, ( + f"raw {field} tensor for step {step} is still alive after " + f"hf_output.{field} was cleared" + ) + + +async def test_post_processing_clears_all_tensor_fields_when_caching_disabled(): + """When `_use_caches=False`, `post_processing` clears tensor-bearing fields. + + The HF output `sequences`, `scores`, `logits`, and `past_key_values` + must be released after post_processing. + """ + backend = _make_backend(1) + backend._use_caches = False # triggers the no-caching clearing path + + mot, refs = await _post_process_holding_only_weakrefs(backend, n_steps=2) + + gc.collect() + gc.collect() + + assert mot.raw.response is None, ( + "raw.response should be dropped on the no-caching path" + ) + + for field in ("sequences", "scores", "logits", "past_key_values"): + for i, ref in enumerate(refs[field]): + assert ref() is None, ( + f"{field} tensor {i} is still alive after post_processing " + "with _use_caches=False" + ) diff --git a/test/backends/test_huggingface_unit.py b/test/backends/test_huggingface_unit.py index beaf52187..7e1d37836 100644 --- a/test/backends/test_huggingface_unit.py +++ b/test/backends/test_huggingface_unit.py @@ -5,6 +5,7 @@ import asyncio from types import SimpleNamespace +from typing import Any, cast from unittest.mock import MagicMock, patch import pytest @@ -18,9 +19,15 @@ ) import base64 +import gc import struct +import weakref -from transformers.generation.utils import GenerateDecoderOnlyOutput +from transformers.cache_utils import CacheLayerMixin, DynamicCache +from transformers.generation.utils import ( + GenerateBeamDecoderOnlyOutput, + GenerateDecoderOnlyOutput, +) from mellea.backends import ModelOption from mellea.backends.adapters import AdapterMixin, IntrinsicAdapter @@ -920,6 +927,178 @@ def transform(self, chunk, rewritten): assert all(t.shape == (vocab_size,) for t in output.generation.logits) +@pytest.mark.asyncio +async def test_intrinsic_closure_cell_and_kv_cache_released_after_post_processing( + stub_backend, +): + """Holding only the MOT after post_processing must not pin intrinsic HF output. + + Two related retention paths are exercised: + + 1. `raw_hf_output_cell` — the closure captured by `_gen.process` (a + `functools.partial` that outlives the call). The cell must be cleared + after its value is transferred to `mot.raw.response`, otherwise the + held MOT retains the full GenerateDecoderOnlyOutput. + + 2. `past_key_values` — the KV cache inside the HF output. On the no-cache + path, post_processing must remove it from the GenerateDecoderOnlyOutput + before clearing `mot.raw.response`. + + The test deliberately retains the MOT while dropping all independent test + references to the HF output, DynamicCache, and KV tensors. + """ + backend = _make_intrinsic_backend_stub(stub_backend) + backend.processing = lambda *args, **kwargs: LocalHFBackend.processing( + backend, *args, **kwargs + ) + backend.post_processing = lambda *args, **kwargs: LocalHFBackend.post_processing( + backend, *args, **kwargs + ) + backend._surface_logits = lambda mot, hf_out: LocalHFBackend._surface_logits( + backend, mot, hf_out + ) + backend._use_caches = False + backend.cache_put = MagicMock() + backend._tokenizer = MagicMock(eos_token_id=0) + backend.model_id = "stub-model" + + # Build a small KV cache with real tensors so weakrefs can verify that the + # cache and its allocations are released. + kv_cache = DynamicCache() + kv_cache.update( + key_states=torch.zeros(1, 1, 1, 4), + value_states=torch.zeros(1, 1, 1, 4), + layer_idx=0, + ) + + fake_scores = (torch.zeros(1, 32000),) + fake_hf_output = GenerateDecoderOnlyOutput( + sequences=torch.tensor([[1, 2]]), + scores=fake_scores, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=kv_cache, + ) + + # Take weakrefs before dropping all direct strong references. + ref_container = weakref.ref(fake_hf_output) + ref_kv_cache = weakref.ref(kv_cache) + ref_kv_tensors = [ + weakref.ref(t) + for layer in kv_cache.layers + if isinstance(layer, CacheLayerMixin) + for t in (layer.keys, layer.values) + if t is not None + ] + + adapter = _make_intrinsic_adapter_stub() + backend._added_adapters = {adapter.qualified_name: adapter} + + class _FakeChatCompletionResponse: + class _Choice: + class _Message: + content = "0.9" + + message = _Message() + + choices = [_Choice()] + + class _FakeResultProcessorPassthrough: + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + def transform(self, chunk: Any, rewritten: Any) -> Any: + return chunk + + def fake_transformers_inputs( + rewritten: Any, tokenizer: Any, model: Any, ll_tokenizer: Any = None + ) -> tuple[dict, dict]: + return {"input_tokens": torch.tensor([[1]])}, {} + + def fake_generate_with_transformers( + tokenizer: Any, model: Any, generate_input: Any, other_input: Any + ) -> Any: + model.generate(inputs=generate_input["input_tokens"]) + return _FakeChatCompletionResponse() + + mock_model = MagicMock() + mock_model.generate = MagicMock(return_value=fake_hf_output) + backend._model = mock_model + + with ( + patch( + "mellea.backends.huggingface.granite_formatters.IntrinsicsRewriter", + _FakeRewriter, + ), + patch( + "mellea.backends.huggingface.granite_formatters.IntrinsicsResultProcessor", + _FakeResultProcessorPassthrough, + ), + patch( + "mellea.formatters.granite.base.util.chat_completion_request_to_transformers_inputs", + side_effect=fake_transformers_inputs, + ), + patch( + "mellea.formatters.granite.base.util.generate_with_transformers", + side_effect=fake_generate_with_transformers, + ), + ): + output = await LocalHFBackend._generate_from_intrinsic( + backend, + Intrinsic("answerability"), + ChatContext().add(Message("user", "Is the sky blue?")), + model_options={ModelOption.LOGITS: True}, + ) + + assert output._gen.generate is not None + await output._gen.generate + + while not output._gen.queue.empty(): + item = output._gen.queue.get_nowait() + if item is not None: + await output._gen.process(output, item) + + output._computed = True + + # MagicMock artificially retains its return_value. A real HF model does not + # retain the result of generate(), so remove this test-only retention path. + mock_model.generate.return_value = None + + # Drop the test's direct strong references. From here onward, the retained + # MOT should be the only object graph capable of keeping the HF output alive. + del fake_hf_output + del kv_cache + del fake_scores + + await backend.post_processing( + output, [], None, False, {}, None, torch.tensor([[1]]) + ) + + gc.collect() + gc.collect() + + assert output.raw.response is None, ( + "raw.response should be None on the no-caching path" + ) + + assert ref_container() is None, ( + "GenerateDecoderOnlyOutput is still alive while the MOT is held; " + "the _gen.process closure or another MOT-owned path is pinning it" + ) + + assert ref_kv_cache() is None, ( + "DynamicCache is still alive while the MOT is held; " + "past_key_values was not fully released on the no-cache path" + ) + + for i, ref in enumerate(ref_kv_tensors): + assert ref() is None, ( + f"KV-cache tensor {i} is still alive while the MOT is held; " + "past_key_values or another reference path is retaining it" + ) + + @pytest.mark.parametrize("images,audio", _MULTIMODAL_CASES) @pytest.mark.asyncio async def test_multimodal_blocks_raise_error(images, audio): @@ -1393,3 +1572,286 @@ def _capture_grammar(schema, overrides=None): assert captured[0].get("whitespace_pattern") == r"[\x20\x0A\x0D\x09]{0,20}", ( f"Expected bounded whitespace_pattern to override False in {path_name}" ) + + +def _make_raw_fake_setup( + batch_size: int, vocab_size: int, n_tokens: int, prompt_len: int +): + """Return (backend, fake_encoding, fake_outputs, actions) for generate_from_raw tests.""" + backend = _make_backend() + fake_input_ids = torch.zeros(batch_size, prompt_len, dtype=torch.long) + fake_encoding = MagicMock() + fake_encoding.__getitem__ = lambda self, k: ( + fake_input_ids + if k == "input_ids" + else torch.ones(batch_size, prompt_len, dtype=torch.long) + ) + fake_encoding.to = MagicMock(return_value=fake_encoding) + backend._tokenizer = MagicMock(eos_token_id=0, vocab_size=vocab_size) + backend._tokenizer.__len__ = MagicMock(return_value=vocab_size) + backend._tokenizer.return_value = fake_encoding + decode_values = [f"result_{chr(ord('a') + i)}" for i in range(batch_size)] + backend._tokenizer.batch_decode = MagicMock(return_value=decode_values) + return backend, fake_encoding, fake_input_ids + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_set_per_mot(): + """Every MOT from generate_from_raw has raw.response set to a GenerateDecoderOnlyOutput. + + Asserts: + - raw.response is not None for each MOT. + - raw.response.sequences.shape == (1, full_seq_len). + - raw.response.sequences shares storage with the original batch sequences tensor (Clone, not view). + - raw.response.past_key_values is None. + - raw.response.attentions is None. + - raw.response.hidden_states is None. + """ + batch_size = 2 + vocab_size = 32000 + n_tokens = 3 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + sequences = torch.zeros(batch_size, full_seq_len, dtype=torch.long) + fake_outputs = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + actions = [Message("user", "hello"), Message("user", "world")] + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + actions, MagicMock(), model_options={} + ) + + assert len(results) == batch_size + for item_idx, result in enumerate(results): + assert result.raw.response is not None, ( + f"item {item_idx}: raw.response must be set" + ) + assert isinstance(result.raw.response, GenerateDecoderOnlyOutput), ( + f"item {item_idx}: raw.response must be GenerateDecoderOnlyOutput" + ) + assert result.raw.response.sequences.shape == (1, full_seq_len), ( + f"item {item_idx}: sequences shape must be (1, {full_seq_len})" + ) + # Clone - must NOT share storage with the original batch tensor. + assert ( + result.raw.response.sequences.untyped_storage().data_ptr() + != sequences.untyped_storage().data_ptr() + ), f"item {item_idx}: sequences must be a clone, not a view" + assert result.raw.response.past_key_values is None, ( + f"item {item_idx}: past_key_values must be None" + ) + assert result.raw.response.attentions is None, ( + f"item {item_idx}: attentions must be None" + ) + assert result.raw.response.hidden_states is None, ( + f"item {item_idx}: hidden_states must be None" + ) + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_scores_are_clones_when_logits_requested(): + """raw.response.scores is a tuple of clones when ModelOption.LOGITS is set. + + Each tensor in raw.response.scores must own compact per-row storage and must + not share storage with the corresponding batch step tensor — consistent with + generation.logits which also holds clones. + """ + batch_size = 2 + vocab_size = 32000 + n_tokens = 3 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + sequences = torch.zeros(batch_size, full_seq_len, dtype=torch.long) + fake_scores = tuple(torch.randn(batch_size, vocab_size) for _ in range(n_tokens)) + fake_outputs = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=fake_scores, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + actions = [Message("user", "hello"), Message("user", "world")] + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + actions, MagicMock(), model_options={ModelOption.LOGITS: True} + ) + + for item_idx, result in enumerate(results): + assert result.raw.response.scores is not None, ( + f"item {item_idx}: raw.response.scores must be set when LOGITS=True" + ) + assert len(result.raw.response.scores) == n_tokens, ( + f"item {item_idx}: one scores tensor per generation step" + ) + for tok_idx, t in enumerate(result.raw.response.scores): + assert t.shape == (1, vocab_size), ( + f"item {item_idx} token {tok_idx}: shape must be (1, vocab_size)" + ) + # Clone - must NOT share storage with the original batch step tensor. + assert ( + t.untyped_storage().data_ptr() + != fake_scores[tok_idx].untyped_storage().data_ptr() + ), f"item {item_idx} token {tok_idx}: raw.response.scores must be a clone" + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_scores_none_when_logits_not_requested(): + """raw.response.scores is None when ModelOption.LOGITS is not set.""" + batch_size = 1 + vocab_size = 32000 + n_tokens = 2 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + # When LOGITS is not set, model.generate() is called without output_scores=True, + # so outputs.scores will be None — simulate that here. + sequences = torch.zeros(batch_size, full_seq_len, dtype=torch.long) + fake_outputs = GenerateDecoderOnlyOutput( + sequences=sequences, + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ) + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + [Message("user", "hi")], MagicMock(), model_options={} + ) + + assert results[0].raw.response.scores is None, ( + "raw.response.scores must be None when model.generate() returns no scores" + ) + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_none_for_non_tensor_sequences(): + """raw.response stays None when cached raw output sequences are not a tensor.""" + batch_size = 2 + vocab_size = 32000 + n_tokens = 3 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + backend._use_caches = True + fake_outputs = cast( + Any, + SimpleNamespace( + sequences=[[0] * full_seq_len for _ in range(batch_size)], + scores=None, + logits=None, + attentions=None, + hidden_states=None, + past_key_values=None, + ), + ) + actions = [Message("user", "hello"), Message("user", "world")] + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + cast(Any, actions), MagicMock(), model_options={} + ) + + assert len(results) == batch_size + assert "raw_batch_non_tensor_sequences_unsupported" in backend._warned_about + for item_idx, result in enumerate(results): + assert result.raw.response is None, ( + f"item {item_idx}: raw.response must stay None for non-tensor sequences" + ) + + +@pytest.mark.asyncio +async def test_generate_from_raw_raw_response_none_for_beam_outputs(): + """raw.response stays None for cached beam-search outputs.""" + batch_size = 2 + vocab_size = 32000 + n_tokens = 3 + prompt_len = 1 + full_seq_len = prompt_len + n_tokens + + backend, _fake_encoding, _fake_input_ids = _make_raw_fake_setup( + batch_size, vocab_size, n_tokens, prompt_len + ) + backend._use_caches = True + fake_outputs = GenerateBeamDecoderOnlyOutput( + sequences=torch.zeros(batch_size, full_seq_len, dtype=torch.long), + sequences_scores=None, + scores=None, + logits=None, + beam_indices=torch.zeros(batch_size, full_seq_len, dtype=torch.long), + attentions=None, + hidden_states=None, + past_key_values=None, + ) + actions = [Message("user", "hello"), Message("user", "world")] + + with ( + patch( + "mellea.backends.huggingface.asyncio.to_thread", return_value=fake_outputs + ), + patch.object(backend, "do_generate_walks"), + patch.object(backend, "formatter") as mock_fmt, + ): + mock_fmt.print = MagicMock(return_value="prompt") + results = await backend.generate_from_raw( + cast(Any, actions), MagicMock(), model_options={} + ) + + assert len(results) == batch_size + assert "raw_batch_beam_search_unsupported" in backend._warned_about + for item_idx, result in enumerate(results): + assert result.raw.response is None, ( + f"item {item_idx}: raw.response must stay None for beam-search outputs" + )