Baladithya Balamurugan

Wave 20: 5-facet AWS-native architecture design (F1-F5)

50e3f0e 29 days ago

25.1 kB

Facet F2 — AWS-Native Dataset Generation (the outer loop)

Account 386931836011 · us-west-2 · Admin · 2026-06-09

The outer (slow) loop of §5 — ingest → multi-teacher replay → invert SWE substrates into FeatureDeletion tasks → normalize to SFT/DPO corpus on S3 — is embarrassingly parallel, bursty, fault-tolerant (report §5, §8). The wrong move is to pick one service for the whole loop. Each of the four stages has a different compute shape, so each gets the right AWS-native primitive, all stitched by Step Functions (the AWS-native analog of the report's Argo controller for the non-cluster path) and all reading/writing one S3 dataset contract.

This note commits to a concrete service per stage and maps every choice to the repo file it backs.

TL;DR — the per-stage verdict

Outer-loop stage	Repo code	Compute shape	AWS service (committed)	Why not the others
(a) Ingest + clean Claude Code traces	`ingestion/claude_code.py`	Light CPU, JSONL → TraceState, IO-bound	Glue Spark ETL job (Glue 5.0)	Trivial scale; one Spark job over a JSONL prefix; Glue's Data Catalog + bookmarks are free here. SM Processing also fine but Glue is the cheaper IO-bound default.
(b) Replay each state across N teachers	`teacher_replay.py`	API-bound fan-out, N×T calls, no local GPU	Bedrock batch inference (`CreateModelInvocationJob`, one job per teacher) + a thin EMR Serverless aggregation step	NOT Glue-Ray (end-of-support). NOT SageMaker-Processing-with-Ray for the calls (you'd pay for idle instances waiting on API latency). Bedrock batch = 50% discount, S3-native, the fan-out IS the job.
(c) Synthesize FeatureDeletionEnv tasks (run tests in sandboxes)	`datagen/substrates.py`, `env.py`, `docker_sandbox.py`, `validator.py`	CPU-heavy, untrusted code, 1 container per task, fault-tolerant	AWS Batch array jobs on EC2 Spot (container = the existing `DockerSandbox` image)	NOT SM Processing (no per-task isolation/retry granularity; 1 job = 1 fleet). NOT Fargate-only (no `--privileged`/gVisor, 4-vCPU cap, no Spot-as-cheap). Batch array index = task index, matches SWE-bench `--max_workers` model exactly.
(d) Normalize → SFT corpus + DPO pairs	`replaysim/normalize.py` (data-juicer), `teacher_replay.extract_dpo_pairs`	CPU Spark/dataframe dedup + the data-juicer op-graph	EMR Serverless (Spark) running data-juicer's Ray/standalone executor in `--mode local` per partition, OR SM Processing for the single-node data-juicer path	EMR Serverless is the Spark-style dataframe normalize/dedup home; data-juicer runs CPU-only here. Glue ETL would also work but EMR Serverless gives finer Spark control + Graviton price/perf.

The recurring principle: API-bound and test-execution work must NOT live on a service that bills for idle wall-clock (Glue/SM Processing keep instances hot while you wait on Bedrock or a 5-minute pytest). Bedrock batch is async-priced; Batch array jobs are per-task Spot. Spark-shaped dataframe work (ingest, dedup, normalize) goes to the Spark services (Glue ETL for the trivial ingest, EMR Serverless for the heavy normalize).

Stage (a) — Ingest + clean traces · Glue 5.0 Spark ETL

ingestion/claude_code.py::ClaudeCodeIngester.ingest() is a pure JSONL→TraceState transform (one TraceState per assistant turn; strip_thinking=False per report §6 — 67% of error-recovery turns are pure thinking). This is IO-bound dataframe work over a prefix of session files.

Service: AWS Glue 5.0 Spark ETL job.

Input: s3://<datagen-bucket>/raw/claude_code/**/*.jsonl (Claude Code session JSONL uploaded from ~/.claude/projects).
The Glue script mapPartitions over the JSONL files; each partition runs ClaudeCodeIngester().ingest(path) unchanged (it already yields TraceState), filters subagent/sidechain records, and writes Parquet.
Output: s3://<datagen-bucket>/traces/v1/run_id=<id>/part-*.parquet partitioned by run_id.
Glue Data Catalog table composer_traces makes the trace store queryable by Athena (debug: "how many error-site turns this run?").

Why Glue over SM Processing here: ingestion is small, periodic, IO-bound; Glue's per-DPU-second billing + job bookmarks (skip already-ingested sessions) fit better than spinning a Processing fleet. ClaudeCodeIngester runs verbatim inside the Spark UDF — zero repo change to the ingester itself.

Live caveat (worth flagging): ClaudeCodeIngester.ingest() is a per-file, single-pass, record-parallel pure-Python iterator with no shuffle/join — i.e. it is not intrinsically Spark-shaped. SageMaker Processing with ProcessingInput.S3DataDistributionType = ShardedByS3Key (custom container, the ingester runs unchanged, slice config from processingjobconfig.json/resourceconfig.json) is the cleaner primitive for this exact shape and avoids any dependence on the Glue stack — relevant because AWS Glue-for-Ray is now closed to new customers (confirmed in the Glue engine docs), so leaning on Glue here means committing to Glue-for-Spark specifically. Verdict stands as Glue ETL for the trivial periodic ingest (catalog + bookmarks are a real convenience), with SM Processing as the equally-valid, Glue-independent fallback. Either way the ingester code is untouched.

Stage (b) — N-teacher replay · Bedrock batch, NOT OpenRouter

Today teacher_replay.py calls OpenRouter over httpx with DEFAULT_TEACHERS = [claude-opus, gpt-5, deepseek-v4] and a max_total_usd cap. The facet question: stay API or move to Bedrock?

Verdict: move the AWS-native default to Bedrock batch inference, keep OpenRouter as a fallback adapter.

Reasoning grounded in research + repo:

Bedrock batch = 50% of on-demand pricing, S3-in/S3-out, ~24h turnaround SLA. The replay is the single most expensive piece of the whole system (report §10: ~$0.98/trace at N=3, ~$64/trace at 8-teacher×1000-step). A 50% cut on the dominant cost line is the highest-leverage AWS-native move in this facet, and replay is exactly the latency-insensitive offline workload Bedrock batch targets.
The fan-out IS one job per teacher. Bedrock batch does not multiplex models in a single job — CreateModelInvocationJob takes one modelId. So the N-teacher fan-out = N batch jobs over the same JSONL, each writing to its own S3 prefix. The flat for state: for teacher: loop in replay_trace becomes "emit one shared JSONL of all states, submit N jobs." This is structurally simpler than the current per-call gather.
Heterogeneity survives on Bedrock — VERIFIED LIVE in this account. aws bedrock list-foundation-models --region us-west-2 returns a genuinely multi-lab pool: anthropic.claude-opus-4-8, anthropic.claude-opus-4-7, anthropic.claude-opus-4-6-v1, anthropic.claude-sonnet-4-6, anthropic.claude-haiku-4-5-..., deepseek.v3.2, deepseek.r1-v1:0, meta.llama4-maverick-17b-instruct, meta.llama3-3-70b-instruct, etc. That is Claude + DeepSeek + Llama = three different families, satisfying the report's N≥3 heterogeneous-population anti-collapse safeguard (§5 #4) and the "drop same-family teachers if Claude is the student" leakage rule (§6).
- CRITICAL batch-eligibility split (read live from the Bedrock "Supported Regions and models for batch inference" table): not every model has a batch ID. Batch-eligible in/through us-west-2: DeepSeek V3.2 (deepseek.v3.2, single-region us-west-2), Claude Sonnet 4.6 / Opus 4.6 (via the us. cross-region inference profile whose destinations include us-west-2), and the Nova/Titan families. On-demand ONLY (no ARN-versioned batch ID, so NOT batchable): Claude Opus 4.7 / 4.8 — reachable via bedrock-runtime Converse/InvokeModel and fanned out with a bounded async/thread pool exactly like today's replay_trace, but you pay on-demand. So the cheap bulk batch pool is {deepseek.v3.2, us.anthropic.claude-sonnet-4-6, us.anthropic.claude-opus-4-6}; if the teacher set must include 4.7/4.8 they ride the on-demand path or are substituted by batchable Opus 4.6.
Bedrock data does not train Anthropic models (governance) — a real reason to prefer it over OpenRouter for a self-improving loop.

Concrete wiring

teacher_replay.py gains a BedrockBatchTeacherPool alongside the OpenRouter path (the TeacherSpec slug becomes a Bedrock modelId or inference-profile id like us.anthropic.claude-haiku-4-5-...):

# new: teacher_replay_bedrock.py  (~180 LOC)
def submit_replay_batch(states, teachers, s3_in, s3_out, role_arn) -> list[jobArn]:
    # 1. write ONE shared abc.jsonl: {recordId: state_id, modelInput: {anthropic_version, messages, max_tokens}}
    #    one record per state (messages = state["messages"]).  recordId == state_id is the join key.
    # 2. for each teacher: bedrock.create_model_invocation_job(
    #        modelId=teacher.model_id, jobName=..., roleArn=role_arn,
    #        inputDataConfig={"s3InputDataConfig":{"s3Uri": s3_in}},
    #        outputDataConfig={"s3OutputDataConfig":{"s3Uri": f"{s3_out}/{teacher.slug}/"}})
    # 3. return jobArns; poll get_model_invocation_job until Completed.

The output .jsonl.out rows carry recordId (=state_id) + modelOutput; an EMR Serverless step joins all N teacher outputs back by state_id into the exact list[TeacherCallResult] shape extract_dpo_pairs already consumes — so extract_dpo_pairs and the entire DPOPair contract are byte-for-byte untouched.

Constraints to honor (from research): Bedrock batch input file ≤ 1 GB, total job ≤ 5 GB, default 100k records/job (adjustable via Service Quotas), ≥ a minimum-records floor per model. The submitter must shard a >100k-state run into multiple JSONL files. A BedrockBatchTeacherPool is the right home for that sharding.

Where OpenRouter stays: when N>3 teachers from labs Bedrock doesn't host (e.g. a brand-new frontier model), or when sub-24h latency is needed for a hot debugging loop. The TeacherSpec TypedDict already abstracts provider; we add a provider: "bedrock"|"openrouter" field and route. This keeps replay_trace's async OpenRouter path as the low-latency escape hatch and Bedrock batch as the cheap bulk default.

Why NOT Glue-Ray / SM-Processing-with-Ray for the fan-out

The facet floats "Glue Ray jobs / SageMaker Processing with Ray for the N-model parallel fan-out." AWS Glue for Ray is end-of-support — AWS now explicitly recommends Ray-on-EKS (KubeRay) instead (docs: AWS Glue for Ray end of support). And running the teacher API calls on a Ray cluster (Glue or SM Processing) means paying for idle CPU/instances while threads block on model latency — the anti-pattern above. Ray's place in this facet is stage (c) orchestration of sandbox fan-out if we ever want Ray semantics, but Batch array jobs are the simpler AWS-native fit there too. So: Bedrock batch for the calls, EMR Serverless to fan-in.

Stage (c) — FeatureDeletion synthesis + test execution · AWS Batch array jobs on Spot

This is the compute-heavy, genuinely-new part (report §8: "per-branch sandbox isolation is the throughput ceiling of the whole idea"). Two sub-steps:

c1. Schema inversion — SweBenchAdapter.to_task(instance) (substrates.py) maps a SWE-bench-shaped row → FeatureDeletionTask (revert gold patch = broken repo; FAIL_TO_PASS = reward target; PASS_TO_PASS = guard). Pure CPU, no Docker. This runs inside the Glue ingest job (a) or a tiny Lambda — it's a dict transform. is_redistributable() license-gate (GPL/AGPL/LGPL filter) runs here.

c2. Materialize + validate the broken repo — this is where it gets heavy. For each task: pull the substrate's frozen image, git apply -R the gold patch, scrub_tree() (the PRIMARY reward-hack control — strip __pycache__/.git/*.pyc), run the test command, confirm FAIL_TO_PASS actually fails and PASS_TO_PASS actually passes (the 4-gate validator.py solvability check). Each task = one untrusted, isolated, ~1-5 minute container run.

Service: AWS Batch array jobs, EC2 Spot compute environment, container = the existing DockerSandbox image.

batch.submit_job(
  jobName=f"fd-validate-{run_id}",
  jobQueue="fd-sandbox-spot",
  jobDefinition="fd-docker-sandbox:N",   # the DockerSandbox image in ECR
  arrayProperties={"size": n_tasks},      # up to 10,000 children
  containerOverrides={"environment": [{"name":"TASK_MANIFEST_S3","value": s3_manifest}]})

Each array child reads AWS_BATCH_JOB_ARRAY_INDEX, looks up its task at that line in the S3 task manifest (tasks/v1/run_id=<id>/manifest.jsonl), boots the broken image, runs validator + FeatureDeletionEnv._grade() against LocalSubprocessSandbox/DockerSandbox, writes one result row to s3://.../task_grades/<run_id>/<task_id>.json.
This is exactly the SWE-bench harness --max_workers model: SWE-bench runs 1 container per instance with N parallel workers; DeepSWE hit Docker-daemon limits at 512 containers/iteration and preloaded images onto NVMe (report §8). Batch array jobs give that fan-out with managed retry (retryStrategy), Spot interruption handling, and per-child CloudWatch logs — without us building a container scheduler.
Isolation tier (report §8 layered posture): the DockerSandbox already bakes the lockdown recipe (network_mode=none, read_only, cap_drop=ALL, no-new-privileges, pids_limit, gVisor runtime='runsc' when available). On Batch EC2 (not Fargate), we can install gVisor on the AMI and run untrusted model-generated fix attempts under runsc by default; Kata+Firecracker for adversarial code requires self-managed node groups (the EKS-MNG CPU-Options gotcha from §8 applies to Batch managed EC2 too — use a self-managed launch template if nested virt is needed).

Why NOT SageMaker Processing for the sandboxes

SM Processing is one job = one fixed instance fleet with ShardedByS3Key splitting input across instances. It has no per-task retry, no Spot-native interruption recovery at task granularity, and no array-index primitive — a single poisoned task (infinite loop, fork bomb) can wedge a whole Processing instance's shard. Batch array jobs isolate each task to its own container with its own timeout (DockerSandbox.exec_timeout_s + Batch timeout), retry, and Spot replacement. For "10,000 untrusted code executions, each independent, some will hang," Batch is the textbook fit and SM Processing is not.

Why NOT plain Fargate

Fargate caps at 4 vCPU / 30 GB without --privileged, cannot run gVisor/Kata, and is pricier than Spot for bursty bulk. Batch-on-EC2-Spot is 60-70% cheaper for this fault-tolerant fan-out (report §10 Spot savings).

Stage (d) — Normalize → SFT corpus + DPO pairs · EMR Serverless (Spark)

After (b) fan-in produces list[TeacherCallResult] and extract_dpo_pairs produces DPOPair rows, replaysim/normalize.py::DJNormalizer runs the data-juicer op-graph (recipes/replaysim/default.yaml: length filter ×2, words-num ×2, special-char ×2, document_deduplicator). ADR-004 chose data-juicer precisely because the op set is CPU-only (no NeMo-Curator GPU dep).

Service: EMR Serverless (Spark) application, Graviton.

The normalize is Spark-shaped dataframe work: read all DPOPair rows for a run, run the data-juicer op-graph, dedup across the corpus, write the partitioned corpus.
data-juicer's DefaultExecutor runs in-process per Spark partition (the repo's DJNormalizer.normalize() already does file-in/file-out via init_configs + DefaultExecutor). On EMR Serverless we mapPartitions → run DJNormalizer(skip_dj=False).normalize(partition_rows) per partition, then a Spark dropDuplicates for cross-partition full-corpus dedup (the repo notes document_deduplicator is per-batch only — Spark closes that gap natively).
EMR Serverless auto-scales workers, bills per vCPU/GB-second on Graviton, and is the AWS-native Spark home for "dataframe normalize + dedup at scale." The repo's replaysim already is the op-graph — EMR Serverless just runs it distributed.

Why EMR Serverless over Glue ETL here: the normalize is the heavier of the two Spark stages, data-juicer wants control over the Python runtime (custom --py-files/wheel), and EMR Serverless gives finer Spark tuning + Graviton price/perf than Glue's DPU model. Glue ETL stays the choice for the light ingest (a). Both are valid Spark; we split by weight.

Two output corpora (the dataset contract below): the SFT corpus (clean winning trajectories — report §5 "SFT-first competence floor") and the DPO pairs (from extract_dpo_pairs + future execution-oracle near-miss rejects).

The S3 dataset contract

One bucket per environment (reuse the existing amazon-sagemaker-386931836011-us-west-2-* or a dedicated composer-datagen-386931836011-us-west-2). Layout is Hive-partitioned by run_id so each outer-loop generation is an immutable, addressable slice (the report's "generation" in the GA framing, §3/§5) and Athena/Glue Catalog can query across generations.

s3://composer-datagen-386931836011-us-west-2/
  raw/claude_code/**/*.jsonl                         # (a) input: uploaded sessions
  traces/v1/run_id=<id>/part-*.parquet               # (a) out: TraceState rows
  tasks/v1/run_id=<id>/manifest.jsonl                # (c1) FeatureDeletionTask rows (1/line; array index = line)
  replay/v1/run_id=<id>/
      input/states.jsonl                             # (b) shared Bedrock batch input (recordId=state_id)
      teacher=<slug>/*.jsonl.out                      # (b) per-teacher Bedrock batch output
  task_grades/v1/run_id=<id>/<task_id>.json          # (c2) validator + _grade() results
  corpus/v1/run_id=<id>/
      sft/part-*.parquet                              # (d) SFT corpus (clean winners)
      dpo/part-*.parquet                              # (d) DPO pairs (normalized DPOPair)
  manifests/run_id=<id>.json                          # run-level manifest: counts, cost, lineage, schema_version
  diloco/rendezvous/round_<NNNNNN>/rank_<RRRR>.pt     # (separate) inner-loop ObjectStoreAllReduce (already exists)

Format decisions:

Parquet for traces/, corpus/sft/, corpus/dpo/ — columnar, compressed, Athena-queryable, the SFT/DPO training read path (TRL load_dataset reads Parquet directly). This is the durable corpus.
JSONL for replay/input/, tasks/manifest, task_grades/ — because that's what Bedrock batch requires (s3InputFormat=JSONL), what AWS Batch array-index line-lookup wants, and what DPOPair/TraceState natively serialize to. JSONL is the wire format between stages; Parquet is the corpus-at-rest format.
recordId == state_id is the universal join key linking trace → replay → DPO pair. Already true: TraceState.state_id is f"{path.stem}::{idx:04d}" and DPOPair.state_id carries it through.
manifests/run_id.json is the run-level contract: {run_id, schema_version, n_traces, n_tasks, n_dpo_pairs, teacher_pool, bedrock_job_arns, total_cost_usd, parent_run_id}. parent_run_id threads the flywheel lineage (report §5: improved student regenerates next round's traces). The held-out-eval guard (safety/HeldoutSplit) reads schema_version + a split column to enforce the disjoint held-out set the report flags as the load-bearing gap (§7 Pushback 4).

Orchestration: Step Functions (the non-cluster Argo analog)

Report §8 uses Argo Workflows on EKS as the outer-loop controller. For the AWS-native, no-persistent-cluster path this facet targets, the equivalent is AWS Step Functions (Standard workflow) driving the DAG:

Ingest(Glue) → InvertSchema(Lambda) → [Bedrock batch ×N (Map)] → FanIn(EMR-Serverless)
  → ExtractDPO+SynthTasks → SandboxValidate(Batch array, .sync) → Normalize(EMR-Serverless)
  → WriteManifest(Lambda)

Step Functions has native .sync integrations for Glue, EMR Serverless, Batch (arn:aws:states:::batch:submitJob.sync — blocks until the whole array completes), and SageMaker, plus a Map state for the N-teacher Bedrock fan-out. This makes the whole outer loop one declarative state machine, retryable, with per-stage IAM. If/when the program moves to the EKS-primary path of §8, the same stage boundaries lift to Argo DAG nodes — the S3 contract is identical, so it's a controller swap, not a rewrite.

Repo delta (what to build in `composer_replication/`)

composer_replication/teacher_replay_bedrock.py (~180 LOC) — BedrockBatchTeacherPool: submit_replay_batch() writes the shared states JSONL, submits one create_model_invocation_job per teacher, polls, and parses .jsonl.out back into list[TeacherCallResult]. Add provider/model_id to TeacherSpec. extract_dpo_pairs untouched.
composer_replication/datagen/aws/batch_validate.py (~120 LOC) — the Batch array-child entrypoint: read AWS_BATCH_JOB_ARRAY_INDEX → task manifest line → boot DockerSandbox/LocalSubprocessSandbox → run validator + _grade() → write task_grades/.../{task_id}.json. Plus a submit_validate_array() helper.
composer_replication/datagen/aws/glue_ingest_job.py (~80 LOC) — Glue Spark entrypoint wrapping ClaudeCodeIngester.ingest in mapPartitions; writes traces/ Parquet + Glue Catalog table.
composer_replication/replaysim/emr_normalize_job.py (~100 LOC) — EMR Serverless Spark entrypoint wrapping DJNormalizer per partition + Spark cross-partition dedup; writes corpus/dpo/ + corpus/sft/ Parquet.
composer_replication/datagen/aws/s3_contract.py (~120 LOC) — the S3 layout constants, RunManifest dataclass, Parquet/JSONL serializers for TraceState/FeatureDeletionTask/DPOPair, the recordId==state_id join helpers, and schema_version/split column injection for the held-out guard.
infra/datagen_stepfunctions.json (+ a thin CDK/infra/datagen_stack.py) — the Step Functions state machine + IAM roles (Bedrock batch service role, Batch Spot compute env, EMR Serverless app, Glue role). ~250 LOC IaC.
pyproject.toml — extend the [aws]/[datagen] extras with boto3 for Bedrock/Batch/Glue/EMR clients (the [serverless] extra already needs s3fs/boto3 per report §9).
Dockerfile — the DockerSandbox ECR image (also the Batch job-definition image), baking gVisor runsc on the AMI/launch-template for default untrusted-code isolation.

The trainer, loss, FeatureDeletionEnv, curriculum, monitor, extract_dpo_pairs, DJNormalizer op-graph, and the DiLoCo/ObjectStoreAllReduce inner loop are all untouched — every AWS-native piece wraps an existing repo entrypoint and reads/writes the S3 contract.

Open questions

Bedrock batch 24h SLA vs flywheel cadence. Report §5 says outer-loop cadence is hours-to-days, so 24h batch turnaround is acceptable for bulk generations — but a fast bootstrap iteration may want the OpenRouter low-latency path. Need a batch_vs_realtime policy knob keyed on max_total_usd + urgency.
Per-branch sandbox cold-start is the §8/§10 explicit falsifier. If Batch array job startup + image pull dominates wall-clock at target fan-out even with gVisor, the report's demotion path applies: keep Batch for control/grading but move bulk sandbox execution to a container-free pool (SWE-MiniSandbox class) or a warm Spot fleet. Must instrument cold-start in batch_validate.py.
Cross-region inference profiles for Bedrock batch (us.anthropic...) route to multiple regions for throughput — confirm the batch service role + S3 bucket policy allow the destination regions, else throttling.
DeepSeek on Bedrock batch in us-west-2 — RESOLVED LIVE: deepseek.v3.2 is batch-eligible single-region in us-west-2; deepseek.r1-v1:0 is in the on-demand catalog. Claude batch IDs are the us. cross-region inference profiles (us.anthropic.claude-sonnet-4-6, us.anthropic.claude-opus-4-6). Still confirm minimum-records floors per model + the records/job & file-size quotas in Service Quotas at build time (adjustable, but the BedrockBatchTeacherPool must shard input JSONL to fit).
Opus 4.7/4.8 are on-demand-only (no batch ID). If the heterogeneity ablation (report §3) wants the newest Claude as a teacher, it costs on-demand pricing — budget for it or substitute the batchable Opus 4.6. GPT-5 (in DEFAULT_TEACHERS) is not on Bedrock at all → that one family member stays on the OpenRouter escape hatch or is replaced by Bedrock Llama 4 / DeepSeek as the third family.
golden_diff must be ACL-isolated. It is repr=False in FeatureDeletionTask and held out of the policy observation; on S3 it must live in a separate deny-by-default prefix (tasks/golden/), never co-located with the policy-visible tasks/v1/... — oracle-cleanliness Gate 1 (report §4).