Benchmarking and tuning guide for writing large JAX arrays to S3-backed Zarr via TensorStore.

For writing 32768×32768 float32 arrays (~4.3 GB each), the following configuration achieves 5–6 GB/s sustained write throughput:

import os
import jax
import tensorstore as ts

size = 2**15  # 32768
chunks_per_dim = 32

spec = {
    "driver": "zarr",
    "kvstore": {
        "driver": "s3",
        "bucket": "ace-usw04a",
        "path": "my-array/",
        "endpoint": "http://ace-usw04a.cwlota.com",
        "host_header": "ace-usw04a.cwlota.com",
        "aws_region": "us-west-04a",
        "aws_credentials": {"type": "environment"},
    },
    "metadata": {
        "dtype": "<f4",
        "shape": [size, size],
        "chunks": [size // chunks_per_dim, size // chunks_per_dim],
        "compressor": None,
    },
    "create": True,
    "open": True,
}

data = jax.random.normal(jax.random.PRNGKey(0), (size, size))
dataset = ts.open(spec).result()
dataset[...] = data

export TENSORSTORE_HTTP_THREADS=64
export TENSORSTORE_S3_REQUEST_CONCURRENCY=512

All experiments write a 32768×32768 float32 array (~4.3 GB) to S3-backed Zarr over 5 iterations. "Steady-state" excludes the first iteration (which pays connection establishment overhead).

The original script used default TensorStore settings with 4 CPU cores, 4×4 chunking (256 MB chunks), and default Blosc compression. No TensorStore environment variables were set.

1. Chunk count + CPU cores are the dominant factors. More chunks create more concurrent S3 PUT operations; more CPU cores feed those operations faster via TensorStore's data_copy_concurrency (which defaults to CPU count).

2. HTTP thread count must match chunk parallelism. TENSORSTORE_HTTP_THREADS caps how many HTTP requests can execute simultaneously. Set it to at least the number of chunks, or higher.

3. S3 request concurrency should exceed chunk count. TENSORSTORE_S3_REQUEST_CONCURRENCY controls TensorStore's S3 upload queue depth. Set it to 2–4× the chunk count.

1. Compression has no effect on random float32 data. Blosc detects incompressible data quickly and short-circuits, so the overhead is negligible. For compressible data (e.g., sparse arrays), compression may still be worthwhile to reduce transfer volume.

2. HTTP threads alone (without more chunks/CPU) don't help. Adding threads to a 4-CPU / 16-chunk workload had no effect — the bottleneck was upstream in chunk encoding, not in the HTTP client.

3. GPU vs CPU source arrays — GPU-resident arrays perform as well or slightly better than CPU arrays. The device-to-host PCIe transfer (~25 GB/s) is not a bottleneck relative to S3 upload throughput.

JAX Array → [data_copy_concurrency: chunk encoding] → [HTTP threads: upload queue] → S3
              ↑ scales with CPU cores                   ↑ TENSORSTORE_HTTP_THREADS
              ↑ more chunks = more parallelism          ↑ TENSORSTORE_S3_REQUEST_CONCURRENCY

The baseline configuration (4 CPU, 16 chunks, 4 HTTP threads) created a bottleneck at every stage. The optimized configuration removes all three:

• 1024 chunks → encoding and upload are highly parallel
• 16 CPU cores → data_copy_concurrency can saturate the pipeline
• 64 HTTP threads / 512 S3 concurrency → no queuing in the HTTP layer

The benchmark script (src/main.py) is parameterized via environment variables. Pod manifests for each experiment are in k8s/experiments/.

# Update the ConfigMap with the benchmark script
kubectl create configmap jaxy-script --from-file=main.py=src/main.py -n <namespace>

# Run all experiments
kubectl apply -f k8s/experiments/

# Collect results
kubectl logs <pod-name> -n <namespace> | grep -E "RESULT|SUMMARY|Best|Worst|Mean|Cumulative"