VibeGame-8B / code /scripts /continue_run.py
khursheed's picture
Release pipeline source, tests, notebook, and licenses
ae36cfb verified
Raw History Blame Contribute Delete
12 kB
#!/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()