code/scripts/continue_run.py

#!/usr/bin/env python3
"""Continue an already running free-Colab training job through evaluation/export.

The saved notebook remains the source for data checks, capacity selection, and
evaluation. Publishing is separate: this runner never receives a credential.
"""
from __future__ import annotations

import argparse
from contextlib import contextmanager
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import threading
import time
import traceback


@contextmanager
def exclusive_export_lock(path):
    import fcntl
    with path.open('a+') as handle:
        # A surviving export process inherits this descriptor and keeps the
        # lock even if its supervisor exits. A new runner waits and rechecks
        # completed artifacts instead of launching a duplicate export.
        fcntl.flock(handle, fcntl.LOCK_EX)
        yield handle.fileno()


def wait_for_training_process(pid, psutil, sleep=time.sleep):
    """A non-child zombie has exited even if its notebook parent has not reaped it."""
    try:
        process = psutil.Process(pid)
        if process.status() == psutil.STATUS_ZOMBIE:
            return
        if 'train.py' not in ' '.join(process.cmdline()):
            raise RuntimeError('The supplied process is not the training command')
        while process.is_running() and process.status() != psutil.STATUS_ZOMBIE:
            sleep(2)
    except psutil.NoSuchProcess:
        return


def evaluator_ready(evaluation, process, launched_at):
    if process is None or process.poll() is not None:
        return False
    path = evaluation / f'.runtime-ready-{process.pid}.json'
    try:
        receipt = json.loads(path.read_text())
    except (OSError, ValueError):
        return False
    return (receipt.get('pid') == process.pid
            and receipt.get('ready_unix_seconds', 0) >= launched_at
            and receipt.get('model_warmup_complete') is True
            and receipt.get('first_batch_processed') is True)


def exports_complete(export_dir, adapter_sha256):
    try:
        merged = json.loads((export_dir / 'merged-fp16/MERGE_COMPLETE.json').read_text())
        gguf = json.loads((export_dir / 'gguf/GGUF_COMPLETE.json').read_text())
        model = export_dir / 'gguf/VibeGame-8B-Q4_K_M.gguf'
        return (merged.get('adapter_sha256') == adapter_sha256
                and gguf.get('source_merge') == merged
                and model.is_file() and model.stat().st_size == gguf.get('bytes'))
    except (OSError, ValueError):
        return False


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--training-pid', type=int, required=True)
    parser.add_argument('--run', type=Path, required=True)
    parser.add_argument('--data', type=Path, required=True)
    parser.add_argument('--godot', required=True)
    args = parser.parse_args()
    import psutil
    root = Path(__file__).resolve().parents[1]
    drive = Path('/content/drive/MyDrive').resolve()
    if not drive.is_dir() or not args.run.resolve().is_relative_to(drive):
        parser.error('Use the mounted free Colab runtime and a persistent Drive run directory')
    os.chdir(root)
    sys.path.insert(0, str(root / 'src'))
    run_dir = args.run.resolve()
    # Prevent two continuation runners from racing on immutable evidence or
    # starting duplicate exports. Child processes do not inherit this lock.
    import fcntl
    import hashlib
    lock_key = hashlib.sha256(str(run_dir).encode()).hexdigest()
    run_lock = (Path('/content') / f'.vibegame-continuation-{lock_key}.lock').open('a+')
    try:
        fcntl.flock(run_lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        raise RuntimeError('A continuation runner is already active for this run')
    evaluation = run_dir / 'evaluation'
    export_dir = Path('/content/vibegame-exports')
    notebook_path = root / 'notebooks/VibeGame-8b.ipynb'
    notebook = json.loads(notebook_path.read_text())
    cells = {index: ''.join(cell['source']) for index, cell in enumerate(notebook['cells'])}
    expected = {10: 'def verify_run_dataset', 12: 'def verify_completed_training',
                14: 'def safe_capacity_recommendation', 16: 'scripts/evaluate.py', 18: 'scripts/build_demo.py'}
    if any(marker not in cells.get(index, '') for index, marker in expected.items()):
        raise RuntimeError('Notebook stages changed; inspect this runner before using it')
    processes = {}
    process_started = {}
    abort = threading.Event()

    def status(filename, **values):
        value = {'utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()), **values}
        destination = run_dir / filename
        temporary = destination.with_suffix('.tmp')
        temporary.write_text(json.dumps(value, indent=2) + '\n')
        temporary.replace(destination)

    def run(*command, **kwargs):
        command = [str(value) for value in command]
        stage = Path(command[1]).stem
        with (run_dir / (stage + '.log')).open('a') as log:
            process_started[stage] = time.time()
            process = subprocess.Popen(command, cwd=root, stdout=log, stderr=subprocess.STDOUT, **kwargs)
            processes[stage] = process
            status('continuation-status.json', stage=stage, pid=process.pid, command=command, status='running')
            code = process.wait()
        status('continuation-status.json', stage=stage, exit_code=code, status='finished' if code == 0 else 'failed')
        if code:
            raise subprocess.CalledProcessError(code, command)

    def run_export_locked(overlap, lock_fd):
        if exports_complete(export_dir, adapter_sha256):
            status('export-status.json', status='complete', reused=True)
            return True
        command = ['nice', '-n', '19', sys.executable, '-u', str(root / 'scripts/export_model.py'),
                   '--adapter', str(run_dir / 'adapter'), '--output', str(export_dir),
                   '--cache', '/content/vibegame-base-cache', '--workspace', '/content/vibegame-export-work']
        with (run_dir / 'export_model.log').open('a') as log:
            process = subprocess.Popen(command, cwd=root, stdout=log, stderr=subprocess.STDOUT,
                                       start_new_session=True, pass_fds=(lock_fd,))
            status('export-status.json', status='running', pid=process.pid, overlap=overlap)
            while process.poll() is None:
                if psutil.virtual_memory().available < 1024**3:
                    try:
                        os.killpg(process.pid, signal.SIGTERM)
                        try:
                            process.wait(timeout=10)
                        except subprocess.TimeoutExpired:
                            os.killpg(process.pid, signal.SIGKILL)
                            process.wait()
                    except ProcessLookupError:
                        process.wait()
                    status('export-status.json', status='stopped_for_host_memory', exit_code=process.returncode)
                    return False
                time.sleep(2)
            complete = process.returncode == 0 and exports_complete(export_dir, adapter_sha256)
            status('export-status.json', status='complete' if complete else 'failed', exit_code=process.returncode)
            return complete

    def run_export(overlap):
        with exclusive_export_lock(export_dir.parent / '.vibegame-export.lock') as lock_fd:
            return run_export_locked(overlap, lock_fd)

    def export_when_ready():
        try:
            if exports_complete(export_dir, adapter_sha256):
                status('export-status.json', status='complete', reused=True)
                return
            # Existing manifests cannot signal readiness for a new process.
            # Wait for its fresh first-batch receipt and available host RAM.
            while not abort.is_set():
                process = processes.get('evaluate')
                if process is not None and process.poll() is not None:
                    return
                if evaluator_ready(evaluation, process, process_started.get('evaluate', float('inf'))) and psutil.virtual_memory().available >= 6 * 1024**3:
                    break
                abort.wait(10)
            if abort.is_set():
                return
            run_export('CPU export overlaps GPU evaluation; timing includes contention')
        except Exception as error:
            status('export-status.json', status='failed', error=type(error).__name__ + ': ' + str(error))
            traceback.print_exc()

    status('continuation-status.json', stage='waiting_for_training', pid=args.training_pid, status='waiting')
    wait_for_training_process(args.training_pid, psutil)
    if not (run_dir / 'adapter/TRAINING_COMPLETE.json').is_file():
        raise RuntimeError('Training exited without a completed adapter; inspect train.log')
    import hashlib
    import shutil
    import torch
    digest = hashlib.sha256()
    with (run_dir / 'adapter/adapter_model.safetensors').open('rb') as handle:
        for block in iter(lambda: handle.read(8 * 1024 * 1024), b''):
            digest.update(block)
    adapter_sha256 = digest.hexdigest()
    namespace = {'__name__': '__vibegame_continuation__', 'PROJECT': root, 'RUN': run_dir,
                 'DATA': args.data.resolve(), 'DRIVE': run_dir.parents[1], 'EVAL': evaluation,
                 'EXPORT': export_dir, 'GODOT': args.godot, 'ADAPTER': run_dir / 'adapter',
                 'DATASET_REPO': '', 'LOCAL_TEACHER_ATTEMPTS': 0, 'Path': Path, 'os': os,
                 'sys': sys, 'subprocess': subprocess, 'json': json, 'hashlib': hashlib,
                 'shutil': shutil, 'torch': torch, 'run': run}
    worker = None
    evaluation_done = False
    demo_done = False
    stage_error = None
    try:
        for index in (10, 12, 14):
            exec(compile(cells[index], f'{notebook_path.name}:cell-{index}', 'exec'), namespace)
        status('continuation-status.json', stage='capacity_selected', status='complete', batch_size=namespace['EVAL_BATCH_SIZE'])
        worker = threading.Thread(target=export_when_ready, name='cpu-export', daemon=False)
        worker.start()
        exec(compile(cells[16], f'{notebook_path.name}:evaluation', 'exec'), namespace)
        evaluation_done = True
        # The runtime saves the video; display happens in the user's notebook.
        capture_source = cells[18].split('from IPython.display import Video, display')[0]
        exec(compile(capture_source, f'{notebook_path.name}:recordings', 'exec'), namespace)
        demo_done = True
    except Exception as error:
        stage_error = type(error).__name__ + ': ' + str(error)
        traceback.print_exc()
    finally:
        abort.set()
        if worker is not None:
            worker.join()
    # The parallel export may have deferred for RAM or stopped under pressure.
    # Retry once with evaluation/capture no longer occupying host memory, even
    # if capture failed. Joining above prevents concurrent duplicate exporters.
    if evaluation_done and not exports_complete(export_dir, adapter_sha256):
        try:
            run_export('Sequential fallback after evaluation; no live evaluator')
        except Exception as error:
            status('export-status.json', status='failed', error=type(error).__name__ + ': ' + str(error))
            traceback.print_exc()
    export_done = exports_complete(export_dir, adapter_sha256)
    complete = evaluation_done and demo_done and export_done
    status('continuation-status.json', stage='finished', status='complete' if complete else 'partial',
           evaluation_complete=evaluation_done, demo_complete=demo_done, export_complete=export_done,
           error=stage_error)
    if not complete:
        raise RuntimeError('Continuation is partial; inspect continuation-status.json and retained logs')


if __name__ == '__main__':
    main()