"""CPU-only tests of batched decoding, timing, and infrastructure accounting.""" from contextlib import nullcontext import json from pathlib import Path import sys import tempfile from types import SimpleNamespace import unittest from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) import generate import evaluate import benchmark_batch class Scalar: def __init__(self, value): self.value = value def item(self): return self.value class Row: def __init__(self, values): self.values = values def sum(self): return Scalar(sum(self.values)) class Matrix: def __init__(self, rows): self.rows = rows self.shape = (len(rows), len(rows[0])) def __getitem__(self, key): if isinstance(key, tuple): row, column = key return Matrix([value[column] for value in self.rows[row]]) return Row(self.rows[key]) def tolist(self): return self.rows class Batch(dict): def __init__(self, ids, masks): super().__init__(input_ids=Matrix(ids), attention_mask=Matrix(masks)) self.input_ids = self["input_ids"] self.attention_mask = self["attention_mask"] def to(self, device): assert device == "cuda" return self class Tokenizer: eos_token_id = 99 padding_side = "right" width = 4 def apply_chat_template(self, messages, **kwargs): assert kwargs == {"tokenize": False, "add_generation_prompt": True} return messages[-1]["content"] def __call__(self, texts, **kwargs): self.padding_observed = self.padding_side self.kwargs_observed = kwargs width = self.width return Batch([[0] * (width - 2) + [10, 11] for _ in texts], [[0] * (width - 2) + [1, 1] for _ in texts]) def decode(self, tokens, **kwargs): return ",".join(str(token) for token in tokens if token not in (98, 99)) class Model: def __init__(self, rows, eos): self.rows = rows self.generation_config = SimpleNamespace(eos_token_id=eos) def generate(self, **kwargs): self.kwargs = kwargs prefix = kwargs["input_ids"].rows return Matrix([head + tail for head, tail in zip(prefix, self.rows)]) def backend(rows, eos=99): instance = generate.UnslothGenerator.__new__(generate.UnslothGenerator) instance.tokenizer = Tokenizer() instance.model = Model(rows, eos) instance.torch = SimpleNamespace(cuda=SimpleNamespace(synchronize=lambda: None), inference_mode=nullcontext) return instance class BatchInferenceTests(unittest.TestCase): def test_left_padding_mask_and_eos_trim(self): generator = backend([[1, 2, 99, 99], [3, 99, 99, 99]]) rows = generator.generate_many(["short", "longer"], 4) self.assertEqual(generator.tokenizer.padding_observed, "left") self.assertEqual(generator.tokenizer.padding_side, "right") self.assertFalse(generator.tokenizer.kwargs_observed["add_special_tokens"]) self.assertIn("attention_mask", generator.model.kwargs) self.assertEqual([row["output"] for row in rows], ["1,2", "3"]) self.assertEqual([row["tokens"] for row in rows], [3, 2]) self.assertEqual([row["prompt_tokens"] for row in rows], [2, 2]) self.assertTrue(all(row["done_reason"] == "stop" for row in rows)) def test_eos_list_and_length_exhaustion(self): generator = backend([[1, 98, 99, 99], [1, 2, 3, 4]], eos=[98, 99]) rows = generator.generate_many(["first", "second"], 4) self.assertEqual(rows[0]["tokens"], 2) self.assertEqual(rows[0]["output"], "1") self.assertEqual(rows[1]["done_reason"], "length") self.assertEqual(rows[1]["tokens"], 4) def test_whole_batch_latency_never_divided(self): generator = backend([[1, 99], [2, 99]]) with patch.object(generate.time, "monotonic", side_effect=[100.0, 101.0, 109.0, 111.0]): rows = generator.generate_many(["first", "second"], 2) self.assertEqual([row["generation_seconds"] for row in rows], [8.0, 8.0]) self.assertEqual([row["generation_call_seconds"] for row in rows], [11.0, 11.0]) self.assertEqual([row["generation_batch_size"] for row in rows], [2, 2]) def test_context_bound_rejects_without_silent_truncation(self): generator = backend([[1, 99]]) generator.tokenizer.width = 4090 with self.assertRaisesRegex(ValueError, "4096-token"): generator.generate_many(["long"], 16) self.assertEqual(generator.tokenizer.padding_side, "right") self.assertFalse(hasattr(generator.model, "kwargs")) def test_batch_limit_enforced(self): generator = backend([]) for prompts in [[], ["x"] * 33]: with self.assertRaises(ValueError): generator.generate_many(prompts, 2) def test_sixteen_prompt_capacity_is_accepted(self): generator = backend([[1, 99]] * 16) rows = generator.generate_many(["training prompt"] * 16, 2) self.assertEqual(len(rows), 16) self.assertTrue(all(row["generation_batch_size"] == 16 for row in rows)) def test_thirty_two_prompt_capacity_is_accepted(self): generator = backend([[1, 99]] * 32) rows = generator.generate_many(["training prompt"] * 32, 2) self.assertEqual(len(rows), 32) self.assertTrue(all(row["generation_batch_size"] == 32 for row in rows)) def test_prompt_to_launch_uses_full_call_with_explicit_legacy_fallback(self): check = {"compile_success": True, "clean_run_60s": True, "validation_queue_seconds": 4.0, "preimport_seconds": 1.0, "import": {"wall_seconds": 2.0}, "runtime": {"startup_seconds": 3.0}} for call_seconds, expected, source in [(11.0, 21.0, "generation_call_seconds"), (None, 18.0, "generation_seconds_legacy_fallback")]: with self.subTest(call_seconds=call_seconds), tempfile.TemporaryDirectory() as folder: root = Path(folder) item = root / "items/0000" item.mkdir(parents=True) result = {"index": 0, "generation_seconds": 8.0, "generation_call_seconds": call_seconds, "prompt_to_launch_seconds": None} with patch("builtins.print"): evaluate.finish_result(item, result, check, root, 1, 60) saved = json.loads((item / "result.json").read_text()) self.assertEqual(saved["prompt_to_launch_seconds"], expected) self.assertEqual(saved["generation_seconds"], 8.0) self.assertEqual(saved["generation_latency_source"], source) def test_benchmark_does_not_select_fast_batch_without_headroom(self): measurements = [ {"status": "success", "batch_size": 1, "sum_tokens_per_second": 10, "peak_allocated_bytes": 5, "peak_reserved_bytes": 6, "full_context_memory_estimate": {"within_limit": True}}, {"status": "success", "batch_size": 4, "sum_tokens_per_second": 35, "peak_allocated_bytes": 8, "peak_reserved_bytes": 9, "full_context_memory_estimate": {"within_limit": True}}, {"status": "success", "batch_size": 16, "sum_tokens_per_second": 90, "peak_allocated_bytes": 10, "peak_reserved_bytes": 13, "full_context_memory_estimate": {"within_limit": True}}, {"status": "success", "batch_size": 32, "sum_tokens_per_second": 130, "peak_allocated_bytes": 9, "peak_reserved_bytes": 10, "full_context_memory_estimate": {"within_limit": False}}, {"status": "failed", "batch_size": 8}, ] best = benchmark_batch.select_candidate(measurements, 12) self.assertEqual(best["batch_size"], 4) self.assertIsNone(benchmark_batch.select_candidate(measurements, 5)) def test_full_context_estimate_uses_actual_qwen_dims_padding_and_gpu_reserve(self): config = {"num_hidden_layers": 28, "num_attention_heads": 28, "num_key_value_heads": 4, "hidden_size": 3584} row = {"batch_size": 16, "per_prompt_input_tokens": [128, 256] * 8, "per_prompt_tokens": [256] * 16, "peak_allocated_bytes": 7_100_000_000, "peak_reserved_bytes": 7_360_000_000} result = benchmark_batch.estimate_full_context_memory(row, config, 16_000_000_000, 12_500_000_000) self.assertEqual(result["bytes_per_sequence_token"], 57344) self.assertEqual(result["observed_cache_tokens"], 511) self.assertEqual(result["additional_cache_tokens"], 4097) self.assertEqual(result["estimated_peak_bytes"], 7_360_000_000 + 4097 * 16 * 57344) self.assertEqual(result["effective_limit_bytes"], 12_500_000_000) self.assertTrue(result["within_limit"]) smaller_gpu = benchmark_batch.estimate_full_context_memory(row, config, 12_000_000_000, 12_500_000_000) self.assertEqual(smaller_gpu["effective_limit_bytes"], 10_000_000_000) self.assertFalse(smaller_gpu["within_limit"]) larger = {**row, "batch_size": 32, "per_prompt_input_tokens": [256] * 32, "per_prompt_tokens": [256] * 32, "peak_reserved_bytes": 9_000_000_000} self.assertFalse(benchmark_batch.estimate_full_context_memory(larger, config, 16_000_000_000, 12_500_000_000)["within_limit"]) def test_memory_gate_rejects_old_report_without_estimate(self): self.assertIsNone(benchmark_batch.select_candidate([ {"status": "success", "batch_size": 16, "sum_tokens_per_second": 80, "peak_allocated_bytes": 7, "peak_reserved_bytes": 8}], 12)) def test_saved_backend_failure_blocks_published_rates(self): with tempfile.TemporaryDirectory() as folder: root = Path(folder) item = root / "items/0000" item.mkdir(parents=True) (item / "result.json").write_text(json.dumps({"index": 0, "compile_success": False, "clean_run_60s": False})) (root / "generation_infrastructure_error.json").write_text(json.dumps({"error": "CUDA out of memory", "indices": [0]})) result = evaluate.summarize(root, 1, 60) self.assertEqual(result["status"], "incomplete") self.assertIsNone(result["compile_success_rate"]) self.assertGreater(result["infrastructure_failures"], 0) def test_single_selected_batch_runs_full_capacity_without_baseline(self): best = {"batch_size": 16} self.assertIsNone(benchmark_batch.full_check_reason(best, None, [16], False)) self.assertIsNotNone(benchmark_batch.full_check_reason(None, None, [16], False)) self.assertIsNotNone(benchmark_batch.full_check_reason(best, None, [16], True)) self.assertIsNotNone(benchmark_batch.full_check_reason(best, None, [4, 16], False)) self.assertIsNone(benchmark_batch.full_check_reason(best, 4.5, [1, 4, 8, 16], False)) def test_malformed_batch_is_a_backend_error_before_persistence(self): for rows in [[], [{"output": "valid"}] * 15, [{"output": None}] * 16, ["not a record"] * 16, {"output": "not a list"}]: with self.assertRaises(RuntimeError): evaluate.check_batch_outputs(rows, 16) evaluate.check_batch_outputs([{"output": ""}] * 16, 16) if __name__ == "__main__": unittest.main()