| |
| |
|
|
| """Eager host contract tests for the TTTv2-native seed lifecycle.""" |
|
|
| import ast |
| import random |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import torch |
|
|
| from models.common.modules.lazy_buffer import LazyBuffer |
| from models.common.modules.sampling.seed_manager_1d import SeedManager1D, SeedState, _hash_request_seed_to_device_seed |
|
|
|
|
| class _EntropySequence: |
| def __init__(self, start=1000): |
| self._next = start |
|
|
| def __call__(self, _bits): |
| value = self._next |
| self._next += 1 |
| return value |
|
|
|
|
| class _TrackingSeedBuffer(LazyBuffer): |
| """LazyBuffer test double that models a materialized stable handle.""" |
|
|
| def __init__(self, defaults): |
| super().__init__(source=defaults) |
| self.handle = object() |
| self.materialize_calls = 0 |
| self.updates = [] |
|
|
| def get_device_buffer(self): |
| self.materialize_calls += 1 |
| return self.handle |
|
|
| def update(self, new_source): |
| self.source = new_source |
| self.updates.append(new_source.detach().clone()) |
|
|
|
|
| def _make_manager(capacity=4, *, salt_duplicate_seeds=True): |
| defaults = torch.arange(capacity, dtype=torch.int32) |
| buffer = _TrackingSeedBuffer(defaults) |
| config = SimpleNamespace(max_batch_size=capacity, seeds=buffer) |
| manager = SeedManager1D( |
| config, |
| entropy_factory=_EntropySequence(), |
| salt_duplicate_seeds=salt_duplicate_seeds, |
| ) |
| return manager, manager.create_state(), buffer, defaults |
|
|
|
|
| def test_seed_state_requires_consistent_caller_owned_slot_storage(expect_error): |
| with expect_error(ValueError, "same capacity"): |
| SeedState( |
| active=[False], |
| request_seeds=[], |
| token_counters=[0], |
| salts=[0], |
| unseeded_rngs=[random.Random(1)], |
| last_absolute_positions=[None], |
| current_device_seeds=[None], |
| ) |
|
|
|
|
| def test_seed_manager_requires_sampling_config_lazy_buffer_and_matching_capacity(expect_error): |
| with expect_error(TypeError, "mutable LazyBuffer-compatible"): |
| SeedManager1D(SimpleNamespace(max_batch_size=4, seeds=torch.arange(4))) |
|
|
| buffer = _TrackingSeedBuffer(torch.arange(3, dtype=torch.int32)) |
| with expect_error(ValueError, "does not match"): |
| SeedManager1D(SimpleNamespace(max_batch_size=4, seeds=buffer)) |
|
|
|
|
| def test_prefill_admission_resets_counters_and_assigns_simultaneous_equal_seed_salts(): |
| manager, state, _, _ = _make_manager() |
|
|
| manager.admit(state, [42, 42, None], [0, 2, 3]) |
| snapshot = state.snapshot() |
|
|
| assert snapshot.active_slots == (0, 2, 3) |
| assert snapshot.request_seeds == (42, None, 42, None) |
| assert snapshot.token_counters == (0, 0, 0, 0) |
| assert snapshot.salts == (0, 0, 1, 0) |
|
|
| manager.refresh(state, [0, 2], positions={0: 8, 2: 8}) |
| manager.admit(state, [42], [2]) |
| assert state.snapshot().token_counters[2] == 0 |
| assert state.snapshot().salts[2] == 1 |
|
|
|
|
| def test_absolute_position_refresh_is_idempotent_and_sequential_refresh_resumes_after_it(): |
| manager, state, buffer, defaults = _make_manager() |
| manager.admit(state, [707], [1]) |
|
|
| first = manager.refresh(state, [1], positions={1: 13}) |
| assert first[1] == _hash_request_seed_to_device_seed(707, 14) |
| assert state.snapshot().token_counters[1] == 15 |
|
|
| repeated = manager.refresh(state, [1], positions={1: 13}) |
| assert repeated == first |
| assert state.snapshot().token_counters[1] == 15 |
|
|
| next_position = manager.refresh(state, [1], positions={1: 14}) |
| assert next_position[1] == _hash_request_seed_to_device_seed(707, 15) |
| assert state.snapshot().token_counters[1] == 16 |
|
|
| sequential = manager.refresh(state, [1]) |
| assert sequential[1] == _hash_request_seed_to_device_seed(707, 16) |
| assert state.snapshot().token_counters[1] == 17 |
| assert torch.equal(buffer.source, defaults) |
|
|
|
|
| def test_equal_request_seeds_produce_distinct_salted_streams(): |
| manager, state, _, _ = _make_manager() |
| manager.admit(state, [1234, 1234, 1234, 1234], [0, 1, 2, 3]) |
|
|
| values = manager.refresh(state, [0, 1, 2, 3], positions=[7, 7, 7, 7]) |
|
|
| assert state.snapshot().salts == (0, 1, 2, 3) |
| assert len(set(values)) == 4 |
|
|
|
|
| def test_equal_request_seeds_share_one_stream_when_salting_is_disabled(): |
| manager, state, _, _ = _make_manager(capacity=32, salt_duplicate_seeds=False) |
| seeds = [1234] * 32 |
| slots = list(range(32)) |
| manager.admit(state, seeds, slots) |
|
|
| values = manager.refresh(state, slots, positions=[7] * 32) |
|
|
| assert state.snapshot().salts == (0,) * 32 |
| assert len(set(values)) == 1 |
| assert values[0] == _hash_request_seed_to_device_seed(1234, 8, 0) |
|
|
|
|
| def test_unseeded_rng_state_is_diverse_and_same_position_trace_refresh_does_not_advance_twice(): |
| manager, state, _, _ = _make_manager() |
| manager.admit(state, [None, None], [0, 1]) |
|
|
| first = manager.refresh(state, [0, 1], positions=[20, 20]) |
| after_first = state.snapshot() |
| repeated = manager.refresh(state, [0, 1], positions=[20, 20]) |
| after_repeated = state.snapshot() |
|
|
| assert first == repeated |
| assert first[0] != first[1] |
| assert after_first.token_counters[:2] == (1, 1) |
| assert after_repeated.token_counters == after_first.token_counters |
| assert after_repeated.unseeded_rng_states == after_first.unseeded_rng_states |
|
|
| second = manager.refresh(state, [0, 1], positions=[21, 21]) |
| assert second[0] != first[0] |
| assert second[1] != first[1] |
| assert state.snapshot().token_counters[:2] == (2, 2) |
|
|
|
|
| def test_synchronize_requires_reset_for_admission_and_preserves_survivor_state(expect_error): |
| manager, state, _, _ = _make_manager() |
| slot_seeds = [42, None, None, None] |
|
|
| with expect_error(RuntimeError, "reset_batch=True"): |
| manager.synchronize(state, slot_seeds, [0], reset_batch=False) |
|
|
| manager.synchronize(state, slot_seeds, [0], reset_batch=True) |
| manager.refresh(state, [0], positions=[9, -1, -1, -1]) |
| survivor = state.snapshot() |
|
|
| manager.synchronize(state, [42, 42, None, None], [0, 1], reset_batch=True) |
| updated = state.snapshot() |
| assert updated.token_counters[0] == survivor.token_counters[0] |
| assert updated.salts[0] == survivor.salts[0] == 0 |
| assert updated.current_device_seeds[0] == survivor.current_device_seeds[0] |
| assert updated.salts[1] == 1 |
|
|
| with expect_error(RuntimeError, "reset_batch=True"): |
| manager.synchronize(state, [77, 42, None, None], [0, 1], reset_batch=False) |
| assert state.snapshot() == updated |
|
|
|
|
| def test_slot_remap_moves_counter_salt_rng_and_current_buffer_value_then_vacates_source(): |
| manager, state, buffer, defaults = _make_manager() |
| manager.admit(state, [55, 55], [0, 3]) |
| manager.refresh(state, [0, 3], positions={0: 8, 3: 8}) |
| source = state.snapshot() |
|
|
| manager.apply_slot_remap(state, torch.tensor([0, 3, 2, 3], dtype=torch.int32)) |
| moved = state.snapshot() |
|
|
| assert moved.request_seeds[1] == source.request_seeds[3] == 55 |
| assert moved.token_counters[1] == source.token_counters[3] |
| assert moved.salts[1] == source.salts[3] == 1 |
| assert moved.unseeded_rng_states[1] == source.unseeded_rng_states[3] |
| assert moved.current_device_seeds[1] == source.current_device_seeds[3] |
| assert moved.active[3] is False |
| assert buffer.updates[-1][1].item() == source.current_device_seeds[3] |
| assert buffer.updates[-1][3].item() == defaults[3].item() |
|
|
|
|
| def test_resume_rejects_reused_salt_when_salting_is_enabled(expect_error): |
| manager, state, _, _ = _make_manager() |
| manager.admit(state, [42, 42], [0, 1]) |
| checkpoint = manager.suspend(state, 1) |
| manager.admit(state, [42], [1]) |
|
|
| with expect_error(RuntimeError, "salt was reused"): |
| manager.resume(state, 2, checkpoint) |
|
|
|
|
| def test_resume_allows_equal_seed_neighbors_when_salting_is_disabled(): |
| manager, state, _, _ = _make_manager(salt_duplicate_seeds=False) |
| manager.admit(state, [42, 42], [0, 1]) |
| checkpoint = manager.suspend(state, 1) |
|
|
| manager.resume(state, 1, checkpoint) |
|
|
| assert state.snapshot().request_seeds[:2] == (42, 42) |
| assert state.snapshot().salts == (0, 0, 0, 0) |
|
|
|
|
| def test_suspend_resume_preserves_unseeded_rng_stream_across_slot_movement(): |
| manager, state, _, _ = _make_manager() |
| manager.admit(state, [None], [3]) |
| manager.refresh(state, [3], positions={3: 10}) |
|
|
| checkpoint = manager.suspend(state, 3) |
| expected_rng = random.Random() |
| expected_rng.setstate(checkpoint.unseeded_rng_state) |
| expected_next = expected_rng.randint(1, 1_000_000) |
|
|
| manager.resume(state, 1, checkpoint) |
| resumed = manager.refresh(state, [1], positions={1: 11}) |
|
|
| assert resumed[1] == expected_next |
| assert state.snapshot().token_counters[1] == checkpoint.token_counter + 1 |
| assert state.snapshot().active_slots == (1,) |
|
|
|
|
| def test_cleanup_removes_ghost_seed_and_restores_default_buffer_rows(): |
| manager, state, buffer, defaults = _make_manager() |
| manager.admit(state, [42], [3]) |
| manager.refresh(state, [3], positions={3: 4}) |
|
|
| manager.cleanup(state, []) |
|
|
| assert state.snapshot().active_slots == () |
| assert state.snapshot().request_seeds == (None, None, None, None) |
| assert torch.equal(buffer.updates[-1], defaults) |
|
|
| manager.admit(state, [42], [1]) |
| assert state.snapshot().salts[1] == 0 |
|
|
|
|
| def test_restore_defaults_keeps_request_state_and_reset_discards_it(): |
| manager, state, buffer, defaults = _make_manager() |
| manager.admit(state, [99], [2]) |
| request_values = manager.refresh(state, [2], positions={2: 6}) |
| before_restore = state.snapshot() |
|
|
| manager.restore_defaults(state) |
|
|
| assert torch.equal(buffer.updates[-1], defaults) |
| assert state.snapshot().active_slots == (2,) |
| assert state.snapshot().token_counters == before_restore.token_counters |
| assert state.snapshot().buffer_is_default is True |
|
|
| replay_values = manager.refresh(state, [2], positions={2: 6}) |
| assert replay_values == request_values |
| assert state.snapshot().token_counters == before_restore.token_counters |
|
|
| manager.reset(state) |
| reset = state.snapshot() |
| assert reset.active_slots == () |
| assert reset.request_seeds == (None, None, None, None) |
| assert reset.token_counters == (0, 0, 0, 0) |
| assert reset.salts == (0, 0, 0, 0) |
| assert reset.current_device_seeds == (None, None, None, None) |
| assert reset.buffer_is_default is True |
| assert torch.equal(buffer.updates[-1], defaults) |
|
|
|
|
| def test_refresh_updates_one_stable_handle_and_never_promotes_request_values_to_defaults(): |
| manager, state, buffer, defaults = _make_manager() |
| handle = manager.get_seed_device_buffer() |
| manager.admit(state, [123], [0]) |
|
|
| manager.refresh(state, [0], positions=[1, -1, -1, -1]) |
| manager.refresh(state, [0], positions=[2, -1, -1, -1]) |
|
|
| assert manager.get_seed_device_buffer() is handle |
| assert buffer.materialize_calls >= 4 |
| assert len(buffer.updates) == 2 |
| assert not torch.equal(buffer.updates[0], buffer.updates[1]) |
| assert torch.equal(buffer.source, defaults) |
|
|
|
|
| def test_slot_validation_rejects_ambiguous_or_out_of_capacity_lifecycle_updates(expect_error): |
| manager, state, _, _ = _make_manager() |
|
|
| with expect_error(ValueError, "unique"): |
| manager.admit(state, [1, 2], [0, 0]) |
| with expect_error(ValueError, "outside"): |
| manager.admit(state, [1], [4]) |
| with expect_error(ValueError, "expected 2 request seeds"): |
| manager.admit(state, [1], [0, 1]) |
| with expect_error(ValueError, "must contain 4"): |
| manager.apply_slot_remap(state, [0, 1]) |
|
|
| manager.admit(state, [1], [3]) |
| with expect_error(ValueError, "do not cover active seed slot 3"): |
| manager.refresh(state, [3], positions=[0]) |
| with expect_error(ValueError, "multiple destinations"): |
| manager.apply_slot_remap(state, [3, 3, 2, 3]) |
|
|
|
|
| def test_seed_manager_1d_imports_no_legacy_sampling_state_or_generator(): |
| module_path = Path(__file__).parents[3] / "modules" / "sampling" / "seed_manager_1d.py" |
| tree = ast.parse(module_path.read_text()) |
| imported_modules = { |
| node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom) and node.module is not None |
| } |
|
|
| assert "models.common.sampling.generator" not in imported_modules |
| assert "models.common.sampling.tt_sampling" not in imported_modules |
| assert "models.common.sampling.tt_penalties" not in imported_modules |
|
|
