"""Execute notebook resume and capacity guards with temporary data and no GPU.""" # Controlled execution of our own generated notebook cells, never external code. # ruff: noqa: S102 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()