diff --git a/.github/markdown-links-config.json b/.github/markdown-links-config.json index e4d81abb5..376a5636f 100644 --- a/.github/markdown-links-config.json +++ b/.github/markdown-links-config.json @@ -8,6 +8,9 @@ }, { "pattern": "^https://codecov.io/gh/Lightning-AI/litData/graph/badge.svg" + }, + { + "pattern": "^https://devblog.pytorchlightning.ai/" } ], "httpHeaders": [ diff --git a/benchmarks/ADAPTIVE_CONCURRENCY.md b/benchmarks/ADAPTIVE_CONCURRENCY.md index cb023aea0..58b5f2ffb 100644 --- a/benchmarks/ADAPTIVE_CONCURRENCY.md +++ b/benchmarks/ADAPTIVE_CONCURRENCY.md @@ -68,6 +68,40 @@ Confirm @ `ba9da13`: interleaved n=3, `max(≥300 batches, ≥30s)`, **w=24 p=0* **Verdict: (a)** — before ≈3.7k (not ~5.5k wrong-tree). Frame high-w as robustness; do **not** headline unverifiable +53%. Artifact: `benchmarks/results/raw_before_vs_after.ba9da13.1785268543.json`. +## Stateful Recovery State Machine & Deadband Control + +The adaptive controller incorporates a 3-state feedback controller to prevent budget jitter: + +| State | Latency Condition | Action | +| :------------ | :--------------------------------------------------------------------------------------- | :----------------------------------------------------------------------------------------------------- | +| **Healthy** | (L\_{\\text{obs}} \\le L\_{\\text{target}}) (40 ms) | Recover backoff factor gradually: (f \\leftarrow \\min(1.0, f + \\alpha\_{\\text{recovery}}(1.0 - f))) | +| **Deadband** | (L\_{\\text{target}} < L\_{\\text{obs}} \\le 1.1 \\times L\_{\\text{target}}) (40–44 ms) | Hold current backoff factor constant (no recovery, no extra backoff) | +| **Congested** | (L\_{\\text{obs}} > 1.1 \\times L\_{\\text{target}}) (>44 ms) | Immediate backoff: (f \\leftarrow \\min(f, L\_{\\text{target}} / L\_{\\text{obs}})) | + +## Environment Variable Configuration Reference + +All controller bounds and operational knobs are configurable via validated environment variables: + +| Environment Variable | Default | Range / Constraints | Description | +| :------------------------------------------- | :--------------------- | :-------------------------------------- | :---------------------------------------------------------- | +| `LITDATA_ASSUMED_BANDWIDTH_BPS` | `104857600` (100 MB/s) | (> 0) | Default aggregate network bandwidth fallback | +| `LITDATA_ASSUMED_REQUEST_LATENCY_S` | `0.040` (40 ms) | (> 0.0) | Target baseline GET request latency | +| `LITDATA_ASSUMED_REQUEST_RATE` | `6000.0` req/s | (\\ge 0.0) | Baseline request rate for Little's law model | +| `LITDATA_SINGLE_PROCESS_CONCURRENCY_CAP` | `128` | (\\ge 1) | Maximum single-process permit cap | +| `LITDATA_AGGREGATE_CONCURRENCY_BUDGET_CAP` | `512` | (\\ge 1) | Maximum aggregate multi-worker permit cap | +| `LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR` | `32` | (1 \\le \\text{floor} \\le \\text{cap}) | Adaptive dynamic concurrency budget floor | +| `LITDATA_BACKOFF_RECOVERY_ALPHA` | `0.1` | (0.0 < \\alpha \\le 1.0) | Recovery EMA factor for healthy state | +| `LITDATA_MIN_EMPIRICAL_SAMPLES` | `5` | (\\ge 1) | Minimum observations required before applying empirical EMA | +| `LITDATA_PERMIT_REFRESH_INTERVAL` | `10` | (\\ge 1) | Iteration interval between dynamic permit recalculations | + +## Adaptive Concurrency Stress Benchmark + +To run the stress benchmark comparing aggregate permits, per-worker permits, BPS EMA, latency EMA, and throughput across worker counts: + +```bash +PYTHONPATH=src .venv/bin/python benchmarks/bench_raw_adaptive_concurrency.py +``` + ## Acceptance (future adaptive) Beats **default** static everywhere; never loses by more than run-to-run noise; removes the w×p tuning matrix from the user’s cognitive load. “Beats tuned static” is the wrong bar — tuned static ties it at best per configuration. diff --git a/benchmarks/bench_raw_adaptive_concurrency.py b/benchmarks/bench_raw_adaptive_concurrency.py new file mode 100644 index 000000000..e749d2753 --- /dev/null +++ b/benchmarks/bench_raw_adaptive_concurrency.py @@ -0,0 +1,401 @@ +#!/usr/bin/env python3 +"""Behavioral Adaptive Concurrency Stress Benchmark. + +Demonstrates cause-and-effect chain: + simulated network condition -> real async delay -> downloader boundary + -> BandwidthTracker observation -> class-gated sample count + -> stateful backoff/recovery -> dynamic semaphore budget -> runtime concurrency. + +Calculates per-worker permit allocations across worker count configurations w in [4, 8, 16, 24]. +""" + +import asyncio +import os +import sys +import tempfile +import time +from typing import Any + +import pytest + +sys.path.insert(0, os.path.abspath("src")) + +from litdata.raw.dataset import ( + StreamingRawDataset, + _aggregate_concurrency_budget, + _DynamicSemaphore, + _effective_concurrency, +) +from litdata.streaming.downloader import Downloader + + +class FakeDownloader(Downloader): + """Deterministic network simulator at the Downloader ABC boundary.""" + + def __init__( + self, + latency_s: float = 0.010, + bandwidth_bps: float = 100 * 1024 * 1024, + error_rate: float = 0.0, + default_size: int = 30 * 1024, + **kwargs: Any, + ): + """Initialize FakeDownloader with latency, bandwidth, and error rates.""" + super().__init__("", "", [], **kwargs) + self.latency_s = latency_s + self.bandwidth_bps = bandwidth_bps + self.error_rate = error_rate + self.default_size = default_size + self._call_count = 0 + + def _extract_requested_size(self, remote_filepath: str) -> int: + """Parse size from remote_filepath query parameter or return default_size.""" + if "?size=" in remote_filepath: + try: + return int(remote_filepath.split("?size=")[1]) + except ValueError: + pass + return self.default_size + + async def adownload_fileobj(self, remote_filepath: str) -> bytes: + """Simulate async object download with dual-component latency and payload transfer delay.""" + self._call_count += 1 + # Deterministic 429 rate-limiting simulation when error_rate > 0 + if self.error_rate > 0 and (self._call_count % 2 == 0): + raise RuntimeError("HTTP 429 Too Many Requests") + req_size = self._extract_requested_size(remote_filepath) + transfer_s = req_size / max(1.0, self.bandwidth_bps) + await asyncio.sleep(self.latency_s + transfer_s) + return b"x" * req_size + + def download_bytes(self, remote_filepath: str, offset: int, length: int, local_chunkpath: str) -> bytes: + """Simulate sync ranged download with dual-component timing delay.""" + self._call_count += 1 + if self.error_rate > 0 and (self._call_count % 2 == 0): + raise RuntimeError("HTTP 429 Too Many Requests") + transfer_s = length / max(1.0, self.bandwidth_bps) + time.sleep(self.latency_s + transfer_s) + data = b"x" * length + with open(local_chunkpath, "wb") as f: + f.write(data) + return data + + +def run_benchmark() -> None: + """Execute behavioral adaptive concurrency stress benchmark across worker matrix.""" + header_title = "BEHAVIORAL ADAPTIVE CONCURRENCY STRESS BENCHMARK" + print("=" * 115) + print(f"{header_title:^115}") + print("=" * 115) + print("Note: BPS (MB/s) = 0.00 in latency-only rows (<64 KiB GETs) is expected class-gated sample isolation.") + print( + f"{'Phase':<12} | {'Workers':<8} | {'Budget':<8} | {'Permits/W':<10} | " + f"{'BPS (MB/s)':<12} | {'Lat (ms)':<10} | {'Tput (MB/s)':<12} | {'Errors':<8}" + ) + print("-" * 115) + + worker_counts = [4, 8, 16, 24] + latency_test_size = 30 * 1024 # <64 KiB (latency population only) + medium_test_size = 128 * 1024 # 64-256 KiB (both latency and bandwidth populations) + bandwidth_test_size = 500 * 1024 # >=256 KiB (bandwidth population only) + + for w in worker_counts: + with tempfile.TemporaryDirectory() as tmp_dir: + sample_file = os.path.join(tmp_dir, "sample.bin") + with open(sample_file, "wb") as f: + f.write(b"x" * 1000) + + ds = StreamingRawDataset( + input_dir=tmp_dir, + cache_dir=tmp_dir, + ) + + # Phase 1: Healthy Baseline (10 ms latency, 100 MB/s bandwidth) + fake_dl_healthy = FakeDownloader( + latency_s=0.010, + bandwidth_bps=100 * 1024 * 1024, + default_size=latency_test_size, + ) + + # Correction 2: Direct timing fidelity check on FakeDownloader using tolerance + t0_direct = time.monotonic() + asyncio.run(fake_dl_healthy.adownload_fileobj(f"s3://mock-bucket/file.bin?size={latency_test_size}")) + direct_dur = time.monotonic() - t0_direct + expected_dur = fake_dl_healthy.latency_s + (latency_test_size / fake_dl_healthy.bandwidth_bps) + tolerance = max(0.015, expected_dur * 0.5) + assert abs(direct_dur - expected_dur) <= tolerance, ( + f"Direct timing ({direct_dur:.4f}s) must be within tolerance of expected ({expected_dur:.4f}s)" + ) + + async def _run_phase1(): + loop = asyncio.get_running_loop() + ds.cache_manager._downloader = fake_dl_healthy + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + tasks = [ + ds.cache_manager._fetch_bytes( + f"s3://mock-bucket/file.bin?size={latency_test_size}", size=latency_test_size + ) + for _ in range(6) + ] + await asyncio.gather(*tasks) + + t0 = time.monotonic() + asyncio.run(_run_phase1()) + elapsed = time.monotonic() - t0 + + tracker = ds.cache_manager._bandwidth_tracker + budget_healthy = _aggregate_concurrency_budget(latency_test_size, tracker=tracker) + permits_w_h = _effective_concurrency( + None, num_workers=w, median_file_bytes=latency_test_size, tracker=tracker + ) + bps_h, lat_h, _, lat_cnt_h = tracker.get_metrics() + bps_mbs_h = (bps_h / (1024 * 1024)) if bps_h else 0.0 + lat_ms_h = (lat_h * 1000) if lat_h else 0.0 + tput_h = (6 * latency_test_size / (1024 * 1024)) / elapsed if elapsed > 0 else 0.0 + + print( + f"{'Healthy':<12} | {w:<8} | {budget_healthy:<8} | {permits_w_h:<10} | " + f"{bps_mbs_h:<12.2f} | {lat_ms_h:<10.2f} | {tput_h:<12.2f} | {0:<8}" + ) + + assert lat_cnt_h >= 5, "Small GETs must update latency sample count" + + # Phase 2: Congested / Slow Link (200 ms latency = 5x target) + fake_dl_congested = FakeDownloader( + latency_s=0.200, + bandwidth_bps=10 * 1024 * 1024, + default_size=latency_test_size, + ) + + async def _run_phase2(): + loop = asyncio.get_running_loop() + ds.cache_manager._downloader = fake_dl_congested + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + tasks = [ + ds.cache_manager._fetch_bytes( + f"s3://mock-bucket/file.bin?size={latency_test_size}", size=latency_test_size + ) + for _ in range(6) + ] + await asyncio.gather(*tasks) + + t0 = time.monotonic() + asyncio.run(_run_phase2()) + elapsed = time.monotonic() - t0 + + budget_congested = _aggregate_concurrency_budget(latency_test_size, tracker=tracker) + permits_w_c = _effective_concurrency( + None, num_workers=w, median_file_bytes=latency_test_size, tracker=tracker + ) + bps_c, lat_c, _, _ = tracker.get_metrics() + bps_mbs_c = (bps_c / (1024 * 1024)) if bps_c else 0.0 + lat_ms_c = (lat_c * 1000) if lat_c else 0.0 + tput_c = (6 * latency_test_size / (1024 * 1024)) / elapsed if elapsed > 0 else 0.0 + + print( + f"{'Congested':<12} | {w:<8} | {budget_congested:<8} | {permits_w_c:<10} | " + f"{bps_mbs_c:<12.2f} | {lat_ms_c:<10.2f} | {tput_c:<12.2f} | {0:<8}" + ) + + assert budget_congested < budget_healthy, ( + f"Congestion budget ({budget_congested}) must be < healthy ({budget_healthy})" + ) + + # Phase 3: Operational Rate-Limiting (Deterministic HTTP 429 Simulation) + fake_dl_429 = FakeDownloader( + latency_s=0.010, + bandwidth_bps=100 * 1024 * 1024, + error_rate=0.5, + default_size=latency_test_size, + ) + + errors_429 = 0 + + async def _run_phase3(): + nonlocal errors_429 + loop = asyncio.get_running_loop() + ds.cache_manager._downloader = fake_dl_429 + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + tasks = [ + ds.cache_manager._fetch_bytes( + f"s3://mock-bucket/file.bin?size={latency_test_size}", size=latency_test_size + ) + for _ in range(10) + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + for r in results: + if isinstance(r, Exception): + errors_429 += 1 + + t0 = time.monotonic() + asyncio.run(_run_phase3()) + elapsed = time.monotonic() - t0 + + budget_429 = _aggregate_concurrency_budget(latency_test_size, tracker=tracker) + permits_w_429 = _effective_concurrency( + None, num_workers=w, median_file_bytes=latency_test_size, tracker=tracker + ) + bps_429, lat_429_v, _, _ = tracker.get_metrics() + bps_mbs_429 = (bps_429 / (1024 * 1024)) if bps_429 else 0.0 + lat_ms_429 = (lat_429_v * 1000) if lat_429_v else 0.0 + tput_429 = (10 * latency_test_size / (1024 * 1024)) / elapsed if elapsed > 0 else 0.0 + + print( + f"{'Rate-Limit':<12} | {w:<8} | {budget_429:<8} | {permits_w_429:<10} | " + f"{bps_mbs_429:<12.2f} | {lat_ms_429:<10.2f} | {tput_429:<12.2f} | {errors_429:<8}" + ) + + # Phase 4: Multi-Step Gradual Recovery (Healthy latency restored) + fake_dl_recovery = FakeDownloader( + latency_s=0.010, + bandwidth_bps=100 * 1024 * 1024, + default_size=latency_test_size, + ) + + recovery_budgets = [budget_congested] + + async def _run_phase4_step(): + loop = asyncio.get_running_loop() + ds.cache_manager._downloader = fake_dl_recovery + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + tasks = [ + ds.cache_manager._fetch_bytes( + f"s3://mock-bucket/file.bin?size={latency_test_size}", size=latency_test_size + ) + for _ in range(2) + ] + await asyncio.gather(*tasks) + + for step in range(3): + asyncio.run(_run_phase4_step()) + b_step = _aggregate_concurrency_budget(latency_test_size, tracker=tracker) + recovery_budgets.append(b_step) + + budget_recovered = recovery_budgets[-1] + permits_w_rec = _effective_concurrency( + None, num_workers=w, median_file_bytes=latency_test_size, tracker=tracker + ) + bps_rec, lat_rec, _, _ = tracker.get_metrics() + bps_mbs_rec = (bps_rec / (1024 * 1024)) if bps_rec else 0.0 + lat_ms_rec = (lat_rec * 1000) if lat_rec else 0.0 + + print( + f"{'Recovered':<12} | {w:<8} | {budget_recovered:<8} | {permits_w_rec:<10} | " + f"{bps_mbs_rec:<12.2f} | {lat_ms_rec:<10.2f} | {'N/A':<12} | {0:<8}" + ) + + for prev_b, curr_b in zip(recovery_budgets, recovery_budgets[1:]): + assert curr_b >= prev_b, f"Recovery sequence must be non-decreasing: {recovery_budgets}" + + assert budget_recovered > budget_congested, ( + f"Recovered budget ({budget_recovered}) must exceed congested budget ({budget_congested})" + ) + + assert budget_recovered <= budget_healthy, ( + f"Recovered budget ({budget_recovered}) must not exceed baseline ({budget_healthy})" + ) + + assert any(b < budget_healthy for b in recovery_budgets[1:-1]), ( + f"Recovery sequence ({recovery_budgets}) must be gradual and not jump immediately " + f"to healthy baseline ({budget_healthy})" + ) + + # Explicit Medium-Object Test (64 KiB <= size < 256 KiB) -> Increments BOTH counters + fake_dl_medium = FakeDownloader( + latency_s=0.010, + bandwidth_bps=50 * 1024 * 1024, + default_size=medium_test_size, + ) + + async def _run_medium_test(): + loop = asyncio.get_running_loop() + ds.cache_manager._downloader = fake_dl_medium + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + tasks = [ + ds.cache_manager._fetch_bytes( + f"s3://mock-bucket/file.bin?size={medium_test_size}", size=medium_test_size + ) + for _ in range(1) + ] + await asyncio.gather(*tasks) + + _, _, bps_before, lat_before = tracker.get_metrics() + asyncio.run(_run_medium_test()) + _, _, bps_after, lat_after = tracker.get_metrics() + + assert bps_after == bps_before + 1, "Medium GET (128 KiB) must increment bps_sample_count" + assert lat_after == lat_before + 1, "Medium GET (128 KiB) must increment lat_sample_count" + + # Large Object Test (>=256 KiB): Bandwidth Population Isolation + fake_dl_large = FakeDownloader( + latency_s=0.010, + bandwidth_bps=50 * 1024 * 1024, + default_size=bandwidth_test_size, + ) + + async def _run_large_test(): + loop = asyncio.get_running_loop() + ds.cache_manager._downloader = fake_dl_large + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + tasks = [ + ds.cache_manager._fetch_bytes( + f"s3://mock-bucket/file.bin?size={bandwidth_test_size}", size=bandwidth_test_size + ) + for _ in range(5) + ] + await asyncio.gather(*tasks) + + asyncio.run(_run_large_test()) + _, _, bps_cnt_large, _ = tracker.get_metrics() + + assert bps_cnt_large >= 5, "Large GETs must update bandwidth sample count" + + ds.cache_manager.reset_runtime_state() + print("-" * 115) + + # Correction 1: Focused in-flight semaphore lifecycle check + # Proves 32 initial holders remain valid on downscale to 16, release 16 -> 17th acquire blocks, + # release remaining -> new acquisitions succeed up to 16 + + async def _run_semaphore_invariant_check(): + sem = _DynamicSemaphore(32) + # Step 1: Acquire all 32 initial permits (32 holders in-flight) + for _ in range(32): + await sem.acquire() + + # Step 2: Target reduced to 16 mid-flight + sem.update_target(16) + assert sem.target_permits == 16 + + # Step 3: Existing 32 holders remain valid, release 16 of them + for _ in range(16): + sem.release() + + # Step 4: 16 holders still remain active — new acquisition attempt must block/timeout + with pytest.raises((asyncio.TimeoutError, TimeoutError)): + await asyncio.wait_for(sem.acquire(), timeout=0.05) + + # Step 5: Release remaining 16 holders + for _ in range(16): + sem.release() + + # Step 6: Verify 16 new acquisitions succeed cleanly + for _ in range(16): + await sem.acquire() + for _ in range(16): + sem.release() + + asyncio.run(_run_semaphore_invariant_check()) + + print("=" * 115) + print("Behavioral adaptive concurrency benchmark completed successfully; all assertions passed.") + + +if __name__ == "__main__": + run_benchmark() diff --git a/src/litdata/CHANGELOG.md b/src/litdata/CHANGELOG.md index b39591fed..3507019af 100644 --- a/src/litdata/CHANGELOG.md +++ b/src/litdata/CHANGELOG.md @@ -48,6 +48,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - `StreamingDataset.subset(indices)` for an index/slice view (same ROI logic as `train_test_split`). - Typed media wrappers matching serializers (`Text`, `Audio`, `Video`, `Image`, `Jpeg`, `JpegArray`, `Pil`, `Tiff`, `File`, `Mesh`, `Pdf`, `Nifti`, `Tensor`, `Graph`) so `optimize` can tell a filepath from a caption. `Tensor(array=)` is a pytree leaf for `TensorSerializer` / `NoHeaderTensorSerializer`; 1-D tensors expose `.shape` for `TokensLoader.encode_data`. Native payloads: `Audio(array=, sampling_rate=)`, `Image(array=, quality=95, format="jpeg")`, `Video(array=, fps=)`, `Mesh(mesh=)`, `Pdf(pdf=)`, `Nifti(array=, affine=)`. Audio/Video also re-encode torchcodec decoders via `_hf_encoded` / `_litdata_encoded`. Image encode downcasts arrays and keeps the PIL native format. Decode is **bytes → tensor** via torchvision (no PIL). EXIF orientation uses PIL only when a JPEG APP1 Exif marker is present. Bare `*.jpg` / `*.png` paths are claimed by `ImageSerializer` so `optimize` can return tensors. Audio decoders support `audio["array"]` / `audio["sampling_rate"]`. - `ParquetReader` / `ParquetLoader` column projection. `ParquetReader` also takes PyArrow `filters` and reshard by row group instead of loading each file. Low-memory `ParquetLoader` returns a row from the Arrow table (no Polars copy of the row group). +- Size-aware adaptive dynamic concurrency control system for `StreamingRawDataset` with process-local EMA feedback (`BandwidthTracker`), Little's Law baseline capacity modeling, stateful congestion backoff with 10% deadband recovery, class-gated sample isolation, transfer-subtracted small-GET latency estimation, dynamic permit management (`_DynamicSemaphore`), and environment variable bounds validation. ### Changed diff --git a/src/litdata/raw/dataset.py b/src/litdata/raw/dataset.py index cf0a3b7de..8a6eaa699 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -108,6 +108,289 @@ # the bandwidth arm alone sizes the budget. Distinct from ``_HEDGE_MAX_BYTES`` # (8 MiB duplicate-GET hedge policy). _LATENCY_MODEL_MAX_MEDIAN_BYTES = 1024 * 1024 +_LATENCY_OBSERVATION_MAX_BYTES = 256 * 1024 +_BANDWIDTH_OBSERVATION_MIN_BYTES = 64 * 1024 +_MIN_EMPIRICAL_SAMPLES = 5 +_LOW_BANDWIDTH_THRESHOLD_BPS = 10 * 1024 * 1024 # 10 MiB/s policy threshold +_HIGH_LATENCY_THRESHOLD_S = 0.100 # 100 ms policy threshold +# Minimum estimated request latency (seconds). Prevents the transfer-subtracted +# latency estimate from collapsing to zero or going negative on fast / cached +# responses. 1 ms is a reasonable floor: well below any real WAN RTT, but +# large enough to keep EMA arithmetic well-conditioned. +_LATENCY_EPSILON_S = 0.001 +_LATENCY_RTT_EPSILON = _LATENCY_EPSILON_S + + +def _env_float(name: str, default: float, min_val: float = 0.0, max_val: float | None = None) -> float: + val_str = os.getenv(name) + if val_str is None: + return default + try: + val = float(val_str) + if val <= min_val or (max_val is not None and val > max_val): + logger.warning("Environment variable %s=%r out of bounds; using default %g", name, val_str, default) + return default + return val + except ValueError: + logger.warning("Environment variable %s=%r invalid float; using default %g", name, val_str, default) + return default + + +def _env_int(name: str, default: int, min_val: int = 1) -> int: + val_str = os.getenv(name) + if val_str is None: + return default + try: + val = int(val_str) + if val < min_val: + logger.warning( + "Environment variable %s=%r must be >= %d; using default %d", + name, + val_str, + min_val, + default, + ) + return default + return val + except ValueError: + logger.warning("Environment variable %s=%r invalid integer; using default %d", name, val_str, default) + return default + + +def _get_assumed_aggregate_bandwidth_bps() -> int: + return _env_int("LITDATA_ASSUMED_BANDWIDTH_BPS", _ASSUMED_AGGREGATE_BANDWIDTH_BPS, min_val=1) + + +def _get_assumed_request_rate() -> float: + return _env_float("LITDATA_ASSUMED_REQUEST_RATE", _ASSUMED_REQUEST_RATE, min_val=0.0) + + +def _get_assumed_request_latency_s() -> float: + return _env_float("LITDATA_ASSUMED_REQUEST_LATENCY_S", _ASSUMED_REQUEST_LATENCY_S, min_val=0.0) + + +def _get_default_median_file_bytes() -> int: + return _env_int("LITDATA_DEFAULT_MEDIAN_FILE_BYTES", _DEFAULT_MEDIAN_FILE_BYTES, min_val=1) + + +def _get_single_process_concurrency_cap() -> int: + return _env_int("LITDATA_SINGLE_PROCESS_CONCURRENCY_CAP", _SINGLE_PROCESS_CONCURRENCY_CAP, min_val=1) + + +def _get_aggregate_concurrency_budget_cap() -> int: + return _env_int("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_CAP", _AGGREGATE_CONCURRENCY_BUDGET_CAP, min_val=1) + + +def _get_aggregate_concurrency_budget_floor() -> int: + floor = _env_int("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR", _AGGREGATE_CONCURRENCY_BUDGET_FLOOR, min_val=1) + cap = _get_aggregate_concurrency_budget_cap() + if floor > cap: + logger.warning("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR=%d exceeds cap=%d; using floor=%d", floor, cap, cap) + return cap + return floor + + +def _get_backoff_recovery_alpha() -> float: + return _env_float("LITDATA_BACKOFF_RECOVERY_ALPHA", 0.1, min_val=0.0, max_val=1.0) + + +def _get_min_empirical_samples() -> int: + return _env_int("LITDATA_MIN_EMPIRICAL_SAMPLES", _MIN_EMPIRICAL_SAMPLES, min_val=1) + + +def _get_permit_refresh_interval() -> int: + return _env_int("LITDATA_PERMIT_REFRESH_INTERVAL", 10, min_val=1) + + +class BandwidthTracker: + """Thread-safe process-local empirical bandwidth and request latency tracker using EMA.""" + + def __init__(self, alpha: float = 0.2) -> None: + self.alpha = max(0.01, min(1.0, alpha)) + self.bandwidth_bps_ema: float | None = None + self.request_latency_s_ema: float | None = None + self.bps_sample_count: int = 0 + self.lat_sample_count: int = 0 + self._sample_count: int = 0 + self._backoff_factor: float = 1.0 + self._lock = threading.Lock() + + def get_backoff_factor(self, target_lat: float) -> float: + """Return stateful backoff factor with 10% deadband and gradual recovery. + + - Congested (obs_lat > 1.1 * target_lat): Immediate backoff. + - Deadband (target_lat < obs_lat <= 1.1 * target_lat): Hold current backoff factor. + - Healthy (obs_lat <= target_lat): Recover slowly toward 1.0. + """ + with self._lock: + min_samples = _get_min_empirical_samples() + if self.request_latency_s_ema is None or self.lat_sample_count < min_samples: + return 1.0 + obs_lat = self.request_latency_s_ema + deadband_lat = 1.1 * target_lat + alpha_rec = _get_backoff_recovery_alpha() + + if obs_lat > deadband_lat: + # Immediate backoff under congestion + target_factor = min(1.0, target_lat / obs_lat) + self._backoff_factor = min(self._backoff_factor, target_factor) + elif obs_lat <= target_lat: + # Gradual recovery under healthy latency + self._backoff_factor = min(1.0, self._backoff_factor + alpha_rec * (1.0 - self._backoff_factor)) + # In deadband (target_lat < obs_lat <= deadband_lat), hold current self._backoff_factor + + return self._backoff_factor + + @property + def sample_count(self) -> int: + """Total GET requests observed by the tracker.""" + with self._lock: + return self._sample_count + + def record_observation(self, size_bytes: int, duration_s: float) -> tuple[int, int] | None: + """Record an empirical GET observation and update EMA estimates. + + Returns (prev_sample_count, new_sample_count) tuple on success, or None for invalid observation. + + Evaluation order is intentional and matters for correctness: + + 1. Read the **current** (pre-update) bandwidth EMA. + 2. Estimate transfer-subtracted request latency using that prior estimate + (if empirical bandwidth sample threshold is met) and update ``request_latency_s_ema``. + 3. Update ``bandwidth_bps_ema`` with the current observation. + + This ordering avoids a circular estimation problem: if the bandwidth EMA + were updated first, the freshly-observed ``size / duration`` would be used + to explain its own transfer time, which collapses the latency estimate + towards ``_LATENCY_EPSILON_S`` on every sample — defeating the purpose + of the subtraction entirely. + + The latency stored in ``request_latency_s_ema`` is *not* TCP RTT; it is + the estimated request latency after removing the payload-transfer + component from wall-clock GET time (connection setup, TLS, server + processing, scheduling, and proxy overhead are all included). + """ + if size_bytes <= 0 or duration_s <= 0: + return None + with self._lock: + prev_sample_count = self._sample_count + self._sample_count += 1 + new_sample_count = self._sample_count + + # Step 1: capture the PREVIOUS bandwidth estimate before modifying it. + # This is the key invariant: the current sample must not be used to + # explain its own transfer time. + prev_bps_ema = self.bandwidth_bps_ema + + # Step 2: update transfer-subtracted latency EMA for small objects. + if size_bytes < _LATENCY_OBSERVATION_MAX_BYTES: + self.lat_sample_count += 1 + # Fall back to assumed aggregate bandwidth when empirical BPS sample + # threshold (_MIN_EMPIRICAL_SAMPLES = 5) has not yet been met. + effective_bps = ( + prev_bps_ema + if self.bps_sample_count >= _MIN_EMPIRICAL_SAMPLES and prev_bps_ema is not None and prev_bps_ema > 0 + else _get_assumed_aggregate_bandwidth_bps() + ) + transfer_time_s = float(size_bytes) / effective_bps + # Estimated request latency: wall-clock minus payload-transfer time. + # Clamped to _LATENCY_EPSILON_S to stay well-conditioned. + estimated_request_latency_s = max(_LATENCY_EPSILON_S, duration_s - transfer_time_s) + if self.request_latency_s_ema is None: + self.request_latency_s_ema = estimated_request_latency_s + else: + self.request_latency_s_ema = ( + self.alpha * estimated_request_latency_s + (1.0 - self.alpha) * self.request_latency_s_ema + ) + + # Step 3: now update bandwidth EMA with the current observation. + if size_bytes >= _BANDWIDTH_OBSERVATION_MIN_BYTES: + self.bps_sample_count += 1 + obs_bps = float(size_bytes) / duration_s + if self.bandwidth_bps_ema is None: + self.bandwidth_bps_ema = obs_bps + else: + self.bandwidth_bps_ema = self.alpha * obs_bps + (1.0 - self.alpha) * self.bandwidth_bps_ema + + return (prev_sample_count, new_sample_count) + + def get_metrics(self) -> tuple[float | None, float | None, int, int]: + """Returns (bandwidth_bps_ema, request_latency_s_ema, bps_sample_count, lat_sample_count).""" + with self._lock: + return ( + self.bandwidth_bps_ema, + self.request_latency_s_ema, + self.bps_sample_count, + self.lat_sample_count, + ) + + def __getstate__(self) -> dict[str, Any]: + with self._lock: + return { + "alpha": self.alpha, + "bandwidth_bps_ema": self.bandwidth_bps_ema, + "request_latency_s_ema": self.request_latency_s_ema, + "bps_sample_count": self.bps_sample_count, + "lat_sample_count": self.lat_sample_count, + "sample_count": self._sample_count, + "backoff_factor": self._backoff_factor, + } + + def __setstate__(self, state: dict[str, Any]) -> None: + self.alpha = state.get("alpha", 0.2) + self.bandwidth_bps_ema = state.get("bandwidth_bps_ema") + self.request_latency_s_ema = state.get("request_latency_s_ema") + self.bps_sample_count = state.get("bps_sample_count", 0) + self.lat_sample_count = state.get("lat_sample_count", 0) + self._sample_count = state.get("sample_count", self.bps_sample_count + self.lat_sample_count) + self._backoff_factor = state.get("backoff_factor", 1.0) + if "bps_sample_count" not in state and "sample_count" in state: + legacy_count = state.get("sample_count", 0) + if self.bandwidth_bps_ema is not None: + self.bps_sample_count = legacy_count + if self.request_latency_s_ema is not None: + self.lat_sample_count = legacy_count + self._lock = threading.Lock() + + +class _DynamicSemaphore: + """An asyncio-compatible semaphore supporting dynamic permit quota updates. + + Replaces standard un-coordinated ``asyncio.Semaphore`` instantiation when + budget updates happen mid-flight. When downscaling (e.g. 32 -> 16), excess + permits are acquired/withheld so new tasks must wait until active in-flight + downloads drain to the new quota. + """ + + def __init__(self, value: int = 1) -> None: + self.target_permits = max(1, value) + self._sem = asyncio.Semaphore(self.target_permits) + + def update_target(self, new_target: int) -> None: + new_target = max(1, new_target) + diff = new_target - self.target_permits + self.target_permits = new_target + if diff > 0: + for _ in range(diff): + self._sem.release() + elif diff < 0: + self._sem._value -= abs(diff) + + def locked(self) -> bool: + return self._sem.locked() + + async def acquire(self) -> None: + await self._sem.acquire() + + def release(self) -> None: + self._sem.release() + + async def __aenter__(self) -> None: + await self.acquire() + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.release() + _RUNNER_LOCK = threading.Lock() _RUNNER: _LoopRunner | None = None @@ -191,7 +474,7 @@ class _LoopRunner: def __init__(self) -> None: self._pid = os.getpid() self.loop: asyncio.AbstractEventLoop = _create_event_loop() - self._executor = ThreadPoolExecutor(max_workers=32, thread_name_prefix="litdata-raw-pool") + self._executor = ThreadPoolExecutor(max_workers=32, thread_name_prefix="asyncio_litdata-raw-pool") self.loop.set_default_executor(self._executor) if _RAW_DEBUG: logger.warning( @@ -413,52 +696,117 @@ def _median_file_bytes(files: Sequence[FileMetadata]) -> int | None: return int(statistics.median(sizes)) -def _aggregate_concurrency_budget(median_file_bytes: int | None) -> int: +def _aggregate_concurrency_budget( + median_file_bytes: int | None, + tracker: BandwidthTracker | None = None, +) -> int: """Aggregate in-flight download slots across all workers (size-aware, clamped). - Takes the max of two models then clamps to ``[floor, cap]``: - - - **bandwidth**: ``(aggregate_bps × pipeline_s) // median_file_bytes`` — keep - ~50 MiB moving for large objects. - - **latency / Little's law**: ``target_rate × assumed_latency`` (~6000×0.040 ≈ - 240) **only when** ``median < _LATENCY_MODEL_MAX_MEDIAN_BYTES`` (1 MiB) so - tiny-object paths are not request-starved. Medians ≥1 MiB stay - bandwidth-bounded (avoids pinning at 240 slots → multi-GB in flight). + Calculates unconstrained baseline capacity: max(bandwidth model, Little's-law model). + `max` is used intentionally so the baseline is not constrained by single-model + underestimation. - Per-worker floor of 8 means realized aggregate is ``max(budget, 8 × num_workers)``. + If empirical measurements indicate congestion (latency > target), applies + stateless backoff factor (L_target / L_obs). """ - median = median_file_bytes if median_file_bytes and median_file_bytes > 0 else _DEFAULT_MEDIAN_FILE_BYTES - target_bytes = int(_ASSUMED_AGGREGATE_BANDWIDTH_BPS * _CONCURRENCY_PIPELINE_SECONDS) + default_median = _get_default_median_file_bytes() + median = median_file_bytes if median_file_bytes and median_file_bytes > 0 else default_median + + floor = _get_aggregate_concurrency_budget_floor() + cap = _get_aggregate_concurrency_budget_cap() + + obs_bps: float | None = None + obs_lat: float | None = None + bps_samples: int = 0 + lat_samples: int = 0 + if tracker is not None: + bps_ema, lat_ema, bps_samples, lat_samples = tracker.get_metrics() + if bps_samples >= _MIN_EMPIRICAL_SAMPLES: + obs_bps = bps_ema + if lat_samples >= _MIN_EMPIRICAL_SAMPLES: + obs_lat = lat_ema + + bandwidth_bps = obs_bps if obs_bps is not None and obs_bps > 0 else _get_assumed_aggregate_bandwidth_bps() + target_bytes = int(bandwidth_bps * _CONCURRENCY_PIPELINE_SECONDS) bandwidth_model = max(1, target_bytes // median) + # Size-gate: Little's-law arm is for request-overhead-bound tiny objects only. + # Fixed baseline: target_rate * target_latency. Observed latency is NEVER multiplied in. + target_lat = _get_assumed_request_latency_s() if median < _LATENCY_MODEL_MAX_MEDIAN_BYTES: - latency_model = max(1, int(_ASSUMED_REQUEST_RATE * _ASSUMED_REQUEST_LATENCY_S)) + req_rate = _get_assumed_request_rate() + latency_model = max(1, int(req_rate * target_lat)) else: latency_model = 0 - raw = max(bandwidth_model, latency_model) - return max(_AGGREGATE_CONCURRENCY_BUDGET_FLOOR, min(_AGGREGATE_CONCURRENCY_BUDGET_CAP, raw)) + + # max is intentional to avoid underestimating baseline capacity + baseline_budget = max(bandwidth_model, latency_model) + + # Stateful congestion control backoff with 10% deadband and gradual recovery + if tracker is not None: + backoff_factor = tracker.get_backoff_factor(target_lat) + elif obs_lat is not None and obs_lat > target_lat: + backoff_factor = min(1.0, target_lat / obs_lat) + else: + backoff_factor = 1.0 + + computed_budget = max(1, int(baseline_budget * backoff_factor)) + + # Guarded adaptive floor reduction: require high-confidence combined evidence + if ( + bps_samples >= _MIN_EMPIRICAL_SAMPLES + and lat_samples >= _MIN_EMPIRICAL_SAMPLES + and obs_bps is not None + and obs_bps < _LOW_BANDWIDTH_THRESHOLD_BPS + and obs_lat is not None + and obs_lat > _HIGH_LATENCY_THRESHOLD_S + ): + effective_floor = 1 + else: + effective_floor = floor + + # Enforce authoritative MAX cap (512) and effective MIN floor + return max(effective_floor, min(cap, computed_budget)) def _effective_concurrency( max_concurrent_downloads: int | None, num_workers: int, median_file_bytes: int | None = None, + tracker: BandwidthTracker | None = None, + worker_id: int | None = None, ) -> int: - """Per-worker download permits for the Stage 1 static clamp. + """Per-worker download permits for the Stage 1 static/adaptive clamp. - ``max_concurrent_downloads is None`` (default): adaptive — - ``max(floor, budget // num_workers)`` with ``budget`` from - :func:`_aggregate_concurrency_budget`. When ``num_workers <= 1``, returns - ``min(budget, _SINGLE_PROCESS_CONCURRENCY_CAP)`` (unbenchmarked path). + budget split across workers while strictly maintaining aggregate budget + invariant: sum(worker_permits) <= aggregate_budget. - Explicit ``int``: **exactly** that many permits (no silent clamp). ``<= 0`` collapses to 1. """ if max_concurrent_downloads is not None: return 1 if max_concurrent_downloads <= 0 else max_concurrent_downloads - budget = _aggregate_concurrency_budget(median_file_bytes) + budget = _aggregate_concurrency_budget(median_file_bytes, tracker=tracker) if num_workers <= 1: - return min(budget, _SINGLE_PROCESS_CONCURRENCY_CAP) - return max(_MIN_CONCURRENCY_PER_WORKER, budget // num_workers) + return min(budget, _get_single_process_concurrency_cap()) + + if budget >= num_workers: + return budget // num_workers + + # When aggregate budget < num_workers, allocate 1 permit to ranks < budget + w_id = worker_id + if w_id is None: + try: + from torch.utils.data import get_worker_info + + info = get_worker_info() + if info is not None: + w_id = info.id + except ImportError: + pass + if w_id is not None and w_id >= budget: + return 0 + return 1 def _num_dataloader_workers() -> int: @@ -557,6 +905,7 @@ def __init__( self.storage_options = storage_options or {} # Index median size (bytes); set by StreamingRawDataset after discovery. self._median_file_bytes: int | None = None + self._bandwidth_tracker = BandwidthTracker() self._downloader: Downloader | None = None self._downloader_pid: int | None = None self._downloader_loop: asyncio.AbstractEventLoop | None = None @@ -610,6 +959,7 @@ def __getstate__(self) -> dict[str, Any]: "cache_dir": self.cache_dir, "storage_options": self.storage_options, "_median_file_bytes": self._median_file_bytes, + "_bandwidth_tracker": self._bandwidth_tracker, # Runtime — always fresh in the child. "_downloader": None, "_downloader_pid": None, @@ -645,6 +995,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: self._range_executor_pid = None self._hedge_fired = 0 self._median_file_bytes = state.get("_median_file_bytes") + self._bandwidth_tracker = state.get("_bandwidth_tracker") or BandwidthTracker() def _shutdown_range_executor(self) -> None: if self._range_executor is not None: @@ -720,6 +1071,19 @@ def downloader(self) -> Downloader: self._downloader_loop = loop return self._downloader + def _record_download_observation(self, size_bytes: int, duration_s: float) -> None: + """Record an empirical GET transfer observation and refresh cached permits when needed.""" + res = self._bandwidth_tracker.record_observation(size_bytes, duration_s) + if res is None: + return + prev_count, new_count = res + min_samples = _get_min_empirical_samples() + refresh_interval = _get_permit_refresh_interval() + if (prev_count < min_samples and new_count >= min_samples) or ( + new_count >= min_samples and new_count % refresh_interval == 0 + ): + self._cached_permits = None + def _effective_download_permits(self) -> int: """Worker-aware permit count for the download semaphore (Stage 1 static clamp). @@ -733,24 +1097,25 @@ def _effective_download_permits(self) -> int: self.max_concurrent_downloads, _num_dataloader_workers(), self._median_file_bytes, + tracker=self._bandwidth_tracker, ) self._cached_permits = permits self._cached_permits_pid = pid return permits - def _get_semaphore(self) -> asyncio.Semaphore: - """Return a semaphore bound to the current event loop with effective permits. + def _get_semaphore(self) -> _DynamicSemaphore | asyncio.Semaphore: + """Return a dynamic semaphore bound to the current event loop with effective permits. Permit count comes from :meth:`_effective_download_permits` (cached per - process). Loop-keyed like other runtime clients; cleared by - ``reset_runtime_state``. + process). Reuses active dynamic semaphore and adjusts target quota when + permits change instead of replacing in-flight semaphores. """ loop = asyncio.get_running_loop() permits = self._effective_download_permits() - if self._semaphore is None or self._semaphore_loop is not loop or self._semaphore_permits != permits: + if self._semaphore is None or self._semaphore_loop is not loop: n_workers = _num_dataloader_workers() budget = ( - _aggregate_concurrency_budget(self._median_file_bytes) + _aggregate_concurrency_budget(self._median_file_bytes, tracker=self._bandwidth_tracker) if self.max_concurrent_downloads is None else None ) @@ -761,9 +1126,15 @@ def _get_semaphore(self) -> asyncio.Semaphore: n_workers, permits, ) - self._semaphore = asyncio.Semaphore(permits) + self._semaphore = _DynamicSemaphore(permits) self._semaphore_loop = loop self._semaphore_permits = permits + elif self._semaphore_permits != permits: + if isinstance(self._semaphore, _DynamicSemaphore): + self._semaphore.update_target(permits) + else: + self._semaphore = _DynamicSemaphore(permits) + self._semaphore_permits = permits return self._semaphore @asynccontextmanager @@ -998,7 +1369,8 @@ async def fetch() -> bytes: scratch = f"{base_scratch}.{offset}.{uuid4().hex}" try: async with self._permit(gated): - data = await asyncio.get_running_loop().run_in_executor( + t0_c = time.monotonic() + res = await asyncio.get_running_loop().run_in_executor( executor, downloader.download_bytes, file_path, @@ -1006,11 +1378,20 @@ async def fetch() -> bytes: length, scratch, ) + data = ( + res + if isinstance(res, (bytes, bytearray)) + else (Path(scratch).read_bytes() if os.path.exists(scratch) else b"") + ) + dur_c = time.monotonic() - t0_c if len(data) != length: raise RuntimeError( f"Ranged GET short read for {file_path}: offset={offset} expected={length} got={len(data)}" ) + if len(data) > 0: + self._record_download_observation(len(data), dur_c) return data + finally: with contextlib.suppress(OSError): os.remove(scratch) @@ -1049,11 +1430,21 @@ async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: # Pay-per-use: hedging off/ineligible → bare permit + download (batch enforces timeout). if delay is None: async with self._permit(gated): - return await self.downloader.adownload_fileobj(file_path) + t0 = time.monotonic() + data = await self.downloader.adownload_fileobj(file_path) + dur = time.monotonic() - t0 + if len(data) > 0: + self._record_download_observation(len(data), dur) + return data async def once() -> bytes: async with self._permit(gated): - return await self.downloader.adownload_fileobj(file_path) + t0_once = time.monotonic() + res = await self.downloader.adownload_fileobj(file_path) + dur_once = time.monotonic() - t0_once + if len(res) > 0: + self._record_download_observation(len(res), dur_once) + return res return await self._hedged(once, delay) @@ -1098,9 +1489,16 @@ async def _download_owned(self, file_path: str, local_path: str, size: int | Non try: if self._path_is_cached(local_path): return local_path + t0 = time.monotonic() try: # Hang protection is batch-level; keep the owned path bare. await self.downloader.adownload_file(file_path, tmp_path) + dur = time.monotonic() - t0 + st_size = ( + size if size and size > 0 else (os.path.getsize(tmp_path) if os.path.exists(tmp_path) else None) + ) + if st_size and st_size > 0: + self._record_download_observation(st_size, dur) except Exception as first_exc: if self._is_non_retryable_download_error(first_exc): raise @@ -1112,6 +1510,7 @@ async def _download_owned(self, file_path: str, local_path: str, size: int | Non with contextlib.suppress(OSError): os.remove(tmp_path) # Caller already holds the download semaphore — avoid nested acquire. + # Note: _fetch_bytes records download observation internally if successful. data = await self._fetch_bytes(file_path, size=size, gated=False) await asyncio.to_thread(Path(tmp_path).write_bytes, data) self._verify_tmp_size(tmp_path, size) diff --git a/tests/raw/test_dataset.py b/tests/raw/test_dataset.py index bead4c069..3cb6158a3 100644 --- a/tests/raw/test_dataset.py +++ b/tests/raw/test_dataset.py @@ -91,7 +91,7 @@ def test_effective_prefetch_vs_num_workers(num_workers, max_prefetch, expected): (32, None, 100_000, 16), # 512//32 # Large objects (≥1 MiB): bandwidth-only (no Little's-law pin at 240) (4, None, 10 * 1024 * 1024, 8), # budget=floor 32, 32//4=8 - (16, None, 10 * 1024 * 1024, 8), # 32//16=2 → floor 8 + (16, None, 10 * 1024 * 1024, 2), # budget=floor 32, 32//16=2 # Unknown size uses default median (256KiB) → latency arm (240) (8, None, None, 30), # 240//8 ], @@ -539,3 +539,729 @@ def transform(data): gds = GroupedDS(input_dir=str(tmp_path), transform=transform) gds.cache_manager.download_file_async = mock_download_file_async assert gds[0] == b"abc" + + +def test_bandwidth_tracker_basic(): + import pytest + + from litdata.raw.dataset import ( + _ASSUMED_AGGREGATE_BANDWIDTH_BPS, + _LATENCY_RTT_EPSILON, + BandwidthTracker, + ) + + assumed_bps = float(_ASSUMED_AGGREGATE_BANDWIDTH_BPS) + tracker = BandwidthTracker(alpha=0.2) + assert tracker.sample_count == 0 + + # --- Tiny GET (50 KB < 64 KB) --- + # Only updates latency EMA; no bandwidth observation recorded. + # prev_bps_ema = None → fallback to assumed bandwidth for transfer estimate. + tracker.record_observation(50_000, 0.010) + bps, lat, bps_count, lat_count = tracker.get_metrics() + assert tracker.sample_count == 1 + assert bps_count == 0 + assert lat_count == 1 + assert bps is None + transfer1 = 50_000 / assumed_bps + expected_lat1 = max(_LATENCY_RTT_EPSILON, 0.010 - transfer1) + assert pytest.approx(lat, abs=1e-6) == expected_lat1 + # Sanity: transfer-subtracted latency must be strictly less than raw duration. + assert lat < 0.010 + + # --- Small/Medium GET (100 KB: 64 KB <= size < 256 KB) --- + # Updates BOTH latency and bandwidth EMAs. + # prev_bps_ema is still None (previous obs was below bandwidth threshold) → fallback. + tracker.record_observation(100_000, 0.020) + bps, lat, bps_count, lat_count = tracker.get_metrics() + assert tracker.sample_count == 2 + assert bps_count == 1 + assert lat_count == 2 + assert bps is not None + assert pytest.approx(bps, abs=1.0) == 5_000_000.0 + transfer2 = 100_000 / assumed_bps # prev_bps_ema still None before this obs + est_lat2 = max(_LATENCY_RTT_EPSILON, 0.020 - transfer2) + expected_lat2 = 0.2 * est_lat2 + 0.8 * expected_lat1 + assert pytest.approx(lat, abs=1e-6) == expected_lat2 + + # --- Large GET (10 MiB >= 256 KB) --- + # Updates bandwidth EMA only; size >= _LATENCY_OBSERVATION_MAX_BYTES. + tracker.record_observation(10 * 1024 * 1024, 0.118) + bps, lat, bps_count, lat_count = tracker.get_metrics() + assert tracker.sample_count == 3 + assert bps_count == 2 + assert lat_count == 2 # unchanged — large GETs do not update latency EMA + assert bps is not None + + +def test_bandwidth_tracker_pickle(): + import pickle + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker() + tracker.record_observation(100_000, 0.020) + tracker.record_observation(10 * 1024 * 1024, 0.200) + blob = pickle.dumps(tracker) + restored = pickle.loads(blob) # noqa: S301 + assert restored.sample_count == tracker.sample_count + assert restored.bps_sample_count == tracker.bps_sample_count + assert restored.lat_sample_count == tracker.lat_sample_count + assert restored.bandwidth_bps_ema == tracker.bandwidth_bps_ema + assert restored.request_latency_s_ema == tracker.request_latency_s_ema + + +def test_concurrency_budget_warmup_gating(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget + + tracker = BandwidthTracker() + # 4 observations (under threshold of 5) -> uses default static budget + for _ in range(4): + tracker.record_observation(10 * 1024 * 1024, 0.001) + + # Median 10MB -> default budget: (100MB/s * 0.5s) // 10MB = 5 -> floor 32 + assert _aggregate_concurrency_budget(10 * 1024 * 1024, tracker=tracker) == 32 + + # 5th observation -> warm-up gate unlocks empirical EMA + tracker.record_observation(10 * 1024 * 1024, 0.001) + # Measured bandwidth is huge -> dynamic budget scales up from 32 to 500 + assert _aggregate_concurrency_budget(10 * 1024 * 1024, tracker=tracker) == 500 + + +def test_concurrency_budget_high_and_low_bandwidth_adaptation(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget + + # High bandwidth scenario + high_tracker = BandwidthTracker() + for _ in range(5): + high_tracker.record_observation(10 * 1024 * 1024, 0.002) + budget_high = _aggregate_concurrency_budget(1 * 1024 * 1024, tracker=high_tracker) + assert budget_high == 512 + + # Low bandwidth scenario + low_tracker = BandwidthTracker() + for _ in range(5): + low_tracker.record_observation(10 * 1024 * 1024, 5.0) + budget_low = _aggregate_concurrency_budget(10 * 1024 * 1024, tracker=low_tracker) + assert budget_low == 32 + + +def test_class_gated_observation_isolation(): + from litdata.raw.dataset import ( + BandwidthTracker, + _aggregate_concurrency_budget, + ) + + # Scenario 1: 5 small GETs (< 64 KB) -> lat_sample_count = 5, bps_sample_count = 0 + tracker_small = BandwidthTracker() + for _ in range(5): + tracker_small.record_observation(10_000, 0.010) + + bps, lat, bps_cnt, lat_cnt = tracker_small.get_metrics() + assert bps_cnt == 0 + assert lat_cnt == 5 + assert bps is None + + # Budget for large 10 MB objects must STILL use static default aggregate bandwidth because bps_sample_count < 5 + # Clamped to floor 32 + assert _aggregate_concurrency_budget(10 * 1024 * 1024, tracker=tracker_small) == 32 + + # Scenario 2: 5 large GETs (>= 256 KB) -> bps_sample_count = 5, lat_sample_count = 0 + tracker_large = BandwidthTracker() + for _ in range(5): + tracker_large.record_observation(10 * 1024 * 1024, 0.001) + + bps, lat, bps_cnt, lat_cnt = tracker_large.get_metrics() + assert bps_cnt == 5 + assert lat_cnt == 0 + assert lat is None + + +def test_environment_variable_overrides(monkeypatch): + from litdata.raw.dataset import ( + _aggregate_concurrency_budget, + _effective_concurrency, + ) + + monkeypatch.setenv("LITDATA_ASSUMED_BANDWIDTH_BPS", str(500 * 1024 * 1024)) + monkeypatch.setenv("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_CAP", "1024") + monkeypatch.setenv("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR", "16") + monkeypatch.setenv("LITDATA_SINGLE_PROCESS_CONCURRENCY_CAP", "256") + + assert _aggregate_concurrency_budget(100_000) == 1024 + assert _effective_concurrency(None, num_workers=1, median_file_bytes=100_000) == 256 + + +def test_no_latency_concurrency_inflation(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget + + tracker = BandwidthTracker() + # Record 5 latency observations with high latency (200ms vs assumed 40ms) + for _ in range(5): + tracker.record_observation(100_000, 0.200) + + # Budget with high latency must not inflate above baseline (240 for sub-1MB) + budget = _aggregate_concurrency_budget(100_000, tracker=tracker) + assert budget <= 240 + + +def test_concurrency_latency_monotonicity(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget + + latencies = [0.040, 0.080, 0.200, 0.500] + budgets = [] + + for lat in latencies: + tracker = BandwidthTracker() + for _ in range(5): + tracker.record_observation(100_000, lat) + budgets.append(_aggregate_concurrency_budget(100_000, tracker=tracker)) + + # For latencies above target (40ms), increasing latency must not increase computed budget + for i in range(len(budgets) - 1): + assert budgets[i + 1] <= budgets[i] + + +def test_concurrency_latency_recovery(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget + + tracker = BandwidthTracker() + # Inject high latency + for _ in range(5): + tracker.record_observation(100_000, 0.200) + budget_degraded = _aggregate_concurrency_budget(100_000, tracker=tracker) + + # Now inject healthy latency + for _ in range(10): + tracker.record_observation(100_000, 0.040) + budget_recovered = _aggregate_concurrency_budget(100_000, tracker=tracker) + + assert budget_recovered > budget_degraded + + +def test_imagenet_bandwidth_observation(): + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker() + # ImageNet JPEG size ~150 KB — straddles both latency (<256 KB) and bandwidth (>=64 KB) thresholds. + imagenet_file_size = 150 * 1024 + for _ in range(5): + tracker.record_observation(imagenet_file_size, 0.015) + + bps, lat, bps_cnt, lat_cnt = tracker.get_metrics() + assert bps_cnt == 5 + assert lat_cnt == 5 + + # ImageNet files MUST record both EMA estimates. + assert bps is not None + assert bps > 0 + assert lat is not None + assert lat > 0 + + # Transfer-subtracted latency must be strictly less than raw wall-clock duration. + # (transfer time is non-zero for a 150 KB file) + assert lat < 0.015 + + # Epsilon floor must hold: estimated latency must not go below _LATENCY_RTT_EPSILON. + assert lat >= 0.001 + + +def test_guarded_adaptive_floor_reduction(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget + + tracker = BandwidthTracker() + # Low bandwidth (<10 MiB/s) AND high latency (>100 ms) + for _ in range(5): + tracker.record_observation(128 * 1024, 0.500) + + # Under severe evidence (low bandwidth + high latency), budget can drop below default 32 floor + budget = _aggregate_concurrency_budget(10 * 1024 * 1024, tracker=tracker) + assert budget < 32 + + +def test_worker_allocation_respects_aggregate_budget(): + from litdata.raw.dataset import BandwidthTracker, _aggregate_concurrency_budget, _effective_concurrency + + test_cases = [ + (512, 8, 100_000), + (32, 8, 10_000_000), + (32, 16, 10_000_000), + (32, 64, 10_000_000), + ] + + for expected_budget_cap, workers, median_bytes in test_cases: + tracker = BandwidthTracker() + budget = _aggregate_concurrency_budget(median_bytes, tracker=tracker) + total_permits = sum( + _effective_concurrency( + None, num_workers=workers, median_file_bytes=median_bytes, tracker=tracker, worker_id=w + ) + for w in range(workers) + ) + # Sum of per-worker permits across all workers must not exceed aggregate budget + assert total_permits <= budget + + +def test_transfer_subtracted_latency(): + """Pre-seed the tracker with 5 large GETs to reach sample threshold. + + Establishes a known bandwidth EMA (20 MB/s), then records a 200 KB request taking 30 ms. + Verify latency EMA is approximately 30 ms - transfer_time (10 ms) = 20 ms. + """ + import pytest + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + prior_bps = 20 * 1024 * 1024 # 20 MB/s + large_size = 10 * 1024 * 1024 # 10 MiB (bandwidth-only GET) + + # Record 5 observations to pass the _MIN_EMPIRICAL_SAMPLES = 5 threshold + for _ in range(5): + tracker.record_observation(large_size, large_size / prior_bps) + + bps, lat, bps_cnt, lat_cnt = tracker.get_metrics() + assert bps_cnt == 5 + assert lat_cnt == 0 + assert pytest.approx(bps, rel=1e-5) == prior_bps + + # Now record 200 KB GET taking 30 ms + size = 200 * 1024 + duration = 0.030 + tracker.record_observation(size, duration) + + _, lat_after, _, _ = tracker.get_metrics() + expected_lat = duration - (size / prior_bps) # 30 ms - 10 ms = 20 ms + assert lat_after is not None + assert pytest.approx(lat_after, abs=1e-4) == expected_lat + + +def test_transfer_subtracted_latency_bootstrap(): + """When no empirical bandwidth exists or bps_sample_count < 5, verify default bandwidth is used.""" + import pytest + + from litdata.raw.dataset import _ASSUMED_AGGREGATE_BANDWIDTH_BPS, BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + assert tracker.bandwidth_bps_ema is None + + # Record a 50 KB GET taking 10 ms (sample 1, below 5 threshold) + size = 50 * 1024 + duration = 0.010 + tracker.record_observation(size, duration) + + _, lat, _, _ = tracker.get_metrics() + expected_transfer = size / float(_ASSUMED_AGGREGATE_BANDWIDTH_BPS) + expected_lat = max(0.001, duration - expected_transfer) + assert lat is not None + assert pytest.approx(lat, abs=1e-6) == expected_lat + + +def test_transfer_subtracted_latency_clamped(): + """Provide an observation where estimated transfer time > observed duration. + + Verify resulting latency is clamped to epsilon (0.001s) rather than becoming zero or negative. + """ + import pytest + + from litdata.raw.dataset import _LATENCY_EPSILON_S, BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + # Seed 5 large GETs at a slow prior BPS (1 MB/s) + slow_bps = 1 * 1024 * 1024 + large_size = 10 * 1024 * 1024 + for _ in range(5): + tracker.record_observation(large_size, large_size / slow_bps) + + # Now record a 100 KB GET arriving in only 5 ms (faster than 100 ms transfer estimate) + tracker.record_observation(100 * 1024, 0.005) + _, lat, _, _ = tracker.get_metrics() + + assert lat is not None + assert pytest.approx(lat, abs=1e-9) == _LATENCY_EPSILON_S + + +def test_current_bandwidth_sample_does_not_explain_itself(): + """Verify bandwidth calculated from current observation is NOT used to calculate transfer time.""" + from litdata.raw.dataset import _LATENCY_EPSILON_S, BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + # Record a single 200 KB GET taking 30 ms (bps_sample_count = 0 before this sample). + # If circular, it would use 200 KB / 30 ms to calculate transfer = 30 ms -> lat = 0 -> clamped to epsilon (1 ms). + # Since it correctly uses fallback default (100 MB/s), transfer = 2 ms -> lat = 28 ms. + tracker.record_observation(200 * 1024, 0.030) + _, lat, _, _ = tracker.get_metrics() + + assert lat is not None + assert lat > 10 * _LATENCY_EPSILON_S # 28 ms is >> 1 ms epsilon + + +def test_256k_to_1m_bandwidth_observation(): + """Verify 256 KiB - 1 MiB objects update BPS EMA without updating latency EMA.""" + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker() + size_500k = 500 * 1024 # 500 KiB + tracker.record_observation(size_500k, 0.050) + + bps, lat, bps_cnt, lat_cnt = tracker.get_metrics() + assert bps_cnt == 1 + assert lat_cnt == 0 + assert bps is not None + assert bps > 0 + assert lat is None + + +def test_dynamic_semaphore_scale_down_target(): + """Verify _DynamicSemaphore updates target permits on downscale and retains permit bounds.""" + import asyncio + + import pytest + + from litdata.raw.dataset import _DynamicSemaphore + + @pytest.mark.asyncio + async def _run(): + dyn_sem = _DynamicSemaphore(32) + assert dyn_sem.target_permits == 32 + + # Downscale to 16 + dyn_sem.update_target(16) + assert dyn_sem.target_permits == 16 + + # Acquire 16 permits + for _ in range(16): + await dyn_sem.acquire() + + # Releasing permits works cleanly + for _ in range(16): + dyn_sem.release() + + asyncio.run(_run()) + + +def test_ranged_gather_individual_chunk_observations(tmp_path): + """Verify that ranged downloads record observations per chunk rather than one aggregate blob.""" + import os + from unittest.mock import MagicMock + + from litdata.raw.dataset import StreamingRawDataset + + (tmp_path / "sample.bin").write_bytes(b"x" * 100) + ds = StreamingRawDataset( + input_dir=str(tmp_path), + cache_dir=str(tmp_path), + range_parallel_threshold=100 * 1024, + range_chunk_size=100 * 1024, + ) + + import asyncio + + async def _run(): + loop = asyncio.get_running_loop() + mock_dl = MagicMock() + + def _mock_write(f_path, off, length, scratch): + from pathlib import Path + + Path(scratch).write_bytes(b"x" * length) + + mock_dl.download_bytes.side_effect = _mock_write + ds.cache_manager._downloader = mock_dl + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + + data = await ds.cache_manager._ranged_download_bytes("s3://mock-bucket/file.bin", size=200 * 1024) + assert len(data) == 200 * 1024 + + asyncio.run(_run()) + # 2 chunks of 100 KiB (>= 64 KiB min) -> 2 bps observations recorded + _, _, bps_cnt, _ = ds.cache_manager._bandwidth_tracker.get_metrics() + ds.cache_manager.reset_runtime_state() + assert bps_cnt == 2 + + +def test_hedged_get_excludes_hedge_delay(tmp_path): + """Verify that hedged GET timing is recorded inside the winning attempt.""" + import os + from unittest.mock import AsyncMock, MagicMock + + from litdata.raw.dataset import StreamingRawDataset + + (tmp_path / "sample.bin").write_bytes(b"y" * 100_000) + ds = StreamingRawDataset( + input_dir=str(tmp_path), + cache_dir=str(tmp_path), + hedge_delay=0.1, + ) + + import asyncio + + async def _run(): + loop = asyncio.get_running_loop() + mock_dl = MagicMock() + mock_dl.adownload_fileobj = AsyncMock(return_value=b"y" * 100_000) + ds.cache_manager._downloader = mock_dl + ds.cache_manager._downloader_pid = os.getpid() + ds.cache_manager._downloader_loop = loop + + data = await ds.cache_manager._fetch_bytes("s3://mock-bucket/file.bin", size=100_000) + assert len(data) == 100_000 + + asyncio.run(_run()) + bps, _, bps_cnt, _ = ds.cache_manager._bandwidth_tracker.get_metrics() + ds.cache_manager.reset_runtime_state() + assert bps_cnt == 1 + assert bps is not None + + +def test_backoff_is_immediate(): + """Verify high latency (>1.1 * target) triggers immediate backoff factor reduction.""" + import pytest + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + # 5 samples with 200 ms latency (target = 40 ms -> factor = 0.20) + for _ in range(5): + tracker.record_observation(10 * 1024, 0.200) + + factor = tracker.get_backoff_factor(0.040) + assert pytest.approx(factor, abs=1e-3) == 0.20 + + +def test_no_recovery_in_deadband(): + """Verify latency inside deadband (target < L <= 1.1 * target) holds current backoff factor.""" + import pytest + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + # Trigger backoff down to 0.50 (80 ms latency) + for _ in range(5): + tracker.record_observation(10 * 1024, 0.080) + assert pytest.approx(tracker.get_backoff_factor(0.040), abs=1e-3) == 0.50 + + # Deadband sample (42 ms: 40 ms < 42 ms <= 44 ms) + tracker.record_observation(10 * 1024, 0.042) + factor = tracker.get_backoff_factor(0.040) + assert pytest.approx(factor, abs=1e-3) == 0.50 + + +def test_gradual_recovery(): + """Verify healthy latency (<= target) triggers gradual recovery towards 1.0.""" + import pytest + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + for _ in range(5): + tracker.record_observation(10 * 1024, 0.080) + factor_initial = tracker.get_backoff_factor(0.040) + assert pytest.approx(factor_initial, abs=1e-3) == 0.50 + + # Healthy sample (20 ms <= 40 ms target) -> recovers by alpha (0.1 * (1.0 - 0.5) = +0.05) + tracker.record_observation(10 * 1024, 0.020) + factor_recovered = tracker.get_backoff_factor(0.040) + assert factor_recovered > factor_initial + assert pytest.approx(factor_recovered, abs=1e-3) == 0.55 + + +def test_congestion_interrupts_recovery(): + """Verify new congestion immediately interrupts gradual recovery.""" + import pytest + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + for _ in range(5): + tracker.record_observation(10 * 1024, 0.080) + tracker.get_backoff_factor(0.040) # 0.50 + + # 1 healthy sample -> recovers to 0.55 + tracker.record_observation(10 * 1024, 0.020) + tracker.get_backoff_factor(0.040) + + # Congestion sample (200 ms) -> immediate drop to 0.20 + tracker.record_observation(10 * 1024, 0.200) + factor = tracker.get_backoff_factor(0.040) + assert pytest.approx(factor, abs=1e-3) == 0.20 + + +def test_recovery_reaches_one(): + """Verify sustained healthy observations restore backoff factor to 1.0.""" + import pytest + + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=1.0) + for _ in range(5): + tracker.record_observation(10 * 1024, 0.200) + tracker.get_backoff_factor(0.040) + + # Repeated healthy samples + for _ in range(100): + tracker.record_observation(10 * 1024, 0.010) + tracker.get_backoff_factor(0.040) + + factor = tracker.get_backoff_factor(0.040) + assert pytest.approx(factor, abs=1e-2) == 1.0 + + +def test_env_invalid_value_falls_back(monkeypatch): + """Verify invalid string in environment variable falls back to safe default.""" + from litdata.raw.dataset import _get_assumed_aggregate_bandwidth_bps + + monkeypatch.setenv("LITDATA_ASSUMED_BANDWIDTH_BPS", "not_a_number") + assert _get_assumed_aggregate_bandwidth_bps() > 0 + + +def test_env_zero_value_falls_back(monkeypatch): + """Verify zero value in float environment variable falls back to safe default.""" + from litdata.raw.dataset import _get_assumed_request_latency_s + + monkeypatch.setenv("LITDATA_ASSUMED_REQUEST_LATENCY_S", "0.0") + assert _get_assumed_request_latency_s() > 0.0 + + +def test_env_negative_value_falls_back(monkeypatch): + """Verify negative integer in environment variable falls back to default.""" + from litdata.raw.dataset import _get_single_process_concurrency_cap + + monkeypatch.setenv("LITDATA_SINGLE_PROCESS_CONCURRENCY_CAP", "-10") + assert _get_single_process_concurrency_cap() > 0 + + +def test_env_floor_cannot_exceed_cap(monkeypatch): + """Verify aggregate floor is clamped to cap when floor > cap.""" + from litdata.raw.dataset import _get_aggregate_concurrency_budget_cap, _get_aggregate_concurrency_budget_floor + + monkeypatch.setenv("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR", "1000") + monkeypatch.setenv("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_CAP", "512") + cap = _get_aggregate_concurrency_budget_cap() + floor = _get_aggregate_concurrency_budget_floor() + assert floor <= cap + + +def test_permit_refresh_respects_reduced_budget_for_new_acquisitions(): + """Verify _DynamicSemaphore target reduction prevents additional permits beyond reduced target.""" + import asyncio + + import pytest + + from litdata.raw.dataset import _DynamicSemaphore + + async def _run(): + sem = _DynamicSemaphore(32) + sem.update_target(16) + assert sem.target_permits == 16 + + # Acquire 16 permits cleanly + for _ in range(16): + await sem.acquire() + + # Target 16 is reached — next acquire without release must fail non-blocking + with pytest.raises((asyncio.TimeoutError, TimeoutError)): + await asyncio.wait_for(sem.acquire(), timeout=0.05) + + for _ in range(16): + sem.release() + + asyncio.run(_run()) + + +def test_record_observation_returns_atomic_sample_counts(): + """Verify record_observation returns (prev_sample_count, new_sample_count) tuple atomically.""" + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker() + assert tracker.record_observation(-1, 0.1) is None + assert tracker.record_observation(100, -0.5) is None + + res1 = tracker.record_observation(100_000, 0.05) + assert res1 == (0, 1) + + res2 = tracker.record_observation(200_000, 0.04) + assert res2 == (1, 2) + + assert tracker.sample_count == 2 + + +def test_record_download_observation_uses_configured_sample_and_refresh_thresholds(tmp_path, monkeypatch): + """Verify permit cache invalidation uses min_empirical_samples and permit_refresh_interval configs.""" + monkeypatch.setenv("LITDATA_MIN_EMPIRICAL_SAMPLES", "3") + monkeypatch.setenv("LITDATA_PERMIT_REFRESH_INTERVAL", "4") + + (tmp_path / "file1.jpg").write_bytes(b"x") + from litdata.raw.dataset import StreamingRawDataset + + ds = StreamingRawDataset(input_dir=str(tmp_path), max_prefetch=0) + cm = ds.cache_manager + + cm._cached_permits = 64 + cm._cached_permits_pid = os.getpid() + + # Sample 1 & 2 (< min_samples 3) -> should not invalidate permits + cm._record_download_observation(100_000, 0.01) + assert cm._cached_permits == 64 + cm._record_download_observation(100_000, 0.01) + assert cm._cached_permits == 64 + + # Sample 3 (== min_samples 3) -> should invalidate permits + cm._record_download_observation(100_000, 0.01) + assert cm._cached_permits is None + + cm._cached_permits = 64 + cm._cached_permits_pid = os.getpid() + + # Sample 4 (% refresh 4 == 0) -> should invalidate permits + cm._record_download_observation(100_000, 0.01) + assert cm._cached_permits is None + + +@pytest.mark.asyncio +async def test_telemetry_duration_excludes_permit_queue_wait(tmp_path): + """Verify semaphore queue wait time is excluded from recorded telemetry duration.""" + import asyncio + from unittest.mock import AsyncMock + + from litdata.raw.dataset import StreamingRawDataset + + (tmp_path / "file.jpg").write_bytes(b"12345") + ds = StreamingRawDataset(input_dir=str(tmp_path), max_prefetch=0) + cm = ds.cache_manager + + # Force 1 permit max + cm.max_concurrent_downloads = 1 + cm._downloader = AsyncMock() + cm._downloader_pid = os.getpid() + cm._downloader_loop = asyncio.get_running_loop() + + async def mock_download(file_path): + await asyncio.sleep(0.02) # actual download takes 20ms + return b"12345" + + cm.downloader.adownload_fileobj.side_effect = mock_download + + recorded_durations = [] + + def mock_record(size, dur): + recorded_durations.append(dur) + + cm._record_download_observation = mock_record + + # Task 1 holds permit for 0.1s + async def task1(): + async with cm._permit(True): + await asyncio.sleep(0.1) + + # Task 2 waits for permit (>= 0.1s wait) then performs fetch + async def task2(): + await asyncio.sleep(0.01) # start after task1 holds permit + await cm._fetch_bytes(str(tmp_path / "file.jpg"), size=5, gated=True) + + await asyncio.gather(task1(), task2()) + + assert len(recorded_durations) == 1 + # Duration recorded must measure only task2's actual fetch (20ms), excluding the ~100ms permit queue wait + assert recorded_durations[0] < 0.08