castle-batch-v2/batch_jobs.py

#!/usr/bin/env python3
"""Render short HF Batch workers. No submission unless --execute is supplied."""
import argparse
import json
import re
import shlex
import subprocess
import sys


def build_command(*, role, code_volume, output_volume, name, project, state_uri,
                  pipeline_args, credential_secret, image='python:3.12-slim',
                  flavor='cpu-basic', timeout=None, location='global'):
    if role not in {'submit','hourly-tick','tick'}:
        raise ValueError('Unknown worker role')
    for value in (code_volume,output_volume):
        if not value.startswith('hf://buckets/') or any(c in value[5:] for c in ':\r\n') or '..' in value.split('/'):
            raise ValueError('Use explicit HF bucket prefixes')
    if code_volume.rstrip('/')==output_volume.rstrip('/'):
        raise ValueError('Code and output must be distinct')
    if not re.fullmatch(r'[A-Z][A-Z0-9_]*',credential_secret):
        raise ValueError('Supply the environment variable name containing ADC JSON, never its value')
    if not project or not state_uri.startswith('gs://'):
        raise ValueError('Explicit Google project and GCS state URI are required')
    managed={'--project','--location','--state-uri','--output-dir','--scratch-dir','--submit'}
    if any(arg.split('=')[0] in managed for arg in pipeline_args):
        raise ValueError('Launcher owns project/location/state/output/scratch/submit flags')
    if role!='submit' and pipeline_args:
        raise ValueError('Tick scope comes entirely from persistent state')
    cmd=['hf','jobs']
    if role=='hourly-tick':
        cmd+=['scheduled','run','hourly','--no-concurrency']
    else:
        cmd+=['run','--detach']
    cmd+=['--flavor',flavor,'--timeout',timeout or ('3h' if role=='submit' else '30m'),
          '--name',name,'--volume',code_volume.rstrip('/')+':/workspace:ro',
          '--volume',output_volume.rstrip('/')+':/output',
          '--secrets',credential_secret,'--env','CASTLE_ADC_SECRET_NAME='+credential_secret,
          '--env','GOOGLE_APPLICATION_CREDENTIALS=/scratch/google-adc.json',
          '--env','PYTHONUNBUFFERED=1','--env','PYTHONDONTWRITEBYTECODE=1',
          '--env','HF_HOME=/scratch/hf','--env','HF_XET_CHUNK_CACHE_SIZE_BYTES=0',
          '--env','HF_XET_HIGH_PERFORMANCE=0','--env','OMP_NUM_THREADS=1']
    runtime=['python','/workspace/batch_pipeline.py','prepare' if role=='submit' else 'tick',
             '--project',project,'--location',location,'--state-uri',state_uri,
             '--output-dir','/output','--scratch-dir','/scratch/castle',*pipeline_args]
    if role=='submit':runtime+=['--submit']
    install_media=('if ! command -v ffmpeg >/dev/null || ! command -v ffprobe >/dev/null; then '
                   'apt-get update && apt-get install -y --no-install-recommends ffmpeg; fi; ') if role=='submit' else ''
    bootstrap=('set -eu; mkdir -p /scratch/castle; '+install_media+
               'python -m pip install --no-cache-dir -r /workspace/requirements-batch.txt; '
               'exec python /workspace/bootstrap_batch.py '+shlex.join(runtime))
    return cmd+['--',image,'bash','-lc',bootstrap]


def main(argv=None):
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('role',choices=['submit','hourly-tick','tick'])
    for field in ['code-volume','output-volume','name','project','state-uri']:
        p.add_argument('--'+field,required=True)
    p.add_argument('--credential-secret',default='GOOGLE_ADC_JSON')
    p.add_argument('--location',default='global')
    p.add_argument('--image',default='python:3.12-slim')
    p.add_argument('--flavor',default='cpu-basic')
    p.add_argument('--timeout')
    p.add_argument('--execute',action='store_true')
    argv=list(sys.argv[1:] if argv is None else argv)
    split=argv.index('--') if '--' in argv else len(argv)
    args=p.parse_args(argv[:split]);forwarded=argv[split+1:]
    values=vars(args).copy();execute=values.pop('execute');values['pipeline_args']=forwarded
    command=build_command(**values)
    if args.role=='submit':
        from batch_pipeline import build_parser
        build_parser().parse_args(['prepare',*forwarded,'--project',args.project,'--location',args.location,
                                  '--state-uri',args.state_uri,'--output-dir','/output','--scratch-dir','/scratch/castle'])
    if not execute:
        print(json.dumps({'submitted':False,'argv':command},indent=2));return 0
    return subprocess.run(command,check=False).returncode


if __name__=='__main__':
    raise SystemExit(main())