#!/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()