Capture the training configuration

"""Capture and verify an existing trainer's setup without executing the trainer."""
import argparse
import ast
import hashlib
import json
import platform
from pathlib import Path
import struct
import sys

HERE = Path(__file__).resolve().parent
SUPPORTED_TRAINER_SHA256 = '41b40a9ac8cc6b5b53b45f7b2f696f0fe9488b7fd583da1f20be53b293af97a9'
FILES = ('trainer/baseline.py', 'trainer/rice.csv', 'job-config.json', 'requirements-lock.txt')


def sha256(raw):
    return hashlib.sha256(raw).hexdigest()


def identity(value):
    raw = json.dumps(value, sort_keys=True, separators=(',', ':'), allow_nan=False).encode()
    return sha256(raw)


def parse_config(raw):
    def pairs(items):
        result = {}
        for key, value in items:
            if key in result:
                raise ValueError('duplicate configuration key')
            result[key] = value
        return result
    def constant(value):
        raise ValueError('nonstandard JSON constant: ' + value)
    return json.loads(raw, object_pairs_hook=pairs, parse_constant=constant)


def trainer_contract(raw):
    if sha256(raw) != SUPPORTED_TRAINER_SHA256:
        raise ValueError('unsupported trainer revision; review the setup adapter before recapturing')
    tree = ast.parse(raw)
    names = {'EPOCHS', 'LEARNING_RATE', 'SALT', 'POSITIVE', 'DATA_SHA256'}
    constants = {}
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id in names:
                    constants[target.id] = ast.literal_eval(node.value)
    function = next(node for node in tree.body if isinstance(node, ast.FunctionDef)
                    and node.name == 'train_logistic')
    # Read the actual definition-time defaults, rather than assuming later global
    # assignment would change already-bound Python function defaults.
    defaults = {argument.arg: constants[value.id]
                for argument, value in zip(function.args.args[-2:], function.args.defaults)
                if isinstance(value, ast.Name)}
    if defaults != {'learning_rate': constants['LEARNING_RATE'], 'epochs': constants['EPOCHS']}:
        raise ValueError('trainer defaults differ from the reviewed adapter')
    imports = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            imports.update(alias.name.split('.')[0] for alias in node.names)
        elif isinstance(node, ast.ImportFrom) and node.module:
            imports.add(node.module.split('.')[0])
    if not imports <= sys.stdlib_module_names:
        raise ValueError('trainer declares imports outside the standard library')
    return constants, defaults, sorted(imports)


def describe_runtime():
    return {'implementation': platform.python_implementation(),
            'python_version': platform.python_version(),
            'compiler': platform.python_compiler(),
            'system': platform.system(), 'machine': platform.machine(),
            'pointer_bits': struct.calcsize('P') * 8}


def capture(root=HERE):
    root = Path(root)
    raw = {name: (root / name).read_bytes() for name in FILES}
    constants, defaults, imports = trainer_contract(raw['trainer/baseline.py'])
    config = parse_config(raw['job-config.json'])
    expected = {'mode': config.get('mode') if isinstance(config, dict) else None,
                'epochs': defaults['epochs'], 'learning_rate': defaults['learning_rate'],
                'split_salt': constants['SALT'], 'positive_class': constants['POSITIVE'],
                'scaler_fitted_on': 'training'}
    if (not isinstance(config, dict) or set(config) != set(expected)
            or config['mode'] not in ('validation', 'final')
            or type(config['epochs']) is not int
            or type(config['learning_rate']) not in (int, float) or config != expected):
        raise ValueError('configuration does not describe the existing trainer behavior')
    if sha256(raw['trainer/rice.csv']) != constants['DATA_SHA256']:
        raise ValueError('dataset differs from the trainer data version')
    argv = ['python', 'trainer/baseline.py']
    if config['mode'] == 'final':
        argv.append('--final')
    return {'format': 1, 'mode': 'setup capture only; trainer not executed',
            'files': {name: sha256(blob) for name, blob in raw.items()},
            'capture_tool_sha256': sha256((HERE / 'capture_job.py').read_bytes()),
            'effective_config': config,
            'dependencies': {'trainer_direct_imports': imports, 'third_party': []},
            'runtime': describe_runtime(), 'planned_argv': argv}


def verify(snapshot, expected_id, root=HERE):
    if identity(snapshot) != expected_id:
        raise ValueError('snapshot differs from the separately expected identity')
    current = capture(root)
    if not isinstance(snapshot, dict) or identity(current) != expected_id:
        differences = sorted(key for key in set(current) | set(snapshot)
                             if current.get(key) != snapshot.get(key)) if isinstance(snapshot, dict) else []
        raise ValueError('setup differs: ' + (', '.join(differences) or 'canonical values'))
    return {'setup_matches': True, 'training_started': False, 'recipe_id': expected_id}


def main():
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest='command', required=True)
    save = commands.add_parser('capture')
    save.add_argument('--root', type=Path, default=HERE)
    save.add_argument('--out', type=Path, required=True)
    check = commands.add_parser('verify')
    check.add_argument('--root', type=Path, default=HERE)
    check.add_argument('--snapshot', type=Path, required=True)
    check.add_argument('--expected-id', required=True)
    args = parser.parse_args()
    if args.command == 'capture':
        snapshot = capture(args.root)
        with args.out.open('x') as file:
            file.write(json.dumps(snapshot, indent=2) + '\n')
        print(identity(snapshot))
    else:
        result = verify(parse_config(args.snapshot.read_bytes()), args.expected_id, args.root)
        print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()

"""Exercise setup capture and drift refusal; never start model training."""
import argparse
import copy
import json
from pathlib import Path
import shutil
import tempfile
from unittest.mock import patch
from capture_job import HERE, FILES, capture, verify, identity


def record():
    baseline = capture()
    matched = verify(baseline, identity(baseline))
    probes = []
    for name in ('configuration', 'dataset', 'trainer'):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            for file in FILES:
                target = root / file
                target.parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(HERE / file, target)
            if name == 'configuration':
                path = root / 'job-config.json'
                config = json.loads(path.read_text()); config['epochs'] = 201
                path.write_text(json.dumps(config) + '\n')
            else:
                path = root / ('trainer/rice.csv' if name == 'dataset' else 'trainer/baseline.py')
                path.write_bytes(path.read_bytes() + b'\n')
            try:
                verify(baseline, identity(baseline), root)
            except ValueError as error:
                probes.append({'changed': name, 'refused': True, 'reason': str(error)})
            else:
                raise AssertionError('changed setup was accepted')
    # An explicitly simulated runtime metadata difference exercises the verifier.
    # This does not claim that a second Python interpreter was installed or used.
    changed = dict(baseline['runtime'], python_version='simulated-other-version')
    with patch('capture_job.describe_runtime', return_value=changed):
        try:
            verify(baseline, identity(baseline))
        except ValueError as error:
            probes.append({'changed': 'simulated runtime metadata', 'refused': True, 'reason': str(error)})
        else:
            raise AssertionError('runtime metadata difference was accepted')
    return {'mode': 'local setup inspection; no training, metric computation or model calls',
            'snapshot': baseline, 'verified': matched, 'probes': probes}


def portable_checks(result):
    """Check portable behavior; runtime enforcement is tested by verify(), not waived."""
    stable = copy.deepcopy(result)
    stable['snapshot'].pop('runtime')
    stable['verified'].pop('recipe_id')  # Derived from the runtime as well as the files.
    return stable


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--check', action='store_true')
    parser.add_argument('--save', type=Path)
    args = parser.parse_args(); result = record()
    if args.check and portable_checks(result) != portable_checks(json.loads((HERE / 'result.json').read_text())):
        raise SystemExit('portable setup checks differ from the recorded result')
    if args.save:
        args.save.write_text(json.dumps(result, indent=2) + '\n')
    snapshot = result['snapshot']
    print('Setup capture only: training_started=False; no new metrics')
    print('Trainer imports:', ', '.join(snapshot['dependencies']['trainer_direct_imports']))
    print('Third-party packages:', len(snapshot['dependencies']['third_party']))
    config = snapshot['effective_config']
    print(f"Effective settings: epochs={config['epochs']}; learning_rate={config['learning_rate']}; mode={config['mode']}")
    print('Current captured setup verifies:', result['verified']['setup_matches'])
    for row in result['probes']:
        print(row['changed'] + ': refused=' + str(row['refused']))
    print('Runtime is recorded in result.json; --check compares portable checks only.')
    print('Use capture_job.py verify to enforce your own captured runtime and input identity.')


if __name__ == '__main__':
    main()