批量启动沙箱运行 RL Rollout

python3 rl-batch-demo.py run --count 10 --output rl-parallel-10

python3 rl-batch-demo.py run --count 10 \\
  --runtime YOUR_RUNTIME_CLASS \\
  --image YOUR_PYTHON_IMAGE \\
  --output rl-parallel-custom

cat rl-parallel-10/results.json

python3 rl-batch-demo.py cleanup --output rl-parallel-10
python3 rl-batch-demo.py run --count 10 --parallelism 1 --output rl-sequential-10
cat rl-sequential-10/results.json

python3 rl-batch-demo.py cleanup --output rl-sequential-10

kubectl -n YOUR_TEST_NAMESPACE get job,pods
kubectl -n YOUR_TEST_NAMESPACE get events --sort-by=.metadata.creationTimestamp
kubectl -n YOUR_TEST_NAMESPACE logs YOUR_POD_NAME -c episode

#!/usr/bin/env python3
"""Reproducible, bounded RL rollout environment demo; requires kubectl."""
import argparse
import concurrent.futures
import datetime
import json
import math
import os
from pathlib import Path
import subprocess
import sys
import time
import uuid

WORKER = r'''
import json, os, pathlib, subprocess, time
p = pathlib.Path('/workspace')
os.chdir(p)
p.joinpath('shipping.py').write_text('def shipping(total):\\n    return 0 if total > 99 else 10\\n')
p.joinpath('test_shipping.py').write_text("""import unittest
from shipping import shipping
class ShippingTest(unittest.TestCase):
    def test_below(self): self.assertEqual(shipping(98), 10)
    def test_boundary(self): self.assertEqual(shipping(99), 0)
    def test_above(self): self.assertEqual(shipping(100), 0)
if __name__ == '__main__': unittest.main()
""")
def emit(event, **data):
    print(json.dumps(dict(event=event, at=time.time(),
          episode_id=os.environ['EPISODE_ID'], **data)), flush=True)
def test():
    result = subprocess.run(['python', '-B', 'test_shipping.py'],
                            text=True, capture_output=True, timeout=20)
    return result.returncode, result.stderr
emit('ready')
before, before_output = test()
if before == 0:
    raise RuntimeError('baseline must fail the boundary test')
# Deterministic stand-in for a coding Agent's action; no external model call.
p.joinpath('shipping.py').write_text(p.joinpath('shipping.py').read_text().replace('total > 99', 'total >= 99'))
after, after_output = test()
emit('reward', reward=int(after == 0), baseline_exit=before,
     tests_exit=after, tests_output=after_output)
if after != 0:
    raise RuntimeError('patched code failed its tests')
'''


def percentile(values, q):
    values = sorted(values)
    i = (len(values) - 1) * q
    lo, hi = math.floor(i), math.ceil(i)
    return values[lo] + (values[hi] - values[lo]) * (i - lo)


def collect(pods, logs, job, count):
    origin = datetime.datetime.fromisoformat(job['metadata']['creationTimestamp'].replace('Z', '+00:00')).timestamp()
    rows = []
    if len(pods) != count:
        raise ValueError('Pod count differs from requested episodes; inspect failures/replacements')
    for pod, log in zip(pods, logs):
        if pod['status']['phase'] != 'Succeeded':
            raise ValueError('Episode did not succeed: ' + pod['metadata']['name'])
        events = [json.loads(line) for line in log.splitlines() if line.strip()]
        ready = [x for x in events if x['event'] == 'ready']
        reward = [x for x in events if x['event'] == 'reward']
        if len(ready) != 1 or len(reward) != 1:
            raise ValueError('Expected one ready and one reward event per episode')
        r, w = ready[0], reward[0]
        if r['episode_id'] != pod['metadata']['name'] or w['episode_id'] != r['episode_id']:
            raise ValueError('Episode identity mismatch')
        if w['reward'] != 1 or w['tests_exit'] != 0 or w['baseline_exit'] == 0:
            raise ValueError('Reward or baseline verification failed')
        created = datetime.datetime.fromisoformat(pod['metadata']['creationTimestamp'].replace('Z', '+00:00')).timestamp()
        if not origin <= created <= r['at'] <= w['at']:
            raise ValueError('Invalid timestamp order; check clock synchronization')
        status = pod['status']['containerStatuses'][0]
        if status['restartCount'] != 0 or status['state']['terminated']['exitCode'] != 0:
            raise ValueError('Unexpected restart or exit code')
        rows.append(dict(pod=r['episode_id'], node=pod['spec']['nodeName'],
            image_id=status.get('imageID'), reward=w['reward'],
            pod_to_ready_seconds=r['at']-created,
            batch_to_ready_seconds=r['at']-origin,
            batch_to_reward_seconds=w['at']-origin))
    startup = [x['pod_to_ready_seconds'] for x in rows]
    return dict(episodes=count, successful_episodes=len(rows), reward_sum=sum(x['reward'] for x in rows),
        batch_to_all_ready_seconds=max(x['batch_to_ready_seconds'] for x in rows),
        batch_to_all_rewards_seconds=max(x['batch_to_reward_seconds'] for x in rows),
        pod_to_ready_p50_seconds=percentile(startup, .5),
        pod_to_ready_p95_seconds=percentile(startup, .95),
        pod_to_ready_max_seconds=max(startup), pods=rows)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=['run', 'cleanup'])
    parser.add_argument('--output', required=True)
    parser.add_argument('--count', type=int, default=10)
    parser.add_argument('--parallelism', type=int)
    parser.add_argument('--image', default='python:3.12-alpine')
    parser.add_argument('--runtime', default='cube')
    args = parser.parse_args()
    output = Path(args.output)
    parallel = args.count if args.parallelism is None else args.parallelism
    if not 1 <= parallel <= args.count <= 20:
        parser.error('require 1 <= parallelism <= count <= 20')
    if args.action == 'cleanup':
        receipt = json.loads((output/'receipt.json').read_text())
        context = receipt['context']
    else:
        context = subprocess.check_output(['kubectl', 'config', 'current-context'], text=True).strip()
    def kubectl(*parts, data=None):
        return subprocess.check_output(['kubectl', '--context', context, '--request-timeout=30s', *parts],
              input=json.dumps(data) if data is not None else None, text=True, timeout=60)
    def create(obj):
        return json.loads(kubectl('create', '-f', '-', '-o', 'json', data=obj))
    if args.action == 'cleanup':
        ns = receipt['namespace']
        raw = kubectl('get', 'namespace', ns, '--ignore-not-found', '-o', 'json')
        if not raw.strip():
            print('Namespace already absent'); return
        live = json.loads(raw)
        if live['metadata']['uid'] != receipt['uid'] or live['metadata'].get('labels', {}).get('agc-demo-run') != receipt['run_id']:
            raise RuntimeError('Ownership check failed; refusing cleanup')
        raw_job = kubectl('-n', ns, 'get', 'job', 'rollout', '--ignore-not-found', '-o', 'json')
        if raw_job.strip():
            expected_uid = receipt.get('job_uid')
            if expected_uid is None and (output/'job.json').exists():
                expected_uid = json.loads((output/'job.json').read_text())['metadata']['uid']
            if json.loads(raw_job)['metadata']['uid'] != expected_uid:
                raise RuntimeError('Job ownership check failed; refusing cleanup')
            kubectl('-n', ns, 'delete', 'job', 'rollout', '--cascade=foreground', '--wait=false')
            deadline = time.monotonic() + 120
            while kubectl('-n', ns, 'get', 'job', 'rollout', '--ignore-not-found', '-o', 'name').strip():
                if time.monotonic() > deadline: raise TimeoutError('Job or Pods still terminating')
                time.sleep(2)
        kubectl('delete', 'namespace', ns, '--wait=false')
        deadline = time.monotonic() + 120
        while kubectl('get', 'namespace', ns, '--ignore-not-found', '-o', 'name').strip():
            if time.monotonic() > deadline: raise TimeoutError('Namespace still terminating; inspect finalizers')
            time.sleep(2)
        print('Deleted test namespace:', ns); return
    output.mkdir(parents=True, exist_ok=False)
    run_id = uuid.uuid4().hex[:12]
    ns = 'agc-rl-' + run_id
    receipt = dict(context=context, namespace=ns, run_id=run_id)
    # Save intended namespace before any cluster mutation so failures remain traceable.
    (output/'receipt.json').write_text(json.dumps(receipt, indent=2))
    kubectl('get', 'runtimeclass', args.runtime)
    created = create(dict(apiVersion='v1', kind='Namespace', metadata=dict(name=ns, labels={'agc-demo-run':run_id})))
    receipt['uid'] = created['metadata']['uid']
    (output/'receipt.json').write_text(json.dumps(receipt, indent=2))
    print('Test namespace:', ns, flush=True)
    create(dict(apiVersion='v1', kind='ConfigMap', metadata=dict(name='episode', namespace=ns), data={'episode.py':WORKER}))
    job = dict(apiVersion='batch/v1', kind='Job', metadata=dict(name='rollout', namespace=ns), spec=dict(
        completions=args.count, parallelism=parallel, backoffLimit=0, activeDeadlineSeconds=240,
        template=dict(metadata=dict(labels={'app':'rl-batch-demo'}), spec=dict(
            runtimeClassName=args.runtime, restartPolicy='Never', automountServiceAccountToken=False,
            securityContext=dict(runAsUser=1000, runAsGroup=1000, fsGroup=1000, runAsNonRoot=True),
            containers=[dict(name='episode', image=args.image, imagePullPolicy='IfNotPresent',
                command=['python','-B','/demo/episode.py'],
                env=[dict(name='EPISODE_ID',valueFrom=dict(fieldRef=dict(fieldPath='metadata.name')))],
                securityContext=dict(allowPrivilegeEscalation=False, capabilities=dict(drop=['ALL']), seccompProfile=dict(type='RuntimeDefault')),
                resources=dict(requests=dict(cpu='100m',memory='64Mi'),limits=dict(cpu='500m',memory='128Mi')),
                volumeMounts=[dict(name='script',mountPath='/demo',readOnly=True),dict(name='workspace',mountPath='/workspace')])],
            volumes=[dict(name='script',configMap=dict(name='episode')),dict(name='workspace',emptyDir={})]))))
    (output/'job-submitted.json').write_text(json.dumps(job, indent=2))
    start = time.monotonic()
    created_job = create(job)
    receipt['job_uid'] = created_job['metadata']['uid']
    (output/'receipt.json').write_text(json.dumps(receipt, indent=2))
    deadline = start + 270
    while True:
        live = json.loads(kubectl('-n',ns,'get','job','rollout','-o','json'))
        conditions = {c['type']:c['status'] for c in live.get('status',{}).get('conditions',[])}
        if conditions.get('Complete')=='True' or conditions.get('Failed')=='True': break
        if time.monotonic()>deadline: break
        time.sleep(1)
    observed = time.monotonic()-start
    pods = json.loads(kubectl('-n',ns,'get','pods','-l','job-name=rollout','-o','json'))['items']
    (output/'job.json').write_text(json.dumps(live, indent=2))
    (output/'pods.json').write_text(json.dumps(pods, indent=2))
    (output/'events.txt').write_text(kubectl('-n',ns,'get','events','--sort-by=.metadata.creationTimestamp'))
    def get_log(pod):
        name = pod['metadata']['name']
        try: log = kubectl('-n',ns,'logs',name,'-c','episode')
        except subprocess.SubprocessError: log = ''
        (output/(name+'.log')).write_text(log)
        return log
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
        logs = list(pool.map(get_log,pods))
    if conditions.get('Complete')!='True':
        raise RuntimeError('Job failed or timed out; inspect saved status/events/logs, then cleanup')
    report = collect(pods, logs, live, args.count)
    report.update(context=context, namespace=ns, image=args.image, runtime=args.runtime,
                  parallelism=parallel, cache_state='uncontrolled', client_job_complete_seconds=observed)
    (output/'results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps({k:v for k,v in report.items() if k!='pods'},indent=2))

if __name__=='__main__':
    main()