s3 improvement

[karank2512]

Summary by CodeRabbit

• New Features

  • Added batch S3 snapshot download/upload support with concurrent processing and preserved output order.
  • Improved S3 handling with automatic local caching, duplicate-request deduplication, and reusable client connections.

• Documentation

  • Expanded the Cloud storage guide with cache behavior, tensor-parallel snapshot naming, credential sourcing, and concurrency settings.

• Bug Fixes

  • Reduced repeated S3 client setup overhead and improved reliability for parallel snapshot transfers.

[vercel]

@karank2512 is attempting to deploy a commit to the Nils 's projects Team on Vercel.

A member of the Team first needs to authorize it.

[coderabbitai]

Caution

Review failed

An error occurred during the review process. Please try again later.

📝 Walkthrough

Walkthrough

The PR adds batch snapshot resolution and upload helpers in thaw_common.cloud, shared S3 client caching, package exports and README updates, a new S3 batch benchmark script with receipt output, and tests covering caching and batch transfer behavior.

Changes

S3 snapshot batching and client reuse

Layer / File(s)	Summary
Batch transfer helpers and concurrency knob python/thaw_common/cloud.py	resolve_snapshots and upload_snapshots add bounded batch transfers with ordered results, duplicate-URI deduplication, local-path passthrough, and first-error propagation; the module docstring and THAW_S3_FILE_CONCURRENCY update the S3 transfer model.
Shared S3 client cache python/thaw_common/cloud.py	_s3_client now reuses boto3 clients by process ID and pool size, and reset_s3_client_cache clears the cache under lock.
Public exports and README python/thaw_common/__init__.py, README.md	resolve_snapshots and upload_snapshots are exported from thaw_common, and the Cloud storage README section adds cache, per-rank URI, credentials, and concurrency details.
Benchmark script and receipt benchmarks/bench_s3_batch.py, site/receipts/2026-06-24_s3_batch.json	A new benchmark script measures client reuse and bounded batch concurrency, prints the results, and writes a JSON receipt for one recorded run.
Client cache tests tests/test_cloud.py	Imports, moto cache-reset hooks, and TestSharedClientCache cover reuse, reset, pool-size differences, and PID-keyed rebuilds.
Batch transfer tests tests/test_cloud.py	TestResolveSnapshotsBatch, TestUploadSnapshotsBatch, and TestBatchRoundTripMoto cover ordered resolution, bounded concurrency, error propagation, duplicate URIs, and moto-backed round trips.

Sequence Diagram(s)

Resolve batch flow:

sequenceDiagram
  participant resolve_snapshots
  participant _batch_workers
  participant resolve_snapshot_path
  resolve_snapshots->>_batch_workers: submit distinct remote URIs
  _batch_workers->>resolve_snapshot_path: resolve each URI
  resolve_snapshot_path-->>_batch_workers: resolved path
  _batch_workers-->>resolve_snapshots: ordered results

Loading

Upload batch flow:

sequenceDiagram
  participant upload_snapshots
  participant _batch_workers
  participant upload_snapshot
  upload_snapshots->>_batch_workers: submit local_path and uri pairs
  _batch_workers->>upload_snapshot: upload each pair
  upload_snapshot-->>_batch_workers: upload complete
  _batch_workers-->>upload_snapshots: completion or first failure

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Poem

/) /)
( •‿•) I hopped through caches, neat and bright,
Batch burrows ran in tidy flight.
Reused clients blinked, then sped away,
While S3 carrots danced today. 🐰

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.93% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title is too vague to convey the actual change; it only says "s3 improvement" and doesn't mention caching, batch transfers, or benchmarks.	Use a more specific title such as "Add S3 snapshot batching, caching, and benchmarks" to reflect the main change.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[matteso1]

yo this is sick — the client caching + the URL dedup before fan-out are exactly right, and the (pid, pool_size) key handling the TP-worker fork case is clean. tests are solid too.

two small things before i merge:

1. the README line "N keys in parallel give ~N×135 MB/s" — we can't claim that one yet. the bench injects a synthetic sleep(), so it's measuring latency-hiding, not real bandwidth (the receipt has no MB/s in it). can you soften it to like "should scale toward ~N×135 MB/s (projected from the single-key ceiling, not yet measured end to end)"? we got burned before citing numbers we hadn't actually run so we keep that bar high. the single 135/key number is great to keep, that one's got a real receipt.

2. the receipt json has a local path baked in (C:\Users\kkapu\...) — mind scrubbing it to a relative path? just don't want personal paths in committed receipts.

one thing for later, not blocking: the shared client pool is 36 connections but the batch can ask for 8×32=256 in flight, so most of the per-file ranged concurrency isn't actually staying warm in the batch path. correctness is fine, it just means the real fan-out win is bounded by the pool size — we can size the pool for file×per-file when we do the shard-at-freeze work.

fix those two and let's merge. great start on the cloud path 🙏