| """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() |
|
|
