From 8acf1ed92fc0453ad1c03122ccb191afa210cab8 Mon Sep 17 00:00:00 2001 From: hillhack <2jyotihill@gmail.com> Date: Sun, 9 Aug 2026 18:38:56 +0530 Subject: [PATCH 1/6] feat(raw): make dataset concurrency budget dynamic with process-local EMA feedback and env overrides --- src/litdata/raw/dataset.py | 188 ++++++++++++++++++++++++++++++++++--- tests/raw/test_dataset.py | 92 ++++++++++++++++++ 2 files changed, 268 insertions(+), 12 deletions(-) diff --git a/src/litdata/raw/dataset.py b/src/litdata/raw/dataset.py index cf0a3b7de..b3a12da9a 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -108,6 +108,113 @@ # 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 = 1024 * 1024 + + +def _env_float(name: str, default: float) -> float: + val = os.getenv(name) + if val is None: + return default + try: + return float(val) + except ValueError: + return default + + +def _env_int(name: str, default: int) -> int: + val = os.getenv(name) + if val is None: + return default + try: + return int(val) + except ValueError: + return default + + +def _get_assumed_aggregate_bandwidth_bps() -> int: + return _env_int("LITDATA_ASSUMED_BANDWIDTH_BPS", _ASSUMED_AGGREGATE_BANDWIDTH_BPS) + + +def _get_assumed_request_rate() -> float: + return _env_float("LITDATA_ASSUMED_REQUEST_RATE", _ASSUMED_REQUEST_RATE) + + +def _get_assumed_request_latency_s() -> float: + return _env_float("LITDATA_ASSUMED_REQUEST_LATENCY_S", _ASSUMED_REQUEST_LATENCY_S) + + +def _get_default_median_file_bytes() -> int: + return _env_int("LITDATA_DEFAULT_MEDIAN_FILE_BYTES", _DEFAULT_MEDIAN_FILE_BYTES) + + +def _get_single_process_concurrency_cap() -> int: + return _env_int("LITDATA_SINGLE_PROCESS_CONCURRENCY_CAP", _SINGLE_PROCESS_CONCURRENCY_CAP) + + +def _get_aggregate_concurrency_budget_cap() -> int: + return _env_int("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_CAP", _AGGREGATE_CONCURRENCY_BUDGET_CAP) + + +def _get_aggregate_concurrency_budget_floor() -> int: + return _env_int("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR", _AGGREGATE_CONCURRENCY_BUDGET_FLOOR) + + +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.sample_count: int = 0 + self._lock = threading.Lock() + + def record_observation(self, size_bytes: int, duration_s: float) -> None: + if size_bytes <= 0 or duration_s <= 0: + return + with self._lock: + self.sample_count += 1 + if size_bytes < _LATENCY_OBSERVATION_MAX_BYTES: + if self.request_latency_s_ema is None: + self.request_latency_s_ema = duration_s + else: + self.request_latency_s_ema = ( + self.alpha * duration_s + (1.0 - self.alpha) * self.request_latency_s_ema + ) + elif size_bytes >= _BANDWIDTH_OBSERVATION_MIN_BYTES: + lat_est = ( + self.request_latency_s_ema + if self.request_latency_s_ema is not None + else _get_assumed_request_latency_s() + ) + transfer_time = max(0.001, duration_s - lat_est) + obs_bps = float(size_bytes) / transfer_time + 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 + + def get_metrics(self) -> tuple[float | None, float | None, int]: + with self._lock: + return self.bandwidth_bps_ema, self.request_latency_s_ema, self.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, + "sample_count": self.sample_count, + } + + 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.sample_count = state.get("sample_count", 0) + self._lock = threading.Lock() + _RUNNER_LOCK = threading.Lock() _RUNNER: _LoopRunner | None = None @@ -191,7 +298,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,7 +520,10 @@ 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]``: @@ -425,26 +535,46 @@ def _aggregate_concurrency_budget(median_file_bytes: int | None) -> int: tiny-object paths are not request-starved. Medians ≥1 MiB stay bandwidth-bounded (avoids pinning at 240 slots → multi-GB in flight). + Uses empirical EMA metrics from ``tracker`` when ``sample_count >= 5``. Per-worker floor of 8 means realized aggregate is ``max(budget, 8 × num_workers)``. """ - 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 + if tracker is not None: + bps_ema, lat_ema, samples = tracker.get_metrics() + if samples >= 5: + obs_bps = bps_ema + 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. 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() + req_lat = obs_lat if obs_lat is not None and obs_lat > 0 else _get_assumed_request_latency_s() + latency_model = max(1, int(req_rate * req_lat)) else: latency_model = 0 + raw = max(bandwidth_model, latency_model) - return max(_AGGREGATE_CONCURRENCY_BUDGET_FLOOR, min(_AGGREGATE_CONCURRENCY_BUDGET_CAP, raw)) + return max(floor, min(cap, raw)) def _effective_concurrency( max_concurrent_downloads: int | None, num_workers: int, median_file_bytes: int | None = None, + tracker: BandwidthTracker | 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 @@ -455,9 +585,9 @@ def _effective_concurrency( """ 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 min(budget, _get_single_process_concurrency_cap()) return max(_MIN_CONCURRENCY_PER_WORKER, budget // num_workers) @@ -557,6 +687,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 +741,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 +777,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 +853,14 @@ 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.""" + prev_count = self._bandwidth_tracker.sample_count + self._bandwidth_tracker.record_observation(size_bytes, duration_s) + new_count = self._bandwidth_tracker.sample_count + if (prev_count < 5 and new_count >= 5) or (new_count >= 5 and new_count % 10 == 0): + self._cached_permits = None + def _effective_download_permits(self) -> int: """Worker-aware permit count for the download semaphore (Stage 1 static clamp). @@ -733,6 +874,7 @@ 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 @@ -750,7 +892,7 @@ def _get_semaphore(self) -> asyncio.Semaphore: if self._semaphore is None or self._semaphore_loop is not loop or self._semaphore_permits != permits: 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 ) @@ -1021,11 +1163,15 @@ async def fetch() -> bytes: data = await fetch() return offset, data + t0 = time.monotonic() parts = await asyncio.gather(*(one(o, n) for o, n in ranges)) + dur = time.monotonic() - t0 parts.sort(key=lambda x: x[0]) joined = b"".join(data for _, data in parts) if len(joined) != size: raise RuntimeError(f"Ranged download size mismatch for {file_path}: expected={size} got={len(joined)}") + if len(joined) > 0: + self._record_download_observation(len(joined), dur) return joined async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: bool = True) -> bytes: @@ -1048,14 +1194,24 @@ async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: delay = _effective_hedge_delay(self.hedge_delay, size) if self._is_remote_object(file_path) else None # Pay-per-use: hedging off/ineligible → bare permit + download (batch enforces timeout). if delay is None: + t0 = time.monotonic() async with self._permit(gated): - return await self.downloader.adownload_fileobj(file_path) + 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) - return await self._hedged(once, delay) + t0_h = time.monotonic() + data = await self._hedged(once, delay) + dur_h = time.monotonic() - t0_h + if len(data) > 0: + self._record_download_observation(len(data), dur_h) + return data def _schedule_write_behind(self, local_path: str, data: bytes) -> None: """Atomically publish ``data`` to ``local_path`` on a worker thread.""" @@ -1098,9 +1254,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 +1275,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..825fd9445 100644 --- a/tests/raw/test_dataset.py +++ b/tests/raw/test_dataset.py @@ -539,3 +539,95 @@ 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_ema_and_partitioning(): + from litdata.raw.dataset import BandwidthTracker + + tracker = BandwidthTracker(alpha=0.2) + assert tracker.sample_count == 0 + + # Small GET (< 256 KB) updates latency EMA + tracker.record_observation(100_000, 0.020) + bps, lat, count = tracker.get_metrics() + assert count == 1 + assert bps is None + assert pytest.approx(lat, abs=1e-6) == 0.020 + + tracker.record_observation(50_000, 0.010) + _, lat, count = tracker.get_metrics() + assert count == 2 + # EMA: 0.2 * 0.010 + 0.8 * 0.020 = 0.018 + assert pytest.approx(lat, abs=1e-6) == 0.018 + + # Large GET (>= 1 MiB) updates bandwidth EMA + tracker.record_observation(10 * 1024 * 1024, 0.118) + bps, lat, count = tracker.get_metrics() + assert count == 3 + assert bps is not None + assert pytest.approx(bps, abs=1.0) == 104_857_600.0 + + +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.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_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 From ea5e043adf932324646863ff6534bfa5ed0460c6 Mon Sep 17 00:00:00 2001 From: hillhack <2jyotihill@gmail.com> Date: Sun, 16 Aug 2026 09:28:10 +0530 Subject: [PATCH 2/6] fix(raw): decouple baseline capacity from congestion backoff and fix worker allocation --- .github/markdown-links-config.json | 3 + src/litdata/raw/dataset.py | 132 +++++++++++++++------ tests/raw/test_dataset.py | 177 ++++++++++++++++++++++++++--- 3 files changed, 261 insertions(+), 51 deletions(-) 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/src/litdata/raw/dataset.py b/src/litdata/raw/dataset.py index b3a12da9a..7b466b889 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -109,7 +109,10 @@ # (8 MiB duplicate-GET hedge policy). _LATENCY_MODEL_MAX_MEDIAN_BYTES = 1024 * 1024 _LATENCY_OBSERVATION_MAX_BYTES = 256 * 1024 -_BANDWIDTH_OBSERVATION_MIN_BYTES = 1024 * 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 def _env_float(name: str, default: float) -> float: @@ -167,37 +170,47 @@ 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.sample_count: int = 0 + self.bps_sample_count: int = 0 + self.lat_sample_count: int = 0 + self._sample_count: int = 0 self._lock = threading.Lock() + @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) -> None: if size_bytes <= 0 or duration_s <= 0: return with self._lock: - self.sample_count += 1 + self._sample_count += 1 if size_bytes < _LATENCY_OBSERVATION_MAX_BYTES: + self.lat_sample_count += 1 if self.request_latency_s_ema is None: self.request_latency_s_ema = duration_s else: self.request_latency_s_ema = ( self.alpha * duration_s + (1.0 - self.alpha) * self.request_latency_s_ema ) - elif size_bytes >= _BANDWIDTH_OBSERVATION_MIN_BYTES: - lat_est = ( - self.request_latency_s_ema - if self.request_latency_s_ema is not None - else _get_assumed_request_latency_s() - ) - transfer_time = max(0.001, duration_s - lat_est) - obs_bps = float(size_bytes) / transfer_time + 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 - def get_metrics(self) -> tuple[float | None, float | None, int]: + 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.sample_count + 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: @@ -205,14 +218,24 @@ def __getstate__(self) -> dict[str, Any]: "alpha": self.alpha, "bandwidth_bps_ema": self.bandwidth_bps_ema, "request_latency_s_ema": self.request_latency_s_ema, - "sample_count": self.sample_count, + "bps_sample_count": self.bps_sample_count, + "lat_sample_count": self.lat_sample_count, + "sample_count": self._sample_count, } 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.sample_count = state.get("sample_count", 0) + 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) + 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() @@ -526,17 +549,12 @@ def _aggregate_concurrency_budget( ) -> int: """Aggregate in-flight download slots across all workers (size-aware, clamped). - Takes the max of two models then clamps to ``[floor, cap]``: + 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. - - **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). - - Uses empirical EMA metrics from ``tracker`` when ``sample_count >= 5``. - 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). """ default_median = _get_default_median_file_bytes() median = median_file_bytes if median_file_bytes and median_file_bytes > 0 else default_median @@ -546,10 +564,13 @@ def _aggregate_concurrency_budget( 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, samples = tracker.get_metrics() - if samples >= 5: + 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() @@ -557,15 +578,39 @@ def _aggregate_concurrency_budget( 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: req_rate = _get_assumed_request_rate() - req_lat = obs_lat if obs_lat is not None and obs_lat > 0 else _get_assumed_request_latency_s() - latency_model = max(1, int(req_rate * req_lat)) + latency_model = max(1, int(req_rate * target_lat)) else: latency_model = 0 - raw = max(bandwidth_model, latency_model) - return max(floor, min(cap, raw)) + # max is intentional to avoid underestimating baseline capacity + baseline_budget = max(bandwidth_model, latency_model) + + # Stateless congestion control backoff: latency > target => reduce budget + if obs_lat is not None and obs_lat > target_lat: + backoff_factor = min(1.0, target_lat / obs_lat) + computed_budget = max(1, int(baseline_budget * backoff_factor)) + else: + computed_budget = baseline_budget + + # 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( @@ -573,13 +618,13 @@ def _effective_concurrency( 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/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. """ @@ -588,7 +633,24 @@ def _effective_concurrency( budget = _aggregate_concurrency_budget(median_file_bytes, tracker=tracker) if num_workers <= 1: return min(budget, _get_single_process_concurrency_cap()) - return max(_MIN_CONCURRENCY_PER_WORKER, budget // num_workers) + + 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: diff --git a/tests/raw/test_dataset.py b/tests/raw/test_dataset.py index 825fd9445..eacd2e2c3 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 ], @@ -541,31 +541,41 @@ def transform(data): assert gds[0] == b"abc" -def test_bandwidth_tracker_ema_and_partitioning(): +def test_bandwidth_tracker_basic(): + import pytest + from litdata.raw.dataset import BandwidthTracker tracker = BandwidthTracker(alpha=0.2) assert tracker.sample_count == 0 - # Small GET (< 256 KB) updates latency EMA - tracker.record_observation(100_000, 0.020) - bps, lat, count = tracker.get_metrics() - assert count == 1 + # Tiny GET (< 64 KB) updates latency EMA only + 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 - assert pytest.approx(lat, abs=1e-6) == 0.020 + assert pytest.approx(lat, abs=1e-6) == 0.010 - tracker.record_observation(50_000, 0.010) - _, lat, count = tracker.get_metrics() - assert count == 2 - # EMA: 0.2 * 0.010 + 0.8 * 0.020 = 0.018 - assert pytest.approx(lat, abs=1e-6) == 0.018 + # Small/Medium GET (64 KB <= size < 256 KB) updates BOTH latency & bandwidth EMA + 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 + # EMA for latency: 0.2 * 0.020 + 0.8 * 0.010 = 0.012 + assert pytest.approx(lat, abs=1e-6) == 0.012 - # Large GET (>= 1 MiB) updates bandwidth EMA + # Large GET (>= 256 KB) updates bandwidth EMA tracker.record_observation(10 * 1024 * 1024, 0.118) - bps, lat, count = tracker.get_metrics() - assert count == 3 + bps, lat, bps_count, lat_count = tracker.get_metrics() + assert tracker.sample_count == 3 + assert bps_count == 2 + assert lat_count == 2 assert bps is not None - assert pytest.approx(bps, abs=1.0) == 104_857_600.0 def test_bandwidth_tracker_pickle(): @@ -579,6 +589,8 @@ def test_bandwidth_tracker_pickle(): 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 @@ -618,6 +630,37 @@ def test_concurrency_budget_high_and_low_bandwidth_adaptation(): 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, @@ -631,3 +674,105 @@ def test_environment_variable_overrides(monkeypatch): 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 + 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 + assert lat is not None + assert abs(lat - 0.015) < 1e-4 + # ImageNet files (150 KB) MUST record bandwidth observation (bps != None) + assert bps is not None + assert bps > 0 + + +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 From 29d7a138614e2223fe7b5b644aeb2a24cabecba6 Mon Sep 17 00:00:00 2001 From: hillhack <2jyotihill@gmail.com> Date: Mon, 17 Aug 2026 23:02:59 +0530 Subject: [PATCH 3/6] fix(raw): use transfer-subtracted latency for small GETs avoiding circular self-estimation --- src/litdata/raw/dataset.py | 53 ++++++++++++- tests/raw/test_dataset.py | 149 +++++++++++++++++++++++++++++++++---- 2 files changed, 185 insertions(+), 17 deletions(-) diff --git a/src/litdata/raw/dataset.py b/src/litdata/raw/dataset.py index 7b466b889..3c1e55457 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -113,6 +113,12 @@ _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) -> float: @@ -182,18 +188,61 @@ def sample_count(self) -> int: return self._sample_count def record_observation(self, size_bytes: int, duration_s: float) -> None: + """Record an empirical GET observation and update EMA estimates. + + 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 with self._lock: self._sample_count += 1 + + # 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 = duration_s + self.request_latency_s_ema = estimated_request_latency_s else: self.request_latency_s_ema = ( - self.alpha * duration_s + (1.0 - self.alpha) * 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 diff --git a/tests/raw/test_dataset.py b/tests/raw/test_dataset.py index eacd2e2c3..e95f66a63 100644 --- a/tests/raw/test_dataset.py +++ b/tests/raw/test_dataset.py @@ -544,21 +544,34 @@ def transform(data): def test_bandwidth_tracker_basic(): import pytest - from litdata.raw.dataset import BandwidthTracker + 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 (< 64 KB) updates latency EMA only + # --- 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 - assert pytest.approx(lat, abs=1e-6) == 0.010 - - # Small/Medium GET (64 KB <= size < 256 KB) updates BOTH latency & bandwidth EMA + 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 @@ -566,15 +579,18 @@ def test_bandwidth_tracker_basic(): assert lat_count == 2 assert bps is not None assert pytest.approx(bps, abs=1.0) == 5_000_000.0 - # EMA for latency: 0.2 * 0.020 + 0.8 * 0.010 = 0.012 - assert pytest.approx(lat, abs=1e-6) == 0.012 + 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 (>= 256 KB) updates bandwidth EMA + # --- 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 + assert lat_count == 2 # unchanged — large GETs do not update latency EMA assert bps is not None @@ -724,10 +740,12 @@ def test_concurrency_latency_recovery(): def test_imagenet_bandwidth_observation(): + import pytest + from litdata.raw.dataset import BandwidthTracker tracker = BandwidthTracker() - # ImageNet JPEG size ~150 KB + # 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) @@ -735,11 +753,17 @@ def test_imagenet_bandwidth_observation(): bps, lat, bps_cnt, lat_cnt = tracker.get_metrics() assert bps_cnt == 5 assert lat_cnt == 5 - assert lat is not None - assert abs(lat - 0.015) < 1e-4 - # ImageNet files (150 KB) MUST record bandwidth observation (bps != None) - assert bps is not None - assert bps > 0 + + # ImageNet files MUST record both EMA estimates. + assert bps is not None and bps > 0 + assert lat is not None and 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(): @@ -776,3 +800,98 @@ def test_worker_allocation_respects_aggregate_budget(): ) # 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 and establish a known bandwidth EMA (20 MB/s), + then record 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 the configured/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 (e.g. cached/fast GET). + 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 that the bandwidth calculated from the current observation is NOT used to calculate + that observation's own transfer time. + """ + import pytest + + 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 From aca6daaa87b7a8eb95f6583a672b57665ecb0151 Mon Sep 17 00:00:00 2001 From: hillhack <2jyotihill@gmail.com> Date: Tue, 18 Aug 2026 08:44:11 +0530 Subject: [PATCH 4/6] feat(raw): refine dynamic concurrency budget logic and tests --- benchmarks/bench_raw_adaptive_concurrency.py | 401 +++++++++++++++++++ src/litdata/raw/dataset.py | 188 +++++++-- tests/raw/test_dataset.py | 299 +++++++++++++- 3 files changed, 833 insertions(+), 55 deletions(-) create mode 100644 benchmarks/bench_raw_adaptive_concurrency.py 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/raw/dataset.py b/src/litdata/raw/dataset.py index 3c1e55457..ab87c9a6a 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -121,52 +121,85 @@ _LATENCY_RTT_EPSILON = _LATENCY_EPSILON_S -def _env_float(name: str, default: float) -> float: - val = os.getenv(name) - if val is None: +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: - return float(val) + 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) -> int: - val = os.getenv(name) - if val is None: +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: - return int(val) + 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) + 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) + 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) + 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) + 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) + 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) + return _env_int("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_CAP", _AGGREGATE_CONCURRENCY_BUDGET_CAP, min_val=1) def _get_aggregate_concurrency_budget_floor() -> int: - return _env_int("LITDATA_AGGREGATE_CONCURRENCY_BUDGET_FLOOR", _AGGREGATE_CONCURRENCY_BUDGET_FLOOR) + 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: @@ -179,8 +212,35 @@ def __init__(self, alpha: float = 0.2) -> 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.""" @@ -225,9 +285,7 @@ def record_observation(self, size_bytes: int, duration_s: float) -> None: # 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 + 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 @@ -238,8 +296,7 @@ def record_observation(self, size_bytes: int, duration_s: float) -> 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 + 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. @@ -288,6 +345,42 @@ def __setstate__(self, state: dict[str, Any]) -> None: 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) + + 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 @@ -638,12 +731,15 @@ def _aggregate_concurrency_budget( # max is intentional to avoid underestimating baseline capacity baseline_budget = max(bandwidth_model, latency_model) - # Stateless congestion control backoff: latency > target => reduce budget - if obs_lat is not None and obs_lat > target_lat: + # 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) - computed_budget = max(1, int(baseline_budget * backoff_factor)) else: - computed_budget = baseline_budget + backoff_factor = 1.0 + + computed_budget = max(1, int(baseline_budget * backoff_factor)) # Guarded adaptive floor reduction: require high-confidence combined evidence if ( @@ -991,16 +1087,16 @@ def _effective_download_permits(self) -> int: 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, tracker=self._bandwidth_tracker) @@ -1014,9 +1110,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 @@ -1250,8 +1352,9 @@ async def fetch() -> bytes: # Unique scratch per attempt so first/hedge never share a path. scratch = f"{base_scratch}.{offset}.{uuid4().hex}" try: + t0_c = time.monotonic() async with self._permit(gated): - data = await asyncio.get_running_loop().run_in_executor( + await asyncio.get_running_loop().run_in_executor( executor, downloader.download_bytes, file_path, @@ -1259,11 +1362,16 @@ async def fetch() -> bytes: length, scratch, ) + data = Path(scratch).read_bytes() + 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) @@ -1274,15 +1382,11 @@ async def fetch() -> bytes: data = await fetch() return offset, data - t0 = time.monotonic() parts = await asyncio.gather(*(one(o, n) for o, n in ranges)) - dur = time.monotonic() - t0 parts.sort(key=lambda x: x[0]) joined = b"".join(data for _, data in parts) if len(joined) != size: raise RuntimeError(f"Ranged download size mismatch for {file_path}: expected={size} got={len(joined)}") - if len(joined) > 0: - self._record_download_observation(len(joined), dur) return joined async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: bool = True) -> bytes: @@ -1314,15 +1418,15 @@ async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: return data async def once() -> bytes: + t0_once = time.monotonic() async with self._permit(gated): - return await self.downloader.adownload_fileobj(file_path) + 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 - t0_h = time.monotonic() - data = await self._hedged(once, delay) - dur_h = time.monotonic() - t0_h - if len(data) > 0: - self._record_download_observation(len(data), dur_h) - return data + return await self._hedged(once, delay) def _schedule_write_behind(self, local_path: str, data: bytes) -> None: """Atomically publish ``data`` to ``local_path`` on a worker thread.""" diff --git a/tests/raw/test_dataset.py b/tests/raw/test_dataset.py index e95f66a63..98e84b52b 100644 --- a/tests/raw/test_dataset.py +++ b/tests/raw/test_dataset.py @@ -740,8 +740,6 @@ def test_concurrency_latency_recovery(): def test_imagenet_bandwidth_observation(): - import pytest - from litdata.raw.dataset import BandwidthTracker tracker = BandwidthTracker() @@ -755,8 +753,10 @@ def test_imagenet_bandwidth_observation(): assert lat_cnt == 5 # ImageNet files MUST record both EMA estimates. - assert bps is not None and bps > 0 - assert lat is not None and lat > 0 + 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) @@ -803,8 +803,10 @@ def test_worker_allocation_respects_aggregate_budget(): def test_transfer_subtracted_latency(): - """Pre-seed the tracker with 5 large GETs to reach sample threshold and establish a known bandwidth EMA (20 MB/s), - then record a 200 KB request taking 30 ms. Verify latency EMA is approximately 30 ms - transfer_time (10 ms) = 20 ms. + """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 @@ -835,7 +837,7 @@ def test_transfer_subtracted_latency(): def test_transfer_subtracted_latency_bootstrap(): - """When no empirical bandwidth exists or bps_sample_count < 5, verify the configured/default bandwidth is used.""" + """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 @@ -856,7 +858,8 @@ def test_transfer_subtracted_latency_bootstrap(): def test_transfer_subtracted_latency_clamped(): - """Provide an observation where estimated transfer time > observed duration (e.g. cached/fast GET). + """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 @@ -879,11 +882,7 @@ def test_transfer_subtracted_latency_clamped(): def test_current_bandwidth_sample_does_not_explain_itself(): - """Verify that the bandwidth calculated from the current observation is NOT used to calculate - that observation's own transfer time. - """ - import pytest - + """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) @@ -895,3 +894,277 @@ def test_current_bandwidth_sample_does_not_explain_itself(): 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()) From 0fbf393034ac75cc804123164bd403d3754476e4 Mon Sep 17 00:00:00 2001 From: hillhack <2jyotihill@gmail.com> Date: Tue, 18 Aug 2026 08:55:28 +0530 Subject: [PATCH 5/6] docs(raw): update CHANGELOG and ADAPTIVE_CONCURRENCY reference --- benchmarks/ADAPTIVE_CONCURRENCY.md | 34 ++++++++++++++++++++++++++++++ src/litdata/CHANGELOG.md | 1 + 2 files changed, 35 insertions(+) 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/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 From 765dcf1fac8d6ad24bbfafc8f33cb2d8e566ee94 Mon Sep 17 00:00:00 2001 From: hillhack <2jyotihill@gmail.com> Date: Tue, 18 Aug 2026 11:26:45 +0530 Subject: [PATCH 6/6] fix(raw): refine telemetry duration measurement and atomic sample count transitions - Update BandwidthTracker.record_observation to return (prev_sample_count, new_sample_count) tuple atomically under lock. - Update CacheManager._record_download_observation to consume atomic return tuple and use _get_min_empirical_samples() and _get_permit_refresh_interval() instead of magic constants. - Move timer start inside async with self._permit(gated) in _fetch_bytes and _ranged_download_bytes to exclude semaphore queuing wait time from network telemetry. - Add locked() method delegation to _DynamicSemaphore. - Add unit tests verifying atomic sample transitions, permit refresh threshold compliance, and telemetry duration permit wait exclusion. --- src/litdata/raw/dataset.py | 48 +++++++++++++------ tests/raw/test_dataset.py | 97 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 131 insertions(+), 14 deletions(-) diff --git a/src/litdata/raw/dataset.py b/src/litdata/raw/dataset.py index ab87c9a6a..8a6eaa699 100644 --- a/src/litdata/raw/dataset.py +++ b/src/litdata/raw/dataset.py @@ -247,9 +247,11 @@ def sample_count(self) -> int: with self._lock: return self._sample_count - def record_observation(self, size_bytes: int, duration_s: float) -> None: + 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. @@ -269,9 +271,11 @@ def record_observation(self, size_bytes: int, duration_s: float) -> None: processing, scheduling, and proxy overhead are all included). """ if size_bytes <= 0 or duration_s <= 0: - return + 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 @@ -308,6 +312,8 @@ def record_observation(self, size_bytes: int, duration_s: float) -> None: 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: @@ -327,6 +333,7 @@ def __getstate__(self) -> dict[str, Any]: "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: @@ -336,6 +343,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: 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: @@ -368,6 +376,9 @@ def update_target(self, new_target: int) -> None: 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() @@ -1062,10 +1073,15 @@ def downloader(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.""" - prev_count = self._bandwidth_tracker.sample_count - self._bandwidth_tracker.record_observation(size_bytes, duration_s) - new_count = self._bandwidth_tracker.sample_count - if (prev_count < 5 and new_count >= 5) or (new_count >= 5 and new_count % 10 == 0): + 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: @@ -1352,9 +1368,9 @@ async def fetch() -> bytes: # Unique scratch per attempt so first/hedge never share a path. scratch = f"{base_scratch}.{offset}.{uuid4().hex}" try: - t0_c = time.monotonic() async with self._permit(gated): - 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, @@ -1362,8 +1378,12 @@ async def fetch() -> bytes: length, scratch, ) - data = Path(scratch).read_bytes() - dur_c = time.monotonic() - t0_c + 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)}" @@ -1409,19 +1429,19 @@ async def _fetch_bytes(self, file_path: str, size: int | None = None, *, gated: delay = _effective_hedge_delay(self.hedge_delay, size) if self._is_remote_object(file_path) else None # Pay-per-use: hedging off/ineligible → bare permit + download (batch enforces timeout). if delay is None: - t0 = time.monotonic() async with self._permit(gated): + t0 = time.monotonic() data = await self.downloader.adownload_fileobj(file_path) - dur = time.monotonic() - t0 + dur = time.monotonic() - t0 if len(data) > 0: self._record_download_observation(len(data), dur) return data async def once() -> bytes: - t0_once = time.monotonic() async with self._permit(gated): + t0_once = time.monotonic() res = await self.downloader.adownload_fileobj(file_path) - dur_once = time.monotonic() - t0_once + dur_once = time.monotonic() - t0_once if len(res) > 0: self._record_download_observation(len(res), dur_once) return res diff --git a/tests/raw/test_dataset.py b/tests/raw/test_dataset.py index 98e84b52b..3cb6158a3 100644 --- a/tests/raw/test_dataset.py +++ b/tests/raw/test_dataset.py @@ -1168,3 +1168,100 @@ async def _run(): 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