Add codameter-bench: shardable config-sweep CLI for Fargate/Batch

README.md

@@ -198,6 +198,25 @@ against the known truth). Export the JSONL with `pixi run frugalmind-export`; th
 drop-in FrugalMind suite lives in
 [`integrations/frugalmind/`](integrations/frugalmind).
 
+### Scaling the sweep (Fargate / AWS Batch)
+
+`codameter-bench` scores a grid of processing configs against every golden case
+and records recovery RMS. The work is deterministic and embarrassingly parallel,
+so it shards cleanly across an array of tasks:
+
+```bash
+codameter-bench plan  --grid multiverse --shard 0/64      # size the job (no compute)
+codameter-bench sweep --grid multiverse --shard 0/64 --jobs 4 --out s3://bucket/run/
+codameter-bench aggregate --src s3://bucket/run/ --out s3://bucket/run/agg/
+```
+
+Each task writes one `shard-<k>-of-<N>.jsonl`, so retries are idempotent and no
+coordination is needed. On AWS Batch the shard is read from
+`AWS_BATCH_JOB_ARRAY_INDEX` + `CODAMETER_SHARDS`, so one image runs every array
+task. Container image and a Batch job definition are in
+[`docker/`](docker); `s3://` output needs the `aws` extra
+(`pip install "codameter[aws]"`).
+
 ---
 
 ## Models, hyperparameters, and physical bounds

docker/Dockerfile

@@ -0,0 +1,24 @@
+# Container for the codameter config-sweep benchmark (codameter-bench).
+# Built to run as one task of an array job (AWS Batch on Fargate, or any array
+# runner): each task computes one --shard of the case x config grid and writes a
+# shard-<k>-of-<N>.jsonl to the --out path (local or s3://).
+FROM python:3.11-slim
+
+# Non-interactive, no .pyc, unbuffered logs (so CloudWatch sees progress live).
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1
+
+WORKDIR /app
+
+# Install the package (with the aws extra for s3:// output). Only the bits the
+# sweep needs are imported at runtime; wheels cover the scientific stack.
+COPY pyproject.toml README.md ./
+COPY src ./src
+COPY tests/data/golden/manifest.json ./tests/data/golden/manifest.json
+RUN pip install ".[aws]"
+
+# The array runner sets the shard: pass --shard k/N explicitly, or set
+# CODAMETER_SHARDS=N and let the CLI read AWS_BATCH_JOB_ARRAY_INDEX.
+ENTRYPOINT ["codameter-bench"]
+CMD ["--help"]

docker/README.md

@@ -0,0 +1,75 @@
+# Scaling the codameter sweep on Fargate / AWS Batch
+
+`codameter-bench` scores a grid of processing configs against every golden case
+and records recovery RMS. The work is deterministic and embarrassingly parallel,
+so it fans out cleanly across an array of Fargate tasks: each task runs one
+`--shard k/N` and writes one `shard-<k>-of-<N>.jsonl` to S3.
+
+## Size the job first
+
+```bash
+codameter-bench plan --grid multiverse --shard 0/64
+# grid=multiverse  cases=10  configs/case=648
+# total cells=6480  shards=64  cells/shard: min=101 max=102
+```
+
+`plan` needs no compute; use it to pick the array size and vCPUs per task. Grids:
+`compact` (smoke), `multiverse` (default), `wide` (overnight).
+
+## Run locally (one machine, all vCPUs)
+
+```bash
+codameter-bench sweep --grid multiverse --out ./out --jobs 8 --shard 0/1
+codameter-bench aggregate --src ./out --out ./agg
+```
+
+`--jobs` spreads one shard's cells across worker processes. Splitting into shards
+*and* using `--jobs` composes: shards across tasks, jobs across a task's vCPUs.
+
+## Run on AWS Batch (Fargate)
+
+1. **Build and push** the image:
+
+   ```bash
+   docker build -f docker/Dockerfile -t codameter-bench .
+   aws ecr get-login-password --region <REGION> | docker login --username AWS \
+     --password-stdin <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com
+   docker tag codameter-bench <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/codameter-bench:latest
+   docker push <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/codameter-bench:latest
+   ```
+
+2. **Register** the job definition (edit placeholders + `CODAMETER_SHARDS` first):
+
+   ```bash
+   aws batch register-job-definition --cli-input-json file://docker/aws-batch-job-def.json
+   ```
+
+3. **Submit** an array of N tasks (N must equal `CODAMETER_SHARDS`):
+
+   ```bash
+   aws batch submit-job --job-name codameter-sweep --job-queue <QUEUE> \
+     --job-definition codameter-sweep --array-properties size=64
+   ```
+
+   Batch sets `AWS_BATCH_JOB_ARRAY_INDEX` on each task; the CLI reads it together
+   with `CODAMETER_SHARDS` to select the shard, so every task runs the same image
+   and command. The `jobRoleArn` needs `s3:PutObject` on the output prefix.
+
+4. **Aggregate** when the array finishes (locally or as a final one-off task):
+
+   ```bash
+   codameter-bench aggregate --src s3://<BUCKET>/codameter-sweep/run-001/ \
+                             --out  s3://<BUCKET>/codameter-sweep/run-001/agg/
+   ```
+
+## Notes
+
+- **Idempotent retries.** Each shard writes exactly its own file, named by shard,
+  so a Batch retry overwrites only that shard. No coordination, no partial merges.
+- **Even load.** Sharding is round-robin (`items[k::N]`), so the slow configs
+  (moving / inversion references, which rerun the estimator per epoch) spread
+  evenly instead of piling into one task.
+- **s3://** output needs the `aws` extra (`pip install "codameter[aws]"`, already
+  in the image). Local paths need nothing.
+- **Cost knob.** vCPUs per task x array size sets throughput; `plan` tells you the
+  cell count so you can trade array size against per-task `--jobs`.

docker/aws-batch-job-def.json

@@ -0,0 +1,28 @@
+{
+  "_comment": "Example AWS Batch job definition for the codameter sweep on Fargate. Replace <ACCOUNT>, <REGION>, <BUCKET>, and the role ARNs. Register with: aws batch register-job-definition --cli-input-json file://aws-batch-job-def.json. Submit an array of N tasks with: aws batch submit-job --job-name codameter-sweep --job-queue <QUEUE> --job-definition codameter-sweep --array-properties size=64 . Batch sets AWS_BATCH_JOB_ARRAY_INDEX on each task; CODAMETER_SHARDS below must equal the array size.",
+  "jobDefinitionName": "codameter-sweep",
+  "type": "container",
+  "platformCapabilities": ["FARGATE"],
+  "containerProperties": {
+    "image": "<ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/codameter-bench:latest",
+    "command": [
+      "sweep",
+      "--grid", "multiverse",
+      "--out", "s3://<BUCKET>/codameter-sweep/run-001/",
+      "--jobs", "4"
+    ],
+    "environment": [
+      { "name": "CODAMETER_SHARDS", "value": "64" }
+    ],
+    "resourceRequirements": [
+      { "type": "VCPU", "value": "4" },
+      { "type": "MEMORY", "value": "8192" }
+    ],
+    "executionRoleArn": "arn:aws:iam::<ACCOUNT>:role/ecsTaskExecutionRole",
+    "jobRoleArn": "arn:aws:iam::<ACCOUNT>:role/codameter-sweep-task-role",
+    "fargatePlatformConfiguration": { "platformVersion": "LATEST" },
+    "networkConfiguration": { "assignPublicIp": "ENABLED" }
+  },
+  "retryStrategy": { "attempts": 2 },
+  "timeout": { "attemptDurationSeconds": 3600 }
+}

pixi.toml

@@ -71,6 +71,9 @@ golden = { cmd = "python -m codameter.golden", description = "Regenerate tests/d
 # Export the golden dv/v datasets as FrugalMind-compatible JSONL + manifest
 frugalmind-export = { cmd = "python -m codameter.frugalmind --out datasets", description = "Export dv/v benchmark rows as JSONL (FrugalMind format)" }
 
+# Shardable config-sweep benchmark (plan / sweep / aggregate); scales to Fargate
+bench = { cmd = "python -m codameter.bench", description = "Config-sweep benchmark over case x config grids" }
+
 # Run only the Parkfield notebook-friendly example
 parkfield = { cmd = "python examples/01_parkfield_synthetic.py", description = "Synthetic Parkfield demo" }
 

pyproject.toml

@@ -62,10 +62,12 @@ dev = [
     "mypy>=1.6",
     "pre-commit>=3.5",
 ]
+aws = ["boto3>=1.28"]  # s3:// output for codameter-bench (Fargate/Batch)
 all = ["codameter[kernels,mcmc]"]
 
 [project.scripts]
 codameter = "codameter.cli:main"
+codameter-bench = "codameter.bench:main"
 
 [project.urls]
 Homepage = "https://github.com/Denolle-Lab/codameter"