diff --git a/Makefile b/Makefile index 93e0953..2dd6b36 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,18 @@ build: sync ## Build sdist and wheel packages .PHONY: test test: sync ## Run tests - DAFT_ANALYTICS_ENABLED=0 uv run pytest tests/ -v $(EXTRA_ARGS) + # RAY_ENABLE_UV_RUN_RUNTIME_ENV=0: Ray walks the driver's process tree, sees + # the `uv run` below, and then replicates that environment to every worker -- + # packaging the whole working directory and rebuilding the venv per worker. + # Under DAFT_RUNNER=ray that makes worker startup fail once the suite has + # created enough actors: 277 tests died with ActorUnschedulableError, and the + # same suite passes in 46s with it off. Workers inherit this venv's + # interpreter from the raylet regardless. + DAFT_ANALYTICS_ENABLED=0 RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 uv run pytest tests/ -v $(EXTRA_ARGS) + +.PHONY: bench +bench: sync ## Measure lance session memory in a real Ray run + DAFT_ANALYTICS_ENABLED=0 uv run --group ray python benchmarks/lance_session_memory.py $(EXTRA_ARGS) .PHONY: format format: sync ## Format Python code diff --git a/README.md b/README.md index ee6d1b5..59a9519 100644 --- a/README.md +++ b/README.md @@ -325,6 +325,51 @@ Native namespace support in `daft.read_lance` / `DataFrame.write_lance` is track [Eventual-Inc/Daft#7282](https://github.com/Eventual-Inc/Daft/issues/7282); until that lands, use the `daft_lance` entry points shown above for namespace-addressed tables. +## Environment Variables + +| Variable | Default | Effect | +| --- | --- | --- | +| `DAFT_LANCE_NAMESPACE_CACHE_SIZE` | `16` | Size of the per-(implementation, properties) namespace client cache. Read once at import time. | +| `DAFT_LANCE_SHARE_SESSION` | `1` | Every dataset this extension opens in a process shares one lance `Session`, so a table's manifest and lance's caches are held once instead of once per scan task, maintenance UDF and writer. Set to `0` to give each open its own session (lance's own default); on a many-fragment table that costs one copy of the fragment list per concurrently open dataset. `benchmarks/lance_session_memory.py` (`make bench`) measures the difference in a real Ray worker. | + +### What sharing a session actually saves, per path + +Measured, because the three paths differ and only one of them multiplies by +concurrency. The numbers come from a `read_lance` -> `write_parquet` pipeline in +a Ray worker, which is the shape this matters most in: a task holds its parquet +writer's buffers at the same time as its manifest, so what sharing frees becomes +headroom the writer can use. + +| Path | Opens per worker process | What sharing changes | +| --- | --- | --- | +| Scans, `read_blobs`, writes | `scantask_max_parallel` concurrent opens in one process | Collapses them to one manifest. Measured: 49 MB less resident on a 4000-fragment table at concurrency 24, and 479 MB on a 100k-fragment one at concurrency 8 -- where the pipeline also finished **9.8x sooner** (610s to 62s), because each unshared task re-parsed the same 8 MB manifest. At that scale the cost of not sharing is mostly time, not memory. | +| Indexed (`nearest`, point lookups) | One index-driven task, whatever the concurrency | The index is loaded once per process instead of once per query. Measured at 200k rows: repeated queries 1.7-2.0x faster. | +| Maintenance (`compact_files`, `update_columns`, `merge_columns`) | One, always | Nothing to collapse: Daft gives each `daft.cls` instance a process of its own, so a maintenance worker opens the table exactly once however high the concurrency. | + +An index does not change the maintenance row. A maintenance task's open brings +in the index's metadata entry and nothing else; the index data is read once, on +the driver, when the commit remaps every entry onto the new fragment ids. The +tests that pin both halves down are +`tests/io/lance/test_ray_session_memory.py::test_maintenance_instances_each_get_their_own_process` +and +`tests/io/lance/test_shared_session.py::test_compaction_loads_the_index_in_its_commit_not_in_its_tasks`. + +Two knobs keep a large table affordable without weakening any of that, because +a manifest is the table's whole fragment list and an open costs the same +whatever its task goes on to read: `read_lance(..., fragment_group_size=N)` +puts several fragments in one scan task, and reading only a prefix of the table +prunes whole fragments. At 100k fragments, reading 1% through groups of ten +turns a billion-row scan into ten million rows and 100k tasks into a hundred, +while each of the concurrent opens still carries the same 47 MB fragment list. +What has to survive is the overlap: enough tasks to keep the concurrency in +flight, which `benchmarks/lance_session_memory.py` checks rather than assumes. + +One caveat if you measure this yourself: a `Session` reports both `size_bytes()` +and `index_cache_size_bytes()`, and they are not the same accounting. Only +`size_bytes()` tracks resident memory — on a 200k-row IVF_PQ index it moves by +about 5 MB per loaded copy, matching the process's own RSS, where +`index_cache_size_bytes()` reads 37 MB. + ## Migration The migration only requires replacing `daft.io.lance` with `daft_lance`. diff --git a/benchmarks/lance_session_memory.py b/benchmarks/lance_session_memory.py new file mode 100644 index 0000000..99146c9 --- /dev/null +++ b/benchmarks/lance_session_memory.py @@ -0,0 +1,1298 @@ +"""Measure what sharing one lance ``Session`` saves inside a real Ray worker. + +``daft_lance`` passes one process-level ``lance.Session`` to every dataset it +opens (``daft_lance.namespace.get_shared_session``), so a table's fragment list +and lance's caches are held once per process instead of once per scan task, +maintenance UDF and writer. The numbers that motivated the change were measured +in a driver process, but under Ray the memory is spent in the worker, and a +driver-side monkeypatch cannot reach it. So this script runs the real pipeline +twice in two subprocesses that differ only in ``DAFT_LANCE_SHARE_SESSION``, and +reads the worker's memory back out of the worker itself: + +* a probe UDF, fused into the scan's own ``SwordfishTask`` (no ``concurrency=``, + no ``use_process=True``, or Daft would give it a process of its own), reports + ``getrusage(RUSAGE_SELF).ru_maxrss`` -- the worker process's high-water mark -- + plus ``Session.size_bytes()`` and the flag value it actually sees; +* Daft's own sampler (``DAFT_PROCESS_MONITOR_ENABLED=1`` + ``DAFT_EVENT_LOG_DIR``) + writes ``process.memory.rss`` every 200ms from inside each worker to + ``//events.jsonl``, which is an independent second opinion. + +The main stage is the shape that matters most in practice: **read Lance, write +Parquet**. It is the one the thresholds are derived from, and it is measured as +a write rather than an aggregation on purpose -- a task holds its parquet +writer's buffers at the same time as its manifest, so the two compete for the +same memory limit, which is what a real backfill does. An aggregation keeps +almost nothing resident and would flatter the worker. + +Three more stages follow, and they make different claims. ``read_blobs`` and the +lance write share the main stage's shape: ``--parallel`` concurrent tasks in one +worker process, so their manifests multiply by concurrency and sharing collapses +them. The indexed stage multiplies by nothing -- a ``nearest`` read plans a +single index-driven task -- and what sharing buys there is loading the index +once per process instead of once per query. The maintenance stage (compaction +and ``update_columns``) is the one that turns out *not* to multiply: Daft gives +each concurrent class instance a process of its own, so a maintenance worker +opens the table exactly once however high the concurrency, and an indexed +table's real maintenance cost is the index remap in the commit, on the driver. + +Usage:: + + make bench + make bench EXTRA_ARGS="--fragments 2000 --parallel 16" + uv run --group ray python benchmarks/lance_session_memory.py --mode shared +""" + +from __future__ import annotations + +import argparse +import json +import os +import pickle +import resource +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +import lance +import pyarrow as pa +import ray +from lance.lance import _Session as LanceSession + +import daft +import daft_lance +from daft import method +from daft_lance._blob import read_blobs +from daft_lance.namespace import get_shared_session +from daft_lance.utils import construct_lance_dataset_handle + +MB = 1024 * 1024 +RESULT_PREFIX = "RESULT_JSON " +# The indexed stage. A vector index costs roughly 26 bytes per row in a +# session's index cache (measured: 5.1 MB at 200k rows), which dwarfs the +# manifest on any normally shaped table -- so a benchmark that only ever reads +# an unindexed table understates what sharing a session saves. +INDEX_ROWS = 200_000 +INDEX_DIM = 128 +INDEX_PARTITIONS = 256 +INDEX_SUB_VECTORS = 16 +INDEX_QUERIES = 6 +RSS_KEY = "process.memory.rss" +WRITE_ROWS = 8 +# Rows per fragment in the table under test. Also load-bearing for the +# measurement, not just for realism -- see build_table. +ROWS_PER_FRAGMENT = 10_000 +# Worker RSS wobbles by a megabyte or so between runs (Ray's own buffers, the +# arrow allocator). Below this the manifest saving is inside the noise, and the +# table under test is simply too small to make the claim with RSS. +RSS_NOISE_FLOOR = 4 * MB +# How long one side of the A/B may take before something is clearly wrong. +CHILD_TIMEOUT = 1800 +# The maintenance stage. Fewer fragments than the read table and no vector +# column: what is being measured there is where a maintenance *worker's* memory +# goes, and a worker opens the table once whatever the fragment count. The index +# is a scalar one for the same reason -- cheap to build, and the compaction +# commit still has to remap it. +MAINT_FRAGMENTS = 200 +MAINT_CONCURRENCY = 6 +# Rows in each small fragment the maintenance stage appends to give compaction +# something to merge. Compaction's target is a few of these, so the table's own +# large fragments stay above the target and are never rewritten -- the stage +# costs seconds whatever the table's size. +MAINT_ROWS_PER_APPEND = 1_000 +MAINT_TARGET_FRAGMENT_ROWS = MAINT_ROWS_PER_APPEND * 4 +# How many extra opens calibrate_per_open_bytes may hold at once. +CALIBRATION_OPENS = 4 + + +def maxrss_bytes() -> int: + """``ru_maxrss`` normalized to bytes: macOS reports bytes, Linux kilobytes.""" + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return int(rss) if sys.platform == "darwin" else int(rss) * 1024 + + +def worker_parallelism() -> dict[str, Any]: + """What the executing process thinks it may run in parallel. + + The scheduling ceiling is min(the Ray node's CPUs, scantask_max_parallel), + but a container's real ceiling is its cgroup quota, and that is what sizes + Daft's execution thread pool. On a laptop these agree; in a pod with + ``num-cpus: 24`` and ``limits.cpu: 2`` they do not, so both get recorded. + """ + info: dict[str, Any] = {"cpu_count": os.cpu_count()} + affinity = getattr(os, "sched_getaffinity", None) + if affinity is not None: + info["affinity"] = len(affinity(0)) + try: + quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split() + info["cgroup_cpu_max"] = None if quota == "max" else round(int(quota) / int(period), 2) + except (OSError, ValueError): + info["cgroup_cpu_max"] = None + return info + + +@daft.func.batch(return_dtype=daft.DataType.string()) +def probe(column: daft.Series) -> list[str]: + """Report the executing process's memory and session state, one row at a time. + + Deliberately a plain batch UDF: Daft fuses those into the scan's task + instead of giving them their own actor, so what this returns describes the + very process that opened the lance datasets. + """ + session = get_shared_session() + payload = json.dumps( + { + "pid": os.getpid(), + "maxrss": maxrss_bytes(), + "session_size": session.size_bytes(), + "session_obj_id": id(session), + "same_session": bool(session.is_same_as(get_shared_session())), + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + "parallelism": worker_parallelism(), + } + ) + return [payload] * len(column) + + +def build_table(uri: str, fragments: int, rows_per_fragment: int = ROWS_PER_FRAGMENT) -> None: + """``fragments`` files of ``rows_per_fragment`` rows each. + + A real lance table holds tens of thousands of rows per fragment, and the row + count matters to this benchmark beyond realism: with one row per fragment a + scan task opens its dataset, reads a row and drops it, so the opens barely + overlap and far fewer than ``--parallel`` manifests are ever resident at + once. Filling the fragments keeps each task busy long enough for the + concurrency to actually materialize. + """ + # Written in chunks of about ten million rows so the whole table is never in + # memory at once, and each chunk is a whole number of fragments so no chunk + # boundary leaves a half-filled one. Chunks are not just a memory knob: each + # one is an append, and an append rewrites the whole manifest, so chunking a + # 100k-fragment table a million rows at a time spends most of its time + # rewriting an ever-growing fragment list. Ten million keeps the arrow + # buffer under a hundred megabytes and cuts the commits tenfold. + rows_per_chunk = max(rows_per_fragment, (10_000_000 // max(rows_per_fragment, 1)) * rows_per_fragment) + total = fragments * rows_per_fragment + start = 0 + while start < total: + count = min(rows_per_chunk, total - start) + table = pa.table({"i": pa.array(range(start, start + count), type=pa.int64())}) + lance.write_dataset( + table, + uri, + mode="create" if start == 0 else "append", + max_rows_per_file=rows_per_fragment, + ) + start += count + + +def build_blob_table(uri: str, fragments: int) -> None: + rows = max(fragments, 2) + table = pa.table( + { + "id": pa.array(range(rows), type=pa.int64()), + "blob": lance.blob_array([b"payload-%d" % i for i in range(rows)]), + } + ) + lance.write_dataset(table, uri, mode="create", data_storage_version="2.2", max_rows_per_file=1) + + +def build_indexed_table(uri: str, rows: int = INDEX_ROWS, dim: int = INDEX_DIM) -> None: + """A normally shaped table -- few fragments, many rows -- carrying a vector index. + + The point of this one is the index, not the fragment list: index size scales + with rows while a manifest scales with fragments, so this is the mirror image + of the one-row-per-fragment table above. + """ + import numpy as np + + rng = np.random.default_rng(0) + chunk = 50_000 + for start in range(0, rows, chunk): + count = min(chunk, rows - start) + vectors = rng.random((count, dim), dtype=np.float32) + table = pa.table( + { + "id": pa.array(range(start, start + count), type=pa.int64()), + "v": pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), type=pa.float32()), dim), + } + ) + lance.write_dataset(table, uri, mode="create" if start == 0 else "append", max_rows_per_file=25_000) + dataset = lance.dataset(uri) + dataset.create_index( + "v", + index_type="IVF_PQ", + num_partitions=INDEX_PARTITIONS, + num_sub_vectors=INDEX_SUB_VECTORS, + ) + + +def indexed_query_stage(uri: str) -> dict[str, Any]: + """Run the same vector search repeatedly and time each one. + + ``read_lance`` with ``nearest`` plans a *single* index-driven scan task, so + this is not about concurrency: it is about whether the index is loaded once + per process or once per query. Shared, the first query pays for the load and + the rest are served from the session's index cache; unshared, every query + re-reads and re-decodes the index. + """ + import numpy as np + + rng = np.random.default_rng(7) + durations: list[float] = [] + sessions: list[int] = [] + for _ in range(INDEX_QUERIES): + query = rng.random(INDEX_DIM, dtype=np.float32) + df = daft_lance.read_lance(uri, default_scan_options={"nearest": {"column": "v", "q": query, "k": 10}}) + started = time.monotonic() + rows = df.to_arrow().num_rows + durations.append((time.monotonic() - started) * 1000) + if rows != 10: + raise SystemExit(f"the indexed query returned {rows} rows, expected 10") + report = report_from_worker() + sessions.append(max(int(row["session_size"]) for row in report.values())) + return { + "durations_ms": durations, + # Steady state, i.e. every query but the first: the first one is where a + # shared session does its one load, and including it would hide the + # difference the rest of them show. + "steady_median_ms": statistics.median(durations[1:]), + "first_ms": durations[0], + "worker_session_sizes": sessions, + } + + +def build_maintenance_table(uri: str, fragments: int, rows_per_fragment: int) -> None: + """A many-fragment table with a scalar index, for the maintenance stage. + + Compaction rewrites every row it touches, so this one carries no vector + column: the index is a BTREE over a string key, which still has to be + remapped at commit but costs seconds to build rather than gigabytes to + rewrite. Written with pyarrow rather than through Daft because Daft emits + ``large_string``, which lance will not index. + """ + total = fragments * rows_per_fragment + rows_per_chunk = max(rows_per_fragment, (1_000_000 // max(rows_per_fragment, 1)) * rows_per_fragment) + start = 0 + while start < total: + count = min(rows_per_chunk, total - start) + table = pa.table( + { + "i": pa.array(range(start, start + count), type=pa.int64()), + "k": pa.array([f"k{value:09d}" for value in range(start, start + count)], type=pa.string()), + } + ) + lance.write_dataset( + table, + uri, + mode="create" if start == 0 else "append", + max_rows_per_file=rows_per_fragment, + ) + start += count + lance.dataset(uri).create_scalar_index("k", index_type="BTREE") + + +def append_compaction_work(uri: str, groups: int, rows_per_append: int = MAINT_ROWS_PER_APPEND) -> int: + """Append small fragments so there is always something to compact. + + The data volume is deliberately trivial. Compaction consumes its own input, + so a stage that compacts whatever it finds works exactly once per table and + reports "nothing to do" on every later run -- which is how a reused table + turned a passing A/B into a failure. Appending is also the honest shape of + the problem: small fragments arriving from streaming writes are what + compaction exists for. + + None of this touches what the stage measures. A maintenance instance's cost + is one open of the table, and an open carries the whole fragment list + whatever the task then rewrites. + """ + dataset = lance.dataset(uri) + start = dataset.count_rows() + # Enough fragments that every instance in the pool gets work, with slack so + # the plan does not degenerate to one task. + count = max(groups, 2) * 4 * rows_per_append + table = pa.table( + { + "i": pa.array(range(start, start + count), type=pa.int64()), + "k": pa.array([f"k{value:09d}" for value in range(start, start + count)], type=pa.string()), + } + ) + lance.write_dataset(table, uri, mode="append", max_rows_per_file=rows_per_append) + return count // rows_per_append + + +class OpenCensus: + """Open the table the way a maintenance worker does, and report who did it. + + Decorated at the call site exactly like ``CompactionTaskUDF`` and the + ``update_columns`` handlers -- ``max_concurrency`` set, ``use_process`` + left alone -- because the question this answers is about placement: if + Daft puts the N instances in one process, the N opens collapse into one + shared session and sharing saves N-1 manifests; if it gives each its own + process, there is exactly one open per process and nothing to collapse. + The real maintenance UDFs cannot be asked this (they return rewrite + results, not pids), so this stands in for them. + """ + + def __init__(self, open_context: Any) -> None: + self.open_context = open_context + self._lance_ds: lance.LanceDataset | None = None + + @method.batch(return_dtype=daft.DataType.string()) + def __call__(self, column: daft.Series) -> list[str]: + if self._lance_ds is None: + # Once per instance, which is what the maintenance UDFs do. + self._lance_ds = self.open_context.open_pinned() + session = get_shared_session() + payload = json.dumps( + { + "pid": os.getpid(), + "instance": id(self), + "session_obj_id": id(session), + "session_size": session.size_bytes(), + "fragments": len(self._lance_ds.get_fragments()), + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + } + ) + return [payload] * len(column) + + +def maintenance_census(open_context: Any, concurrency: int) -> dict[str, Any]: + """Run ``concurrency`` maintenance-shaped instances and count processes.""" + census = daft.cls(OpenCensus, max_concurrency=concurrency) + partitions = concurrency * 8 + df = daft.from_pydict({"x": list(range(partitions))}).into_partitions(partitions) + rows = [json.loads(value) for value in df.select(census(open_context)(daft.col("x")).alias("p")).to_pydict()["p"]] + instances = {(row["pid"], row["instance"]): row for row in rows} + return { + "requested": concurrency, + "instances": len(instances), + "processes": len({pid for pid, _ in instances}), + "sessions": len({row["session_obj_id"] for row in instances.values()}), + "flags": sorted({row["flag"] for row in instances.values()}, key=str), + "fragments": sorted({row["fragments"] for row in instances.values()}), + # Per process, so it is comparable between the sides regardless of how + # many processes each got. + "session_size_max": max(int(row["session_size"]) for row in instances.values()), + } + + +def maintenance_stage(data_dir: str, side: str, fragments: int, concurrency: int) -> dict[str, Any]: + """Compaction and ``update_columns`` on an indexed table, with the driver watched too. + + Two separate things are measured here, and they answer different halves of + "does an index make maintenance expensive": + + * the worker side, via :func:`maintenance_census`: how many processes the + concurrent instances land in, and what one of them holds; + * the driver side, directly, because under Ray the driver *is* this process: + what the compaction leaves in it. Compaction has to remap every index + entry onto the new fragment ids, and that happens in the commit, on the + driver, not in the tasks -- so this is where an indexed table's + maintenance cost shows up at all. The figure is manifest and index + together; ``test_shared_session.py`` is where the two are separated, + which needs the tasks and the commit pulled apart. + + Each side of the A/B gets its own copy of the table: compaction consumes its + own input, so whoever ran second would otherwise find nothing to compact. + The stage also makes its own work rather than compacting whatever happens to + be there, because the tables outlive a run -- see :func:`append_compaction_work`. + """ + uri = str(Path(data_dir) / f"maint-{side}-{fragments}.lance") + appended = append_compaction_work(uri, concurrency) + session = get_shared_session() + + def snapshot() -> dict[str, int]: + # ``size_bytes()`` only. A session also reports + # ``index_cache_size_bytes()``, but that is the accounting lance + # enforces its index budget with, not resident bytes: measured on a + # 200k-row IVF_PQ index it reads 37 MB where both ``size_bytes()`` and + # the process's own RSS move by about 5 MB per loaded copy. Printing it + # as memory would overstate an index's cost sevenfold. + return {"total": session.size_bytes()} + + driver: dict[str, dict[str, int]] = {"start": snapshot()} + handle = construct_lance_dataset_handle(uri) + driver["after_open"] = snapshot() + result: dict[str, Any] = { + "uri": uri, + "fragments_appended": appended, + "fragments_before": len(handle.dataset.get_fragments()), + "rows": handle.dataset.count_rows(), + "indices": [index["name"] for index in handle.dataset.list_indices()], + "concurrency": concurrency, + "census": maintenance_census(handle.worker_open_context(), concurrency), + } + driver["after_census"] = snapshot() + + started = time.monotonic() + metrics = daft_lance.compact_files( + uri, + # Keyed to the appended fragments, not to the table's own: the large + # fragments already sit above this and are not candidates, so only the + # freshly appended ones are rewritten. + compaction_options={"target_rows_per_fragment": MAINT_TARGET_FRAGMENT_ROWS}, + partition_num=concurrency, + concurrency=concurrency, + ) + result["compact_seconds"] = time.monotonic() - started + driver["after_compact"] = snapshot() + if metrics is None: + raise SystemExit( + f"compaction found nothing to do on {uri} even after appending {appended} small fragments; " + "check that target_rows_per_fragment is above the appended fragments' size" + ) + result["compaction"] = { + "fragments_removed": metrics.fragments_removed, + "fragments_added": metrics.fragments_added, + } + + started = time.monotonic() + updated = daft_lance.update_columns( + uri, + transform={"i": "i + 0"}, + columns=["i"], + max_concurrency=concurrency, + ) + result["update_seconds"] = time.monotonic() - started + driver["after_update"] = snapshot() + if updated.rows_updated != result["rows"]: + raise SystemExit(f"update_columns touched {updated.rows_updated} rows, expected {result['rows']}") + result["rows_updated"] = updated.rows_updated + result["driver_session"] = driver + return result + + +def report_from_worker() -> dict[int, dict[str, Any]]: + """Run a one-row query so the worker reports its current state back.""" + df = daft.from_pydict({"x": [0]}).with_column("p", probe(daft.col("x"))) + rows = [json.loads(value) for value in df.to_pydict()["p"]] + return {row["pid"]: row for row in rows} + + +def sampled_rss_peaks(event_log_dir: str) -> dict[int, dict[str, Any]]: + """Peak ``process.memory.rss`` per pid, from Daft's own 200ms sampler.""" + peaks: dict[int, dict[str, Any]] = {} + for path in sorted(Path(event_log_dir).glob("*/events.jsonl")): + for line in path.read_text().splitlines(): + try: + record = json.loads(line) + except ValueError: + continue + if record.get("event") != "process_stats": + continue + rss = record.get("metrics", {}).get(RSS_KEY) + pid = record.get("pid") + if rss is None or pid is None: + continue + entry = peaks.setdefault(int(pid), {"component": record.get("component"), "peak": 0}) + entry["peak"] = max(entry["peak"], int(rss)) + return peaks + + +def payload_sizes(uri: str) -> dict[str, int]: + """Bytes a scan task used to carry versus what it carries now. + + ``LanceDataset.__reduce__`` ships the re-encoded manifest, so closing over a + dataset -- which ``read_blobs`` used to do -- cost that much per task. The + replacement is a uri plus the pinned open kwargs. + """ + dataset = lance.dataset(uri) + open_kwargs = {"version": dataset.version, "storage_options": dataset._storage_options} + return { + "pickled_dataset": len(pickle.dumps(dataset)), + "pickled_open_kwargs": len(pickle.dumps((uri, open_kwargs))), + } + + +def run_child(args: argparse.Namespace) -> int: + """One half of the A/B: a full read + read_blobs + write under Ray. + + Either starts a cluster of its own (the local ``--mode both`` path, which + spawns this twice) or attaches to one that is already there, which is how it + runs as a RayJob entrypoint: there the two sides are separate jobs and the + flag arrives through the pod spec instead of process inheritance. + """ + address = os.environ.get("RAY_ADDRESS") + if address: + # An attached run must not ask for CPUs: the cluster already decided, + # and Ray ignores num_cpus with a warning when it is only connecting. + # Concurrency comes from the worker's ``ray start --num-cpus``, which is + # a logical count and need not match the pod's cpu limit. + ray.init() + else: + ray.init(num_cpus=args.parallel, include_dashboard=False) + daft.set_runner_ray(noop_if_initialized=True) + daft.set_execution_config(scantask_max_parallel=args.parallel) + + uri = str(Path(args.data_dir) / f"table-{args.fragments}.lance") + blob_uri = str(Path(args.data_dir) / "blobs.lance") + dataset = lance.dataset(uri) + result: dict[str, Any] = { + "mode": args.mode, + "parallel": args.parallel, + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + "ray_address": address, + "rows_per_fragment": args.rows_per_fragment, + "cluster_cpus": ray.cluster_resources().get("CPU"), + "fragments_at_read": len(dataset.get_fragments()), + "manifest_bytes": len(dataset._ds.serialized_manifest()), + } + + durations: dict[str, float] = {} + before = report_from_worker() + + # The pipeline under test: read Lance, write Parquet. One task per fragment + # group, `parallel` of them open at once, and each task holds a parquet + # writer's buffers at the same time as its manifest -- which is the point of + # measuring this shape rather than an aggregation. An aggregation keeps + # almost nothing resident, so it flatters the worker; a write makes the + # manifests compete with the writer for the same memory limit, which is what + # a real backfill does. + # + # Deliberately not count_rows(): count pushdown would collapse the whole + # plan into a single task that opens no fragment at all. + # + # count_rows() here is the dataset's own metadata, read on the driver before + # the scan: cheap, and it says how much the scan below had to touch. The two + # sides see slightly different totals because each appends in its lance + # write stage, which is why this is recorded rather than asserted against a + # formula. + rows_total = dataset.count_rows() + result["rows_total"] = rows_total + # Two knobs for keeping a large table affordable, and neither weakens the + # measurement: a manifest is the table's whole fragment list, so an open + # costs the same whether its task reads one fragment or fifty, one row or + # ten thousand. --read-fraction prunes whole fragments (``i`` is written + # strictly increasing, so a cutoff is a fragment prefix) and + # --fragment-group-size puts several fragments in one task. + # + # What does have to survive is the overlap: enough tasks to keep + # --parallel of them in flight, each long enough that their opens coexist. + # Both counts are reported so that is checkable rather than hoped for. + result["read_fraction"] = args.read_fraction + result["fragment_group_size"] = args.fragment_group_size + reader = daft_lance.read_lance(uri, fragment_group_size=args.fragment_group_size) + if args.read_fraction < 1.0: + cutoff = int(rows_total * args.read_fraction) + result["read_cutoff"] = cutoff + # Arithmetic, not a second pass over 100k fragments: build_table fills + # fragments in order, so a cutoff of N rows is the first N/rows + # fragments. + touched = -(-cutoff // args.rows_per_fragment) + result["fragments_touched"] = touched + result["tasks_estimated"] = -(-touched // max(args.fragment_group_size or 1, 1)) + reader = reader.where(daft.col("i") < cutoff) + else: + result["fragments_touched"] = result["fragments_at_read"] + result["tasks_estimated"] = -(-result["fragments_at_read"] // max(args.fragment_group_size or 1, 1)) + if result["tasks_estimated"] < 2 * args.parallel: + raise SystemExit( + f"lance -> parquet would plan about {result['tasks_estimated']} tasks at concurrency " + f"{args.parallel}: too few for the opens to overlap, so raise --read-fraction or lower " + "--fragment-group-size" + ) + # One output directory per side. Overwrite rather than append so a rerun + # does not measure the previous run's files, and so the two sides never read + # each other's output back as input to anything. + parquet_dir = Path(args.data_dir) / f"out-parquet-{args.mode}" + started = time.monotonic() + written = reader.write_parquet(str(parquet_dir), write_mode="overwrite").to_pydict() + durations["read_to_parquet"] = time.monotonic() - started + files = written.get("path") or [] + if not files: + raise SystemExit("the pipeline wrote no parquet files, so it never reached the data") + result["parquet_files"] = len(files) + result["parquet_bytes"] = sum(path.stat().st_size for path in parquet_dir.rglob("*.parquet")) + after_read = report_from_worker() + + started = time.monotonic() + blob_ds = lance.dataset(blob_uri) + blob_df = daft.read_lance(blob_uri, default_scan_options={"with_row_id": True}) + result["blobs_read"] = len(read_blobs(blob_df, blob_ds, "blob").to_pydict()["blob"]) + durations["read_blobs"] = time.monotonic() - started + after_blobs = report_from_worker() + + # A small append with tiny files, so the writer flushes `write_fragments` + # several times; each flush is a manifest load when no session is passed. + started = time.monotonic() + daft.from_pydict({"i": list(range(WRITE_ROWS))}).write_lance(uri, mode="append", max_rows_per_file=1).collect() + durations["write"] = time.monotonic() - started + after_write = report_from_worker() + + # Last, because it loads an index into whatever session it gets and would + # otherwise contaminate the manifest measurements above. + indexed_uri = str(Path(args.data_dir) / f"indexed-{INDEX_ROWS}.lance") + started = time.monotonic() + result["indexed"] = indexed_query_stage(indexed_uri) + durations["indexed_queries"] = time.monotonic() - started + + # After everything above, and on a table of its own, because compaction + # rewrites its input: it must not disturb the read table the stages above + # measured, and its own driver-side numbers are read as deltas so the + # earlier stages' caches do not matter. + started = time.monotonic() + result["maintenance"] = maintenance_stage( + args.data_dir, + args.mode, + args.maint_fragments, + args.maint_concurrency, + ) + durations["maintenance"] = time.monotonic() - started + result["durations"] = durations + + # Calibrated here rather than by whoever orchestrates the A/B: the threshold + # then comes from the same architecture, the same lance build and the same + # allocator as the measurement it gates. Last, not first: it holds several + # unshared copies of the manifest, which at 100k fragments is a few hundred + # megabytes, and the driver needs that headroom while it is planning one + # scan task per fragment. Nothing above depends on the number. + result["per_open_bytes"] = calibrate_per_open_bytes(uri, args.parallel) + + workers: dict[str, Any] = {} + for pid, row in after_write.items(): + workers[str(pid)] = { + "flag": row["flag"], + "same_session": row["same_session"], + "parallelism": row.get("parallelism"), + "session_obj_ids": sorted( + { + stage[pid]["session_obj_id"] + for stage in (before, after_read, after_blobs, after_write) + if pid in stage + } + ), + "maxrss": { + "before": before.get(pid, {}).get("maxrss"), + "after_read": after_read.get(pid, {}).get("maxrss"), + "after_blobs": after_blobs.get(pid, {}).get("maxrss"), + "after_write": after_write.get(pid, {}).get("maxrss"), + }, + "session_size": { + "before": before.get(pid, {}).get("session_size"), + "after_read": after_read.get(pid, {}).get("session_size"), + "after_blobs": after_blobs.get(pid, {}).get("session_size"), + "after_write": after_write.get(pid, {}).get("session_size"), + }, + } + result["workers"] = workers + result["worker_pids_per_stage"] = { + "before": sorted(before), + "after_read": sorted(after_read), + "after_blobs": sorted(after_blobs), + "after_write": sorted(after_write), + } + event_log_dir = os.environ.get("DAFT_EVENT_LOG_DIR") + result["sampled_rss_peak"] = ( + {str(pid): entry for pid, entry in sampled_rss_peaks(event_log_dir).items()} if event_log_dir else {} + ) + print(RESULT_PREFIX + json.dumps(result)) + return 0 + + +def calibrate_per_open_bytes(uri: str, opens: int) -> float: + """What one *additional* unshared open of this table costs, measured here and now. + + Thresholds are derived from this rather than hardcoded, so the script stays + honest on a different table or a different lance version. + + Marginal rather than average, and capped at a handful of opens, for two + reasons. The cap: at 100k fragments one open is around 64 MB, so holding + ``--parallel`` of them puts half a gigabyte in the driver -- a calibration + must not be able to push over the process it is calibrating, and at 100k + fragments on a 3Gi head it did. Marginal: the first open also pays one-time + costs (the object-store client, lance's own statics) that do not repeat per + open, so averaging them in overstates the per-open cost and with it the + threshold. + """ + first = [lance.dataset(uri, session=LanceSession())] + extra = max(min(opens - 1, CALIBRATION_OPENS), 1) + before = maxrss_bytes() + held = [lance.dataset(uri, session=LanceSession()) for _ in range(extra)] + after = maxrss_bytes() + assert len(first) == 1 and len(held) == extra + return (after - before) / extra + + +def spawn_child(mode: str, args: argparse.Namespace, event_log_dir: Path) -> dict[str, Any]: + env = dict(os.environ) + env.update( + { + "DAFT_LANCE_SHARE_SESSION": "1" if mode == "shared" else "0", + "DAFT_RUNNER": "ray", + "DAFT_PROGRESS_BAR": "0", + "DAFT_ANALYTICS_ENABLED": "0", + "DAFT_PROCESS_MONITOR_ENABLED": "1", + "DAFT_EVENT_LOG_DIR": str(event_log_dir), + # Ray notices when the driver's process tree was started by `uv run` + # -- which it is, via `make bench` -- and then replicates that + # environment to the workers: it packages the whole working + # directory and reinstalls it in a fresh venv per worker. That costs + # minutes here and has hung outright. The workers already inherit + # this venv's interpreter from the raylet, so turn it off. + "RAY_ENABLE_UV_RUN_RUNTIME_ENV": "0", + } + ) + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--mode", + mode, + "--fragments", + str(args.fragments), + "--parallel", + str(args.parallel), + "--data-dir", + str(args.data_dir), + "--rows-per-fragment", + str(args.rows_per_fragment), + "--maint-fragments", + str(args.maint_fragments), + "--maint-concurrency", + str(args.maint_concurrency), + "--read-fraction", + str(args.read_fraction), + ] + if args.fragment_group_size: + command += ["--fragment-group-size", str(args.fragment_group_size)] + print(f"\n=== running {mode} (DAFT_LANCE_SHARE_SESSION={env['DAFT_LANCE_SHARE_SESSION']}) ===", flush=True) + try: + process = subprocess.run(command, env=env, capture_output=True, text=True, check=False, timeout=CHILD_TIMEOUT) + except subprocess.TimeoutExpired as expired: + sys.stderr.write(expired.stderr or "") + raise SystemExit( + f"{mode} child did not finish within {CHILD_TIMEOUT}s. The read is one task per fragment, so a large " + "--fragments takes a long time; lower it, or raise CHILD_TIMEOUT." + ) from expired + payload = None + for line in process.stdout.splitlines(): + if line.startswith(RESULT_PREFIX): + payload = json.loads(line[len(RESULT_PREFIX) :]) + if payload is None: + sys.stdout.write(process.stdout) + sys.stderr.write(process.stderr) + raise SystemExit(f"{mode} child produced no result (exit code {process.returncode})") + return payload + + +def worker_row(result: dict[str, Any]) -> tuple[str, dict[str, Any]]: + workers = result["workers"] + if len(workers) != 1: + raise SystemExit(f"expected one worker process, got {sorted(workers)}; run this on a single-node cluster") + return next(iter(workers.items())) + + +def compare(baseline: dict[str, Any], shared: dict[str, Any], per_open_bytes: float) -> list[str]: + """Print the comparison and return the failures, if any.""" + failures: list[str] = [] + baseline_pid, baseline_worker = worker_row(baseline) + shared_pid, shared_worker = worker_row(shared) + + # Before anything else: the A/B only means something if the flag actually + # reached the worker. Ray actors do not inherit the driver's environment by + # contract, so this must be observed, never assumed. + for name, result, worker, expected in ( + ("baseline", baseline, baseline_worker, "0"), + ("shared", shared, shared_worker, "1"), + ): + if worker["flag"] != expected: + failures.append( + f"{name}: worker saw DAFT_LANCE_SHARE_SESSION={worker['flag']!r}, driver set {expected!r} -- " + "the environment did not reach the worker, so these numbers compare nothing" + ) + if result["flag"] != expected: + failures.append(f"{name}: driver saw flag {result['flag']!r}, expected {expected!r}") + + # Concurrency decides how many manifests can be resident at once, so the two + # sides have to have had the same ceiling, and the reader has to be able to + # see what it was. + for name, result, worker in ( + ("baseline", baseline, baseline_worker), + ("shared", shared, shared_worker), + ): + limits = worker.get("parallelism") or {} + print( + f"{name:10} rows/fragment {result.get('rows_per_fragment')}, " + f"scantask_max_parallel {result['parallel']}, ray cluster CPUs {result.get('cluster_cpus')}, " + f"worker sees cpu_count={limits.get('cpu_count')} affinity={limits.get('affinity')} " + f"cgroup_cpu_max={limits.get('cgroup_cpu_max')}" + ) + for key in ("parallel", "rows_per_fragment", "read_fraction", "fragment_group_size"): + if baseline.get(key) != shared.get(key): + failures.append( + f"the two sides ran with different {key} ({baseline.get(key)!r} vs {shared.get(key)!r}), " + "so their memory is not comparable" + ) + + def delta(worker: dict[str, Any], field: str, stage: str) -> int: + return int(worker[field][stage]) - int(worker[field]["before"]) + + rows = [] + for name, result, worker, pid in ( + ("baseline", baseline, baseline_worker, baseline_pid), + ("shared", shared, shared_worker, shared_pid), + ): + sampled = result["sampled_rss_peak"].get(pid, {}) + rows.append( + { + "name": name, + "pid": pid, + "fragments": result["fragments_at_read"], + "rss_delta_read": delta(worker, "maxrss", "after_read"), + "rss_delta_total": delta(worker, "maxrss", "after_write"), + "sampled_peak": sampled.get("peak"), + "sampled_component": sampled.get("component"), + "session_delta_read": delta(worker, "session_size", "after_read"), + "session_delta_blobs": delta(worker, "session_size", "after_blobs"), + "session_delta_total": delta(worker, "session_size", "after_write"), + "same_session": worker["same_session"], + "sessions_seen": len(worker["session_obj_ids"]), + } + ) + + print(f"\nper-open cost of this table, measured here: {per_open_bytes / MB:.2f} MB") + for name, result in (("baseline", baseline), ("shared", shared)): + stages = ", ".join(f"{stage} {seconds:.1f}s" for stage, seconds in result["durations"].items()) + print(f"{name:10} {stages}") + print( + f"manifest: {baseline['manifest_bytes'] / 1024:.1f} KiB, table rows: {baseline.get('rows_total')}, " + f"parquet written: {baseline.get('parquet_files')} files / " + f"{(baseline.get('parquet_bytes') or 0) / MB:.1f} MB" + ) + print( + f"lance -> parquet: {baseline.get('read_fraction')} of the table, fragment_group_size " + f"{baseline.get('fragment_group_size')}, so about {baseline.get('fragments_touched')} fragments in " + f"{baseline.get('tasks_estimated')} tasks at concurrency {baseline['parallel']} " + "(a manifest costs the same per open whatever the task reads)" + ) + print( + f"{'':10} {'pid':>7} {'frags':>7} {'rssDpq':>9} {'rssDall':>9} {'rssPeak':>9} " + f"{'sessDpq':>10} {'sessDblob':>10} {'sessDall':>10}" + ) + for row in rows: + sampled = "-" if row["sampled_peak"] is None else f"{row['sampled_peak'] / MB:.1f}" + print( + f"{row['name']:10} {row['pid']:>7} {row['fragments']:>7} " + f"{row['rss_delta_read'] / MB:>9.1f} {row['rss_delta_total'] / MB:>9.1f} {sampled:>9} " + f"{row['session_delta_read'] / MB:>10.3f} " + f"{(row['session_delta_blobs'] - row['session_delta_read']) / MB:>10.3f} " + f"{row['session_delta_total'] / MB:>10.3f}" + ) + print( + "(MB. rssDpq: the worker's ru_maxrss growth across the lance-to-parquet stage. rssPeak: Daft's\n" + " own sampled rss peak for that pid, an independent reading of the same process. sessD: growth of\n" + " the process-wide lance session, which is zero whenever an open did not go through it.)" + ) + + baseline_row, shared_row = rows[0], rows[1] + + # Mechanism, with no RSS noise in it: with sharing on, a scan task's open + # lands in the process-wide session, so its size grows. With sharing off + # the global session is never touched at all. + if baseline_row["session_delta_total"] != 0: + failures.append( + f"baseline: process session grew by {baseline_row['session_delta_total']} bytes; " + "with sharing off no open should touch it" + ) + if shared_row["session_delta_read"] <= 0: + failures.append( + "shared: the lance-to-parquet stage did not grow the process session, so its scan tasks bypassed it" + ) + if shared_row["session_delta_total"] <= shared_row["session_delta_read"]: + failures.append("shared: read_blobs and the write did not add to the process session") + if not shared_row["same_session"] or shared_row["sessions_seen"] != 1: + failures.append( + f"shared: worker reported {shared_row['sessions_seen']} distinct sessions, expected one per process" + ) + if baseline_row["same_session"]: + failures.append("baseline: the kill switch did not take effect in the worker") + + # The memory claim itself. Only the per-open cost beyond the first is + # saved, and never all of it: a task holds its manifest just while its + # dataset is open, so the opens overlap only partly and the peak never + # reaches the full K-way sum. Measured at --parallel 24, the saving came in + # at 33-41 MB against a 67 MB theoretical maximum, so demanding 30% of it + # keeps roughly a 2x margin over the observed value while still failing + # outright if sharing stops working, which takes the saving to zero. + # + # Measured over the lance-to-parquet stage alone, not the whole run. + # ``ru_maxrss`` is a high-water mark, and the later stages allocate about ten + # megabytes of their own: on the shared side, whose peak was lower, that + # pushes the mark back up to roughly where the baseline already was, hiding + # the very difference being measured. That stage is also the only one whose + # concurrency is pinned, and the one the pipeline this benchmark exists for + # actually runs, so it is the one the threshold is derived from. + # The indexed stage, which is a different claim from everything above: a + # `nearest` read plans one index-driven scan task, so nothing multiplies by + # concurrency. What sharing changes is that the index is loaded once per + # process instead of once per query. + baseline_index, shared_index = baseline.get("indexed"), shared.get("indexed") + if baseline_index and shared_index: + print( + f"\nindexed vector search, {INDEX_QUERIES} queries of the same table:\n" + f"{'':10} {'first':>9} {'steady':>9} {'sessionEnd':>12}" + ) + for name, side in (("baseline", baseline_index), ("shared", shared_index)): + print( + f"{name:10} {side['first_ms']:>8.1f}ms {side['steady_median_ms']:>8.1f}ms " + f"{side['worker_session_sizes'][-1] / 1024:>10.1f}KiB" + ) + print( + "(first: the query that pays for loading the index. steady: the median of the rest. sessionEnd: what\n" + " the worker's process-wide session holds at the end, so zero means every query loaded the index\n" + " into a session of its own and threw it away.)" + ) + baseline_end = baseline_index["worker_session_sizes"][-1] + shared_end = shared_index["worker_session_sizes"][-1] + if shared_end <= 0: + failures.append("shared: the indexed queries left nothing in the process session, so none of them used it") + # Compared against the shared side rather than against zero: an empty + # lance Session is not zero bytes, it is about 152 of bookkeeping, so + # "holds no index" means "holds orders of magnitude less than one". + elif baseline_end > shared_end // 10: + failures.append( + f"baseline: the process session holds {baseline_end / 1024:.1f} KiB after the indexed queries, " + f"against {shared_end / 1024:.1f} KiB shared; with sharing off no query should have filled it" + ) + # Deliberately not a hard threshold on the ratio: it depends on index + # size and on how far the storage is, and on a local disk with a warm + # page cache it is the smallest it will ever be. Measured locally at + # 200k rows: 5.6-7.0ms shared against 14.0-14.7ms unshared. + ratio = baseline_index["steady_median_ms"] / max(shared_index["steady_median_ms"], 1e-9) + print( + f"steady-state queries are {ratio:.2f}x slower without a shared session " + f"({baseline_index['steady_median_ms']:.1f}ms vs {shared_index['steady_median_ms']:.1f}ms); " + "the gap grows with index size and storage latency" + ) + if ratio < 1.0: + failures.append( + f"the unshared side was not slower on repeated indexed queries (ratio {ratio:.2f}x), so reloading " + "the index cost nothing measurable here -- check that the index exists and that the queries used it" + ) + + # The maintenance path, which is a different shape of claim again. A scan + # task shares a process with `parallel`-1 others, so its manifest is one of + # many; a maintenance instance is given a process to itself, so its open is + # the only one there. Both are measured, because only measuring says which. + baseline_maint, shared_maint = baseline.get("maintenance"), shared.get("maintenance") + if baseline_maint and shared_maint: + print( + f"\nmaintenance on an indexed table " + f"({shared_maint['fragments_before']} fragments, {shared_maint['rows']} rows, " + f"indices {shared_maint['indices']}), concurrency {shared_maint['concurrency']}:\n" + f"{'':10} {'insts':>6} {'procs':>6} {'sessPerProc':>12} {'driverCommit':>13} " + f"{'compact':>9} {'update':>9}" + ) + for name, side in (("baseline", baseline_maint), ("shared", shared_maint)): + census = side["census"] + driver = side["driver_session"] + driver_commit = driver["after_compact"]["total"] - driver["after_census"]["total"] + print( + f"{name:10} {census['instances']:>6} {census['processes']:>6} " + f"{census['session_size_max'] / 1024:>10.1f}KiB {driver_commit / 1024:>11.1f}KiB " + f"{side['compact_seconds']:>8.1f}s {side['update_seconds']:>8.1f}s" + ) + print( + "(insts/procs: concurrent maintenance instances, and how many processes they landed in. sessPerProc:\n" + " the largest process-wide session any one of them held. driverCommit: what compaction added to the\n" + " *driver's* session, manifest and remapped index together -- remapping an index onto the new fragment\n" + " ids happens in the commit, on the driver, not in the tasks.)" + ) + for name, side, expected in (("baseline", baseline_maint, "0"), ("shared", shared_maint, "1")): + census = side["census"] + if census["flags"] != [expected]: + failures.append( + f"{name}: the maintenance instances saw DAFT_LANCE_SHARE_SESSION={census['flags']}, " + f"expected [{expected!r}] -- the flag did not reach them" + ) + if census["instances"] < 2: + failures.append( + f"{name}: only {census['instances']} maintenance instance ran, so concurrency was never " + "reached and this stage measured nothing" + ) + shared_census, baseline_census = shared_maint["census"], baseline_maint["census"] + if shared_census["session_size_max"] <= baseline_census["session_size_max"]: + failures.append( + "shared: a maintenance instance's process session is no larger than the baseline's, so the " + "maintenance open did not go through it" + ) + if baseline_census["sessions"] != baseline_census["instances"]: + failures.append( + f"baseline: {baseline_census['instances']} maintenance instances reported " + f"{baseline_census['sessions']} sessions; with sharing off each open gets its own" + ) + shared_driver = shared_maint["driver_session"] + shared_driver_commit = shared_driver["after_compact"]["total"] - shared_driver["after_census"]["total"] + if shared_driver_commit <= 0: + failures.append( + "shared: compaction left nothing in the driver's session, so the driver's planning and commit " + "opened outside the process session" + ) + baseline_driver = baseline_maint["driver_session"] + if baseline_driver["after_compact"]["total"] != baseline_driver["start"]["total"]: + failures.append( + "baseline: the driver's process session grew during maintenance; with sharing off every open " + "should get a session of its own" + ) + # Not a failure, a finding: if Daft ever fuses these instances into one + # process, the maintenance path acquires a saving to measure, and this + # stage has to be rewritten to measure it rather than quietly passing. + if shared_census["processes"] < shared_census["instances"]: + print( + f"NOTE: {shared_census['instances']} instances shared {shared_census['processes']} processes. " + "Daft used to give each one its own, which is why this stage asserts nothing about concurrency " + "multiplying manifests. It now does multiply them -- this stage should be extended to measure it." + ) + else: + print( + f"Each of the {shared_census['instances']} instances got its own process, so a maintenance worker " + f"opens the table exactly once ({shared_census['session_size_max'] / 1024:.1f} KiB of session) no " + "matter the concurrency: nothing here multiplies by it, and the index cost sits in the driver's " + "commit instead." + ) + + extra_opens = max(shared["parallel"] - 1, 1) + expected_saving = 0.3 * extra_opens * per_open_bytes + observed_saving = baseline_row["rss_delta_read"] - shared_row["rss_delta_read"] + print( + f"\nsaving over lance -> parquet: {observed_saving / MB:.1f} MB observed, " + f"{expected_saving / MB:.1f} MB required " + f"(30% of {extra_opens} extra opens x {per_open_bytes / MB:.2f} MB). " + f"Over the whole run: {(baseline_row['rss_delta_total'] - shared_row['rss_delta_total']) / MB:.1f} MB, " + "which the later stages' own allocations dilute." + ) + if expected_saving < RSS_NOISE_FLOOR: + print( + f"NOT MEASURED: {extra_opens} extra opens of this table only amount to " + f"{extra_opens * per_open_bytes / MB:.1f} MB, which RSS cannot resolve. The deterministic checks above " + "still hold; raise --fragments (or --parallel) to put a number on the saving." + ) + elif observed_saving < expected_saving: + failures.append( + f"the shared run saved {observed_saving / MB:.1f} MB, expected at least {expected_saving / MB:.1f} MB" + ) + return failures + + +def run_prepare(args: argparse.Namespace) -> int: + """Write the tables and nothing else, for a one-off prep job. + + Both sides of a Kubernetes A/B read the *same* copy, so the data has to be + built once into a shared volume: lance names its data files with random + uuids, so a per-pod copy would leave the driver's manifest pointing at file + names the worker's copy does not have. + """ + data_dir = Path(args.data_dir) + data_dir.mkdir(parents=True, exist_ok=True) + uri = data_dir / f"table-{args.fragments}.lance" + blob_uri = data_dir / "blobs.lance" + if uri.exists(): + print(f"{uri} already exists, keeping it") + else: + build_table(str(uri), args.fragments, args.rows_per_fragment) + print(f"wrote {args.fragments} fragments x {args.rows_per_fragment} rows to {uri}") + if blob_uri.exists(): + print(f"{blob_uri} already exists, keeping it") + else: + build_blob_table(str(blob_uri), min(args.fragments, 64)) + print(f"wrote {blob_uri}") + indexed_uri = data_dir / f"indexed-{INDEX_ROWS}.lance" + if indexed_uri.exists(): + print(f"{indexed_uri} already exists, keeping it") + else: + build_indexed_table(str(indexed_uri)) + print(f"wrote {INDEX_ROWS} rows with an IVF_PQ index to {indexed_uri}") + # One per side: the maintenance stage compacts its table, so a shared copy + # would leave whichever side ran second with nothing to compact. They are + # built identically, which is what makes the two sides comparable. + for side in ("baseline", "shared"): + maint_uri = data_dir / f"maint-{side}-{args.maint_fragments}.lance" + if maint_uri.exists(): + print(f"{maint_uri} already exists, keeping it") + else: + build_maintenance_table(str(maint_uri), args.maint_fragments, args.rows_per_fragment) + print( + f"wrote {args.maint_fragments} fragments x {args.rows_per_fragment} rows " + f"with a BTREE index to {maint_uri}" + ) + return 0 + + +def run_compare(args: argparse.Namespace) -> int: + """Compare two result files produced by separate runs of ``--mode``.""" + baseline = json.loads(Path(args.baseline).read_text()) + shared = json.loads(Path(args.shared).read_text()) + failures = compare(baseline, shared, float(baseline["per_open_bytes"])) + if failures: + print("\nFAILED:") + for failure in failures: + print(f" - {failure}") + return 1 + print("\nOK: sharing one session cut the worker's memory as expected.") + return 0 + + +def run_parent(args: argparse.Namespace) -> int: + address = os.environ.get("RAY_ADDRESS") + if address and not address.startswith("local"): + raise SystemExit( + f"RAY_ADDRESS={address!r} points at an existing cluster. The A/B needs the worker to inherit this " + "process's environment, which only holds for a cluster started here, so refusing to report numbers." + ) + + data_dir = Path(args.data_dir) + data_dir.mkdir(parents=True, exist_ok=True) + uri = str(data_dir / f"table-{args.fragments}.lance") + blob_uri = str(data_dir / "blobs.lance") + if not Path(uri).exists(): + print(f"writing a {args.fragments}-fragment table to {uri} (reused by later runs)", flush=True) + build_table(uri, args.fragments, args.rows_per_fragment) + if not Path(blob_uri).exists(): + build_blob_table(blob_uri, min(args.fragments, 64)) + indexed_uri = str(data_dir / f"indexed-{INDEX_ROWS}.lance") + if not Path(indexed_uri).exists(): + print(f"writing a {INDEX_ROWS}-row table with an IVF_PQ index to {indexed_uri}", flush=True) + build_indexed_table(indexed_uri) + for side in ("baseline", "shared"): + maint_uri = data_dir / f"maint-{side}-{args.maint_fragments}.lance" + if not maint_uri.exists(): + print(f"writing the {side} maintenance table to {maint_uri}", flush=True) + build_maintenance_table(str(maint_uri), args.maint_fragments, args.rows_per_fragment) + + # Each child gets its own event-log directory: the sampler keys records by + # pid, and pids are not unique across runs. + with tempfile.TemporaryDirectory(prefix="lance-session-events-") as events: + baseline = spawn_child("baseline", args, Path(events) / "baseline") + shared = spawn_child("shared", args, Path(events) / "shared") + + payload = payload_sizes(uri) + small = payload_sizes(blob_uri) + print( + f"\nper-task payload: pickling the dataset costs {payload['pickled_dataset'] / 1024:.1f} KiB at " + f"{args.fragments} fragments versus {small['pickled_dataset'] / 1024:.1f} KiB at " + f"{min(args.fragments, 64)}; the uri plus pinned open kwargs that replaced it costs " + f"{payload['pickled_open_kwargs']} and {small['pickled_open_kwargs']} bytes." + ) + + per_open_bytes = float(baseline["per_open_bytes"]) + failures = compare(baseline, shared, per_open_bytes) + if args.json: + print( + json.dumps( + {"per_open_bytes": per_open_bytes, "baseline": baseline, "shared": shared, "failures": failures}, + indent=2, + ) + ) + if failures: + print("\nFAILED:") + for failure in failures: + print(f" - {failure}") + return 1 + print("\nOK: sharing one session cut the worker's memory as expected.") + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--fragments", + type=int, + default=4000, + help=( + "fragments in the table under test. The per-open cost grows with this, and so does the saving, but " + "the read is one task per fragment, so the run time grows with it too" + ), + ) + parser.add_argument( + "--rows-per-fragment", + type=int, + default=ROWS_PER_FRAGMENT, + help=( + "rows in each fragment. Below a few thousand a scan task finishes too fast for the opens to overlap, " + "so fewer manifests are resident at once than --parallel suggests" + ), + ) + parser.add_argument( + "--parallel", + type=int, + default=24, + help=( + "Ray CPUs and concurrent scan tasks. This, not --fragments, is the lever that makes the saving " + "stand out of the RSS noise: it multiplies how many manifests are held at once without adding " + "tasks, so the run takes no longer. It may exceed the machine's core count" + ), + ) + parser.add_argument( + "--read-fraction", + type=float, + default=1.0, + help=( + "fraction of the table the read stage actually reads, as a prefix of the sorted key. A manifest " + "costs the same per open whatever the task reads, so this cuts the run time without touching the " + "measurement -- it is how a 100k-fragment table stays affordable" + ), + ) + parser.add_argument( + "--fragment-group-size", + type=int, + default=None, + help=( + "fragments per scan task (read_lance's own option). Fewer, longer tasks; each still opens the " + "dataset once and still carries the whole fragment list, so the per-open cost is unchanged" + ), + ) + parser.add_argument( + "--maint-fragments", + type=int, + default=MAINT_FRAGMENTS, + help=( + "fragments in the maintenance table. Compaction rewrites all of them, so this is the stage's whole " + "cost; it does not change what the stage concludes, since a maintenance worker opens the table once" + ), + ) + parser.add_argument( + "--maint-concurrency", + type=int, + default=MAINT_CONCURRENCY, + help=( + "concurrent compaction / update_columns instances. Each one is a separate process, so this is a " + "process count, not a thread count -- keep it near the pod's real cpu limit" + ), + ) + parser.add_argument( + "--data-dir", + default=str(Path(tempfile.gettempdir()) / "daft-lance-session-bench"), + help="where the test tables live; reused across runs", + ) + parser.add_argument( + "--mode", + choices=["both", "shared", "baseline", "prepare", "compare"], + default="both", + help=( + "both: run the whole A/B locally in two subprocesses. shared/baseline: one side only, for an " + "orchestrator that runs each side separately (a RayJob per side). prepare: only write the tables. " + "compare: read two sides' result files and judge them" + ), + ) + parser.add_argument("--baseline", help="--mode compare: the baseline side's result JSON") + parser.add_argument("--shared", help="--mode compare: the shared side's result JSON") + parser.add_argument("--json", action="store_true", help="also dump the raw measurements") + args = parser.parse_args() + if args.mode == "compare": + if not args.baseline or not args.shared: + raise SystemExit("--mode compare needs both --baseline and --shared") + return run_compare(args) + if args.parallel < 2: + raise SystemExit("--parallel must be at least 2 for the comparison to mean anything") + if not 0 < args.read_fraction <= 1.0: + raise SystemExit("--read-fraction must be in (0, 1]") + if args.mode == "prepare": + return run_prepare(args) + return run_parent(args) if args.mode == "both" else run_child(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/daft_lance/_blob.py b/daft_lance/_blob.py index 3ce44b8..560fc2f 100644 --- a/daft_lance/_blob.py +++ b/daft_lance/_blob.py @@ -12,6 +12,7 @@ from daft.dependencies import pa from daft.series import Series from daft_lance._metadata import LANCE_BLOB_DESCRIPTOR_TYPE +from daft_lance.namespace import open_dataset_from_open_kwargs if TYPE_CHECKING: from collections.abc import Iterable @@ -340,13 +341,26 @@ def read_blobs( if batch_size < 1: raise ValueError(f"batch_size must be at least 1, got {batch_size}") - # 2. Create a UDF closure over the lance dataset so we can read the blobs. The dataset - # pickles with its version and manifest, so workers read the same snapshot. + # 2. Create a UDF closure that reopens the pinned snapshot on the worker. + # Closing over ``ds`` would pickle it, and ``LanceDataset.__reduce__`` + # carries the whole re-encoded manifest -- one copy of the table's entire + # fragment list per task, which also bypasses the shared session's + # manifest cache. The uri, version and storage options are enough. + uri = ds.uri + open_kwargs: dict[str, Any] = {"version": ds.version, "storage_options": ds._storage_options} + worker_ds: list[lance.LanceDataset] = [] + + def _worker_dataset() -> lance.LanceDataset: + # Opened once per worker process, then reused for every batch. + if not worker_ds: + worker_ds.append(open_dataset_from_open_kwargs(uri, open_kwargs)) + return worker_ds[0] + @daft.func.batch(return_dtype=DataType.binary(), batch_size=batch_size) def read_blobs_udf(row_ids: Series) -> Series: ids = row_ids.to_pylist() valid = [i for i, row_id in enumerate(ids) if row_id is not None] - blobs = ds.read_blobs( + blobs = _worker_dataset().read_blobs( column, ids=[ids[i] for i in valid], io_buffer_size=io_buffer_size, diff --git a/daft_lance/lance_data_sink.py b/daft_lance/lance_data_sink.py index c67cbc8..a854ea3 100644 --- a/daft_lance/lance_data_sink.py +++ b/daft_lance/lance_data_sink.py @@ -30,6 +30,7 @@ ResolvedNamespaceTable, get_namespace_commit_kwargs, get_namespace_kwargs, + get_shared_session, get_write_fragments_kwargs, merge_storage_options, resolve_namespace_table, @@ -208,6 +209,19 @@ def _validate_overwrite_where_against_table(self, dataset: lance.LanceDataset) - def _namespace_kwargs(self) -> dict[str, Any]: return get_namespace_kwargs(self._namespace_impl, self._namespace_properties, self._table_id) + @property + def _dataset_open_kwargs(self) -> dict[str, Any]: + """Namespace wiring plus this process's shared lance session. + + Every dataset this sink opens goes through here. Without the session + each open builds its own ``Arc`` holding the table's whole + fragment list, so a many-fragment table costs that again per open -- + including the one the Rust writer does inside ``write_fragments`` on + every buffer flush. Commits keep using ``_namespace_commit_kwargs``: + ``LanceDataset.commit()`` takes no session. + """ + return {"session": get_shared_session(), **self._namespace_kwargs} + @property def _namespace_commit_kwargs(self) -> dict[str, Any]: return get_namespace_commit_kwargs( @@ -347,7 +361,7 @@ def _absorb_existing_dataset(self) -> lance.LanceDataset | None: dataset: lance.LanceDataset | None try: dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs ) except (ValueError, FileNotFoundError, OSError) as e: # Pinned to the Rust message format; lance has no typed exception. See test_lance_message_format_unchanged. @@ -430,6 +444,9 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada fragments = lance.fragment.write_fragments( wrapped, dataset_uri=self._table_uri, + # A uri destination makes the Rust writer load the dataset on every + # flush; the shared session turns the repeats into cache hits. + session=get_shared_session(), mode=self._write_mode, storage_options=self._storage_options, max_rows_per_file=self._max_rows_per_file, @@ -450,7 +467,9 @@ def _write_arrow_table(self, table: pa.Table) -> WriteResult[list[FragmentMetada def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: assert self._effective_pyarrow_schema is not None, "LanceDataSink.start() must run before writes" try: - ds = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + ds = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) except (ValueError, FileNotFoundError, OSError): ds = None @@ -465,7 +484,7 @@ def _ensure_mem_wal_dataset(self) -> lance.LanceDataset: storage_options=self._storage_options, data_storage_version=self._data_storage_version, use_legacy_format=self._use_legacy_format, - **self._namespace_kwargs, + **self._dataset_open_kwargs, ) details = ds.mem_wal_index_details() @@ -631,7 +650,9 @@ def checkpoint_commit_exists(self) -> bool: if self._checkpoint is None: raise RuntimeError("checking a Lance checkpoint commit requires the driver-side CheckpointStore") self._ensure_started() - dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) return self._idempotence_key_exists(dataset) def _finalize_uncheckpointed(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: @@ -679,7 +700,7 @@ def _finalize_insert_overwrite(self, fragments: list[FragmentMetadata]) -> Micro self._dataset_uri_arg, version=self._version, storage_options=self._storage_options, - **self._namespace_kwargs, + **self._dataset_open_kwargs, ) open_context = DatasetOpenContext.from_dataset( pinned, @@ -701,12 +722,14 @@ def _finalize_insert_overwrite(self, fragments: list[FragmentMetadata]) -> Micro self._overwrite_where, ) dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs ) return _dataset_stats(dataset) def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadata]]]) -> MicroPartition: - dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) if self._compact_after_write: logger.info( @@ -728,7 +751,7 @@ def _finalize_mem_wal(self, write_results: list[WriteResult[list[FragmentMetadat ), ) dataset = lance.dataset( - self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs ) return _dataset_stats(dataset) @@ -743,7 +766,9 @@ def _finalize_checkpointed(self, write_results: list[WriteResult[list[FragmentMe assert self._checkpoint is not None self._ensure_started() - dataset = lance.dataset(self._dataset_uri_arg, storage_options=self._storage_options, **self._namespace_kwargs) + dataset = lance.dataset( + self._dataset_uri_arg, storage_options=self._storage_options, **self._dataset_open_kwargs + ) pending_ids = self._pending_checkpoint_ids() # Check the Lance history again at finalize time. The pre-check in diff --git a/daft_lance/namespace.py b/daft_lance/namespace.py index ea7c8a4..dcd3463 100644 --- a/daft_lance/namespace.py +++ b/daft_lance/namespace.py @@ -28,16 +28,80 @@ from __future__ import annotations import os +import threading from dataclasses import dataclass, field from functools import lru_cache from typing import Any from urllib.parse import unquote, urlparse import lance +from lance.lance import _Session as LanceSession from lance_namespace import DeclareTableResponse, DescribeTableResponse, LanceNamespace _NAMESPACE_CACHE_SIZE = int(os.environ.get("DAFT_LANCE_NAMESPACE_CACHE_SIZE", "16")) +_SHARED_SESSION: LanceSession | None = None +_SHARED_SESSION_LOCK = threading.Lock() + + +def _sharing_enabled() -> bool: + """Whether session sharing is on; ``DAFT_LANCE_SHARE_SESSION=0`` turns it off. + + An escape hatch for a process that must not pool lance's caches, and the + only way to get a pre-sharing baseline inside a Ray worker: a driver-side + monkeypatch never reaches the worker process. Read on every call rather + than once at import, which costs nothing next to a dataset open and lets a + test flip it in-process. + """ + return os.environ.get("DAFT_LANCE_SHARE_SESSION", "1").strip().lower() not in {"0", "false", "no"} + + +def get_shared_session() -> LanceSession: + """Return this process's one lance ``Session``, creating it on first use. + + Every ``lance.dataset()`` that is opened without a session builds its own: + its own ``Arc`` holding the table's whole fragment list, its own + index and metadata cache budgets (6 GiB and 1 GiB each, not a shared + process total) and its own object-store client. A table with many fragments + therefore costs that much again for every concurrently open dataset, and + Daft opens one per scan task, per maintenance UDF instance and per writer. + + Only some of those overlap inside one process, which is where the saving + comes from. Scan tasks and writers do: ``scantask_max_parallel`` of them run + in the same worker, so their manifests add up. Maintenance instances do not + -- Daft gives each ``daft.cls`` instance a process of its own, measured on + both runners -- so there the shared session keeps the cost at one manifest + per process rather than collapsing many. An indexed table does not change + that: a maintenance task's open brings in the index's metadata entry only, + and the index data is read once on the driver, when the commit remaps it + onto the new fragment ids. + + Sharing one session collapses that: the manifest is a metadata-cache entry + keyed by (uri, version, e_tag), so every open of the snapshot the driver + pinned hands back a clone of the same ``Arc``. Measured on a + 20k-fragment table, 64 opens cost 859 MB unshared and 6 MB shared. + + Created lazily, never at import: a session carries an object-store client, + which must not be inherited across a fork. + """ + if not _sharing_enabled(): + # One fresh session per open, which is byte-for-byte what lance does + # when no session is passed at all. + return LanceSession() + global _SHARED_SESSION + if _SHARED_SESSION is None: + with _SHARED_SESSION_LOCK: + if _SHARED_SESSION is None: + _SHARED_SESSION = LanceSession() + return _SHARED_SESSION + + +def reset_shared_session() -> None: + """Drop the process's shared session. For tests only.""" + global _SHARED_SESSION + with _SHARED_SESSION_LOCK: + _SHARED_SESSION = None + def has_namespace_params(namespace_impl: str | None, table_id: list[str] | None) -> bool: """Whether namespace addressing is in effect (``namespace_properties`` stays optional).""" @@ -273,14 +337,20 @@ class DatasetOpenContext: commit as if the table were an unmanaged uri table. This context is the serializable substitute. It deliberately excludes the - live dataset, the namespace client, ``serialized_manifest``, ``session``, - ``commit_lock`` and ``asof``: + live dataset, the namespace client, ``serialized_manifest``, ``session`` + (none of which pickle usefully), ``commit_lock`` and ``asof``: * the client is rebuilt per process from the triple via the ``lru_cache``; * the manifest is omitted so the payload stays independent of fragment - count -- workers pay one pinned-manifest read instead; + count -- workers pay one pinned-manifest read instead, and + :func:`get_shared_session` makes the whole process share that one read; * ``asof``/tags are already resolved by the driver into :attr:`version`, so workers cannot drift onto a different snapshot. + + Because the session is shared, :attr:`index_cache_size` and + :attr:`metadata_cache_size_bytes` no longer size a per-dataset cache: lance + ignores them once a session is supplied, and the session's own budgets + apply to every dataset in the process. """ uri: str @@ -362,6 +432,7 @@ def _open(self, version: int | None) -> lance.LanceDataset: namespace_client=get_or_create_namespace(self.namespace_impl, self.namespace_properties), table_id=self.table_id, namespace_client_managed_versioning=self.managed_versioning, + session=get_shared_session(), ) def open_pinned(self) -> lance.LanceDataset: @@ -387,6 +458,7 @@ def open_dataset_from_open_kwargs(ds_uri: str | None, open_kwargs: dict[str, Any namespace_impl, namespace_properties, table_id = pop_namespace_params(open_kwargs) return lance.dataset( None if has_namespace_params(namespace_impl, table_id) else ds_uri, + session=get_shared_session(), **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), **open_kwargs, ) diff --git a/daft_lance/utils.py b/daft_lance/utils.py index c3cf79d..2742bc2 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -13,6 +13,7 @@ DatasetOpenContext, get_namespace_commit_kwargs, get_namespace_kwargs, + get_shared_session, has_namespace_params, merge_storage_options, resolve_namespace_table, @@ -214,6 +215,9 @@ def construct_lance_dataset_handle( dataset_uri, storage_options=merged_storage_options, version=version, + # The native runner plans and executes in one process, so the driver's + # manifest is the one every scan task reuses. + session=get_shared_session(), **get_namespace_kwargs(namespace_impl, namespace_properties, table_id), **kwargs, ) diff --git a/tests/io/lance/test_ray_session_memory.py b/tests/io/lance/test_ray_session_memory.py new file mode 100644 index 0000000..c05c26a --- /dev/null +++ b/tests/io/lance/test_ray_session_memory.py @@ -0,0 +1,357 @@ +"""Under Ray, the dataset opens happen in a worker process, not in the driver. + +``daft_lance`` hands one process-level ``lance.Session`` to every open, so a +table's fragment list is held once per process instead of once per scan task, +blob reader and writer. None of that is observable from the driver: a +monkeypatch or a spy only ever patches the driver, which is why the spy tests in +``test_shared_session.py`` are native-only. + +So these ask the worker itself. A plain batch UDF is fused into the scan's own +``SwordfishTask``, so a UDF that reports ``Session.size_bytes()`` reports it from +the very process that opened the datasets. ``size_bytes()`` is used rather than +RSS on purpose: it is exact, it needs no large table, and it is zero for any open +that did not go through the process session -- which is exactly what the +pre-sharing behavior looks like. ``benchmarks/lance_session_memory.py`` is the +counterpart that puts megabytes on it. + +The last two tests cover the maintenance path, which behaves differently enough +to be worth stating: its opens do not multiply by concurrency, because Daft runs +each ``daft.cls`` instance in a process of its own. They assert that, so the day +it stops being true is the day the tests say so. +""" + +from __future__ import annotations + +import json +import os +from typing import Any + +import lance +import numpy as np +import pyarrow as pa +import pytest +from pytest import TempPathFactory + +import daft +import daft_lance +from daft_lance._blob import read_blobs +from daft_lance.utils import construct_lance_dataset_handle + +ray = pytest.importorskip("ray") + +pytestmark = pytest.mark.skipif( + os.environ.get("DAFT_RUNNER") != "ray", + reason="a worker process's lance session is only observable under the Ray runner", +) + +NUM_FRAGMENTS = 32 +DRIVER_FLAG = os.environ.get("DAFT_LANCE_SHARE_SESSION") + + +def report_from_worker() -> dict[int, dict[str, Any]]: + """Ask every worker that runs a task for its lance session state. + + The probe is defined here, not at module scope, so cloudpickle ships it by + value and the worker does not have to import this test module. It must stay + a plain batch UDF: ``concurrency=`` or ``use_process=True`` would move it + into a process of its own, and it would then describe the wrong process. + """ + + @daft.func.batch(return_dtype=daft.DataType.string()) + def probe(column: daft.Series) -> list[str]: + from daft_lance.namespace import get_shared_session + + session = get_shared_session() + payload = json.dumps( + { + "pid": os.getpid(), + "session_size": session.size_bytes(), + "session_obj_id": id(session), + # ``is_same_as`` exists at runtime; pylance 12's stub omits it. + "same_session": bool(session.is_same_as(get_shared_session())), # type: ignore[attr-defined] + "flag": os.environ.get("DAFT_LANCE_SHARE_SESSION"), + } + ) + return [payload] * len(column) + + df = daft.from_pydict({"x": [0]}).with_column("p", probe(daft.col("x"))) + rows = [json.loads(value) for value in df.to_pydict()["p"]] + return {row["pid"]: row for row in rows} + + +def one_worker(report: dict[int, dict[str, Any]]) -> dict[str, Any]: + assert len(report) == 1, f"expected a single worker process, got pids {sorted(report)}" + return next(iter(report.values())) + + +def session_growth(before: dict[int, dict[str, Any]], after: dict[int, dict[str, Any]]) -> int: + """How much the worker's process-wide session grew between two reports.""" + assert sorted(before) == sorted(after), ( + f"the reports ran on different workers ({sorted(before)} then {sorted(after)}), so nothing can be compared" + ) + pid = next(iter(before)) + return int(after[pid]["session_size"]) - int(before[pid]["session_size"]) + + +def write_table(uri: str, fragments: int = NUM_FRAGMENTS) -> lance.LanceDataset: + """One row per fragment, so the manifest is mostly its fragment list.""" + table = pa.table({"i": pa.array(range(fragments), type=pa.int64())}) + return lance.write_dataset(table, uri, mode="create", max_rows_per_file=1) + + +def test_worker_sees_the_drivers_share_session_flag() -> None: + """Every other test here, and the A/B benchmark, rests on this. + + Ray actors do not inherit the driver's environment by contract, so if the + kill switch stopped reaching the worker, the benchmark would quietly compare + two identical runs and report a flattering number. + """ + worker = one_worker(report_from_worker()) + + assert worker["flag"] == DRIVER_FLAG + + +def test_scan_tasks_open_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + uri = str(tmp_path_factory.mktemp("read") / "table.lance") + dataset = write_table(uri) + + before = report_from_worker() + # Never count_rows(): count pushdown collapses the plan to a single task + # that opens no fragment at all. + rows = daft_lance.read_lance(uri).to_arrow().num_rows + after = report_from_worker() + + assert rows == NUM_FRAGMENTS + growth = session_growth(before, after) + assert growth >= len(dataset._ds.serialized_manifest()), ( + f"the worker's session grew by only {growth} bytes, which cannot hold this table's manifest: " + "the scan tasks opened their datasets outside the process session" + ) + assert one_worker(after)["same_session"] + + +def test_scan_tasks_reuse_one_session_per_process(tmp_path_factory: TempPathFactory) -> None: + """Two queries in a row must land on the same session object, not two.""" + uri = str(tmp_path_factory.mktemp("reuse") / "table.lance") + write_table(uri) + + daft_lance.read_lance(uri).to_arrow() + first = one_worker(report_from_worker()) + daft_lance.read_lance(uri).to_arrow() + second = one_worker(report_from_worker()) + + assert first["same_session"] and second["same_session"], ( + "the worker handed out a fresh session per call, so nothing is shared between its tasks" + ) + assert first["session_obj_id"] == second["session_obj_id"], ( + "the two queries saw different session objects, so the session is per-query, not per-process" + ) + # The second read of the same pinned snapshot is served from the manifest + # cache, so it adds nothing. + assert second["session_size"] == first["session_size"], ( + "re-reading the same pinned snapshot grew the session again: the manifest cache was not hit" + ) + + +def test_writes_open_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + """``write_fragments`` reloads the manifest per flush unless it gets a session.""" + uri = str(tmp_path_factory.mktemp("write") / "table.lance") + write_table(uri) + + before = report_from_worker() + daft.from_pydict({"i": list(range(8))}).write_lance(uri, mode="append", max_rows_per_file=1).collect() + after = report_from_worker() + + assert session_growth(before, after) > 0, ( + "the worker's session did not grow across the write: the writer opened outside it" + ) + assert lance.dataset(uri).count_rows() == NUM_FRAGMENTS + 8 + + +def test_indexed_queries_load_the_index_once_per_process(tmp_path_factory: TempPathFactory) -> None: + """A vector index is the other half of what a session holds, and the costly half. + + A session carries an index cache as well as a metadata cache, and index size + scales with rows where a manifest scales with fragments: measured on a + 200k-row table, one IVF_PQ index costs about 5 MB per session, against 9 KB + of manifest. A ``nearest`` read plans a single index-driven scan task, so + nothing here multiplies by concurrency -- what sharing changes is that the + index is loaded once per worker process instead of once per query, which is + a storage round trip and a decode every query without it. + """ + uri = str(tmp_path_factory.mktemp("indexed") / "table.lance") + rows, dim = 512, 8 + rng = np.random.default_rng(0) + vectors = rng.random((rows, dim), dtype=np.float32) + table = pa.table( + { + "id": pa.array(range(rows), type=pa.int64()), + "v": pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), type=pa.float32()), dim), + } + ) + dataset = lance.write_dataset(table, uri, max_rows_per_file=128) + dataset.create_index("v", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=2) + + # The same vector every time, on purpose: a different one probes different + # IVF partitions and legitimately loads more of the index, which would make + # "was it reloaded?" unanswerable on a table this small. + query = rng.random(dim, dtype=np.float32) + + def search() -> int: + nearest = {"column": "v", "q": query, "k": 4} + return daft_lance.read_lance(uri, default_scan_options={"nearest": nearest}).to_arrow().num_rows + + before = report_from_worker() + assert search() == 4 + after_first = report_from_worker() + assert search() == 4 + after_second = report_from_worker() + + first_growth = session_growth(before, after_first) + assert first_growth > 0, ( + "the worker's session did not grow across an indexed query: the index was loaded into a session of its own" + ) + # The second query must be served from the index cache the first one filled. + assert session_growth(after_first, after_second) == 0, ( + "the second indexed query grew the session again, so it reloaded the index instead of reusing it" + ) + + +def test_read_blobs_opens_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + """The blob reader reopens the pinned snapshot on the worker. + + That the reopen uses a uri plus open kwargs -- rather than closing over the + driver's ``LanceDataset``, whose ``__reduce__`` ships the whole re-encoded + manifest per task -- is asserted by the spy test in ``test_shared_session``, + which can only run natively. What is checked here is the other half: under + Ray the payload really does travel and the open really does land in the + worker's process-wide session. + """ + uri = str(tmp_path_factory.mktemp("blobs") / "table.lance") + rows = 4 + table = pa.table( + { + "id": pa.array(range(rows), type=pa.int64()), + "blob": lance.blob_array([b"payload-%d" % i for i in range(rows)]), + } + ) + dataset = lance.write_dataset(table, uri, data_storage_version="2.2", max_rows_per_file=1) + + before = report_from_worker() + df = daft.read_lance(uri, default_scan_options={"with_row_id": True}) + blobs = read_blobs(df, dataset, "blob").sort("id").to_pydict()["blob"] + after = report_from_worker() + + assert blobs == [b"payload-%d" % i for i in range(rows)] + assert session_growth(before, after) > 0, ( + "the worker's session did not grow across the blob read: it opened outside the process session" + ) + + +MAINT_CONCURRENCY = 2 + + +def maintenance_census(open_context: Any, concurrency: int) -> list[dict[str, Any]]: + """Open the table the way a maintenance worker does, and report who did it. + + Daft's maintenance operators are ``daft.cls`` instances -- one open per + instance, cached on ``self`` -- so what their memory costs turns entirely on + where Daft puts those instances. The real operators cannot be asked (they + return rewrite results, not pids), so this stands in for them, decorated the + same way: ``max_concurrency`` set, ``use_process`` left alone. + + Defined inside the function so cloudpickle ships it by value, exactly as in + :func:`report_from_worker`. + """ + + class OpenCensus: + def __init__(self, context: Any) -> None: + self.context = context + self._lance_ds: lance.LanceDataset | None = None + + @daft.method.batch(return_dtype=daft.DataType.string()) + def __call__(self, column: daft.Series) -> list[str]: + import json as json_ + import os as os_ + + from daft_lance.namespace import get_shared_session + + if self._lance_ds is None: + self._lance_ds = self.context.open_pinned() + session = get_shared_session() + payload = json_.dumps( + { + "pid": os_.getpid(), + "instance": id(self), + "session_obj_id": id(session), + "session_size": session.size_bytes(), + "fragments": len(self._lance_ds.get_fragments()), + } + ) + return [payload] * len(column) + + census = daft.cls(OpenCensus, max_concurrency=concurrency) + # One partition per instance several times over, so every instance in the + # pool is given work and reports back. + partitions = concurrency * 8 + df = daft.from_pydict({"x": list(range(partitions))}).into_partitions(partitions) + rows = [json.loads(value) for value in df.select(census(open_context)(daft.col("x")).alias("p")).to_pydict()["p"]] + instances = {(row["pid"], row["instance"]): row for row in rows} + assert len(instances) >= 2, ( + f"only {len(instances)} instance ran at concurrency {concurrency}, so this measured nothing about concurrency" + ) + return list(instances.values()) + + +def test_maintenance_instances_each_get_their_own_process(tmp_path_factory: TempPathFactory) -> None: + """Why concurrency does not multiply a maintenance worker's manifests. + + A scan task shares its process with ``scantask_max_parallel``-1 others, so + without a shared session their manifests add up and sharing collapses them. + A maintenance instance does not: Daft runs each one in a process of its own, + so it opens the table exactly once however high the concurrency, and there + is nothing for a session to collapse. This is asserted rather than assumed + because it is the whole reason the maintenance path needs no saving -- if + Daft ever fuses these instances into one process, the manifests start + multiplying and both this file and the A/B benchmark need a stage for it. + """ + uri = str(tmp_path_factory.mktemp("maint_procs") / "table.lance") + write_table(uri) + open_context = construct_lance_dataset_handle(uri).worker_open_context() + + instances = maintenance_census(open_context, MAINT_CONCURRENCY) + + pids = {row["pid"] for row in instances} + assert len(pids) == len(instances), ( + f"{len(instances)} maintenance instances shared {len(pids)} processes. Daft used to give each its own; " + "now that they overlap, their opens multiply within a process and that saving should be measured" + ) + assert {row["fragments"] for row in instances} == {NUM_FRAGMENTS}, ( + "the instances disagreed about the table's fragment count, so they did not all open the pinned snapshot" + ) + + +def test_maintenance_workers_open_through_the_process_session(tmp_path_factory: TempPathFactory) -> None: + """The open itself still has to go through the session, per process. + + One open per process is the floor, not zero: the worker pays one manifest + read either way. What this checks is that it lands in the process-wide + session rather than beside it, which is what keeps the floor at one -- a + maintenance process that also scans, or that Daft later gives a second + instance, then shares the one copy. + """ + uri = str(tmp_path_factory.mktemp("maint_session") / "table.lance") + dataset = write_table(uri) + manifest_bytes = len(dataset._ds.serialized_manifest()) + open_context = construct_lance_dataset_handle(uri).worker_open_context() + + instances = maintenance_census(open_context, MAINT_CONCURRENCY) + + for row in instances: + assert row["session_size"] >= manifest_bytes, ( + f"a maintenance instance's process session holds {row['session_size']} bytes, too little for this " + f"table's {manifest_bytes}-byte manifest: it opened outside the process session" + ) + assert len({row["session_obj_id"] for row in instances}) == len(instances), ( + "two instances in different processes reported the same session object id, which cannot be" + ) diff --git a/tests/io/lance/test_shared_session.py b/tests/io/lance/test_shared_session.py new file mode 100644 index 0000000..bba23e8 --- /dev/null +++ b/tests/io/lance/test_shared_session.py @@ -0,0 +1,208 @@ +"""Every dataset a process opens must share one lance Session. + +Without it each ``lance.dataset()`` builds its own ``Arc`` holding +the table's whole fragment list, plus its own index/metadata cache budgets and +object-store client. Daft opens one dataset per scan task, per maintenance UDF +instance and per writer, so on a many-fragment table that cost is paid again +for every concurrently open dataset. + +How much that saves depends on how many of those opens share a process, which +differs by path and is measured rather than assumed: see +``test_ray_session_memory.py`` for the scan and maintenance paths, and +``test_compaction_loads_the_index_in_its_commit_not_in_its_tasks`` below for +where an indexed table's maintenance cost actually sits. +""" + +from __future__ import annotations + +import os +import pickle + +import lance +import numpy as np +import pyarrow as pa +import pytest +from pytest import TempPathFactory + +import daft +from daft_lance._blob import read_blobs +from daft_lance.namespace import ( + DatasetOpenContext, + get_shared_session, + open_dataset_from_open_kwargs, + reset_shared_session, +) +from daft_lance.utils import construct_lance_dataset_handle + +NUM_FRAGMENTS = 32 + +ON_RAY = os.environ.get("DAFT_RUNNER") == "ray" + + +def same_session(left: object, right: object) -> bool: + """``Session.is_same_as`` exists at runtime; pylance 12's stub omits it.""" + return bool(left.is_same_as(right)) # type: ignore[attr-defined] + + +@pytest.fixture +def many_fragment_uri(tmp_path_factory: TempPathFactory) -> str: + """A table with one fragment per row, so the manifest is dominated by fragments.""" + uri = str(tmp_path_factory.mktemp("many_fragments") / "table.lance") + table = pa.table({"i": pa.array(range(NUM_FRAGMENTS), type=pa.int64())}) + lance.write_dataset(table, uri, mode="create", max_rows_per_file=1) + return uri + + +@pytest.fixture(autouse=True) +def fresh_session() -> None: + reset_shared_session() + + +def test_get_shared_session_returns_one_session_per_process() -> None: + first = get_shared_session() + assert get_shared_session() is first + assert same_session(first, get_shared_session()) + + +def test_reset_shared_session_builds_a_new_one() -> None: + first = get_shared_session() + reset_shared_session() + assert not same_session(get_shared_session(), first) + + +def test_kill_switch_gives_every_open_its_own_session(monkeypatch: pytest.MonkeyPatch) -> None: + """``DAFT_LANCE_SHARE_SESSION=0`` restores lance's per-open sessions. + + This is the only way to get pre-sharing behavior inside a Ray worker, so + the A/B benchmark depends on it: a driver-side patch never reaches there. + """ + monkeypatch.setenv("DAFT_LANCE_SHARE_SESSION", "0") + unshared = get_shared_session() + assert not same_session(get_shared_session(), unshared) + + monkeypatch.delenv("DAFT_LANCE_SHARE_SESSION") + shared = get_shared_session() + assert same_session(get_shared_session(), shared) + assert not same_session(shared, unshared) + + +def test_scan_task_opens_share_the_session(many_fragment_uri: str) -> None: + """``open_dataset_from_open_kwargs`` is what every scan task calls.""" + shared = get_shared_session() + open_kwargs = {"version": 1} + datasets = [open_dataset_from_open_kwargs(many_fragment_uri, dict(open_kwargs)) for _ in range(NUM_FRAGMENTS)] + + assert all(same_session(ds.session(), shared) for ds in datasets) + assert all(len(ds.get_fragments()) == NUM_FRAGMENTS for ds in datasets) + + +def test_driver_and_worker_opens_share_the_session(many_fragment_uri: str) -> None: + """The native runner plans and executes in one process, so these must match.""" + handle = construct_lance_dataset_handle(many_fragment_uri) + worker_ds = handle.worker_open_context().open_pinned() + + assert same_session(handle.dataset.session(), worker_ds.session()) + assert worker_ds.version == handle.dataset.version + + +def test_open_context_still_pickles(many_fragment_uri: str) -> None: + """The session must never become a field: it does not pickle, the context must.""" + context = DatasetOpenContext.from_dataset(lance.dataset(many_fragment_uri), many_fragment_uri) + + restored = pickle.loads(pickle.dumps(context)) + + assert restored == context + assert same_session(restored.open_pinned().session(), get_shared_session()) + + +def test_lance_session_itself_does_not_pickle() -> None: + with pytest.raises(TypeError): + pickle.dumps(get_shared_session()) + + +@pytest.mark.skipif(ON_RAY, reason="the spy patches the driver process, which is not where a Ray worker runs") +def test_read_blobs_reopens_on_the_worker(tmp_path_factory: TempPathFactory, monkeypatch: pytest.MonkeyPatch) -> None: + """The UDF must reopen the snapshot, not close over the driver's dataset. + + Closing over it pickles ``LanceDataset``, whose ``__reduce__`` carries the + whole re-encoded manifest -- one copy per task -- and a dataset rebuilt + from a serialized manifest also bypasses the session's manifest cache. + """ + uri = str(tmp_path_factory.mktemp("blobs")) + table = pa.table({"id": pa.array([1, 2], pa.int64()), "blob": lance.blob_array([b"first", b"second"])}) + ds = lance.write_dataset(table, uri, data_storage_version="2.2", max_rows_per_file=1) + + calls: list[dict[str, object]] = [] + + def spy(ds_uri: str | None, open_kwargs: dict[str, object] | None) -> lance.LanceDataset: + calls.append({"uri": ds_uri, **(open_kwargs or {})}) + return open_dataset_from_open_kwargs(ds_uri, open_kwargs) + + monkeypatch.setattr("daft_lance._blob.open_dataset_from_open_kwargs", spy) + + df = daft.read_lance(ds.uri, default_scan_options={"with_row_id": True}) + blobs = read_blobs(df, ds, "blob").sort("id").to_pydict()["blob"] + + assert blobs == [b"first", b"second"] + assert calls, "read_blobs did not reopen the dataset on the worker" + assert all(call["uri"] == ds.uri and call["version"] == ds.version for call in calls) + + +def test_compaction_loads_the_index_in_its_commit_not_in_its_tasks(tmp_path_factory: TempPathFactory) -> None: + """On an indexed table, compaction's index cost lands on the driver. + + The worry this settles: a maintenance task opens the table itself, so on an + indexed table every concurrent task might load the index and the cost would + multiply by the concurrency. It does not. A task's open brings in the + index's *metadata* entry and nothing more; the index data is read when the + rewritten fragments are committed, because that is when every index entry + has to be remapped onto the new fragment ids -- and the commit runs once, + on the driver, on ``lance_ds`` rather than on the workers' reopened copy. + + Run in-process on purpose: it needs the tasks and the commit to be + separable, which ``compact_files`` does not expose. ``CompactionTaskUDF`` is + exactly what Daft instantiates on a worker, so calling it directly measures + the same opens. + + The numbers come from ``index_cache_size_bytes()``, which is the accounting + lance caps its index budget with rather than resident bytes (it reads about + seven times what ``size_bytes()`` and the process's own RSS move by). Only + growth within that one accounting is compared here, never against a size in + bytes. + """ + from lance.optimize import Compaction, CompactionOptions + + from daft_lance.lance_compaction import CompactionTaskUDF + + uri = str(tmp_path_factory.mktemp("compact_index") / "table.lance") + rows, dim = 512, 8 + rng = np.random.default_rng(0) + vectors = rng.random((rows, dim), dtype=np.float32) + table = pa.table( + { + "i": pa.array(range(rows), type=pa.int64()), + "v": pa.FixedSizeListArray.from_arrays(pa.array(vectors.reshape(-1), type=pa.float32()), dim), + } + ) + dataset = lance.write_dataset(table, uri, max_rows_per_file=rows // 16) + dataset.create_index("v", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=2) + + session = get_shared_session() + handle = construct_lance_dataset_handle(uri) + plan = Compaction.plan(handle.dataset, CompactionOptions(target_rows_per_fragment=rows // 4)) + assert plan.num_tasks() > 1, "nothing to compact, so this measured no task opens at all" + + before_tasks = session.index_cache_size_bytes() + worker = CompactionTaskUDF(handle.worker_open_context()) + rewrites = [worker(task) for task in plan.tasks] + after_tasks = session.index_cache_size_bytes() + Compaction.commit(handle.dataset, rewrites) + after_commit = session.index_cache_size_bytes() + + task_growth = after_tasks - before_tasks + commit_growth = after_commit - after_tasks + assert commit_growth > 3 * task_growth, ( + f"the {plan.num_tasks()} compaction tasks grew the index cache by {task_growth} bytes and the commit by " + f"{commit_growth}; the index is no longer remapped in the commit, so a maintenance task's index cost may " + "now multiply by the concurrency and benchmarks/lance_session_memory.py needs revisiting" + )