diff --git a/docs/README.md b/docs/README.md index aae75ea8..0fa67e51 100644 --- a/docs/README.md +++ b/docs/README.md @@ -44,6 +44,7 @@ The built documentation will be in the `site/` directory. - `write.md` - Write operations guide - `distributed-indexing.md` - Distributed indexing guide - `compaction.md` - Distributed compaction guide + - `cleanup.md` - Distributed cleanup guide - `data-evolution.md` - Data evolution guide - `examples.md` - Usage examples - `.pages` - Navigation configuration diff --git a/docs/src/.pages b/docs/src/.pages index 4e3fde89..fcc6c4c9 100644 --- a/docs/src/.pages +++ b/docs/src/.pages @@ -5,5 +5,6 @@ nav: - Write: write.md - Index: distributed-indexing.md - Compaction: compaction.md + - Cleanup: cleanup.md - Data Evolution: data-evolution.md - Examples: examples.md diff --git a/docs/src/cleanup.md b/docs/src/cleanup.md new file mode 100644 index 00000000..d61e4241 --- /dev/null +++ b/docs/src/cleanup.md @@ -0,0 +1,87 @@ +# Distributed Cleanup + +Overwrites, compactions and index rebuilds leave old versions behind, together with the data, deletion and index files only those versions reference. `cleanup_old_versions` removes them. On a large table most of its time goes into listing the dataset's directories, which a single process does one page at a time, and deleting the files it finds one request at a time. + +Lance-Ray splits that work across Ray workers. Lance plans the cleanup on the driver: it reads every manifest, removes the old ones, and splits the listing of each directory into key ranges. Each range is swept by a worker, which lists it and removes the files there that no retained version references. + +Which files are removed is decided by Lance with the same rules as `lance.LanceDataset.cleanup_old_versions`. Lance-Ray only schedules the ranges. + +Requires a pylance release with `LanceDataset.plan_cleanup_old_versions`. With older releases the cleanup runs on the driver, and a warning is logged. + +## `cleanup_old_versions` + +```python +cleanup_old_versions( + uri=None, + *, + table_id=None, + older_than=None, + retain_versions=None, + versions=None, + delete_unverified=False, + error_if_tagged_old_versions=True, + delete_rate_limit=None, + delete_concurrency=None, + num_workers=4, + max_tasks=None, + max_task_retries=2, + storage_options=None, + namespace_impl=None, + namespace_properties=None, + ray_remote_args=None, +) +``` + +**Parameters:** + +- `uri`: Dataset URI (either `uri` OR `namespace_impl` + `table_id` required) +- `table_id`: Table identifier as a list of strings (requires `namespace_impl`) +- `older_than`, `retain_versions`, `versions`, `delete_unverified`, `error_if_tagged_old_versions`: As for `lance.LanceDataset.cleanup_old_versions`. Without `older_than`, `retain_versions` or `versions`, versions older than two weeks are removed. +- `delete_rate_limit`: Approximate maximum delete requests per second for the whole cleanup +- `delete_concurrency`: Maximum in-flight deletes per worker +- `num_workers`: Number of Ray workers (default: 4) +- `max_tasks`: Maximum number of key ranges (default: four per worker). Every listed directory gets at least one. +- `max_task_retries`: Retries per failed range (default: 2) +- `storage_options`: Optional storage configuration dictionary +- `namespace_impl`: Namespace implementation type (e.g., `"rest"`, `"dir"`) +- `namespace_properties`: Properties for connecting to the namespace +- `ray_remote_args`: Optional kwargs for the Ray workers + +**Returns:** `CleanupStats` for the whole cleanup. + +The driver logs a line as each range finishes, with the running number of files and bytes removed. + +## `explain_cleanup_old_versions` + +Takes the same arguments, without the delete limits, plus `max_files` (default: 1000), and returns the `CleanupExplanation` of what `cleanup_old_versions` would remove. Nothing is deleted. + +## Safety + +- A file is removed only if no retained version references it. Before listing its range, each worker reads the versions committed since the plan and keeps their files. +- A file that no version references at all, for example one left by a failed write, is removed only once it is seven days old, counted from when the plan was created. `delete_unverified=True` removes such files regardless of age. Only set it when nothing else writes to the dataset until the cleanup ends, or the files of an in-progress write can be removed. +- The old versions are removed while planning, before any file is swept, so none of them can be restored, tagged or branched from while its files are being removed. If a tag or branch is created on one of them while the plan is being made, the cleanup fails before removing anything. +- If a range still fails after its retries, a `RuntimeError` is raised. The old versions are already gone, and no version references the files the range left, so later cleanups only remove them as unverified files. +- Every worker must run the same pylance version as the driver. +- Cleaning referenced branches is not supported. + +## Limits + +- `delete_rate_limit` is divided evenly between the workers, at least one request per second each, so the total is approximate. +- On a store that does not list keys in lexical order (for example S3 Express One Zone), a range would have to list its whole directory, so each directory gets a single range. +- Reading the manifests and removing the old ones still happens on the driver. + +## Example + +```python +from datetime import timedelta + +import lance_ray as lr + +stats = lr.cleanup_old_versions( + "s3://bucket/path/to/table.lance", + older_than=timedelta(days=7), + num_workers=16, + delete_rate_limit=2000, +) +print(stats.old_versions, stats.bytes_removed) +``` diff --git a/lance_ray/__init__.py b/lance_ray/__init__.py index cea756a3..68104502 100644 --- a/lance_ray/__init__.py +++ b/lance_ray/__init__.py @@ -8,6 +8,7 @@ __version__ = "0.5.0" __author__ = "LanceDB Devs" __email__ = "dev@lancedb.com" +from .cleanup import cleanup_old_versions, explain_cleanup_old_versions from .compaction import compact_database, compact_files # Main imports @@ -42,6 +43,8 @@ "optimize_indices", "compact_files", "compact_database", + "cleanup_old_versions", + "explain_cleanup_old_versions", "LanceFragmentWriter", "LanceFragmentCommitter", ] diff --git a/lance_ray/_dispatch.py b/lance_ray/_dispatch.py new file mode 100644 index 00000000..d569efd4 --- /dev/null +++ b/lance_ray/_dispatch.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Run independent, idempotent work items on a set of Ray actors.""" + +from __future__ import annotations + +import logging +import time +from collections import deque +from collections.abc import Callable +from typing import Any, Optional + +import ray +from ray.exceptions import ActorUnavailableError + +logger = logging.getLogger(__name__) + +Handler = Callable[[Any], dict[str, Any]] + + +class _BatchWorker: + """Ray actor that runs one handler call at a time.""" + + def run(self, handler: Handler, item: Any) -> dict[str, Any]: + return handler(item) + + +def _map_async_with_pool( + create_handler: Callable[[], Handler], + items: list[Any], + *, + num_workers: int, + ray_remote_args: Optional[dict[str, Any]], + error_prefix: str, + max_task_retries: int = 0, + on_result: Optional[Callable[[int, dict[str, Any]], None]] = None, +) -> list[dict[str, Any]]: + """Run items on a pool of Ray actors, one item per idle worker. + + Each item is handed to whichever worker is idle at that moment, so a slow + node only delays the items it is actually running. (Ray's + ``multiprocessing.Pool`` assigns items round-robin at submission time, + which makes a map of N items wait for the slowest worker's fixed share + of N.) Items whose handler returned ``{"status": "error"}`` or whose + Ray task/actor failed are resubmitted immediately, up to + ``max_task_retries`` times each, rather than after every other item has + finished; handlers must therefore be idempotent. + + ``on_result(index, result)`` is called once per item, with its final result. + + Workers are actors created with ``ray_remote_args`` (``num_cpus=0`` unless + given, like Ray's Pool). A worker whose process dies is dropped from the + pool unless ``{"max_restarts": ...}`` is set, in which case Ray restarts + it and it keeps receiving items. + """ + remote_args = {"num_cpus": 0, **(ray_remote_args or {})} + worker_cls: Any = ray.remote(**remote_args)(_BatchWorker) + workers: list[Any] = [worker_cls.remote() for _ in range(num_workers)] + try: + handler = create_handler() + results: dict[int, dict[str, Any]] = {} + pending: deque[tuple[int, int]] = deque( + (index, 0) for index in range(len(items)) + ) + idle: list[Any] = [] + # A worker joins the idle list once Ray has actually started it, so an + # item never queues on an actor still waiting for cluster resources. + starting: dict[Any, Any] = {w.__ray_ready__.remote(): w for w in workers} + in_flight: dict[Any, tuple[Any, int, int]] = {} + while pending or in_flight: + while pending and idle: + index, attempt = pending.popleft() + worker = idle.pop() + ref = worker.run.remote(handler, items[index]) + in_flight[ref] = (worker, index, attempt) + if not in_flight and not starting: + raise RuntimeError("every worker has died") + [ref], _ = ray.wait(list(in_flight) + list(starting), num_returns=1) + if ref in starting: + worker = starting.pop(ref) + try: + ray.get(ref) + idle.append(worker) + except ActorUnavailableError: + # Restarting under max_restarts: probe again shortly. + time.sleep(0.5) + starting[worker.__ray_ready__.remote()] = worker + except Exception as exc: + logger.warning( + "Worker unavailable (%s); dropping it. Pass max_restarts " + "in ray_remote_args to restart dead workers", + exc, + ) + continue + worker, index, attempt = in_flight.pop(ref) + try: + result = ray.get(ref) + idle.append(worker) + except Exception as exc: + result = {"status": "error", "error": f"{type(exc).__name__}: {exc}"} + # A live worker is ready again at once; a dead one is dropped + # (or, with max_restarts, restarted) by the start path above. + starting[worker.__ray_ready__.remote()] = worker + results[index] = result + if result.get("status") != "success" and attempt < max_task_retries: + logger.warning( + "Batch %d failed (attempt %d of %d): %s; resubmitting", + index, + attempt + 1, + max_task_retries + 1, + result.get("error"), + ) + pending.append((index, attempt + 1)) + elif on_result is not None: + on_result(index, result) + return [results[i] for i in range(len(items))] + except Exception as exc: + raise RuntimeError(f"{error_prefix}: {exc}") from exc + finally: + for worker in workers: + ray.kill(worker) diff --git a/lance_ray/cleanup.py b/lance_ray/cleanup.py new file mode 100644 index 00000000..dd569c1e --- /dev/null +++ b/lance_ray/cleanup.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Distributed ``cleanup_old_versions`` on Ray.""" + +from __future__ import annotations + +import logging +from datetime import timedelta +from functools import lru_cache +from typing import Any, Optional, cast + +import lance +from lance.dataset import LanceDataset + +from ._dispatch import _map_async_with_pool +from .utils import ( + get_namespace_kwargs, + resolve_namespace_table, + validate_uri_or_namespace, +) + +logger = logging.getLogger(__name__) + +# Tasks each worker gets by default, so a slow range does not hold up the rest. +_TASKS_PER_WORKER = 4 + + +def _freeze(value: Optional[dict[str, Any]]) -> Optional[tuple[tuple[str, Any], ...]]: + return tuple(sorted(value.items())) if value else None + + +def _open_dataset( + uri: str, + storage_options: Optional[tuple[tuple[str, Any], ...]], + namespace_impl: Optional[str], + namespace_properties: Optional[tuple[tuple[str, str], ...]], + table_id: Optional[tuple[str, ...]], +) -> LanceDataset: + namespace_kwargs = get_namespace_kwargs( + namespace_impl, + dict(namespace_properties) if namespace_properties else None, + list(table_id) if table_id else None, + ) + return lance.LanceDataset( + uri, + storage_options=dict(storage_options) if storage_options else None, + **namespace_kwargs, + ) + + +@lru_cache(maxsize=8) +def _open_worker_dataset(*dataset_args: Any) -> LanceDataset: + """Opens the dataset once for all of a worker's tasks: a task runs on any + version of it. The driver plans on the latest version, so it does not cache.""" + return _open_dataset(*dataset_args) + + +def _supports_distributed_cleanup() -> bool: + return hasattr(LanceDataset, "plan_cleanup_old_versions") + + +def _sweep_handler( + dataset_args: tuple[Any, ...], + delete_concurrency: Optional[int], + delete_rate_limit: Optional[int], +) -> Any: + def run(task: Any) -> dict[str, Any]: + try: + dataset = _open_worker_dataset(*dataset_args) + result = task.execute( + dataset, + delete_concurrency=delete_concurrency, + delete_rate_limit=delete_rate_limit, + ) + return {"status": "success", "result": result} + except Exception as e: + return {"status": "error", "error": f"{type(e).__name__}: {e}"} + + return run + + +def _cleanup_kwargs( + older_than: Optional[timedelta], + retain_versions: Optional[int], + versions: Optional[list[int]], + delete_unverified: bool, + error_if_tagged_old_versions: bool, +) -> dict[str, Any]: + return { + "older_than": older_than, + "retain_versions": retain_versions, + "versions": versions, + "delete_unverified": delete_unverified, + "error_if_tagged_old_versions": error_if_tagged_old_versions, + } + + +def _run_cleanup( + *, + uri: Optional[str], + table_id: Optional[list[str]], + cleanup_kwargs: dict[str, Any], + explain: bool, + max_files: int, + delete_rate_limit: Optional[int], + delete_concurrency: Optional[int], + num_workers: int, + max_tasks: Optional[int], + max_task_retries: int, + storage_options: Optional[dict[str, str]], + namespace_impl: Optional[str], + namespace_properties: Optional[dict[str, str]], + ray_remote_args: Optional[dict[str, Any]], +) -> Any: + if num_workers <= 0: + raise ValueError(f"num_workers must be positive, got {num_workers}") + if max_tasks is not None and max_tasks <= 0: + raise ValueError(f"max_tasks must be positive, got {max_tasks}") + if max_task_retries < 0: + raise ValueError( + f"max_task_retries must be non-negative, got {max_task_retries}" + ) + validate_uri_or_namespace(uri, namespace_impl, table_id) + uri, merged_storage_options = resolve_namespace_table( + uri, storage_options, namespace_impl, namespace_properties, table_id + ) + dataset_args = ( + uri, + _freeze(merged_storage_options), + namespace_impl, + _freeze(namespace_properties), + tuple(table_id) if table_id else None, + ) + dataset = _open_dataset(*dataset_args) + + if not _supports_distributed_cleanup(): + logger.warning( + "pylance %s cannot distribute cleanup; running it on the driver", + lance.__version__, + ) + if explain: + return dataset.explain_cleanup_old_versions( + **cleanup_kwargs, include_files=True, max_files=max_files + ) + # Older releases do not take every limit, so only pass the ones set. + limits: dict[str, Any] = { + name: value + for name, value in [ + ("delete_rate_limit", delete_rate_limit), + ("delete_concurrency", delete_concurrency), + ] + if value is not None + } + return dataset.cleanup_old_versions(**cleanup_kwargs, **limits) + + # Typed as Any: the pylance this is type-checked against may predate the API. + plan = cast(Any, dataset).plan_cleanup_old_versions( + **cleanup_kwargs, + delete_rate_limit=delete_rate_limit, + delete_concurrency=delete_concurrency, + explain=explain, + max_files=max_files, + ) + tasks = plan.tasks(max_tasks or num_workers * _TASKS_PER_WORKER) + workers = min(num_workers, len(tasks)) + # Each worker deletes at its own rate, so split the budget between them. + worker_rate = max(1, delete_rate_limit // workers) if delete_rate_limit else None + tasks_per_subtree: dict[str, int] = {} + for task in tasks: + tasks_per_subtree[task.subtree] = tasks_per_subtree.get(task.subtree, 0) + 1 + logger.info( + "Planned cleanup of %s at version %d: %d old versions, %d tasks (%s) on %d workers", + uri, + plan.read_version, + plan.num_old_versions, + len(tasks), + ", ".join( + f"{name}: {count}" for name, count in sorted(tasks_per_subtree.items()) + ), + workers, + ) + + progress = {"done": 0, "files": 0, "bytes": 0} + + def on_result(index: int, result: dict[str, Any]) -> None: + task = tasks[index] + if result["status"] != "success": + logger.warning( + "Cleanup task %d (%s) failed: %s", index, task.subtree, result["error"] + ) + return + stats = result["result"].stats + progress["done"] += 1 + progress["files"] += ( + stats.data_files_removed + + stats.transaction_files_removed + + stats.index_files_removed + + stats.deletion_files_removed + ) + progress["bytes"] += stats.bytes_removed + logger.info( + "Cleanup task %d/%d done: %s (%s, %s]; %d files, %d bytes %s so far", + progress["done"], + len(tasks), + task.subtree, + task.start_after or "", + task.end_inclusive or "", + progress["files"], + progress["bytes"], + "found" if explain else "removed", + ) + + results = _map_async_with_pool( + create_handler=lambda: _sweep_handler( + dataset_args, delete_concurrency, worker_rate + ), + items=tasks, + num_workers=workers, + ray_remote_args=ray_remote_args, + error_prefix="Failed to complete distributed cleanup", + max_task_retries=max_task_retries, + on_result=on_result, + ) + failed = [r for r in results if r["status"] != "success"] + if failed: + errors = "; ".join(r["error"] for r in failed) + if explain: + raise RuntimeError( + f"Explaining cleanup failed in {len(failed)} of {len(tasks)} tasks: " + f"{errors}" + ) + raise RuntimeError( + f"Cleanup failed in {len(failed)} of {len(tasks)} tasks; the old versions " + "are removed, and later cleanups only remove the files these tasks left " + f"as unverified files: {errors}" + ) + sweep_results = [r["result"] for r in results] + if explain: + return plan.finish_explain(dataset, sweep_results) + stats = plan.finish(dataset, sweep_results) + logger.info("Finished cleanup of %s: %s", uri, stats) + return stats + + +def cleanup_old_versions( + uri: Optional[str] = None, + *, + table_id: Optional[list[str]] = None, + older_than: Optional[timedelta] = None, + retain_versions: Optional[int] = None, + versions: Optional[list[int]] = None, + delete_unverified: bool = False, + error_if_tagged_old_versions: bool = True, + delete_rate_limit: Optional[int] = None, + delete_concurrency: Optional[int] = None, + num_workers: int = 4, + max_tasks: Optional[int] = None, + max_task_retries: int = 2, + storage_options: Optional[dict[str, str]] = None, + namespace_impl: Optional[str] = None, + namespace_properties: Optional[dict[str, str]] = None, + ray_remote_args: Optional[dict[str, Any]] = None, +) -> Any: + """Remove old versions of a Lance dataset, listing and deleting on Ray workers. + + Lance plans the cleanup on the driver: it removes the old manifests and splits + the listing of the dataset's directories into key ranges. Each range is swept + by a Ray worker, which removes the files there that no retained version + references, using the same rules as + :meth:`lance.LanceDataset.cleanup_old_versions`. Lance-Ray only schedules the + ranges; it never decides which files are removed. + + A failed range is retried up to ``max_task_retries`` times. If a range still + fails, a ``RuntimeError`` is raised. The old versions are already gone by + then, and later cleanups only remove the files the range left as unverified + files. With pylance releases that cannot distribute a cleanup, it runs on the + driver instead. + + Args: + uri: The URI of the dataset. Either uri OR (namespace_impl + table_id) + must be provided. + table_id: The table identifier, used together with namespace_impl. + older_than: Only versions older than this are removed. Defaults to two + weeks when ``retain_versions`` and ``versions`` are not set either. + retain_versions: Retain the last N versions. + versions: Remove only these versions. + delete_unverified: Also remove recent files no version references. Only + safe when nothing else writes to the dataset until the cleanup ends. + error_if_tagged_old_versions: Raise if a tagged version would otherwise + be removed. + delete_rate_limit: Approximate maximum delete requests per second for + the whole cleanup. Each worker gets an equal share, at least one. + delete_concurrency: Maximum in-flight deletes per worker, and on the + driver when removing the old manifests. + num_workers: Number of Ray workers (default: 4). + max_tasks: Maximum number of ranges to split the listing into. Defaults + to four per worker. Every listed directory gets at least one. + max_task_retries: Retries per failed range (default: 2). + storage_options: Storage options for the dataset. + namespace_impl: The namespace implementation type (e.g., "rest", "dir"). + namespace_properties: Properties for connecting to the namespace. + ray_remote_args: Options for the Ray workers (e.g., num_cpus, resources). + + Returns: + The ``CleanupStats`` of the whole cleanup. + """ + return _run_cleanup( + uri=uri, + table_id=table_id, + cleanup_kwargs=_cleanup_kwargs( + older_than, + retain_versions, + versions, + delete_unverified, + error_if_tagged_old_versions, + ), + explain=False, + max_files=1, + delete_rate_limit=delete_rate_limit, + delete_concurrency=delete_concurrency, + num_workers=num_workers, + max_tasks=max_tasks, + max_task_retries=max_task_retries, + storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + ray_remote_args=ray_remote_args, + ) + + +def explain_cleanup_old_versions( + uri: Optional[str] = None, + *, + table_id: Optional[list[str]] = None, + older_than: Optional[timedelta] = None, + retain_versions: Optional[int] = None, + versions: Optional[list[int]] = None, + delete_unverified: bool = False, + error_if_tagged_old_versions: bool = True, + max_files: int = 1000, + num_workers: int = 4, + max_tasks: Optional[int] = None, + max_task_retries: int = 2, + storage_options: Optional[dict[str, str]] = None, + namespace_impl: Optional[str] = None, + namespace_properties: Optional[dict[str, str]] = None, + ray_remote_args: Optional[dict[str, Any]] = None, +) -> Any: + """Explain what :func:`cleanup_old_versions` would remove, without deleting. + + Takes the same arguments as :func:`cleanup_old_versions`, plus ``max_files``, + the maximum number of candidate files to include. The listing is split + across Ray workers in the same way. + + Returns: + A ``CleanupExplanation``, as from + :meth:`lance.LanceDataset.explain_cleanup_old_versions` with + ``include_files=True``. + """ + if max_files <= 0: + raise ValueError("max_files must be positive") + return _run_cleanup( + uri=uri, + table_id=table_id, + cleanup_kwargs=_cleanup_kwargs( + older_than, + retain_versions, + versions, + delete_unverified, + error_if_tagged_old_versions, + ), + explain=True, + max_files=max_files, + delete_rate_limit=None, + delete_concurrency=None, + num_workers=num_workers, + max_tasks=max_tasks, + max_task_retries=max_task_retries, + storage_options=storage_options, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + ray_remote_args=ray_remote_args, + ) diff --git a/lance_ray/index.py b/lance_ray/index.py index 6474531b..099b2055 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -6,9 +6,7 @@ import logging import math import os -import time import uuid -from collections import deque from collections.abc import Callable from typing import Any, Literal, Optional, TypeAlias, cast, get_args @@ -19,8 +17,8 @@ from lance.dataset import Index, IndexConfig, LanceDataset from lance.indices import IndicesBuilder from packaging import version -from ray.exceptions import ActorUnavailableError +from ._dispatch import _map_async_with_pool from .field_path import resolve_arrow_field_path, resolve_dataset_field_path from .utils import ( get_namespace_kwargs, @@ -142,104 +140,6 @@ def _distribute_fragments_balanced( return non_empty_batches -class _BatchWorker: - """Ray actor that runs one fragment-batch handler call at a time.""" - - def run(self, handler: _FragmentHandler, fragment_ids: list[int]) -> dict[str, Any]: - return handler(fragment_ids) - - -def _map_async_with_pool( - create_fragment_handler: Callable[[], _FragmentHandler], - fragment_batches: list[list[int]], - *, - num_workers: int, - ray_remote_args: Optional[dict[str, Any]], - error_prefix: str, - max_task_retries: int = 0, -) -> list[dict[str, Any]]: - """Run fragment batches on a pool of Ray actors, one batch per idle worker. - - Each batch is handed to whichever worker is idle at that moment, so a slow - node only delays the batches it is actually running. (Ray's - ``multiprocessing.Pool`` assigns batches round-robin at submission time, - which makes a map of N batches wait for the slowest worker's fixed share - of N.) Batches whose handler returned ``{"status": "error"}`` or whose - Ray task/actor failed are resubmitted immediately, up to - ``max_task_retries`` times each, rather than after every other batch has - finished; handlers must therefore be idempotent. - - Workers are actors created with ``ray_remote_args`` (``num_cpus=0`` unless - given, like Ray's Pool). A worker whose process dies is dropped from the - pool unless ``{"max_restarts": ...}`` is set, in which case Ray restarts - it and it keeps receiving batches. - """ - remote_args = {"num_cpus": 0, **(ray_remote_args or {})} - worker_cls: Any = ray.remote(**remote_args)(_BatchWorker) - workers: list[Any] = [worker_cls.remote() for _ in range(num_workers)] - try: - fragment_handler = create_fragment_handler() - results: dict[int, dict[str, Any]] = {} - pending: deque[tuple[int, int]] = deque( - (index, 0) for index in range(len(fragment_batches)) - ) - idle: list[Any] = [] - # A worker joins the idle list once Ray has actually started it, so a - # batch never queues on an actor still waiting for cluster resources. - starting: dict[Any, Any] = {w.__ray_ready__.remote(): w for w in workers} - in_flight: dict[Any, tuple[Any, int, int]] = {} - while pending or in_flight: - while pending and idle: - index, attempt = pending.popleft() - worker = idle.pop() - ref = worker.run.remote(fragment_handler, fragment_batches[index]) - in_flight[ref] = (worker, index, attempt) - if not in_flight and not starting: - raise RuntimeError("every worker has died") - [ref], _ = ray.wait(list(in_flight) + list(starting), num_returns=1) - if ref in starting: - worker = starting.pop(ref) - try: - ray.get(ref) - idle.append(worker) - except ActorUnavailableError: - # Restarting under max_restarts: probe again shortly. - time.sleep(0.5) - starting[worker.__ray_ready__.remote()] = worker - except Exception as exc: - logger.warning( - "Worker unavailable (%s); dropping it. Pass max_restarts " - "in ray_remote_args to restart dead workers", - exc, - ) - continue - worker, index, attempt = in_flight.pop(ref) - try: - result = ray.get(ref) - idle.append(worker) - except Exception as exc: - result = {"status": "error", "error": f"{type(exc).__name__}: {exc}"} - # A live worker is ready again at once; a dead one is dropped - # (or, with max_restarts, restarted) by the start path above. - starting[worker.__ray_ready__.remote()] = worker - results[index] = result - if result.get("status") != "success" and attempt < max_task_retries: - logger.warning( - "Batch %d failed (attempt %d of %d): %s; resubmitting", - index, - attempt + 1, - max_task_retries + 1, - result.get("error"), - ) - pending.append((index, attempt + 1)) - return [results[i] for i in range(len(fragment_batches))] - except Exception as exc: - raise RuntimeError(f"{error_prefix}: {exc}") from exc - finally: - for worker in workers: - ray.kill(worker) - - def _resolve_num_segments(num_workers: int, num_segments: Optional[int]) -> int: if num_workers <= 0: raise ValueError(f"num_workers must be positive, got {num_workers}") @@ -792,8 +692,8 @@ def create_fragment_handler() -> Any: ) results = _map_async_with_pool( - create_fragment_handler=create_fragment_handler, - fragment_batches=fragment_batches, + create_handler=create_fragment_handler, + items=fragment_batches, num_workers=pool_workers, ray_remote_args=ray_remote_args, error_prefix="Failed to complete distributed index building", @@ -1661,8 +1561,8 @@ def create_fragment_handler() -> Any: ) results = _map_async_with_pool( - create_fragment_handler=create_fragment_handler, - fragment_batches=fragment_batches, + create_handler=create_fragment_handler, + items=fragment_batches, num_workers=pool_workers, ray_remote_args=ray_remote_args, error_prefix="Failed to complete distributed vector index building", diff --git a/tests/_ray_test_support.py b/tests/_ray_test_support.py index 552a1bda..247f0e9a 100644 --- a/tests/_ray_test_support.py +++ b/tests/_ray_test_support.py @@ -158,3 +158,61 @@ def handler(fragment_ids: list[int]) -> dict[str, Any]: return {"status": "success", "fragment_ids": fragment_ids} return handler + + +# Stand-ins for pylance's cleanup sweep task and result, so the distributed +# cleanup driver can be tested against pylance releases without them. Workers +# unpickle them, so they live here for the same reason as the handlers above. + + +class FakeCleanupStats: + """The fields of ``CleanupStats`` the driver reads.""" + + def __init__(self, data_files_removed: int = 0, bytes_removed: int = 0) -> None: + self.data_files_removed = data_files_removed + self.transaction_files_removed = 0 + self.index_files_removed = 0 + self.deletion_files_removed = 0 + self.bytes_removed = bytes_removed + + +class FakeSweepResult: + def __init__( + self, + index: int, + delete_concurrency: Optional[int], + delete_rate_limit: Optional[int], + ) -> None: + self.index = index + self.delete_concurrency = delete_concurrency + self.delete_rate_limit = delete_rate_limit + self.stats = FakeCleanupStats(data_files_removed=1, bytes_removed=10) + + +class FakeSweepTask: + """Fails its first ``fail_attempts`` attempts, counted with files in + ``marker_dir``; its result carries the limits it ran with.""" + + def __init__(self, index: int, marker_dir: str, fail_attempts: int = 0) -> None: + self.index = index + self.subtree = "data" + self.start_after: Optional[str] = None if index == 0 else f"data/{index}" + self.end_inclusive: Optional[str] = f"data/{index + 1}" + self.marker_dir = marker_dir + self.fail_attempts = fail_attempts + + def execute( + self, + dataset: Any, + *, + delete_concurrency: Optional[int] = None, + delete_rate_limit: Optional[int] = None, + ) -> FakeSweepResult: + prefix = f"{self.index}." + attempts = len([f for f in os.listdir(self.marker_dir) if f.startswith(prefix)]) + open(os.path.join(self.marker_dir, f"{prefix}{time.time_ns()}"), "w").close() + if attempts < self.fail_attempts: + raise RuntimeError(f"injected failure of task {self.index}") + if dataset.count_rows() == 0: + raise RuntimeError("the worker opened an empty dataset") + return FakeSweepResult(self.index, delete_concurrency, delete_rate_limit) diff --git a/tests/test_distributed_cleanup.py b/tests/test_distributed_cleanup.py new file mode 100644 index 00000000..1390220b --- /dev/null +++ b/tests/test_distributed_cleanup.py @@ -0,0 +1,334 @@ +"""Tests for lance_ray.cleanup.""" + +import os +import time +from collections.abc import Callable, Iterator +from datetime import timedelta +from pathlib import Path +from typing import Any, Literal, Optional +from unittest.mock import patch + +import lance +import lance_ray as lr +import pyarrow as pa +import pytest +import ray +from lance.dataset import LanceDataset +from lance_ray import cleanup as cleanup_mod + +from _ray_test_support import FakeCleanupStats, FakeSweepTask + +requires_distributed_cleanup = pytest.mark.skipif( + not cleanup_mod._supports_distributed_cleanup(), + reason="installed pylance cannot distribute a cleanup", +) + +_TABLE = pa.table({"id": list(range(100)), "value": [f"v{i}" for i in range(100)]}) + + +def _write_versions(uri: str) -> None: + """Old versions with an index, deletions and an aged orphan, then one more.""" + lance.write_dataset(_TABLE, uri, max_rows_per_file=25) + lance.write_dataset(_TABLE, uri, mode="overwrite", max_rows_per_file=25) + dataset = lance.dataset(uri) + dataset.create_scalar_index("id", "BTREE") + dataset.delete("id < 10") + lance.write_dataset(_TABLE, uri, mode="overwrite", max_rows_per_file=25) + orphan = Path(uri) / "data" / "orphan.lance" + orphan.write_bytes(b"orphan") + eight_days_ago = time.time() - timedelta(days=8).total_seconds() + os.utime(orphan, (eight_days_ago, eight_days_ago)) + + +def _files(uri: str) -> set[str]: + return {str(path) for path in Path(uri).rglob("*") if path.is_file()} + + +@pytest.fixture +def uri(tmp_path: Path) -> str: + uri = str(tmp_path / "dataset") + _write_versions(uri) + return uri + + +def _expected(uri: str) -> tuple[Any, set[str]]: + """Single-process explanation, and the files left after cleaning by it.""" + explanation = lance.dataset(uri).explain_cleanup_old_versions( + older_than=timedelta(0), include_files=True, max_files=100_000 + ) + removed = {"/" + file.path for file in explanation.candidate_files} + assert removed <= _files(uri) + return explanation, _files(uri) - removed + + +@requires_distributed_cleanup +def test_matches_single_process_cleanup(uri: str) -> None: + expected, remaining = _expected(uri) + assert expected.stats.old_versions == 4 + assert expected.stats.data_files_removed > 0 + assert expected.stats.index_files_removed > 0 + assert expected.stats.deletion_files_removed > 0 + + explanation = lr.explain_cleanup_old_versions( + uri, older_than=timedelta(0), num_workers=2 + ) + assert explanation.stats.bytes_removed == expected.stats.bytes_removed + assert sorted(f.path for f in explanation.candidate_files) == sorted( + f.path for f in expected.candidate_files + ) + + stats = lr.cleanup_old_versions(uri, older_than=timedelta(0), num_workers=2) + assert stats.old_versions == expected.stats.old_versions + assert stats.data_files_removed == expected.stats.data_files_removed + assert stats.bytes_removed == expected.stats.bytes_removed + assert _files(uri) == remaining + assert lance.dataset(uri).to_table() == _TABLE + + +@requires_distributed_cleanup +def test_cleanup_with_directory_namespace(tmp_path: Path) -> None: + import lance_namespace as ln + from lance_namespace import DescribeTableRequest + + table_id = ["cleanup_table"] + namespace: dict[str, Any] = { + "namespace_impl": "dir", + "namespace_properties": {"root": str(tmp_path)}, + } + modes: list[Literal["create", "overwrite"]] = ["create", "overwrite"] + for mode in modes: + lr.write_lance( + ray.data.from_arrow(_TABLE), table_id=table_id, mode=mode, **namespace + ) + location = ( + ln.connect("dir", {"root": str(tmp_path)}) + .describe_table(DescribeTableRequest(id=table_id)) + .location + ) + assert len(lance.dataset(location).versions()) == 2 + + stats = lr.cleanup_old_versions( + table_id=table_id, older_than=timedelta(0), num_workers=2, **namespace + ) + assert stats.old_versions == 1 + assert len(lance.dataset(location).versions()) == 1 + assert lance.dataset(location).to_table().sort_by("id") == _TABLE + + +def _failing_first_attempt(marker_dir: Path) -> Callable[..., Any]: + """Wraps the sweep handler so each task's first attempt fails.""" + real = cleanup_mod._sweep_handler + path = str(marker_dir) + + def create(*args: Any) -> Any: + run = real(*args) + + def handler(task: Any) -> dict[str, Any]: + marker = os.path.join(path, str(task.index)) + if not os.path.exists(marker): + open(marker, "w").close() + return {"status": "error", "error": "injected"} + result: dict[str, Any] = run(task) + return result + + return handler + + return create + + +@requires_distributed_cleanup +def test_failed_task_raises(uri: str, tmp_path: Path) -> None: + _, remaining = _expected(uri) + with ( + patch.object(cleanup_mod, "_sweep_handler", _failing_first_attempt(tmp_path)), + pytest.raises(RuntimeError, match="files these tasks left"), + ): + lr.cleanup_old_versions( + uri, older_than=timedelta(0), num_workers=2, max_task_retries=0 + ) + # The old versions went with the plan; the dataset is intact. + assert len(lance.dataset(uri).versions()) == 1 + assert lance.dataset(uri).to_table() == _TABLE + # No version references what the failed tasks left, so only an unverified + # cleanup removes it before it is seven days old. + stats = lr.cleanup_old_versions( + uri, older_than=timedelta(0), delete_unverified=True + ) + assert stats.old_versions == 0 + assert _files(uri) == remaining + + +class _FakePlan: + """Stands in for pylance's ``CleanupPlan`` and records how it is finished.""" + + def __init__(self, marker_dir: Path, fail_attempts: dict[int, int], count: int): + self.read_version = 3 + self.num_old_versions = 2 + self.max_tasks: list[int] = [] + self.finished: list[list[Any]] = [] + self.explained: list[list[Any]] = [] + self._tasks = [ + FakeSweepTask(index, str(marker_dir), fail_attempts.get(index, 0)) + for index in range(count) + ] + + def tasks(self, max_tasks: int) -> list[FakeSweepTask]: + self.max_tasks.append(max_tasks) + return self._tasks + + def finish(self, dataset: Any, results: list[Any]) -> FakeCleanupStats: + self.finished.append(results) + return FakeCleanupStats( + data_files_removed=sum(r.stats.data_files_removed for r in results) + ) + + def finish_explain(self, dataset: Any, results: list[Any]) -> str: + self.explained.append(results) + return "explanation" + + +@pytest.fixture +def fake_plan(uri: str, tmp_path: Path) -> Iterator[Callable[..., _FakePlan]]: + """Installs a fake ``plan_cleanup_old_versions`` with the real signature, so the + driver and the Ray workers run against any pylance release.""" + marker_dir = tmp_path / "markers" + marker_dir.mkdir() + patches: list[Any] = [] + + def install( + count: int = 5, fail_attempts: Optional[dict[int, int]] = None + ) -> _FakePlan: + plan = _FakePlan(marker_dir, fail_attempts or {}, count) + + def plan_cleanup_old_versions( + self: LanceDataset, + older_than: Optional[timedelta] = None, + retain_versions: Optional[int] = None, + *, + delete_unverified: bool = False, + error_if_tagged_old_versions: bool = True, + delete_rate_limit: Optional[int] = None, + versions: Optional[list[int]] = None, + delete_concurrency: Optional[int] = None, + explain: bool = False, + max_files: int = 1000, + ) -> _FakePlan: + return plan + + patcher = patch.object( + LanceDataset, + "plan_cleanup_old_versions", + plan_cleanup_old_versions, + create=True, + ) + patcher.start() + patches.append(patcher) + return plan + + yield install + for patcher in reversed(patches): + patcher.stop() + + +@pytest.mark.parametrize("explain", [False, True]) +def test_dispatches_every_task_and_finishes( + uri: str, fake_plan: Callable[..., _FakePlan], explain: bool +) -> None: + # Task 0 fails once, so it is retried after the others are dispatched. + plan = fake_plan(count=5, fail_attempts={0: 1}) + run = lr.explain_cleanup_old_versions if explain else lr.cleanup_old_versions + kwargs: dict[str, Any] = {} if explain else {"delete_concurrency": 3} + with patch.object(cleanup_mod, "logger") as log: + result = run( + uri, older_than=timedelta(0), num_workers=2, max_task_retries=1, **kwargs + ) + + assert plan.max_tasks == [2 * cleanup_mod._TASKS_PER_WORKER] + calls, other = ( + (plan.explained, plan.finished) if explain else (plan.finished, plan.explained) + ) + assert other == [] + assert len(calls) == 1 + # One result per task, in task order. + assert [r.index for r in calls[0]] == list(range(5)) + if explain: + assert result == "explanation" + else: + assert result.data_files_removed == 5 + assert {r.delete_concurrency for r in calls[0]} == {3} + # The driver logs progress once per task, with running totals. + progress = [ + call.args + for call in log.info.call_args_list + if call.args[0].startswith("Cleanup task") + ] + assert [args[1] for args in progress] == [1, 2, 3, 4, 5] + assert progress[-1][6] == 5 + + +def test_passes_max_tasks_to_the_plan( + uri: str, fake_plan: Callable[..., _FakePlan] +) -> None: + plan = fake_plan(count=3) + lr.cleanup_old_versions(uri, older_than=timedelta(0), num_workers=2, max_tasks=3) + assert plan.max_tasks == [3] + + +@pytest.mark.parametrize("explain", [False, True]) +def test_failure_after_retries_does_not_finish( + uri: str, fake_plan: Callable[..., _FakePlan], explain: bool +) -> None: + plan = fake_plan(count=3, fail_attempts={1: 2}) + run = lr.explain_cleanup_old_versions if explain else lr.cleanup_old_versions + message = "Explaining cleanup failed" if explain else "files these tasks left" + with pytest.raises(RuntimeError, match=message): + run(uri, older_than=timedelta(0), num_workers=2, max_task_retries=1) + assert plan.finished == [] + assert plan.explained == [] + + +@pytest.mark.parametrize( + ("rate", "num_workers", "num_tasks", "expected"), + [(None, 2, 5, None), (100, 4, 5, 25), (100, 8, 2, 50), (3, 4, 5, 1)], +) +def test_rate_limit_is_split_between_workers( + uri: str, + fake_plan: Callable[..., _FakePlan], + rate: Optional[int], + num_workers: int, + num_tasks: int, + expected: Optional[int], +) -> None: + plan = fake_plan(count=num_tasks) + lr.cleanup_old_versions( + uri, older_than=timedelta(0), num_workers=num_workers, delete_rate_limit=rate + ) + assert [r.delete_rate_limit for r in plan.finished[0]] == [expected] * num_tasks + + +@pytest.mark.parametrize("explain", [False, True]) +def test_falls_back_without_distributed_cleanup(uri: str, explain: bool) -> None: + expected, remaining = _expected(uri) + with patch.object(cleanup_mod, "_supports_distributed_cleanup", return_value=False): + if explain: + result = lr.explain_cleanup_old_versions(uri, older_than=timedelta(0)) + assert result.stats.bytes_removed == expected.stats.bytes_removed + return + stats = lr.cleanup_old_versions(uri, older_than=timedelta(0)) + assert stats.old_versions == expected.stats.old_versions + assert _files(uri) == remaining + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"num_workers": 0}, "num_workers"), + ({"max_tasks": 0}, "max_tasks"), + ({"max_task_retries": -1}, "max_task_retries"), + ], +) +def test_rejects_invalid_arguments( + tmp_path: Path, kwargs: dict[str, Any], message: str +) -> None: + with pytest.raises(ValueError, match=message): + lr.cleanup_old_versions(str(tmp_path / "missing"), **kwargs) diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 7a2fb687..d5bc3944 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -61,6 +61,9 @@ def _load_index_module_with_stubs() -> ModuleType: except ImportError: # pragma: no cover - environment dependent index_mod = _load_index_module_with_stubs() +# Loaded with the index module, from the same (possibly stubbed) package. +from lance_ray import _dispatch as dispatch_mod # noqa: E402 + class _FakeField: def __init__(self, name: str, field_type: Any = None) -> None: @@ -147,9 +150,9 @@ def test_map_async_with_pool_dispatches_to_idle_workers() -> None: submission would put both slow batches on the same worker, one after the other; idle dispatch runs them concurrently on different workers. """ - results = index_mod._map_async_with_pool( - create_fragment_handler=lambda: timed_sleep_handler, - fragment_batches=[[200], [5], [200], [5]], + results = dispatch_mod._map_async_with_pool( + create_handler=lambda: timed_sleep_handler, + items=[[200], [5], [200], [5]], num_workers=2, ray_remote_args={"num_cpus": 1}, error_prefix="failed", @@ -161,6 +164,7 @@ def test_map_async_with_pool_dispatches_to_idle_workers() -> None: def test_map_async_with_pool_retries_failed_tasks(tmp_path: Path) -> None: marker_dir = str(tmp_path) + reported: list[tuple[int, str]] = [] def handler(fragment_ids: list[int]) -> dict[str, Any]: attempt = attempt_marker(marker_dir, fragment_ids) @@ -168,48 +172,54 @@ def handler(fragment_ids: list[int]) -> dict[str, Any]: return {"status": "error", "fragment_ids": fragment_ids, "error": "boom"} return {"status": "success", "fragment_ids": fragment_ids} - results = index_mod._map_async_with_pool( - create_fragment_handler=lambda: handler, - fragment_batches=[[0], [1], [2]], + results = dispatch_mod._map_async_with_pool( + create_handler=lambda: handler, + items=[[0], [1], [2]], num_workers=2, ray_remote_args=None, error_prefix="failed", max_task_retries=1, + on_result=lambda index, result: reported.append((index, result["status"])), ) assert [r["fragment_ids"] for r in results] == [[0], [1], [2]] assert [r["status"] for r in results] == ["success"] * 3 + # Each item is reported once, with its final result. + assert sorted(reported) == [(0, "success"), (1, "success"), (2, "success")] attempts = sorted(name.split(".")[0] for name in os.listdir(marker_dir)) assert attempts == ["0", "1", "1", "2"] def test_map_async_with_pool_reports_exhausted_retries(tmp_path: Path) -> None: marker_dir = str(tmp_path) + reported: list[tuple[int, str]] = [] def handler(fragment_ids: list[int]) -> dict[str, Any]: attempt_marker(marker_dir, fragment_ids) return {"status": "error", "fragment_ids": fragment_ids, "error": "boom"} - results = index_mod._map_async_with_pool( - create_fragment_handler=lambda: handler, - fragment_batches=[[0]], + results = dispatch_mod._map_async_with_pool( + create_handler=lambda: handler, + items=[[0]], num_workers=1, ray_remote_args=None, error_prefix="failed", max_task_retries=2, + on_result=lambda index, result: reported.append((index, result["status"])), ) assert results == [{"status": "error", "fragment_ids": [0], "error": "boom"}] assert len(os.listdir(marker_dir)) == 3 + assert reported == [(0, "error")] def test_map_async_with_pool_retries_on_a_live_worker_after_worker_death( tmp_path: Path, ) -> None: """A batch whose worker dies is retried elsewhere; the dead worker is dropped.""" - results = index_mod._map_async_with_pool( - create_fragment_handler=lambda: die_once_handler(str(tmp_path / "died")), - fragment_batches=[[0], [1]], + results = dispatch_mod._map_async_with_pool( + create_handler=lambda: die_once_handler(str(tmp_path / "died")), + items=[[0], [1]], num_workers=2, ray_remote_args=None, error_prefix="failed", @@ -223,9 +233,9 @@ def test_map_async_with_pool_keeps_restartable_workers_after_death( tmp_path: Path, ) -> None: """With max_restarts, Ray restarts the dead worker and it keeps serving.""" - results = index_mod._map_async_with_pool( - create_fragment_handler=lambda: die_once_handler(str(tmp_path / "died")), - fragment_batches=[[0]], + results = dispatch_mod._map_async_with_pool( + create_handler=lambda: die_once_handler(str(tmp_path / "died")), + items=[[0]], num_workers=1, ray_remote_args={"max_restarts": 1}, error_prefix="failed", @@ -240,9 +250,9 @@ def handler(fragment_ids: list[int]) -> dict[str, Any]: os._exit(1) with pytest.raises(RuntimeError, match="failed: every worker has died"): - index_mod._map_async_with_pool( - create_fragment_handler=lambda: handler, - fragment_batches=[[0]], + dispatch_mod._map_async_with_pool( + create_handler=lambda: handler, + items=[[0]], num_workers=1, ray_remote_args=None, error_prefix="failed", @@ -288,7 +298,7 @@ def fake_put_vector_index_artifacts(ivf_centroids: Any, pq_codebook: Any) -> Any return "ivf_ref", "pq_ref" def fake_map_async_with_pool(**kwargs: Any) -> Any: - fragment_handler = kwargs["create_fragment_handler"]() + fragment_handler = kwargs["create_handler"]() return [fragment_handler([0, 1])] monkeypatch.setattr(index_mod, "_check_pylance_version", lambda: None) @@ -381,7 +391,7 @@ def fake_build_rabitq_model(*, dimension: int, num_bits: int) -> str: def fake_map_async_with_pool(**kwargs: Any) -> Any: captured["map_kwargs"] = kwargs - kwargs["create_fragment_handler"]() + kwargs["create_handler"]() return [ { "status": "success", @@ -469,7 +479,7 @@ def fake_build_rabitq_model(**kwargs: Any) -> None: def fake_map_async_with_pool(**kwargs: Any) -> Any: captured["map_kwargs"] = kwargs - kwargs["create_fragment_handler"]() + kwargs["create_handler"]() return [ { "status": "success", @@ -578,7 +588,7 @@ def fake_handle_scalar_segment_index(**kwargs: Any) -> Any: def fake_map_async_with_pool(**kwargs: Any) -> Any: captured["map_kwargs"] = kwargs - kwargs["create_fragment_handler"]() + kwargs["create_handler"]() return [ { "status": "success", @@ -659,7 +669,7 @@ def fake_put_vector_index_artifacts(ivf_centroids: Any, pq_codebook: Any) -> Any def fake_map_async_with_pool(**kwargs: Any) -> Any: captured["map_kwargs"] = kwargs - kwargs["create_fragment_handler"]() + kwargs["create_handler"]() return [ { "status": "success",