"""Scientific controls for CPU budgeting, paired supervision and saved state. Fixtures are tiny synthetic RV32 programs; no experiment or sealed data is read. """ from copy import deepcopy import gzip import hashlib import io import pytest import torch from lightnest import rv32 from lightnest.relation_patch import RelationConfig, RelationPatchNetwork from experiments.dependency_train import ( ARMS, DEFAULT_SELECTION, PairedBatchStream, build_predictor, json_bytes, load_checkpoint, make_initial_model, promotion_decision, public_path, relation_objectives, run_cpu_budget, save_checkpoint, tensor_inventory, training_step, ) @pytest.fixture(autouse=True) def bounded_threads(): torch.set_num_threads(2) def assemble(words): return b''.join(word.to_bytes(4, 'little') for word in words) def synthetic_rows(): rows = [] for origin in range(3): for nuisance in (0, 1): for variant in (0, 1): for kind in ('edit', 'noop'): binary = assemble([rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), rv32.add('x3', 'x1', 'x2'), rv32.add('x4', 'x1', 'x2'), rv32.addi('x0', 'x0', 10 + nuisance), rv32.addi('x10', 'x3' if variant == 0 else 'x4', 0), rv32.EBREAK]) edits = [{'offset': 11 + 4 * variant, 'value': 64}] if kind == 'edit' else [] target = bytearray(binary) for edit in edits: target[edit['offset']] = edit['value'] row = dict(id=f'g{origin}/n{nuisance}/v{variant}/{kind}', episode_id=f'g{origin}/n{nuisance}/v{variant}/{kind}', source='new', origin_id=f'g{origin}', variant=variant, nuisance_variant=nuisance, state_mask=0, pair_id=f'g{origin}/n{nuisance}/return' if kind == 'edit' else None, pair_state_key=f'g{origin}/n{nuisance}/return#0' if kind == 'edit' else None, invariance_state_key=f'g{origin}/v{variant}/{kind}#0', command='change returned add to sub' if kind == 'edit' else 'keep every byte', binary=binary.hex(), binary_bytes=binary, target=bytes(target).hex(), edits=edits) rows.append(row) for index in range(8): binary = assemble([rv32.addi('x10', 'x0', index), rv32.EBREAK]) rows.append(dict(id=f'replay{index}', episode_id=f'replay{index}', source='replay', command='keep every byte', binary=binary.hex(), binary_bytes=binary, target=binary.hex(), edits=[], pair_state_key=None)) return rows def tiny_model(): return RelationPatchNetwork(RelationConfig(width=32, binary_capacity=64, command_capacity=128)) def test_alternating_streams_preserve_every_new_row_noops_and_identical_prefix(): rows = synthetic_rows() first, control = PairedBatchStream(rows, 104729), PairedBatchStream(rows, 104729) exposure = {'return': [], 'invariance': []} for _ in range(4): indices = first.next_batch() assert indices == control.next_batch() assert first.last_units == control.last_units assert len(indices) == 16 assert all(rows[i]['source'] == 'new' for i in indices[:12]) assert all(rows[i]['source'] == 'replay' for i in indices[12:]) exposure[first.last_stream].extend(indices[:12]) for unit in first.last_units: a, b = [rows[indices[pos]] for pos in unit['positions']] if unit['kind'] == 'return': assert a['pair_state_key'] == b['pair_state_key'] == unit['key'] assert {a['variant'], b['variant']} == {0, 1} elif unit['kind'] == 'invariance': assert a['invariance_state_key'] == b['invariance_state_key'] == unit['key'] assert {a['nuisance_variant'], b['nuisance_variant']} == {0, 1} else: assert not a['edits'] and not b['edits'] expected = [i for i, row in enumerate(rows) if row['source'] == 'new'] for stream in exposure.values(): assert sorted(stream) == expected # Both streams include ALL original no-ops. a, b = first.state_dict(), control.state_dict() assert a == b def test_sampler_resume_reconstructs_next_batches_and_explicit_units(): rows = synthetic_rows() stream = PairedBatchStream(rows, 29) for _ in range(3): stream.next_batch() state = deepcopy(stream.state_dict()) expected = [] for _ in range(9): expected.append((stream.next_batch(), deepcopy(stream.last_units), stream.last_stream)) restored = PairedBatchStream(rows, 29) restored.load_state_dict(state) actual = [] for _ in range(9): actual.append((restored.next_batch(), deepcopy(restored.last_units), restored.last_stream)) assert actual == expected with pytest.raises(ValueError, match='different stream'): PairedBatchStream(rows, 30).load_state_dict(state) def test_pair_metadata_is_validated_not_guessed_from_names(): rows = synthetic_rows() rows[0].pop('invariance_state_key') with pytest.raises(ValueError, match='invariance_state_key'): PairedBatchStream(rows, 1) rows = synthetic_rows() rows[0]['state_mask'] = 7 with pytest.raises(ValueError, match='return-choice'): PairedBatchStream(rows, 1) def test_return_margin_uses_exclusive_sets_two_directions_and_skips_stops(): scores = torch.zeros(4, 9, requires_grad=True) with torch.no_grad(): scores[0, 1] = scores[1, 3] = 4. scores[:2, 2] = 90. # Shared acceptable position must not enter either exclusive mass. group = [dict(binary_bytes=b'12345678', pair_state_key='same', edits=edits) for edits in ( [{'offset': 1}, {'offset': 2}], [{'offset': 2}, {'offset': 3}], [], [])] units = [dict(kind='return', positions=[0, 1], key='same'), dict(kind='return', positions=[2, 3], key='same')] margin, js, counts = relation_objectives(scores, group, units, 8) torch.testing.assert_close(margin, torch.nn.functional.softplus(torch.tensor(-3.))) assert js.item() == 0. assert counts['return_candidate_pairs'] == 2 assert counts['return_effective_pairs'] == 1 assert counts['return_effective_directions'] == 2 margin.backward() assert scores.grad[0, 1] < 0 and scores.grad[0, 3] > 0 assert scores.grad[1, 3] < 0 and scores.grad[1, 1] > 0 assert torch.count_nonzero(scores.grad[:, 2]) == 0 assert torch.count_nonzero(scores.grad[2:]) == 0 def test_js_includes_stop_excludes_padding_and_is_symmetric(): group = [dict(binary_bytes=b'1234', invariance_state_key='inv', edits=[]) for _ in range(2)] unit = [dict(kind='invariance', positions=[0, 1], key='inv')] scores = torch.zeros(2, 9, requires_grad=True) with torch.no_grad(): scores[0, 4:8] = 1000. scores[1, 4:8] = -1000. _, js, counts = relation_objectives(scores, group, unit, 8) torch.testing.assert_close(js, torch.tensor(0.), atol=1e-7, rtol=0.) assert counts['invariance_effective_pairs'] == 1 changed = scores.detach().clone().requires_grad_(True) with torch.no_grad(): changed[0, 8] = 4. _, unequal, _ = relation_objectives(changed, group, unit, 8) _, reversed_js, _ = relation_objectives(changed.flip(0), group, unit, 8) assert unequal.item() > 0. torch.testing.assert_close(unequal, reversed_js) unequal.backward() assert changed.grad[0, 8] != 0 assert torch.count_nonzero(changed.grad[:, 4:8]) == 0 def test_cpu_budget_excludes_evaluation_and_logs_without_equating_updates(): clock = {'cpu': 0., 'wall': 0.} boundaries = [] def step(): clock['cpu'] += 3. clock['wall'] += 2. return {'loss': 1.} def log(_): clock['cpu'] += 11. clock['wall'] += 11. def checkpoint(boundary): boundaries.append(boundary) clock['cpu'] += 100. clock['wall'] += 500. return {'stage': boundary['stage']} result = run_cpu_budget(step, checkpoint, cpu_seconds=10., on_step=log, cpu_clock=lambda: clock['cpu'], wall_clock=lambda: clock['wall']) assert result['actual_updates'] == 4 assert result['training_cpu_seconds'] == 12. assert result['training_wall_seconds'] == 8. assert result['examples_seen'] == 64 assert result['budget_overshoot_seconds'] == 2. assert [row['fraction'] for row in boundaries] == [.25, .5, .75, 1.] assert [row['step'] for row in boundaries] == [1, 2, 3, 4] assert all(row['cpu_overshoot_seconds'] < row['crossing_step_cpu_seconds'] for row in boundaries) def metric_fixture(): return dict(edit_total=64, edit_strict=20, pair_total=32, pair_both_strict=5, noop_total=64, noop_strict=60, false_edit_on_correct=2, invariance_edit_total=32, invariance_edit_both_strict=8, retention_total=128, retention_strict=100, development_loss=1.) def test_promotion_requires_task_gain_and_keeps_original_retention_anchor(): original = metric_fixture() better = dict(original, edit_strict=21, development_loss=2., retention_strict=98) assert promotion_decision(better, original, original, DEFAULT_SELECTION)['accept'] # A lower loss or pure invariance gain cannot substitute for edit/return-pair gain. for candidate in (dict(original, development_loss=.01), dict(original, invariance_edit_both_strict=9), dict(better, invariance_edit_both_strict=7), dict(better, pair_both_strict=4), dict(better, false_edit_on_correct=3)): assert not promotion_decision(candidate, original, original, DEFAULT_SELECTION)['accept'] incumbent = dict(original, retention_strict=99) candidate = dict(better, retention_strict=97) decision = promotion_decision(candidate, incumbent, original, DEFAULT_SELECTION) assert not decision['accept'] assert decision['retention_floor'] == 98 # Not incumbent 99 minus two. with pytest.raises(ValueError, match='denominators'): promotion_decision(dict(better, edit_total=63), original, original, DEFAULT_SELECTION) def test_gzip_checkpoint_replays_optimizer_sampler_and_exact_next_update(tmp_path): rows = synthetic_rows() torch.manual_seed(137) model = tiny_model() sampler = PairedBatchStream(rows, 104729) optimizer = torch.optim.AdamW(model.parameters(), lr=.001, weight_decay=.01) first = training_step(model, 'relation', rows, sampler, optimizer) assert first['loss'] == pytest.approx(first['base_loss'] + .2 * first['return_margin_loss'] + .1 * first['invariance_js_loss'], rel=1e-6) path = tmp_path / 'checkpoint.pt.gz' saved = save_checkpoint(path, model, optimizer, sampler, kind='relation', metadata={'role': 'synthetic-test'}) assert path.read_bytes()[4:8] == b'\0\0\0\0' # Reproducible gzip header, no timestamp. assert saved['model_sha256'] == tensor_inventory(model.state_dict())['sha256'] restored, payload = load_checkpoint(path) assert tensor_inventory(restored.state_dict()) == tensor_inventory(model.state_dict()) restored_optimizer = torch.optim.AdamW(restored.parameters(), lr=.001, weight_decay=.01) restored_optimizer.load_state_dict(payload['optimizer']) restored_sampler = PairedBatchStream(rows, 104729) restored_sampler.load_state_dict(payload['sampler']) expected = training_step(model, 'relation', rows, sampler, optimizer) torch.set_rng_state(payload['torch_rng_state']) actual = training_step(restored, 'relation', rows, restored_sampler, restored_optimizer) assert expected == actual for name, value in model.state_dict().items(): torch.testing.assert_close(value, restored.state_dict()[name], atol=0., rtol=0.) for first_state, second_state in zip(optimizer.state.values(), restored_optimizer.state.values()): for key in first_state: torch.testing.assert_close(first_state[key], second_state[key], atol=0., rtol=0.) def test_initial_relation_import_and_portable_predictor_are_explicit(tmp_path): model = tiny_model() parent = tmp_path / 'parent.pt' model.save(parent) for name, depth in (('relation-d4', 4), ('relation-d8', 8)): child, report = make_initial_model(ARMS[name], parent, 104729) assert child.config.message_steps == depth assert report['new_numel'] == 0 assert report['function_preserving'] == (depth == 4) assert report['parameter_name_map'] == {key: key for key in model.state_dict()} assert tensor_inventory(child.state_dict()) == tensor_inventory(model.state_dict()) path = tmp_path / 'portable.pt.gz' saved = save_checkpoint(path, model, None, None, kind='relation', metadata={}) predictor = build_predictor(dict(id='tiny', kind='relation', path=str(path), sha256=saved['sha256'], message_steps=8)) assert predictor.message_steps == 8 assert tensor_inventory(predictor.model.state_dict()) == tensor_inventory(model.state_dict()) with pytest.raises(ValueError, match='digest'): build_predictor(dict(kind='relation', path=str(path), sha256='0' * 64)) def test_sealed_paths_are_rejected_before_opening(): with pytest.raises(ValueError, match='prohibited'): public_path('work/unopened-round/sealed/final.json') def test_checkpoint_internal_tensor_commitment_rejects_tampering(tmp_path): path = tmp_path / 'tamper.pt.gz' save_checkpoint(path, tiny_model(), None, None, kind='relation', metadata={}) payload = torch.load(io.BytesIO(gzip.decompress(path.read_bytes())), weights_only=True) payload['state_dict']['pointer.bias'][0] += 1. buffer = io.BytesIO() torch.save(payload, buffer) path.write_bytes(gzip.compress(buffer.getvalue(), mtime=0)) with pytest.raises(ValueError, match='digest'): load_checkpoint(path)