castle-batch-v5/batch_tick.py

#!/usr/bin/env python3
"""Run the pinned tick, then mirror collected stage JSON into the HF output mount.

This auxiliary entrypoint deliberately lives outside castle_pipeline/: copying
artifacts must not change the code identity of existing runs. GCS remains the
source of truth; the mirror is never used as a checkpoint or submission input.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import tempfile

from castle_pipeline.batch_cloud import BatchCloud, parse_gs_uri


class MirrorConflict(ValueError):
    """Safe diagnostic fields: no result data or provider error payloads."""

    def __init__(self, relative):
        super().__init__('Stage mirror already has different content: ' + relative)
        self.relative = relative


def mirror_results(cloud, gcs_prefix, output_dir):
    bucket, run = parse_gs_uri(gcs_prefix)
    prefix = run.rstrip('/') + '/results/'
    output = Path(output_dir)
    counts = {'copied': 0, 'unchanged': 0}
    for blob in cloud.storage.list_blobs(bucket, prefix=prefix):
        if not blob.name.startswith(prefix):
            continue
        relative = blob.name[len(prefix):]
        if not re.fullmatch(r'(audio|annotation|review)/[A-Za-z0-9_-]+\.json', relative):
            continue
        target = output / 'results' / relative
        expected = (blob.metadata or {}).get('sha256')
        existing = hashlib.sha256(target.read_bytes()).hexdigest() if target.is_file() else None
        if expected and existing == expected:
            counts['unchanged'] += 1
            continue
        if expected and existing is not None:
            raise MirrorConflict(relative)
        data = blob.download_as_bytes(if_generation_match=int(blob.generation))
        actual = hashlib.sha256(data).hexdigest()
        if expected and actual != expected:
            raise ValueError('GCS stage content hash mismatch: ' + relative)
        if not isinstance(json.loads(data), dict):
            raise ValueError('Stage result must be a JSON object: ' + relative)
        if existing is not None:
            if existing != actual:
                raise MirrorConflict(relative)
            counts['unchanged'] += 1
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        # Publish only complete bytes. This temp file belongs to this invocation.
        temporary = None
        try:
            with tempfile.NamedTemporaryFile(dir=target.parent, prefix='.stage-', suffix='.tmp', delete=False) as handle:
                temporary = Path(handle.name)
                handle.write(data)
            # Catch a destination populated during the download. This is not a
            # cross-mount lock: each output prefix must still have one writer.
            if target.is_file():
                if hashlib.sha256(target.read_bytes()).hexdigest() != actual:
                    raise MirrorConflict(relative)
                counts['unchanged'] += 1
                continue
            os.replace(temporary, target)
        finally:
            if temporary is not None:
                temporary.unlink(missing_ok=True)
        counts['copied'] += 1
    return counts


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=['tick'])
    parser.add_argument('--project', required=True)
    parser.add_argument('--location', default='global')
    parser.add_argument('--state-uri', required=True)
    parser.add_argument('--output-dir', type=Path, required=True)
    parser.add_argument('--scratch-dir', type=Path, required=True)
    parser.add_argument('--allow-missing-state', action='store_true')
    args = parser.parse_args(argv)
    parse_gs_uri(args.state_uri)
    cloud = BatchCloud(args.project, args.location)
    state, _ = cloud.read_state(args.state_uri)
    if state is None:
        if not args.allow_missing_state:
            raise ValueError('No batch state exists')
        print(json.dumps({'event': 'batch_tick_waiting', 'status': 'waiting_for_state',
                          'state_uri': args.state_uri, 'stop_schedule': False}), flush=True)
        return 0
    config = state['config']
    if config['project'] != args.project or config['location'] != args.location:
        raise ValueError('Project/location must match the saved run')
    if args.state_uri != config['gcs_prefix'].rstrip('/') + '/state.json':
        raise ValueError('State URI must belong to the saved GCS run prefix')
    from batch_pipeline import code_hash
    if config.get('code_hash') != code_hash():
        raise ValueError('Use the pinned code version that initialized this run')

    command = [sys.executable, str(Path(__file__).with_name('batch_pipeline.py')), 'tick',
               '--project', args.project, '--location', args.location, '--state-uri', args.state_uri,
               '--output-dir', str(args.output_dir), '--scratch-dir', str(args.scratch_dir)]
    result = subprocess.run(command, check=False)
    # Even a failed tick can have persisted valid stage results before failing.
    # Replaying this mirror never creates model requests or alters GCS state.
    try:
        counts = mirror_results(cloud, config['gcs_prefix'], args.output_dir)
    except Exception as error:
        details = {'file': error.relative, 'reason': 'different_content'} if isinstance(error, MirrorConflict) else {}
        print(json.dumps({'event': 'batch_stage_mirror_failed', 'error_type': type(error).__name__,
                          'message': 'Stage mirror failed; GCS results are preserved. Inspect and retry.', **details}),
              file=sys.stderr, flush=True)
        return result.returncode or 1
    print(json.dumps({'event': 'batch_stage_mirror', **counts}), flush=True)
    return result.returncode


if __name__ == '__main__':
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        raise SystemExit(130)
    except Exception as error:
        print(json.dumps({'ok': False, 'error_type': type(error).__name__,
                          'message': 'Batch tick wrapper failed; inspect configuration and persisted state.'}),
              file=sys.stderr)
        raise SystemExit(1)