diff --git a/support_scripts/bench_validate_shards.py b/support_scripts/bench_validate_shards.py new file mode 100644 index 0000000..0456f31 --- /dev/null +++ b/support_scripts/bench_validate_shards.py @@ -0,0 +1,356 @@ +#!/usr/bin/env python +#%% +"""bench_validate_shards.py — benchmark wsds shard validation strategies.""" + +import os +import time +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from pathlib import Path + +import wsds +from wsds.utils import preload_shard, validate_shards + +# ---- config ---- +DATASET_PATH = "/mnt/weka/data-wsds/data-ar/indices/source" # <-- change me +MAX_SHARDS = None # None = all shards +N_RUNS = 3 # timed runs per strategy +CHUNK_SIZE = 16 # shards per thread for thread_chunked / process_chunked +# hybrid auto-tunes (n_processes, n_threads) from workload size; see _choose_hybrid_params + +# %% + +ds = wsds.WSDataset(str(DATASET_PATH)) +shard_list = ds.get_shard_list() +if MAX_SHARDS: + shard_list = shard_list[:MAX_SHARDS] + +col_dirs = sorted( + {locs[0][0] for locs in ds.fields.values() if locs} - set(ds.computed_columns) +) + +print(f"shards: {len(shard_list):,}") +print(f"col_dirs: {len(col_dirs)} {col_dirs[:5]}{'…' if len(col_dirs) > 5 else ''}") + +# %% +def validate_stat_only(dataset, shards, column_dirs): + """Main-thread os.path.exists() — no file I/O, no subprocess.""" + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + return [ + (shard, all(os.path.exists(dataset.get_shard_path(cd, shard)) for cd in actual)) + for shard in shards + ] + + +def validate_thread(dataset, shards, column_dirs, max_workers=64): + """ThreadPoolExecutor variant — same magic-byte check, avoids fork overhead.""" + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + files = [dataset.get_shard_path(cd, s) for s in shards for cd in actual] + with ThreadPoolExecutor(max_workers=min(len(files), max_workers)) as ex: + results = list(ex.map(preload_shard, files)) + n = len(actual) + return [(s, all(results[i * n:(i + 1) * n])) for i, s in enumerate(shards)] + + +def validate_thread_chunked(dataset, shards, column_dirs, max_workers=64, chunk_size=CHUNK_SIZE): + """ThreadPoolExecutor where each task processes a chunk of shards sequentially. + + Reduces per-task scheduling overhead vs submitting one file at a time. + """ + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + + def check_chunk(chunk): + return [ + (s, all(preload_shard(dataset.get_shard_path(cd, s)) for cd in actual)) + for s in chunk + ] + + chunks = [shards[i:i + chunk_size] for i in range(0, len(shards), chunk_size)] + with ThreadPoolExecutor(max_workers=min(len(chunks), max_workers)) as ex: + chunk_results = list(ex.map(check_chunk, chunks)) + return [item for chunk in chunk_results for item in chunk] + + +def validate_thread_per_chunk(dataset, shards, column_dirs, chunk_size=32): + """Single ThreadPoolExecutor with one thread per chunk of `chunk_size` shards. + + n_threads = ceil(len(shards) / chunk_size). Each thread takes exactly one + chunk and processes it sequentially. No process pool — pure threads. + """ + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + + shard_paths = [ + (s, [dataset.get_shard_path(cd, s) for cd in actual]) + for s in shards + ] + chunks = [shard_paths[i:i + chunk_size] for i in range(0, len(shard_paths), chunk_size)] + + def check_chunk(chunk): + return [(shard, all(preload_shard(p) for p in paths)) for shard, paths in chunk] + + with ThreadPoolExecutor(max_workers=len(chunks)) as ex: + results = list(ex.map(check_chunk, chunks)) + return [item for chunk in results for item in chunk] + + +def _check_shard_chunk_threaded(args): + """Top-level worker: split a process-chunk into per-thread sub-chunks. + + args = (shard_path_chunk, n_threads) + shard_path_chunk = list of (shard_ref, [file_paths]) + Each thread receives a contiguous sub-chunk and validates it sequentially. + """ + shard_path_chunk, n_threads = args + sub_size = max(1, -(-len(shard_path_chunk) // n_threads)) # ceiling division + sub_chunks = [shard_path_chunk[i:i + sub_size] for i in range(0, len(shard_path_chunk), sub_size)] + + def check_sub(sub): + return [(shard, all(preload_shard(p) for p in paths)) for shard, paths in sub] + + with ThreadPoolExecutor(max_workers=n_threads) as ex: + results = list(ex.map(check_sub, sub_chunks)) + return [item for sub in results for item in sub] + + +def _in_debugger() -> bool: + """True when running under a debugger (pdb, debugpy, etc.). Process pools + don't play well with debuggers — child processes aren't traced and can + hang on breakpoints — so the auto-tune avoids them in that case.""" + import sys + return sys.gettrace() is not None or "debugpy" in sys.modules + + +def _choose_hybrid_params(n_ops, cpus=None, verbose=False): + """Auto-tune (n_processes, n_threads) for hybrid validation. + + Heuristics, calibrated empirically: + - in a debugger: 1 process, threads only (process pool breaks tracing) + - n_ops <= 512: 1 process, few threads. Tiny workload — even threadpool setup matters. + - n_ops <= 4096: 1 process, many threads. Process fork cost (~50–100ms each) + still dominates over the parallelism gain. + - n_ops > 4096: spread across cpu_count processes, 32 threads each. + """ + cpus = cpus or os.cpu_count() or 8 + if _in_debugger(): + # Cap threads at 64 — that's already plenty for I/O-bound work and + # avoids spawning hundreds of threads for huge workloads. + n_threads = min(64, max(1, -(-n_ops // 32))) + params = (1, n_threads) + reason = "debugger" + elif n_ops <= 512: + n_threads = min(32, max(1, -(-n_ops // 16))) + params = (1, n_threads) + reason = "tiny" + elif n_ops <= 4096: + n_threads = min(128, max(8, -(-n_ops // 32))) + params = (1, n_threads) + reason = "medium" + else: + n_threads = 32 + desired_threads = max(cpus, -(-n_ops // 32)) + n_processes = min(cpus, max(2, -(-desired_threads // n_threads))) + params = (n_processes, n_threads) + reason = "large" + if verbose: + print(f" auto-tune ({reason}): n_ops={n_ops} → n_processes={params[0]}, n_threads={params[1]}") + return params + + +def validate_hybrid(dataset, shards, column_dirs, + n_processes=None, n_threads=None, verbose=False): + """N processes × K threads. When `n_processes` / `n_threads` are None they + are auto-picked by `_choose_hybrid_params` based on the workload size.""" + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + + shard_paths = [ + (s, [dataset.get_shard_path(cd, s) for cd in actual]) + for s in shards + ] + if n_processes is None or n_threads is None: + auto_p, auto_t = _choose_hybrid_params(len(shards) * len(actual), verbose=verbose) + if n_processes is None: n_processes = auto_p + if n_threads is None: n_threads = auto_t + + # Fast path: skip the process pool when one process is enough + if n_processes <= 1: + return _check_shard_chunk_threaded((shard_paths, n_threads)) + + chunk_size = max(1, -(-len(shard_paths) // n_processes)) + chunks = [shard_paths[i:i + chunk_size] for i in range(0, len(shard_paths), chunk_size)] + with ProcessPoolExecutor(max_workers=n_processes) as ex: + chunk_results = list(ex.map(_check_shard_chunk_threaded, + [(c, n_threads) for c in chunks])) + return [item for chunk in chunk_results for item in chunk] + + +_O_NOATIME = getattr(os, "O_NOATIME", 0) # 0 on non-Linux + + +def _preload_shard_pread(shard_fname): + """Raw-syscall variant of preload_shard: open + lseek + read + close. + + Same semantics (checks 6-byte ARROW1 magic at file end) but bypasses + Python's BufferedReader. Adds O_NOATIME on Linux to skip the per-open + atime write-back (a real cost on networked filesystems). + """ + fname = str(shard_fname) + try: + fd = os.open(fname, os.O_RDONLY | _O_NOATIME) + except PermissionError: + # O_NOATIME requires file ownership; fall back without it. + try: + fd = os.open(fname, os.O_RDONLY) + except OSError: + return False + except OSError: + return False + try: + os.lseek(fd, -6, os.SEEK_END) + return os.read(fd, 6) == b"ARROW1" + except OSError: + return False + finally: + os.close(fd) + + +def _check_shard_chunk_threaded_pread(args): + """Same as _check_shard_chunk_threaded but uses _preload_shard_pread.""" + shard_path_chunk, n_threads = args + sub_size = max(1, -(-len(shard_path_chunk) // n_threads)) + sub_chunks = [shard_path_chunk[i:i + sub_size] for i in range(0, len(shard_path_chunk), sub_size)] + + def check_sub(sub): + return [(shard, all(_preload_shard_pread(p) for p in paths)) for shard, paths in sub] + + with ThreadPoolExecutor(max_workers=n_threads) as ex: + results = list(ex.map(check_sub, sub_chunks)) + return [item for sub in results for item in sub] + + +def validate_hybrid_pread(dataset, shards, column_dirs, + n_processes=None, n_threads=None, verbose=False): + """Hybrid (N processes × K threads) using raw os.open/lseek/read. + Auto-tunes n_processes / n_threads from workload size when not given.""" + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + + shard_paths = [ + (s, [dataset.get_shard_path(cd, s) for cd in actual]) + for s in shards + ] + if n_processes is None or n_threads is None: + auto_p, auto_t = _choose_hybrid_params(len(shards) * len(actual), verbose=verbose) + if n_processes is None: n_processes = auto_p + if n_threads is None: n_threads = auto_t + + if n_processes <= 1: + return _check_shard_chunk_threaded_pread((shard_paths, n_threads)) + + chunk_size = max(1, -(-len(shard_paths) // n_processes)) + chunks = [shard_paths[i:i + chunk_size] for i in range(0, len(shard_paths), chunk_size)] + with ProcessPoolExecutor(max_workers=n_processes) as ex: + chunk_results = list(ex.map(_check_shard_chunk_threaded_pread, + [(c, n_threads) for c in chunks])) + return [item for chunk in chunk_results for item in chunk] + + +def _check_shard_chunk(shard_path_chunk): + """Top-level worker for ProcessPoolExecutor (must be picklable). + + Receives a list of (shard_ref, [file_paths]) and returns [(shard_ref, bool)]. + """ + return [(shard, all(preload_shard(p) for p in paths)) for shard, paths in shard_path_chunk] + + +def validate_process_chunked(dataset, shards, column_dirs, max_workers=64, chunk_size=CHUNK_SIZE): + """ProcessPoolExecutor where each task validates a chunk of shards. + + Pre-computes file paths in the main process (dataset isn't picklable) then + dispatches chunks to worker processes. + """ + actual = [cd for cd in column_dirs if cd not in dataset.computed_columns] + if not actual or not shards: + return [] + + shard_paths = [ + (s, [dataset.get_shard_path(cd, s) for cd in actual]) + for s in shards + ] + chunks = [shard_paths[i:i + chunk_size] for i in range(0, len(shard_paths), chunk_size)] + with ProcessPoolExecutor(max_workers=min(len(chunks), max_workers)) as ex: + chunk_results = list(ex.map(_check_shard_chunk, chunks)) + return [item for chunk in chunk_results for item in chunk] + + +strategies = { + "original": lambda sl: validate_shards(ds, sl, col_dirs), + "thread": lambda sl: validate_thread(ds, sl, col_dirs), + "thread_chunked": lambda sl: validate_thread_chunked(ds, sl, col_dirs), + "process_chunked": lambda sl: validate_process_chunked(ds, sl, col_dirs), + "process_chunk64": lambda sl: validate_process_chunked(ds, sl, col_dirs, chunk_size=64), + "stat_only": lambda sl: validate_stat_only(ds, sl, col_dirs), + "thread_per_chunk": lambda sl: validate_thread_per_chunk(ds, sl, col_dirs), + "hybrid": lambda sl: validate_hybrid(ds, sl, col_dirs, verbose=True), + "hybrid_pread": lambda sl: validate_hybrid_pread(ds, sl, col_dirs, verbose=True), +} + +# %% +SHARD_LEVELS = [None, 300, 100, 10] # None = all shards + +results_by_level = {} +times_by_level = {} + +for level in SHARD_LEVELS: + sl = shard_list if level is None else shard_list[:min(level, len(shard_list))] + if not sl: + continue + print(f"\n=== {len(sl)} shards ===") + results = {} + times = {} + for name, fn in strategies.items(): + run_times = [] + for _ in range(N_RUNS): + t0 = time.perf_counter() + out = fn(sl) + run_times.append(time.perf_counter() - t0) + results[name] = dict(out) + times[name] = run_times + n_ok = sum(1 for ok in results[name].values() if ok) + n_bad = len(results[name]) - n_ok + best_ms = min(run_times) * 1000 + print(f" {name:16s} best={best_ms:.0f}ms shards={len(results[name])} ({n_ok} ok, {n_bad} missing) runs={[f'{t*1000:.0f}ms' for t in run_times]}") + results_by_level[len(sl)] = results + times_by_level[len(sl)] = times + +# %% +# Verify all strategies agree on which shards are valid (per level) +for n, results in results_by_level.items(): + ref_name = next(iter(results)) + ref = results[ref_name] + print(f"\n=== {n} shards: vs {ref_name} ===") + for name, res in results.items(): + match = res == ref + print(f" {name:16s}: {'✓ match' if match else '✗ MISMATCH'}") + +# %% +# Summary: speedup vs first strategy (per level) +for n, times in times_by_level.items(): + base_name = next(iter(times)) + base = min(times[base_name]) + print(f"\n=== {n} shards: vs {base_name} ===") + for name, ts in times.items(): + best = min(ts) + ratio = base / best if best > 0 else float("inf") + print(f" {name:16s} best={best*1000:.0f}ms ({ratio:.1f}x vs {base_name})") + +# %% diff --git a/wsds/utils.py b/wsds/utils.py index 8643ade..b9b496b 100644 --- a/wsds/utils.py +++ b/wsds/utils.py @@ -1,7 +1,8 @@ import os import re +import sys import traceback -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING @@ -248,56 +249,131 @@ def format_duration(duration): return f"{hours:.2f} hours" +_O_NOATIME = getattr(os, "O_NOATIME", 0) # 0 on non-Linux platforms + + def preload_shard(shard_fname): + """Open a shard file and verify the trailing 6-byte ARROW1 magic. + + Uses raw os.open/lseek/read/close (bypassing Python's BufferedReader) and + O_NOATIME on Linux (skips a per-open atime write back to the FS server). + Falls back to plain O_RDONLY on EPERM (file not owned by us). + """ + fname = str(shard_fname) try: - with open(shard_fname, "rb") as f: - f.seek(-6, os.SEEK_END) - if f.read() != b"ARROW1": - print(f"Invalid file format {shard_fname}: ARROW1 magic not found") - return False + try: + fd = os.open(fname, os.O_RDONLY | _O_NOATIME) + except PermissionError: + fd = os.open(fname, os.O_RDONLY) except FileNotFoundError: return False except OSError as err: traceback.print_exc() - print(f"OSError while loading {shard_fname}: {err}") + print(f"OSError while loading {fname}: {err}") return False + try: + os.lseek(fd, -6, os.SEEK_END) + if os.read(fd, 6) != b"ARROW1": + print(f"Invalid file format {fname}: ARROW1 magic not found") + return False + except OSError as err: + traceback.print_exc() + print(f"OSError while loading {fname}: {err}") + return False + finally: + os.close(fd) return True +def _in_debugger() -> bool: + """True under pdb/debugpy. Used to skip ProcessPoolExecutor — child + processes aren't traced and can hang on breakpoints.""" + return sys.gettrace() is not None or "debugpy" in sys.modules + + +def _choose_validation_params(n_ops: int, cpus: int | None = None) -> tuple[int, int]: + """Pick (n_processes, n_threads) for `validate_shards` based on workload size. + + Calibrated empirically on Weka with ~50µs/op effective latency: + - debugger: 1 process, threads only (process pool breaks tracing) + - n_ops <= 512: 1 process, few threads. Threadpool setup itself matters here. + - n_ops <= 4096: 1 process, many threads. Process fork (~50–100ms each) + still dominates over the parallelism gain. + - n_ops > 4096: spread across cpu_count processes, 32 threads each. + """ + cpus = cpus or os.cpu_count() or 8 + if _in_debugger(): + return 1, min(64, max(1, -(-n_ops // 32))) + if n_ops <= 512: + return 1, min(32, max(1, -(-n_ops // 16))) + if n_ops <= 4096: + return 1, min(128, max(8, -(-n_ops // 32))) + n_threads = 32 + desired_threads = max(cpus, -(-n_ops // 32)) + n_processes = min(cpus, max(2, -(-desired_threads // n_threads))) + return n_processes, n_threads + + +def _validate_shard_chunk(args): + """Top-level worker (must be picklable for ProcessPoolExecutor). + + Receives (shard_path_chunk, n_threads). Splits the chunk into per-thread + sub-chunks and runs `preload_shard` on each file sequentially within a + thread. + """ + shard_path_chunk, n_threads = args + sub_size = max(1, -(-len(shard_path_chunk) // n_threads)) + sub_chunks = [shard_path_chunk[i:i + sub_size] for i in range(0, len(shard_path_chunk), sub_size)] + + def check_sub(sub): + return [(shard, all(preload_shard(p) for p in paths)) for shard, paths in sub] + + with ThreadPoolExecutor(max_workers=n_threads) as ex: + results = list(ex.map(check_sub, sub_chunks)) + return [item for sub in results for item in sub] + + def validate_shards( - dataset: "WSDataset", shards: list[tuple[str, str]], column_dirs: list[str], tail_bytes: int = 10240 + dataset: "WSDataset", + shards: list[tuple[str, str]], + column_dirs: list[str], ): """Prefetch and validate shard files for the given shards and column dirs. - Uses a ProcessPoolExecutor to load them concurrently, which helps with network filesystems - where latency is the bottleneck. This is useful before operations that need to read - from multiple shards (like WSSample.__repr__ or SQL queries). + Verifies each shard's ARROW1 magic across every column dir. Auto-tunes + its concurrency (single-process threads for small workloads, process pool + for large ones) via `_choose_validation_params`. - Args: - dataset: The WSDataset instance - shards: List of (partition, shard_name) tuples identifying the shards - column_dirs: List of column directory names to prefetch - tail_bytes: Number of bytes to read from the end of each file (default 10KB) + Already-validated shards on `dataset` (tracked in + `dataset._validated_shards`) are skipped and returned with `ok=True`. - Returns: - List of shards that loaded successfully across all column dirs + Returns a list of `(shard_ref, ok)` for every input shard. """ - - # Filter out computed columns (they don't have actual shard files) - actual_column_dirs = [s for s in column_dirs if s not in dataset.computed_columns] - + actual_column_dirs = [c for c in column_dirs if c not in dataset.computed_columns] if not actual_column_dirs or not shards: - return [] + return [(s, True) for s in shards] - # Create all combinations of shards and column dirs - shard_files = [ - dataset.get_shard_path(column_dir, shard_name) for shard_name in shards for column_dir in actual_column_dirs - ] + cache_key = tuple(sorted(actual_column_dirs)) + cache: set = dataset._validated_shards.setdefault(cache_key, set()) + pending = [s for s in shards if s not in cache] + if not pending: + return [(s, True) for s in shards] - with ProcessPoolExecutor(max_workers=min(len(shard_files), 64)) as executor: - results = list(executor.map(preload_shard, shard_files)) + shard_paths = [ + (s, [dataset.get_shard_path(cd, s) for cd in actual_column_dirs]) + for s in pending + ] + n_processes, n_threads = _choose_validation_params(len(shard_paths) * len(actual_column_dirs)) - # Check that all column dirs loaded successfully for each shard - # Return all given shards with an ok flag - num_column_dirs = len(actual_column_dirs) - return [(shard, all(results[i * num_column_dirs : (i + 1) * num_column_dirs])) for i, shard in enumerate(shards)] + if n_processes <= 1: + verified = _validate_shard_chunk((shard_paths, n_threads)) + else: + chunk_size = max(1, -(-len(shard_paths) // n_processes)) + chunks = [shard_paths[i:i + chunk_size] for i in range(0, len(shard_paths), chunk_size)] + with ProcessPoolExecutor(max_workers=n_processes) as ex: + chunk_results = list(ex.map(_validate_shard_chunk, [(c, n_threads) for c in chunks])) + verified = [item for chunk in chunk_results for item in chunk] + + verified_map = dict(verified) + cache.update(s for s, ok in verified if ok) + return [(s, True) if s in cache else (s, verified_map.get(s, False)) for s in shards] diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py index 580fb57..f6ccce1 100644 --- a/wsds/ws_dataset.py +++ b/wsds/ws_dataset.py @@ -107,6 +107,8 @@ def __init__( self._open_shards = {} self._linked_datasets = {} + # column-dirs tuple -> set of shard refs that already passed validate_shards + self._validated_shards: dict[tuple[str, ...], set] = {} self._register_wsds_links() @@ -118,6 +120,7 @@ def close(self): for ds in self._linked_datasets.values(): ds.close() self._linked_datasets.clear() + self._validated_shards.clear() def enable_filter(self, filter_name: str, filter_df: pl.DataFrame): """