| """Execute notebook resume and capacity guards with temporary data and no GPU.""" |
| |
| |
| import hashlib |
| import importlib.util |
| import json |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| from pathlib import Path |
| from types import SimpleNamespace |
| from unittest import mock |
|
|
| import pytest |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / 'scripts')) |
| from generate import directory_fingerprint |
|
|
| from vibegame import BASE_MODEL, SYSTEM_PROMPT |
|
|
| spec = importlib.util.spec_from_file_location('notebook_builder', ROOT / 'scripts/build_notebook.py') |
| builder = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(builder) |
| CELLS = [cell['source'] for cell in builder.build()['cells'] if cell['cell_type'] == 'code'] |
| DATA_CELL = next(cell for cell in CELLS if 'def verify_run_dataset' in cell) |
| TRAIN_CELL = next(cell for cell in CELLS if 'def verify_completed_training' in cell) |
| CAPACITY_CELL = next(cell for cell in CELLS if 'def capacity_report_matches' in cell) |
|
|
|
|
| def write_json(path, value): |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(value)) |
|
|
|
|
| def fixture(root): |
| data, run, project = [root / part for part in ('data', 'run', 'project')] |
| for path in (data, run, project / 'data'): |
| path.mkdir(parents=True) |
| signature = {'base_model': BASE_MODEL, 'base_revision': 'pinned', |
| 'system_prompt_sha256': hashlib.sha256(SYSTEM_PROMPT.encode()).hexdigest()} |
| for field, name in [('train_sha256', 'train.jsonl'), ('dev_sha256', 'dev.jsonl'), |
| ('blind_sha256', 'eval_prompts.jsonl'), ('validation_manifest_sha256', 'validation_manifest.jsonl')]: |
| (data / name).write_text(name + '\n') |
| signature[field] = hashlib.sha256((data / name).read_bytes()).hexdigest() |
| manifest = {'signature': signature} |
| write_json(run / 'run_manifest.json', manifest) |
| adapter = run / 'adapter' |
| write_json(adapter / 'run_manifest.json', manifest) |
| complete = {'status': 'complete', 'global_step': 14, 'target_steps': 14} |
| write_json(adapter / 'TRAINING_COMPLETE.json', complete) |
| weight = b'unchanged-weights' |
| (adapter / 'adapter_model.safetensors').write_bytes(weight) |
| checkpoint = run / 'checkpoint-14' |
| checkpoint.mkdir() |
| (checkpoint / 'adapter_model.safetensors').write_bytes(weight) |
| write_json(checkpoint / 'COMPLETE.json', {'global_step': 14, 'sha256': { |
| 'adapter_model.safetensors': hashlib.sha256(weight).hexdigest()}}) |
| rows = [{'id': i, 'instruction': f'Training prompt {i}'} for i in range(4)] |
| (project / 'data/teacher_prompts.jsonl').write_text(''.join(json.dumps(row) + '\n' for row in rows)) |
| return {'DATA': data, 'RUN': run, 'PROJECT': project, 'DRIVE': root, |
| 'EVAL': run / 'evaluation', 'ADAPTER': adapter, |
| 'Path': Path, 'hashlib': hashlib, 'json': json, 'sys': sys, |
| 'subprocess': subprocess, 'run': mock.Mock(), |
| 'torch': SimpleNamespace(cuda=SimpleNamespace(get_device_name=lambda _: 'Tesla T4'))} |
|
|
|
|
| def test_existing_run_reuses_dataset_without_touching_admission(): |
| with tempfile.TemporaryDirectory() as directory: |
| env = fixture(Path(directory)) |
| before = {p.name: p.read_bytes() for p in env['DATA'].iterdir()} |
| exec(DATA_CELL, env) |
| env['run'].assert_not_called() |
| assert before == {p.name: p.read_bytes() for p in env['DATA'].iterdir()} |
| (env['DATA'] / 'validation_manifest.jsonl').write_text('changed evidence') |
| with pytest.raises(RuntimeError, match='validation_manifest.jsonl'): |
| exec(DATA_CELL, env) |
| env['run'].assert_not_called() |
|
|
|
|
| def test_completed_adapter_skips_training_and_rejects_false_completion(): |
| with tempfile.TemporaryDirectory() as directory: |
| env = fixture(Path(directory)) |
| exec(DATA_CELL, env) |
| exec(TRAIN_CELL, env) |
| env['run'].assert_not_called() |
| write_json(env['ADAPTER'] / 'TRAINING_COMPLETE.json', {'status': 'complete', 'global_step': 13, 'target_steps': 14}) |
| with pytest.raises(RuntimeError, match='incomplete'): |
| exec(TRAIN_CELL, env) |
| env['run'].assert_not_called() |
|
|
|
|
| def test_fresh_default_excludes_existing_optional_teacher_pilot(): |
| with tempfile.TemporaryDirectory() as directory: |
| env = fixture(Path(directory)) |
| (env['RUN'] / 'run_manifest.json').unlink() |
| authored = env['DATA'] / 'codex_teacher_batch_a.jsonl' |
| authored.write_text('{}\n') |
| pilot = env['DRIVE'] / 'teacher-v2/validated_candidates.jsonl' |
| pilot.parent.mkdir() |
| pilot.write_text('{}\n') |
| calls = [] |
| def admit(*args): |
| calls.append(args) |
| write_json(env['DATA'] / 'dataset_manifest.json', {'train': 1, 'dev': 1}) |
| (env['DATA'] / 'validation_manifest.jsonl').write_text(json.dumps({'accepted': True, 'source': str(authored)}) + '\n') |
| env.update({'run': admit, 'shutil': shutil, 'DATASET_REPO': '', 'LOCAL_TEACHER_ATTEMPTS': 0, 'GODOT': '/never-launched'}) |
| exec(DATA_CELL, env) |
| assert len(calls) == 1 |
| assert calls[0][1] == 'scripts/build_dataset.py' |
| assert authored in calls[0] |
| assert pilot not in calls[0] |
|
|
|
|
| def capacity_report(env, size, success): |
| prompts = [f'Training prompt {i}' for i in range(4)] |
| row = {'batch_size': size, 'status': 'success' if success else 'failed', |
| 'peak_allocated_bytes': 8_000_000_000, 'peak_reserved_bytes': 9_000_000_000, |
| 'full_context_memory_estimate': {'within_limit': success, 'assessed_batch_size': size, |
| 'target_context_tokens': 4096, 'cache_growth_allowance_tokens': 512, |
| 'reserve_bytes': 2_000_000_000, 'estimated_peak_bytes': 10_000_000_000, |
| 'effective_limit_bytes': 12_500_000_000}} |
| signature = json.loads((env['ADAPTER'] / 'run_manifest.json').read_text())['signature'] |
| return {'status': 'complete', 'batch_sizes': [size], 'kind': 'trained_adapter_capacity', |
| 'memory_estimate_policy': 'kv_full_context_v1', 'gpu_memory_bytes': 16_000_000_000, |
| 'blind_prompts_read': False, 'full_max_new_tokens': 3072, 'memory_limit_bytes': 12_500_000_000, |
| 'training_prompts_sha256': hashlib.sha256(json.dumps(prompts, ensure_ascii=False).encode()).hexdigest(), |
| 'versions': {name: 'test' for name in ['torch', 'unsloth', 'transformers', 'bitsandbytes']}, |
| 'gpu': 'Tesla T4', 'base_model': BASE_MODEL, 'base_revision': 'pinned', |
| 'generation': {'system_prompt_sha256': signature['system_prompt_sha256'], |
| 'adapter': directory_fingerprint(env['ADAPTER'])}, |
| 'measurements': [row], 'full_check': {**row, 'max_new_tokens': 3072, 'within_memory_limit': success}, |
| 'recommended_batch_size': size if success else None} |
|
|
|
|
| def test_capacity_preserves_partial_reports_falls_back_and_reuses_results(): |
| with tempfile.TemporaryDirectory() as directory: |
| env = fixture(Path(directory)) |
| exec(DATA_CELL, env) |
| exec(TRAIN_CELL, env) |
| original = env['RUN'] / 'adapter-capacity.json' |
| write_json(original, {'status': 'full_length_check', 'batch_sizes': [32]}) |
| original_bytes = original.read_bytes() |
| calls = [] |
| def capacity_run(*args): |
| size = args[-1] |
| calls.append(size) |
| write_json(args[2], capacity_report(env, size, success=size == 16)) |
| env['run'] = capacity_run |
| with mock.patch('importlib.metadata.version', return_value='test'): |
| exec(CAPACITY_CELL, env) |
| assert calls == [32, 16] |
| assert env['EVAL_BATCH_SIZE'] == 16 |
| assert original.read_bytes() == original_bytes |
| assert (env['RUN'] / 'adapter-capacity-001.json').exists() |
| assert (env['RUN'] / 'adapter-capacity-002.json').exists() |
| exec(CAPACITY_CELL, env) |
| assert calls == [32, 16] |
| altered = capacity_report(env, 32, True) |
| altered['generation']['adapter']['sha256'] = 'changed' |
| assert not env['capacity_report_matches'](altered, 32) |
| legacy = capacity_report(env, 32, True) |
| legacy.pop('memory_estimate_policy') |
| assert not env['capacity_report_matches'](legacy, 32) |
| unsafe = capacity_report(env, 32, True) |
| unsafe['full_check']['full_context_memory_estimate']['within_limit'] = False |
| assert not env['safe_capacity_recommendation'](unsafe, 32) |
|
|
|
|
| def test_blind_resume_keeps_original_batch_without_capacity_jobs(): |
| with tempfile.TemporaryDirectory() as directory: |
| env = fixture(Path(directory)) |
| exec(DATA_CELL, env) |
| exec(TRAIN_CELL, env) |
| write_json(env['EVAL'] / 'manifest.json', {'generation': {'adapter': directory_fingerprint(env['ADAPTER'])}, |
| 'expected_count': 500, 'max_new_tokens': 3072, 'run_seconds': 60, 'smoke_test': False, |
| 'prompts_sha256': env['file_sha256'](env['DATA'] / 'eval_prompts.jsonl'), 'batch_size': 8}) |
| with mock.patch('importlib.metadata.version', return_value='test'): |
| exec(CAPACITY_CELL, env) |
| assert env['EVAL_BATCH_SIZE'] == 8 |
| env['run'].assert_not_called() |
|
|
