julia/router/tests/test_engine.py

import tempfile
import unittest
from unittest.mock import patch
import numpy as np
import torch

from julia.data import Collator, sequence
from julia.inference import TransformerEngine as Engine
from julia.router.engine import FastEngine
from julia.router.native import BendReducer
from synthetic import make_checkpoint, sample_rows


class NativeMathTests(unittest.TestCase):
    def test_transformer_math(self):
        reducer = BendReducer()
        rng = np.random.default_rng(42)
        for width in [1, 7, 64, 384]:
            x = rng.normal(size=(3, width)).astype('float32')
            gamma = rng.normal(size=width).astype('float32')
            beta = rng.normal(size=width).astype('float32')
            expected = torch.nn.functional.layer_norm(torch.from_numpy(x), (width,),
                        torch.from_numpy(gamma), torch.from_numpy(beta)).numpy()
            np.testing.assert_allclose(reducer.layernorm(x, gamma, beta), expected, atol=2e-6, rtol=2e-5)
            for row in x:
                index, probs = reducer.softmax(row)
                self.assertEqual(index, int(row.argmax()))
                np.testing.assert_allclose(probs, torch.from_numpy(row).softmax(-1).numpy(), atol=2e-7)
        _, probabilities = reducer.softmax(np.array([10000, 10001, -10000], np.float32))
        self.assertAlmostEqual(float(probabilities.sum()), 1.0)

    def test_resident_dense_tiles_reuse_and_partial_edges(self):
        from julia.router.native import BendMatrix
        rng = np.random.default_rng(52)
        for inputs, outputs, rows in [(1, 1, 1), (7, 5, 9), (64, 17, 7), (33, 65, 19), (384, 1152, 5)]:
            weights = rng.normal(size=(outputs, inputs)).astype('float32')
            values = rng.normal(size=(rows, inputs)).astype('float32')
            matrix = BendMatrix(weights)
            reference = values @ weights.T
            first = matrix(values)
            np.testing.assert_allclose(first, reference, atol=1e-4, rtol=1e-4)
            for _ in range(8):
                np.testing.assert_array_equal(matrix(values), first)
            with self.assertRaises(ValueError): matrix(np.ones((2, inputs + 1), np.float32))
            with self.assertRaises(ValueError): matrix(np.full((1, inputs), np.nan, np.float32))
            matrix.close()
            matrix.close()
            with self.assertRaises(RuntimeError): matrix(values)
        with self.assertRaises(ValueError): BendMatrix([[float('nan')]])

    def test_tensor_chunks_and_opt_in_native_profile(self):
        from julia.router.native import BendMatrix
        reducer = BendReducer()
        weights = np.random.default_rng(8).normal(size=(257, 3)).astype('float32')
        matrix = BendMatrix(weights, reducer)
        # Noncontiguous input and > 1M output cells require two native chunks.
        values = torch.arange(4090 * 6, dtype=torch.float32).reshape(4090, 6)[:, ::2] / 1000
        try:
            reducer.profile(enabled=True, reset=True)
            actual = matrix.tensor(values)
            measured = reducer.profile(enabled=False)
            self.assertEqual(measured['calls'], 2)
            self.assertTrue(all(measured[k] >= 0 for k in ('pack_ms', 'execute_ms', 'unpack_ms')))
            torch.testing.assert_close(actual, values @ torch.from_numpy(weights).T, rtol=1e-5, atol=1e-5)
            matrix.tensor(values[:1])
            self.assertEqual(reducer.profile(), measured)
            with self.assertRaises(ValueError):
                matrix.tensor(torch.full((1, 3), float('nan')))
        finally:
            reducer.profile(enabled=False, reset=True)
            matrix.close()

    def test_layernorm_parallel_callers_and_large_rows(self):
        import concurrent.futures
        reducer = BendReducer()
        rng = np.random.default_rng(19)
        values = [rng.normal(size=(rows, 384)).astype('float32') for rows in (1, 7, 32, 128)]
        references = [torch.nn.functional.layer_norm(torch.from_numpy(x), (384,)).numpy() for x in values]
        with concurrent.futures.ThreadPoolExecutor(4) as pool:
            actual = list(pool.map(reducer.layernorm, values * 8))
        for i, result in enumerate(actual):
            np.testing.assert_allclose(result, references[i % len(values)], atol=3e-6, rtol=2e-5)

    def test_packed_matrix_concurrent_owners(self):
        from concurrent.futures import ThreadPoolExecutor
        from julia.router.native import BendMatrix
        rng = np.random.default_rng(2026)
        weights = [rng.normal(size=(n, 33)).astype('float32') for n in (5, 17, 128)]
        matrices = [BendMatrix(w) for w in weights]
        values = rng.normal(size=(19, 33)).astype('float32')
        try:
            with ThreadPoolExecutor(6) as pool:
                results = list(pool.map(lambda i: matrices[i % 3](values), range(96)))
            for i, actual in enumerate(results):
                np.testing.assert_allclose(actual, values @ weights[i % 3].T, atol=2e-5, rtol=2e-5)
        finally:
            for matrix in matrices:
                matrix.close()

    def test_rejected_native_math(self):
        reducer = BendReducer()
        for x in [[], [float('nan')], [float('inf')]]:
            with self.assertRaises(ValueError):
                reducer.softmax(x)
        for kwargs in [dict(epsilon=0), dict(gamma=[1]), dict(beta=[0]), dict(epsilon=float('nan'))]:
            with self.assertRaises(ValueError):
                reducer.layernorm([[1, 2, 3]], **kwargs)


class ProbabilityDisplayTests(unittest.TestCase):
    def test_rounds_decisive_answer_to_certain(self):
        from julia.probabilities import display_probabilities

        self.assertEqual(display_probabilities([0.955, 0.044, 0.001]), [1.0, 0.0, 0.0])
        self.assertNotEqual(display_probabilities([0.951, 0.049]), [1.0, 0.0])

    def test_redistributes_sub_one_percent_values(self):
        from julia.probabilities import display_probabilities

        result = display_probabilities([0.6, 0.395, 0.005])
        self.assertEqual(result[2], 0.0)
        self.assertAlmostEqual(sum(result), 1.0)
        self.assertAlmostEqual(result[0] / result[1], 0.6 / 0.395)


class InputValidationTests(unittest.TestCase):
    def test_public_loader_rejects_invalid_backend(self):
        from julia import load_model
        from julia.inference import Engine

        self.assertIs(load_model, Engine)
        with self.assertRaisesRegex(ValueError, 'backend'):
            load_model('unused', backend='unknown')

    def test_rejects_malformed_requests(self):
        from julia.data import validate_row

        valid = {'state': 'context', 'question': 'Choose', 'options': ['A', 'B']}
        for row in (None, [], {**valid, 'teacher_logits': 'bad'},
                    {**valid, 'teacher_logits': [0.0, float('nan')]},
                    {**valid, 'teacher_logits': [True, 0.0]}):
            with self.subTest(row=row), self.assertRaises(ValueError):
                validate_row(row, 1)


class InferenceTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        torch.set_num_threads(2)
        cls.temp = tempfile.TemporaryDirectory()
        cls.path = make_checkpoint(cls.temp.name, tiny=True)
        cls.baseline = Engine(cls.path, device='cpu')
        with patch('julia.router.native.BendReducer', side_effect=AssertionError('Default CPU must not load Bend')):
            cls.fast = FastEngine(cls.path, device='cpu', batch_size=2)
        cls.rows = sample_rows(5)
        cls.rows[0]['state'] = {'text': 'Olá 世界 [MASK]', 'items': [1, 2]}

    @classmethod
    def tearDownClass(cls):
        del cls.baseline, cls.fast
        cls.temp.cleanup()

    def test_modernbert_decision_encoder_deduplicates_rope(self):
        from transformers import ModernBertConfig, ModernBertModel
        from julia.model import JuliaDecisionModel
        from julia.router.encoder import specialize_decision_encoder
        config = ModernBertConfig(vocab_size=64, pad_token_id=0, bos_token_id=1, eos_token_id=2, hidden_size=32, intermediate_size=64,
            num_hidden_layers=4, num_attention_heads=4, max_position_embeddings=128,
            local_attention=16, global_attn_every_n_layers=3, attention_dropout=0.,
            embedding_dropout=0., mlp_dropout=0.)
        config._attn_implementation = 'sdpa'
        model = JuliaDecisionModel(ModernBertModel(config), head_layers=0).eval()
        ids = torch.arange(24).reshape(2,12) % 64
        mask = torch.ones_like(ids); mask[1,8:] = 0
        calls=[]
        hook=model.encoder.rotary_emb.register_forward_hook(lambda *args: calls.append(1))
        with torch.inference_mode():
            expected=model.encoder(input_ids=ids,attention_mask=mask).last_hidden_state
            self.assertEqual(len(calls), 4)
            calls.clear()
            self.assertTrue(specialize_decision_encoder(model))
            actual=model.encoder(input_ids=ids,attention_mask=mask).last_hidden_state
            self.assertEqual(len(calls), 2)
            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
        hook.remove()

    def test_default_backend_and_padding_budget(self):
        self.assertEqual(self.fast.transformer_backend, 'torch')
        self.assertIsNone(self.fast.bend)
        self.assertFalse(hasattr(self.fast, 'bend_projection_count'))
        rows = [dict(self.rows[0], state='short'), dict(self.rows[1], state='long context ' * 100),
                dict(self.rows[2], state='another short state')]
        encoded = self.fast._encode(rows)
        groups = list(self.fast._batch_indices(encoded))
        self.assertEqual(sorted(i for group in groups for i in group), list(range(3)))
        for group in groups:
            lengths = [len(encoded[i]['ids']) for i in group]
            self.assertLessEqual(max(lengths) * len(group), self.fast.padding_ratio * sum(lengths))
        actual = self.fast.predict(rows)
        expected = self.baseline.predict(rows)
        for a, b in zip(actual, expected):
            np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)

    def test_lfs_pointer_error(self):
        from pathlib import Path
        with tempfile.TemporaryDirectory() as directory:
            (Path(directory) / 'model.safetensors').write_text('version https://git-lfs.github.com/spec/v1\n')
            with self.assertRaisesRegex(ValueError, 'Git LFS pointers'):
                FastEngine(directory)

    def test_serialization_and_pack(self):
        encoded = self.fast._encode(self.rows)
        self.assertEqual(encoded, [sequence(self.baseline.tokenizer, row) for row in self.rows])
        packed = self.fast._pack(encoded)
        reference = Collator(self.baseline.tokenizer, 1024, 256)(self.rows)
        for name in packed:
            self.assertTrue(torch.equal(packed[name], reference[name]), name)
        self.assertEqual(self.fast.predict([]), [])
        self.assertEqual(self.fast.logits([]), [])

    def test_inference_collation_omits_training_tensors(self):
        rows = [dict(row, target=0, teacher_logits=[0.] * len(row['options']))
                for row in self.rows]
        collate = self.baseline.collate
        training = collate(rows)
        inference = collate(rows, include_targets=False)
        self.assertIn('labels', training)
        self.assertIn('teacher_logits', training)
        self.assertEqual(set(training) - set(inference), {'labels', 'teacher_logits'})
        for name, value in inference.items():
            self.assertTrue(torch.equal(value, training[name]), name)
        self.assertEqual(self.baseline.predict(rows), self.baseline.predict(self.rows))

    def test_nonfinite_real_scores_are_rejected(self):
        original = self.fast.forward
        try:
            for invalid in (float('nan'), float('inf'), -float('inf')):
                def forward(**batch):
                    result = torch.zeros_like(batch['marker_pos'], dtype=torch.float32)
                    result[0, 0] = invalid
                    return result
                self.fast.forward = forward
                with self.assertRaises(FloatingPointError):
                    self.fast.predict(self.rows)
        finally:
            self.fast.forward = original

    def test_memory_mapped_checkpoint_dtype_and_logits(self):
        from julia.model import JuliaDecisionModel
        for dtype in (torch.float32, torch.bfloat16):
            with tempfile.TemporaryDirectory() as directory:
                model = JuliaDecisionModel.from_pretrained(self.path).to(dtype=dtype)
                model.save_pretrained(directory)
                loaded = JuliaDecisionModel.from_pretrained(directory, memory_map=True)
                copied = JuliaDecisionModel.from_pretrained(directory, memory_map=False)
                for name, value in loaded.state_dict().items():
                    self.assertEqual(value.dtype, copied.state_dict()[name].dtype)
                    torch.testing.assert_close(value, copied.state_dict()[name], rtol=0, atol=0)
                if dtype == torch.float32:
                    batch = self.baseline.collate(self.rows)
                    with torch.inference_mode():
                        torch.testing.assert_close(loaded.eval()(**batch), copied.eval()(**batch), rtol=0, atol=0)

    def test_selected_head_scores_actions_and_training(self):
        model = self.baseline.model
        batch = self.baseline.collate(self.rows)
        batch['qtype'] = torch.arange(len(self.rows)) % 3
        with torch.inference_mode():
            model.marker_only_head = False
            expected = model(**batch, return_actions=True)
            model.marker_only_head = True
            actual = model(**batch, return_actions=True)
        for a, b in zip(expected, actual):
            torch.testing.assert_close(a, b, atol=2e-6, rtol=2e-5)
        # Training must retain the original full-sequence/dropout path.
        model.train()
        torch.manual_seed(123)
        model.marker_only_head = False
        expected = model(**batch)
        torch.manual_seed(123)
        model.marker_only_head = True
        actual = model(**batch)
        torch.testing.assert_close(expected, actual, atol=0, rtol=0)
        actual.sum().backward()
        self.assertTrue(torch.isfinite(model.scorer[-1].weight.grad).all())
        model.zero_grad(set_to_none=True)
        model.eval()
        model.marker_only_head = False

    def test_strict_encoding_reuses_audit_and_rejects_loss(self):
        engine = FastEngine(self.path, device='cpu', strict_encoding=True)
        row = dict(self.rows[1], state='word1 word2')
        audit = engine.encoding_info([row])[0]
        self.assertFalse(audit['stateTruncated'])
        encoded = next(iter(engine._encoded.values()))
        engine.predict([row])
        self.assertIs(next(iter(engine._encoded.values())), encoded)
        for bad in [dict(row, state='[MASK]'),
                    dict(row, state={'value': '[MASK]'}),
                    dict(row, options=['word1 ' * 49, 'word2']),
                    dict(row, state='word1 ' * 2000)]:
            with self.assertRaises(ValueError):
                engine.predict([bad])
        # Changes in the encoding budget must not reuse stale cached IDs.
        engine.max_length = 32
        with self.assertRaises(ValueError):
            engine.predict([dict(row, state='word1 ' * 100)])
        engine.head_length = 16
        with self.assertRaises(ValueError):
            engine.predict([dict(row, question='word1 ' * 40)])

    def test_inference_parity_and_order(self):
        reference = self.baseline.predict(self.rows)
        actual = self.fast.predict(self.rows)
        self.assertEqual([x['index'] for x in reference], [x['index'] for x in actual])
        for a, b in zip(reference, actual):
            np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
        self.assertEqual(self.fast.predict(self.rows, probabilities=False),
                         [dict(index=x['index']) for x in actual])
        self.assertTrue(self.fast._encoded)
        self.fast.clear_cache()
        self.assertFalse(self.fast._encoded)
        self.assertFalse(self.fast._tokens.cache)

    def test_bend_transformer_parity(self):
        bend = FastEngine(self.path, device='cpu', batch_size=2,
                          transformer_backend='bend', bend_postprocess=True)
        reference = self.baseline.predict(self.rows)
        actual = bend.predict(self.rows)
        # Different reduction orders can break an exact FP32 tie. Require the
        # selected reference probability to be within tolerance of its maximum.
        for a, b in zip(reference, actual):
            p = a['probabilities']
            self.assertLessEqual(max(p) - p[b['index']], 2e-6)
        for a, b in zip(reference, actual):
            np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
        self.assertEqual(bend.predict(self.rows, probabilities=False),
                         [dict(index=x['index']) for x in actual])

    def test_bend_dense_encoder_parity(self):
        dense = FastEngine(self.path, device='cpu', batch_size=2,
                           transformer_backend='bend-dense')
        self.assertGreater(dense.bend_projection_count, 0)
        expected = self.baseline.predict(self.rows)
        actual = dense.predict(self.rows)
        for a, b in zip(expected, actual):
            np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=2e-6)
            self.assertLessEqual(max(a['probabilities']) - a['probabilities'][b['index']], 2e-6)
        from julia.router.transformer import BendLinear
        layer = next(m for m in dense.model.modules() if isinstance(m, BendLinear))
        with torch.no_grad():
            layer.weight.add_(1)
        with self.assertRaisesRegex(RuntimeError, 'resident weights changed'):
            dense.predict(self.rows)

    @unittest.skipUnless(torch.cuda.is_available(), 'CUDA hardware unavailable')
    def test_cuda(self):
        baseline = Engine(self.path, device='cuda')
        fast = FastEngine(self.path, device='cuda', batch_size=2)
        reference, actual = baseline.predict(self.rows), fast.predict(self.rows)
        self.assertEqual([x['index'] for x in reference], [x['index'] for x in actual])
        for a, b in zip(reference, actual):
            np.testing.assert_allclose(a['probabilities'], b['probabilities'], atol=5e-3)


if __name__ == '__main__':
    unittest.main()