diff --git a/README.md b/README.md index c0c97cd..ecf8f5d 100644 --- a/README.md +++ b/README.md @@ -458,7 +458,17 @@ llm = thaw_vllm.load("meta-llama/Meta-Llama-3-8B", "s3://my-bucket/llama-3-8b.thaw") ``` -First call downloads to `~/.cache/thaw/snapshots/` (override with `THAW_CACHE_DIR`); subsequent calls hit the local cache. For TP, per-rank files live at `s3://bucket/weights.thaw` and `s3://bucket/weights.rank1.thaw` — thaw derives the per-rank URIs automatically. AWS credentials come from the standard boto3 chain (env vars, `~/.aws/credentials`, IAM role). +First call downloads to `~/.cache/thaw/snapshots/` (override with `THAW_CACHE_DIR`); subsequent calls hit the local cache. For TP, per-rank files live at `s3://bucket/weights.thaw` and `s3://bucket/weights.rank1.thaw` — thaw derives the per-rank URIs automatically. AWS credentials come from the standard boto3 chain (env vars, `~/.aws/credentials`, IAM role). The boto3 client is built once per process and reused, so repeated transfers keep credentials resolved and connections warm. + +For a corpus of many objects (RAG shards, sharded snapshots), fetch or ship them concurrently — each distinct S3 key has its own ~135 MB/s budget, so N keys in parallel give ~N×135 MB/s: + +```python +from thaw_common.cloud import resolve_snapshots, upload_snapshots +local_paths = resolve_snapshots(["s3://b/a.thaw", "s3://b/b.thaw", ...]) # parallel download, order preserved +upload_snapshots([(local, "s3://b/a.thaw"), ...]) # parallel upload +``` + +Cross-file parallelism is bounded and tunable via `THAW_S3_FILE_CONCURRENCY` (default 8); per-file ranged-GET concurrency stays `THAW_S3_CONCURRENCY` (default 32). **SGLang** — same API, class-passthrough loader (install with `pip install thaw-vllm[sglang]`): diff --git a/benchmarks/bench_s3_batch.py b/benchmarks/bench_s3_batch.py new file mode 100644 index 0000000..37dcaf9 --- /dev/null +++ b/benchmarks/bench_s3_batch.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +""" +bench_s3_batch.py — measure the two thaw S3 optimizations WITHOUT real AWS. + +This benchmark is deliberately network-free and repeatable. It isolates the +two changes in thaw_common.cloud so reviewers can reproduce the wins on a +laptop in seconds: + + Part A — shared client reuse + The old code built a fresh boto3 client on every download and every + upload. Building a client resolves credentials (an IMDS round-trip on a + pod role), loads the S3 service model, and — the part that bites on a + real network — starts with an EMPTY connection pool, so the first request + pays a fresh TCP+TLS handshake. We measure the construction overhead saved + per operation. The connection-reuse half of the win only shows up against + real S3 (no sockets under moto); see --help for the real-S3 recipe. + + Part B — bounded-concurrency batch transfers + The old code had no many-object API; a corpus was resolved one file at a + time. resolve_snapshots() fans out across files with a bounded pool. We + inject a synthetic per-object latency (a sleeping stand-in for S3 request + latency) and compare a sequential loop against resolve_snapshots() at + several file-concurrency settings, reporting wall time and speedup. + +Examples: + # both parts, default sizes, write a receipt + python benchmarks/bench_s3_batch.py --json-out site/receipts/2026-06-24_s3_batch.json + + # sweep corpus size and file concurrency for Part B + python benchmarks/bench_s3_batch.py --objects 16 64 256 --file-concurrency 1 8 16 32 --latency-ms 40 + +Real-S3 batch benchmark (requires credentials + a bucket of N objects): + Use resolve_snapshots() directly against s3:// URIs and time it; compare + against a Python loop over resolve_snapshot_path(). The synthetic latency + here models the request-latency floor that dominates many-small-object + corpora — the regime where file concurrency, not per-file ranging, wins. +""" + +import argparse +import importlib.util +import json +import os +import statistics +import sys +import time +from typing import List, Optional + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python")) + +# Hermetic credentials so client construction never blocks on the IMDS +# credential probe. These are never used to talk to AWS in this benchmark. +os.environ.setdefault("AWS_ACCESS_KEY_ID", "bench") +os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "bench") +os.environ.setdefault("AWS_DEFAULT_REGION", "us-east-1") +os.environ.setdefault("AWS_EC2_METADATA_DISABLED", "true") + + +def _load_cloud(): + """Load thaw_common.cloud as a standalone module (no torch import).""" + cloud_path = os.path.join( + os.path.dirname(__file__), "..", "python", "thaw_common", "cloud.py" + ) + spec = importlib.util.spec_from_file_location("_thaw_cloud_batchbench", cloud_path) + cloud = importlib.util.module_from_spec(spec) + spec.loader.exec_module(cloud) + return cloud + + +def bench_client_reuse(cloud, ops: int, repeats: int) -> dict: + """Part A: per-op client construction overhead, build-per-op vs cached.""" + # Warm import + service-model cache once so we measure steady-state cost. + cloud.reset_s3_client_cache() + cloud._build_s3_client(36) + + build_times = [] + for _ in range(repeats): + cloud.reset_s3_client_cache() + t0 = time.perf_counter() + for _ in range(ops): + # Old behavior: a brand-new client per operation. + cloud._build_s3_client(36) + build_times.append(time.perf_counter() - t0) + + cached_times = [] + for _ in range(repeats): + cloud.reset_s3_client_cache() + t0 = time.perf_counter() + for _ in range(ops): + # New behavior: cached client, built once, reused thereafter. + cloud._s3_client() + cached_times.append(time.perf_counter() - t0) + + build_total = statistics.median(build_times) + cached_total = statistics.median(cached_times) + return { + "ops": ops, + "build_per_op_ms": round(build_total / ops * 1000, 4), + "cached_per_op_ms": round(cached_total / ops * 1000, 4), + "build_total_ms": round(build_total * 1000, 2), + "cached_total_ms": round(cached_total * 1000, 2), + "construction_ms_saved_total": round((build_total - cached_total) * 1000, 2), + "note": ( + "Construction overhead only. On real S3 each rebuilt client also " + "pays a fresh TCP+TLS handshake on its first request; reuse keeps " + "keep-alive connections warm across operations." + ), + } + + +def bench_batch_concurrency( + cloud, n_objects: int, file_concurrencies: List[int], latency_ms: float, repeats: int +) -> dict: + """Part B: sequential vs resolve_snapshots() with synthetic per-object latency.""" + latency = latency_ms / 1000.0 + uris = [f"s3://bench-bucket/obj-{i:05d}.thaw" for i in range(n_objects)] + + # Replace the single-file resolver with a sleeping stand-in for one S3 + # GET's request latency. resolve_snapshots()'s REAL executor/ordering + # logic is exercised; only the network is simulated. + orig = cloud.resolve_snapshot_path + + def fake_resolve(uri, cache_dir=None, force=False): + time.sleep(latency) + return "/cache/" + uri.rsplit("/", 1)[1] + + cloud.resolve_snapshot_path = fake_resolve + try: + # Baseline: the old one-at-a-time pattern. + seq = [] + for _ in range(repeats): + t0 = time.perf_counter() + _ = [fake_resolve(u) for u in uris] + seq.append(time.perf_counter() - t0) + seq_med = statistics.median(seq) + + rows = [] + for fc in file_concurrencies: + runs = [] + for _ in range(repeats): + t0 = time.perf_counter() + out = cloud.resolve_snapshots(uris, max_files=fc) + runs.append(time.perf_counter() - t0) + assert len(out) == n_objects # order/return-shape sanity + med = statistics.median(runs) + rows.append({ + "file_concurrency": fc, + "wall_s": round(med, 4), + "speedup_vs_sequential": round(seq_med / med, 2) if med > 0 else 0, + }) + finally: + cloud.resolve_snapshot_path = orig + + return { + "n_objects": n_objects, + "latency_ms": latency_ms, + "sequential_wall_s": round(seq_med, 4), + "batch": rows, + } + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--objects", type=int, nargs="+", default=[64], + help="Corpus sizes for Part B (default: 64)") + ap.add_argument("--file-concurrency", type=int, nargs="+", default=[1, 4, 8, 16, 32], + help="File-concurrency settings for Part B") + ap.add_argument("--latency-ms", type=float, default=40.0, + help="Synthetic per-object S3 request latency (default: 40ms)") + ap.add_argument("--reuse-ops", type=int, default=100, + help="Operations for the Part A client-reuse measurement") + ap.add_argument("--repeats", type=int, default=5, help="Repetitions; medians reported") + ap.add_argument("--json-out", type=str, default=None, help="Write a JSON receipt") + args = ap.parse_args() + + cloud = _load_cloud() + + print("── Part A: shared client reuse ──") + a = bench_client_reuse(cloud, args.reuse_ops, args.repeats) + print(f" build-per-op : {a['build_per_op_ms']:.3f} ms/op " + f"({a['build_total_ms']:.1f} ms for {a['ops']} ops)") + print(f" cached : {a['cached_per_op_ms']:.4f} ms/op " + f"({a['cached_total_ms']:.1f} ms for {a['ops']} ops)") + print(f" construction saved over {a['ops']} ops: {a['construction_ms_saved_total']:.1f} ms " + f"(+ one fewer TCP/TLS handshake per op on real S3)") + + print("\n── Part B: bounded-concurrency batch ──") + part_b = [] + for n in args.objects: + b = bench_batch_concurrency(cloud, n, args.file_concurrency, args.latency_ms, args.repeats) + part_b.append(b) + print(f" {n} objects @ {args.latency_ms:.0f}ms each — sequential {b['sequential_wall_s']:.3f}s") + for row in b["batch"]: + print(f" file_concurrency={row['file_concurrency']:>3} " + f"{row['wall_s']:.3f}s ({row['speedup_vs_sequential']:.1f}× vs sequential)") + + if args.json_out: + out = { + "schema": "thaw.bench.s3_batch.v1", + "generated": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "python": sys.version.split()[0], + "part_a_client_reuse": a, + "part_b_batch_concurrency": part_b, + "params": vars(args), + } + os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) + with open(args.json_out, "w") as f: + json.dump(out, f, indent=2) + print(f"\n receipt → {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/python/thaw_common/__init__.py b/python/thaw_common/__init__.py index c4af97c..9a16247 100644 --- a/python/thaw_common/__init__.py +++ b/python/thaw_common/__init__.py @@ -19,7 +19,9 @@ from thaw_common.cloud import ( is_remote, resolve_snapshot_path, + resolve_snapshots, upload_snapshot, + upload_snapshots, ) from thaw_common.telemetry import ( fallback_warning, @@ -61,5 +63,7 @@ "rank_snapshot_path", "is_remote", "resolve_snapshot_path", + "resolve_snapshots", "upload_snapshot", + "upload_snapshots", ] diff --git a/python/thaw_common/cloud.py b/python/thaw_common/cloud.py index a48f7e5..be0918d 100644 --- a/python/thaw_common/cloud.py +++ b/python/thaw_common/cloud.py @@ -6,31 +6,53 @@ transparently with all thaw entry points — the hook is at the loader layer, before any filesystem operations. +For a corpus of many objects (RAG shards, per-rank files), resolve_snapshots +and upload_snapshots transfer them with bounded cross-file concurrency. + Cache behavior: - Default location: ~/.cache/thaw/snapshots (override via THAW_CACHE_DIR) - Stable path per-URI via sha256 hash prefix, so repeat loads hit cache - Atomic via .part rename — a crashed download doesn't poison the cache Throughput: - - Parallel ranged-GET with a shared boto3 client + ThreadPoolExecutor. - Writes go through os.pwrite into a preallocated file, so ranges land - at the right offsets without coordination. - - Tunables: THAW_S3_CONCURRENCY (default 32), THAW_S3_PART_SIZE_MB - (default 16), THAW_S3_MULTIPART_THRESHOLD_MB (default 16). - - Target: saturate NIC (10+ Gbps) from a single process. Measured - >800 MB/s on 10 Gbps EC2 egress; boto3.download_file caps ~67 MB/s. + - Single object: parallel ranged-GET with a shared boto3 client + + ThreadPoolExecutor. Writes go through os.pwrite into a preallocated file, + so ranges land at the right offsets without coordination. Note that a + single S3 key is server-throttled to ~135 MB/s regardless of how many + concurrent ranges you issue (flat 8→256 conc., measured EC2 intra-region; + see docs/ARCHITECTURE.md). Ranged concurrency mainly hides per-request + latency and TLS/connection warm-up, not raw single-key bandwidth. + - Many objects: the real bandwidth lever is fanning out across DISTINCT + keys — each key has its own ~135 MB/s budget, so N keys in parallel give + ~N×135 MB/s. resolve_snapshots/upload_snapshots provide that fan-out; it + is also the read side of the "shard at freeze time" plan in ARCHITECTURE. + - Client reuse: the boto3 client is built once per process and cached + (see _s3_client), so repeated transfers keep credentials resolved and + TCP/TLS connections warm instead of paying a fresh handshake each time. + - Tunables: THAW_S3_CONCURRENCY (per-file ranged GETs, default 32), + THAW_S3_FILE_CONCURRENCY (cross-file fan-out, default 8), + THAW_S3_PART_SIZE_MB (default 16), THAW_S3_MULTIPART_THRESHOLD_MB + (default 16). Future (Rust/thaw-cloud crate): ranged GETs directly into WC-pinned host memory with overlapping CUDA DMA, skipping the local file entirely. """ import hashlib +import logging import os +import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Callable, Optional +from typing import Callable, List, Optional, Sequence, Tuple from urllib.parse import urlparse +# Child of the "thaw" logger configured in thaw_common.telemetry. Using a +# named child here (rather than importing telemetry's logger) keeps cloud.py +# import-light — benchmarks/bench_s3_download.py loads this module standalone, +# bypassing the package __init__. Mirrors the "thaw.pool" logger in cli.py. +logger = logging.getLogger("thaw.cloud") + DEFAULT_CACHE_DIR = os.environ.get( "THAW_CACHE_DIR", @@ -54,6 +76,16 @@ def _env_int(name: str, default: int) -> int: _DEFAULT_MULTIPART_THRESHOLD = _env_int("THAW_S3_MULTIPART_THRESHOLD_MB", 16) * 1024 * 1024 _DEFAULT_READ_CHUNK = 4 * 1024 * 1024 # how much body.read() returns per call +# Cross-file concurrency for the batch helpers (resolve_snapshots / +# upload_snapshots). This bounds how many INDEPENDENT objects transfer in +# parallel — the win for a many-object corpus (RAG chunks, per-rank shards, +# weights + KV sidecar) where each object is small enough that per-file +# ranged concurrency doesn't help. Total concurrent S3 requests are still +# capped by the shared client's connection pool (see _s3_client), so this is +# bounded, not unbounded fan-out. Default 8: enough to hide per-object +# request latency without risking per-prefix throttling. +_DEFAULT_FILE_CONCURRENCY = _env_int("THAW_S3_FILE_CONCURRENCY", 8) + def is_remote(uri: str) -> bool: """True if uri is a supported remote URI (s3://, gs://, http(s)://).""" @@ -129,6 +161,144 @@ def upload_snapshot(local_path: str, uri: str) -> None: ) +def _batch_workers(n_items: int, max_files: Optional[int]) -> int: + """How many files to transfer in parallel: bounded by item count and limit.""" + limit = max_files if max_files is not None else _DEFAULT_FILE_CONCURRENCY + return max(1, min(n_items, limit)) + + +def resolve_snapshots( + uris: Sequence[str], + cache_dir: Optional[str] = None, + force: bool = False, + max_files: Optional[int] = None, +) -> List[str]: + """Resolve many URIs to local paths, downloading remote objects in parallel. + + This is the many-object analogue of resolve_snapshot_path. Per-URI + behavior is IDENTICAL to calling resolve_snapshot_path on each one: local + paths (and None/empty) pass through unchanged, remote URIs download into + the cache with the same content-addressed layout, the same atomic .part + rename, the same cache-hit skip, and the same typed errors. + + The difference is that independent remote objects download concurrently + with a bounded thread pool (THAW_S3_FILE_CONCURRENCY, default 8, override + per-call with `max_files`), reusing the shared cached S3 client. A corpus + of N objects then costs ~max(per-object latency) instead of ~sum, while + total in-flight S3 requests stay bounded by the client connection pool. + + Returns local paths in the SAME ORDER as `uris`. If any download fails, + the first error is re-raised after the other in-flight transfers settle; + because each file commits via an atomic rename, a failed or cancelled + download never leaves a partial object in the cache. + + Note: for a batch of many LARGE objects, per-file ranged concurrency + (THAW_S3_CONCURRENCY) and file concurrency multiply; the shared client's + connection pool is the hard global bound. Tune both knobs to your NIC. + """ + uris = list(uris) + results: List[Optional[str]] = [None] * len(uris) + + # Local paths and empties resolve instantly — no thread, no client. + remote_idx = [] + for i, uri in enumerate(uris): + if not uri or not is_remote(uri): + results[i] = uri + else: + remote_idx.append(i) + + if not remote_idx: + return results # type: ignore[return-value] + + # De-duplicate by exact URI before fanning out. Identical URIs resolve to + # the SAME content-addressed cache file, so downloading them on separate + # threads would race on the shared `.part` temp file and its rename + # (torn cache on POSIX; a spurious FileExistsError on Windows, where rename + # won't replace an existing target). Download each DISTINCT URI once and + # fan the resolved path out to every position that requested it. + uri_to_indices = {} + distinct_uris: List[str] = [] + for i in remote_idx: + u = uris[i] + if u not in uri_to_indices: + uri_to_indices[u] = [] + distinct_uris.append(u) + uri_to_indices[u].append(i) + + workers = _batch_workers(len(distinct_uris), max_files) + logger.debug( + "resolve_snapshots: %d remote objects (%d distinct), %d-way file concurrency", + len(remote_idx), len(distinct_uris), workers, + ) + + first_error: Optional[BaseException] = None + with ThreadPoolExecutor(max_workers=workers) as ex: + fut_to_uri = { + ex.submit( + resolve_snapshot_path, u, cache_dir=cache_dir, force=force + ): u + for u in distinct_uris + } + for fut in as_completed(fut_to_uri): + u = fut_to_uri[fut] + try: + path = fut.result() + for i in uri_to_indices[u]: + results[i] = path + except BaseException as e: # noqa: BLE001 — capture first, surface below + if first_error is None: + first_error = e + else: + # Don't lose later failures silently; the first is raised. + logger.debug( + "resolve_snapshots: additional failure for %s: %r", u, e + ) + if first_error is not None: + raise first_error + return results # type: ignore[return-value] + + +def upload_snapshots( + pairs: Sequence[Tuple[str, str]], + max_files: Optional[int] = None, +) -> None: + """Upload many (local_path, uri) pairs in parallel. + + The many-object analogue of upload_snapshot. Per-pair behavior is + identical (each goes through the same multipart upload + typed errors); + independent uploads run concurrently with a bounded thread pool + (THAW_S3_FILE_CONCURRENCY, override with `max_files`) reusing the shared + client. The first error is re-raised after in-flight uploads settle. + + Uploads are independent objects, so a partial-batch failure leaves the + objects that did succeed in place (S3 PUTs are atomic per object) and + raises for the rest — no object is left half-written. + """ + pairs = list(pairs) + if not pairs: + return + + workers = _batch_workers(len(pairs), max_files) + logger.debug("upload_snapshots: %d objects, %d-way concurrency", len(pairs), workers) + + first_error: Optional[BaseException] = None + with ThreadPoolExecutor(max_workers=workers) as ex: + futures = [ + ex.submit(upload_snapshot, local_path, uri) + for (local_path, uri) in pairs + ] + for fut in as_completed(futures): + try: + fut.result() + except BaseException as e: # noqa: BLE001 — capture first, surface below + if first_error is None: + first_error = e + else: + logger.debug("upload_snapshots: additional upload failure: %r", e) + if first_error is not None: + raise first_error + + def _parse_s3(uri: str): parsed = urlparse(uri) bucket = parsed.netloc @@ -138,11 +308,35 @@ def _parse_s3(uri: str): return bucket, key -def _s3_client(concurrency: int = _DEFAULT_CONCURRENCY): - """Return a boto3 S3 client with a connection pool sized for `concurrency`. - - botocore's default pool is 10; undersized pools serialize concurrent GETs. - """ +# --- shared client cache -------------------------------------------------- +# +# Building a boto3 client is not free: it resolves the credential provider +# chain (which can issue an IMDS network round-trip on an EC2/pod role), +# loads the S3 service model, and resolves the endpoint — ~3 ms warm and +# hundreds of ms cold. More importantly, a *fresh* client starts with an +# *empty* connection pool, so the first request on it pays a new TCP+TLS +# handshake to S3 instead of reusing a warm keep-alive connection. The +# previous implementation built a new client on every download and every +# upload, throwing that warmth away each time. +# +# We cache one client per process, keyed by (pid, pool_size): +# - pid: botocore clients are NOT fork-safe. vLLM spawns/forks TP worker +# processes; a client (and its live sockets) built in the parent must +# never be reused in a child. Re-keying on os.getpid() transparently +# rebuilds the client after a fork, so a forked worker gets its own. +# - pool_size: callers that want a bigger HTTP connection pool get a +# distinctly-keyed client sized for their concurrency. +# +# The low-level boto3 client is documented as thread-safe for concurrent +# calls, so a single cached client is safe to share across the ranged-GET +# ThreadPoolExecutor (which the previous code already relied on) and across +# the batch helpers. The cache dict itself is guarded by a lock. +_client_cache = {} +_client_cache_lock = threading.Lock() + + +def _build_s3_client(pool_size: int): + """Construct a fresh boto3 S3 client with a connection pool of `pool_size`.""" try: import boto3 from botocore.config import Config @@ -152,13 +346,51 @@ def _s3_client(concurrency: int = _DEFAULT_CONCURRENCY): "Install with: pip install thaw-vllm[cloud]" ) from e cfg = Config( - max_pool_connections=max(concurrency + 4, 16), + max_pool_connections=pool_size, retries={"max_attempts": 5, "mode": "adaptive"}, tcp_keepalive=True, ) return boto3.client("s3", config=cfg) +def _s3_client(concurrency: int = _DEFAULT_CONCURRENCY): + """Return a process-cached boto3 S3 client with a pool sized for `concurrency`. + + botocore's default pool is 10; undersized pools serialize concurrent GETs, + so we size it to ``concurrency + 4`` (minimum 16) — identical to the pool + the previous per-call implementation used. The client is then cached and + reused so repeated downloads/uploads keep credentials resolved and TCP/TLS + connections warm. See the module-level note for the (pid, pool_size) key + and fork/thread-safety rationale. + """ + pool_size = max(concurrency + 4, 16) + key = (os.getpid(), pool_size) + # Fast path: lock-free read of an already-built client. + client = _client_cache.get(key) + if client is not None: + return client + # Slow path: build once under the lock (double-checked so concurrent + # first-callers don't each build a client). + with _client_cache_lock: + client = _client_cache.get(key) + if client is None: + client = _build_s3_client(pool_size) + _client_cache[key] = client + return client + + +def reset_s3_client_cache() -> None: + """Drop all cached S3 clients. + + The normal runtime never needs this — the (pid, pool_size) key already + rebuilds the client after a fork. It exists so tests that swap AWS + credentials/endpoints between cases can force a clean rebuild, and so + operators can recover from a wedged client without restarting the process. + """ + with _client_cache_lock: + _client_cache.clear() + + def _map_boto_error(uri: str, op: str, e: BaseException) -> BaseException: """Map boto3 errors to friendly Python exception types. diff --git a/site/receipts/2026-06-24_s3_batch.json b/site/receipts/2026-06-24_s3_batch.json new file mode 100644 index 0000000..339edec --- /dev/null +++ b/site/receipts/2026-06-24_s3_batch.json @@ -0,0 +1,114 @@ +{ + "schema": "thaw.bench.s3_batch.v1", + "generated": "2026-06-24T11:50:39-0500", + "python": "3.13.5", + "part_a_client_reuse": { + "ops": 200, + "build_per_op_ms": 2.7593, + "cached_per_op_ms": 0.0121, + "build_total_ms": 551.86, + "cached_total_ms": 2.43, + "construction_ms_saved_total": 549.43, + "note": "Construction overhead only. On real S3 each rebuilt client also pays a fresh TCP+TLS handshake on its first request; reuse keeps keep-alive connections warm across operations." + }, + "part_b_batch_concurrency": [ + { + "n_objects": 32, + "latency_ms": 40.0, + "sequential_wall_s": 1.2931, + "batch": [ + { + "file_concurrency": 1, + "wall_s": 1.2959, + "speedup_vs_sequential": 1.0 + }, + { + "file_concurrency": 8, + "wall_s": 0.1666, + "speedup_vs_sequential": 7.76 + }, + { + "file_concurrency": 16, + "wall_s": 0.0861, + "speedup_vs_sequential": 15.02 + }, + { + "file_concurrency": 32, + "wall_s": 0.044, + "speedup_vs_sequential": 29.41 + } + ] + }, + { + "n_objects": 128, + "latency_ms": 40.0, + "sequential_wall_s": 5.1629, + "batch": [ + { + "file_concurrency": 1, + "wall_s": 5.1708, + "speedup_vs_sequential": 1.0 + }, + { + "file_concurrency": 8, + "wall_s": 0.6494, + "speedup_vs_sequential": 7.95 + }, + { + "file_concurrency": 16, + "wall_s": 0.33, + "speedup_vs_sequential": 15.65 + }, + { + "file_concurrency": 32, + "wall_s": 0.1693, + "speedup_vs_sequential": 30.49 + } + ] + }, + { + "n_objects": 512, + "latency_ms": 40.0, + "sequential_wall_s": 20.6571, + "batch": [ + { + "file_concurrency": 1, + "wall_s": 20.6739, + "speedup_vs_sequential": 1.0 + }, + { + "file_concurrency": 8, + "wall_s": 2.5874, + "speedup_vs_sequential": 7.98 + }, + { + "file_concurrency": 16, + "wall_s": 1.294, + "speedup_vs_sequential": 15.96 + }, + { + "file_concurrency": 32, + "wall_s": 0.6502, + "speedup_vs_sequential": 31.77 + } + ] + } + ], + "params": { + "objects": [ + 32, + 128, + 512 + ], + "file_concurrency": [ + 1, + 8, + 16, + 32 + ], + "latency_ms": 40.0, + "reuse_ops": 200, + "repeats": 5, + "json_out": "C:\\Users\\kkapu\\OneDrive\\Desktop\\thaw\\site\\receipts\\2026-06-24_s3_batch.json" + } +} \ No newline at end of file diff --git a/tests/test_cloud.py b/tests/test_cloud.py index 18bf772..69dabe6 100644 --- a/tests/test_cloud.py +++ b/tests/test_cloud.py @@ -12,6 +12,7 @@ import os import sys import tempfile +import time import unittest from unittest.mock import MagicMock, patch @@ -37,10 +38,15 @@ def mock_aws(fn): from thaw_common.cloud import ( is_remote, resolve_snapshot_path, + resolve_snapshots, upload_snapshot, + upload_snapshots, + reset_s3_client_cache, _cache_path, _parse_s3, + _s3_client, ) +import thaw_common.cloud as cloud def _mock_s3_with_payload(payload: bytes) -> MagicMock: @@ -291,6 +297,16 @@ class TestS3RoundTripMoto(unittest.TestCase): error mapping are validated for real. """ + def setUp(self): + # Each method runs under its own @mock_aws backend. Drop any cached + # S3 client between methods so a client built under one method's mock + # context can't leak into the next (defense-in-depth for the shared + # client cache; keeps these green regardless of test ordering). + reset_s3_client_cache() + + def tearDown(self): + reset_s3_client_cache() + @mock_aws def test_round_trip_small(self): bucket = "thaw-test-bucket" @@ -390,5 +406,294 @@ def test_register_still_rejects_bad_local(self): pool.register("bad", "/definitely/not/a/real/path.thaw") +class TestSharedClientCache(unittest.TestCase): + """The S3 client is built once and reused, keyed by (pid, pool_size). + + These tests stub _build_s3_client so they need neither boto3 credentials + nor a network — they assert the CACHING contract, not boto3 itself. + """ + + def setUp(self): + reset_s3_client_cache() + + def tearDown(self): + reset_s3_client_cache() + + def test_client_is_reused_across_calls(self): + builds = {"n": 0} + + def fake_build(pool_size): + builds["n"] += 1 + return MagicMock(name=f"client-{builds['n']}") + + with patch.object(cloud, "_build_s3_client", side_effect=fake_build): + c1 = _s3_client() + c2 = _s3_client() + c3 = _s3_client() + + self.assertIs(c1, c2) + self.assertIs(c2, c3) + self.assertEqual(builds["n"], 1, "client should be built exactly once") + + def test_reset_forces_rebuild(self): + with patch.object( + cloud, "_build_s3_client", side_effect=lambda p: MagicMock() + ) as build: + _s3_client() + _s3_client() + self.assertEqual(build.call_count, 1) + reset_s3_client_cache() + _s3_client() + self.assertEqual(build.call_count, 2) + + def test_distinct_pool_size_distinct_client(self): + with patch.object( + cloud, "_build_s3_client", side_effect=lambda p: MagicMock() + ) as build: + _s3_client(concurrency=32) # pool 36 + _s3_client(concurrency=32) # cache hit + _s3_client(concurrency=128) # pool 132 -> distinct key + self.assertEqual(build.call_count, 2) + + def test_pool_size_matches_concurrency(self): + seen = {} + + def fake_build(pool_size): + seen["pool"] = pool_size + return MagicMock() + + with patch.object(cloud, "_build_s3_client", side_effect=fake_build): + _s3_client(concurrency=32) + # Same pool math as the original per-call implementation: conc + 4, min 16. + self.assertEqual(seen["pool"], 36) + + def test_fork_rebuilds_client(self): + """A different pid must yield a different client (botocore isn't fork-safe).""" + with patch.object( + cloud, "_build_s3_client", side_effect=lambda p: MagicMock() + ) as build: + with patch.object(cloud.os, "getpid", return_value=1000): + _s3_client() + _s3_client() + self.assertEqual(build.call_count, 1) + # Simulate being in a forked child: pid changes -> rebuild. + with patch.object(cloud.os, "getpid", return_value=2000): + _s3_client() + self.assertEqual(build.call_count, 2) + + +class TestResolveSnapshotsBatch(unittest.TestCase): + """Bounded-concurrency batch download orchestration. + + Patches the single-file resolve_snapshot_path so these tests exercise the + fan-out/ordering/error-propagation logic without boto3 or POSIX pwrite. + """ + + def test_empty_list(self): + self.assertEqual(resolve_snapshots([]), []) + + def test_all_local_passthrough_no_threads(self): + # Local paths must pass through untouched and must NOT call the S3 path. + with patch.object(cloud, "resolve_snapshot_path") as rsp: + out = resolve_snapshots(["/a.thaw", "", None, "./b.thaw"]) + self.assertEqual(out, ["/a.thaw", "", None, "./b.thaw"]) + rsp.assert_not_called() + + def test_order_preserved_under_concurrency(self): + # Make earlier items take LONGER so completion order != input order; + # the result list must still match input order. + def fake_resolve(uri, cache_dir=None, force=False): + idx = int(uri.rsplit("/", 1)[1].split(".")[0]) + time.sleep(0.05 * (5 - idx)) # item 0 sleeps longest + return f"/cache/{idx}" + + uris = [f"s3://b/{i}.thaw" for i in range(5)] + with patch.object(cloud, "resolve_snapshot_path", side_effect=fake_resolve): + out = resolve_snapshots(uris, max_files=5) + self.assertEqual(out, [f"/cache/{i}" for i in range(5)]) + + def test_runs_concurrently(self): + # 6 items each sleeping 0.2s: sequential would be >1.0s, 6-way ~0.2s. + def slow_resolve(uri, cache_dir=None, force=False): + time.sleep(0.2) + return uri.replace("s3://b/", "/cache/") + + uris = [f"s3://b/{i}.thaw" for i in range(6)] + with patch.object(cloud, "resolve_snapshot_path", side_effect=slow_resolve): + t0 = time.monotonic() + resolve_snapshots(uris, max_files=6) + elapsed = time.monotonic() - t0 + self.assertLess(elapsed, 0.8, f"expected concurrent (~0.2s), took {elapsed:.2f}s") + + def test_respects_max_files_bound(self): + # With max_files=2, 4 items of 0.2s each => 2 waves => ~0.4s, not ~0.2s. + def slow_resolve(uri, cache_dir=None, force=False): + time.sleep(0.2) + return uri + uris = [f"s3://b/{i}.thaw" for i in range(4)] + with patch.object(cloud, "resolve_snapshot_path", side_effect=slow_resolve): + t0 = time.monotonic() + resolve_snapshots(uris, max_files=2) + elapsed = time.monotonic() - t0 + self.assertGreaterEqual(elapsed, 0.35, f"max_files=2 should serialize into waves, took {elapsed:.2f}s") + + def test_mixed_local_and_remote(self): + def fake_resolve(uri, cache_dir=None, force=False): + return "/cache/" + uri.rsplit("/", 1)[1] + items = ["/local/a.thaw", "s3://b/x.thaw", "/local/c.thaw", "s3://b/y.thaw"] + with patch.object(cloud, "resolve_snapshot_path", side_effect=fake_resolve): + out = resolve_snapshots(items) + self.assertEqual( + out, ["/local/a.thaw", "/cache/x.thaw", "/local/c.thaw", "/cache/y.thaw"] + ) + + def test_first_error_propagates(self): + def fake_resolve(uri, cache_dir=None, force=False): + if uri.endswith("2.thaw"): + raise FileNotFoundError("thaw S3 download failed: object not found at " + uri) + time.sleep(0.05) + return "/cache/ok" + uris = [f"s3://b/{i}.thaw" for i in range(4)] + with patch.object(cloud, "resolve_snapshot_path", side_effect=fake_resolve): + with self.assertRaises(FileNotFoundError): + resolve_snapshots(uris, max_files=4) + + def test_passes_cache_dir_and_force(self): + seen = [] + def fake_resolve(uri, cache_dir=None, force=False): + seen.append((cache_dir, force)) + return "/cache/x" + with patch.object(cloud, "resolve_snapshot_path", side_effect=fake_resolve): + resolve_snapshots(["s3://b/x.thaw"], cache_dir="/custom", force=True) + self.assertEqual(seen, [("/custom", True)]) + + def test_duplicate_uris_downloaded_once(self): + # Duplicate URIs must download exactly once (they share one cache file; + # racing two .part writes/renames would corrupt it) and every slot must + # receive the same resolved path, in order. + calls = [] + def fake_resolve(uri, cache_dir=None, force=False): + calls.append(uri) + return "/cache/" + uri.rsplit("/", 1)[1] + uris = [ + "s3://b/dup.thaw", + "s3://b/other.thaw", + "s3://b/dup.thaw", + "s3://b/dup.thaw", + ] + with patch.object(cloud, "resolve_snapshot_path", side_effect=fake_resolve): + out = resolve_snapshots(uris, max_files=8) + self.assertEqual(out, [ + "/cache/dup.thaw", "/cache/other.thaw", + "/cache/dup.thaw", "/cache/dup.thaw", + ]) + self.assertEqual(sorted(calls), ["s3://b/dup.thaw", "s3://b/other.thaw"]) + self.assertEqual(calls.count("s3://b/dup.thaw"), 1) + + +class TestUploadSnapshotsBatch(unittest.TestCase): + def test_empty_noop(self): + # Must not raise and must not touch the S3 path. + with patch.object(cloud, "upload_snapshot") as up: + upload_snapshots([]) + up.assert_not_called() + + def test_all_uploaded(self): + calls = [] + with patch.object(cloud, "upload_snapshot", side_effect=lambda lp, uri: calls.append((lp, uri))): + upload_snapshots([("/a", "s3://b/a"), ("/c", "s3://b/c")]) + self.assertEqual(sorted(calls), [("/a", "s3://b/a"), ("/c", "s3://b/c")]) + + def test_runs_concurrently(self): + def slow_upload(lp, uri): + time.sleep(0.2) + pairs = [(f"/f{i}", f"s3://b/{i}") for i in range(6)] + with patch.object(cloud, "upload_snapshot", side_effect=slow_upload): + t0 = time.monotonic() + upload_snapshots(pairs, max_files=6) + elapsed = time.monotonic() - t0 + self.assertLess(elapsed, 0.8, f"expected concurrent (~0.2s), took {elapsed:.2f}s") + + def test_first_error_propagates(self): + def maybe_fail(lp, uri): + if uri.endswith("bad"): + raise PermissionError("access denied") + pairs = [("/a", "s3://b/ok"), ("/b", "s3://b/bad"), ("/c", "s3://b/ok2")] + with patch.object(cloud, "upload_snapshot", side_effect=maybe_fail): + with self.assertRaises(PermissionError): + upload_snapshots(pairs, max_files=3) + + +@unittest.skipUnless(HAVE_MOTO, "moto not installed") +@unittest.skipUnless(hasattr(os, "pwrite"), "ranged download path requires POSIX os.pwrite") +class TestBatchRoundTripMoto(unittest.TestCase): + """End-to-end batch round trip against moto. POSIX-only: the download + write path uses os.pwrite. Runs in CI (Ubuntu); skipped on Windows.""" + + @mock_aws + def test_resolve_and_upload_many(self): + reset_s3_client_cache() + bucket = "thaw-batch-bucket" + s3 = boto3.client("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket) + + payloads = {f"obj{i}.thaw": os.urandom(2048 + i) for i in range(6)} + srcs = [] + try: + pairs = [] + for name, data in payloads.items(): + fd, p = tempfile.mkstemp(suffix=".thaw") + with os.fdopen(fd, "wb") as f: + f.write(data) + srcs.append(p) + pairs.append((p, f"s3://{bucket}/{name}")) + + # Batch upload, then batch download, verify bit-identical + order. + upload_snapshots(pairs) + with tempfile.TemporaryDirectory() as cache_dir: + uris = [f"s3://{bucket}/{n}" for n in payloads] + locals_ = resolve_snapshots(uris, cache_dir=cache_dir) + self.assertEqual(len(locals_), len(uris)) + for (name, data), local in zip(payloads.items(), locals_): + with open(local, "rb") as f: + self.assertEqual(f.read(), data, f"mismatch for {name}") + finally: + for p in srcs: + if os.path.exists(p): + os.unlink(p) + + @mock_aws + def test_batch_surfaces_missing_object(self): + reset_s3_client_cache() + bucket = "thaw-batch-bucket2" + boto3.client("s3", region_name="us-east-1").create_bucket(Bucket=bucket) + with tempfile.TemporaryDirectory() as cache_dir: + uris = [f"s3://{bucket}/missing-{i}.thaw" for i in range(3)] + with self.assertRaises(FileNotFoundError): + resolve_snapshots(uris, cache_dir=cache_dir) + + @mock_aws + def test_duplicate_uris_no_race(self): + # Real-path regression for the duplicate-URI race: without dedup, two + # threads resolving the same URI race the same .part rename (torn cache + # on POSIX, FileExistsError on Windows). With dedup it downloads once + # and every slot gets the same byte-identical file. + reset_s3_client_cache() + bucket = "thaw-dup-bucket" + key = "shared.thaw" + payload = os.urandom(8192) + s3 = boto3.client("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket) + s3.put_object(Bucket=bucket, Key=key, Body=payload) + + uri = f"s3://{bucket}/{key}" + with tempfile.TemporaryDirectory() as cache_dir: + out = resolve_snapshots([uri] * 5, cache_dir=cache_dir, max_files=8) + self.assertEqual(len(out), 5) + self.assertEqual(len(set(out)), 1) # all slots -> same cache file + with open(out[0], "rb") as f: + self.assertEqual(f.read(), payload) + + if __name__ == "__main__": unittest.main()