From 01166423a08db069228483e03ef35a7293637673 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 22 Jun 2026 20:09:44 +0800 Subject: [PATCH] refactor: reuse the global Ray Pool across distributed operations compact_files, create_scalar_index, create_index and add_columns each created a fresh Ray Pool, ignoring a Pool configured via init_global_pool/ set_global_pool. Route them through get_or_create_pool so a configured global Pool is reused. Warn when an active global Pool causes per-call ray_remote_args to be ignored, and tolerate processes=None. Update test_map_async_with_pool_closes_and_joins_pool to patch the Pool in lance_ray.pool, since the Pool is now constructed there rather than in lance_ray.index. --- lance_ray/compaction.py | 28 ++++++++------------ lance_ray/index.py | 25 +++++++++--------- lance_ray/io.py | 41 +++++++++++++++--------------- lance_ray/pool.py | 17 ++++++++++--- tests/test_pool.py | 38 +++++++++++++++++++++++++++ tests/test_vector_index_options.py | 7 ++++- 6 files changed, 101 insertions(+), 55 deletions(-) diff --git a/lance_ray/compaction.py b/lance_ray/compaction.py index b24e16dd..d4e96be3 100644 --- a/lance_ray/compaction.py +++ b/lance_ray/compaction.py @@ -4,8 +4,8 @@ import lance from lance.lance import CompactionMetrics from lance.optimize import Compaction, CompactionOptions, CompactionTask -from ray.util.multiprocessing import Pool +from .pool import get_or_create_pool from .utils import ( get_namespace_kwargs, get_or_create_namespace, @@ -151,10 +151,8 @@ def compact_files( num_workers = compaction_plan.num_tasks() logger.info(f"Adjusted num_workers to {num_workers} to match task count") - # Step 2: Execute tasks in parallel using Ray Pool - pool = Pool(processes=num_workers, ray_remote_args=ray_remote_args) - - # Create the task handler function + # Step 2: Execute tasks in parallel using a Ray Pool, reusing the global + # Pool when one is configured via init_global_pool/set_global_pool. task_handler = _handle_compaction_task( dataset_uri=uri, storage_options=merged_storage_options, @@ -163,21 +161,17 @@ def compact_files( table_id=table_id, ) - # Submit tasks using Pool.map_async - rst_futures = pool.map_async( - task_handler, - compaction_plan.tasks, - chunksize=1, - ) - - # Wait for results try: - results = rst_futures.get() + with get_or_create_pool( + processes=num_workers, ray_remote_args=ray_remote_args + ) as pool: + results = pool.map_async( + task_handler, + compaction_plan.tasks, + chunksize=1, + ).get() except Exception as e: raise RuntimeError(f"Failed to complete distributed compaction: {e}") from e - finally: - pool.close() - pool.join() # Check for failures failed_results = [r for r in results if r["status"] == "error"] diff --git a/lance_ray/index.py b/lance_ray/index.py index 45acc2f2..3996604c 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -13,9 +13,9 @@ from lance.dataset import Index, IndexConfig, LanceDataset from lance.indices import IndicesBuilder from packaging import version -from ray.util.multiprocessing import Pool from .field_path import resolve_arrow_field_path, resolve_dataset_field_path +from .pool import get_or_create_pool from .utils import ( get_namespace_kwargs, has_namespace_params, @@ -140,22 +140,21 @@ def _map_async_with_pool( This helper encapsulates the common Pool.map_async + get + error wrapping logic so that both scalar and vector distributed index builders can share - the same implementation. + the same implementation. The global Pool is reused when one is configured + via init_global_pool/set_global_pool. """ - pool = Pool(processes=num_workers, ray_remote_args=ray_remote_args) try: - fragment_handler = create_fragment_handler() - rst_futures = pool.map_async( - fragment_handler, - fragment_batches, - chunksize=1, - ) - results = rst_futures.get() + with get_or_create_pool( + processes=num_workers, ray_remote_args=ray_remote_args + ) as pool: + fragment_handler = create_fragment_handler() + results = pool.map_async( + fragment_handler, + fragment_batches, + chunksize=1, + ).get() except Exception as exc: # pragma: no cover - exercised via integration tests raise RuntimeError(f"{error_prefix}: {exc}") from exc - finally: - pool.close() - pool.join() return results diff --git a/lance_ray/io.py b/lance_ray/io.py index 270ab287..a590a5b0 100644 --- a/lance_ray/io.py +++ b/lance_ray/io.py @@ -12,11 +12,11 @@ from lance.dataset import LanceDataset, LanceOperation from lance.udf import BatchUDF from ray.data import Dataset, read_datasource -from ray.util.multiprocessing import Pool from .datasink import LanceDatasink from .datasource import LanceDatasource from .fragment import prepare_fragment_write_options +from .pool import get_or_create_pool from .utils import ( get_namespace_kwargs, has_namespace_params, @@ -576,30 +576,29 @@ def add_columns( **namespace_kwargs, ) fragment_ids = [f.metadata.id for f in lance_ds.get_fragments()] - pool = Pool(processes=concurrency, ray_remote_args=ray_remote_args) - rst_futures = pool.map_async( - _handle_fragment( - uri, - transform, - read_columns, - batch_size, - reader_schema, - read_version, - storage_options, - namespace_impl, - namespace_properties, - table_id, - ), - fragment_ids, - chunksize=1, + fragment_handler = _handle_fragment( + uri, + transform, + read_columns, + batch_size, + reader_schema, + read_version, + storage_options, + namespace_impl, + namespace_properties, + table_id, ) try: - result = rst_futures.get() + with get_or_create_pool( + processes=concurrency, ray_remote_args=ray_remote_args + ) as pool: + result = pool.map_async( + fragment_handler, + fragment_ids, + chunksize=1, + ).get() except Exception as exc: raise RuntimeError(f"Failed to add columns: {exc}") from exc - finally: - pool.close() - pool.join() commit_messages = [] new_schema = None diff --git a/lance_ray/pool.py b/lance_ray/pool.py index 6dc374eb..43fad505 100644 --- a/lance_ray/pool.py +++ b/lance_ray/pool.py @@ -24,8 +24,8 @@ def _read_pool_processes(pool: Any) -> int | None: return None -def _warn_if_process_count_differs(requested_processes: int) -> None: - if _GLOBAL_POOL_PROCESSES is None: +def _warn_if_process_count_differs(requested_processes: Optional[int]) -> None: + if _GLOBAL_POOL_PROCESSES is None or requested_processes is None: return if requested_processes == _GLOBAL_POOL_PROCESSES: return @@ -38,6 +38,16 @@ def _warn_if_process_count_differs(requested_processes: int) -> None: ) +def _warn_if_remote_args_ignored(ray_remote_args: Optional[dict[str, Any]]) -> None: + if ray_remote_args: + logger.warning( + "Reusing global Ray Pool; per-call ray_remote_args %s are ignored " + "while the global Pool is active. Clear the global Pool to apply " + "per-call remote args.", + ray_remote_args, + ) + + def set_global_pool(pool: Any | None) -> None: """Set a process-wide Ray Pool for Lance-Ray operations to reuse. @@ -93,7 +103,7 @@ def clear_global_pool(*, close: bool = False, join: bool = True) -> None: @contextmanager def get_or_create_pool( *, - processes: int, + processes: Optional[int], ray_remote_args: Optional[dict[str, Any]], ) -> Iterator[Any]: """Yield the global Pool if present, otherwise a local close-and-join Pool.""" @@ -101,6 +111,7 @@ def get_or_create_pool( pool = _GLOBAL_POOL if pool is not None: _warn_if_process_count_differs(processes) + _warn_if_remote_args_ignored(ray_remote_args) if pool is not None: yield pool diff --git a/tests/test_pool.py b/tests/test_pool.py index febfe53d..d657a941 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -77,6 +77,44 @@ class FakePool: assert "requested 16 workers will be ignored" in caplog.text +def test_get_or_create_pool_warns_when_remote_args_ignored(caplog): + class FakePool: + processes = 4 + + pool_mod.set_global_pool(FakePool()) + try: + with ( + caplog.at_level(logging.WARNING, logger="lance_ray.pool"), + pool_mod.get_or_create_pool( + processes=4, ray_remote_args={"num_cpus": 2} + ) as pool, + ): + assert pool is pool_mod.get_global_pool() + finally: + pool_mod.clear_global_pool() + + assert "per-call ray_remote_args" in caplog.text + assert "are ignored" in caplog.text + + +def test_get_or_create_pool_handles_none_processes_with_global_pool(caplog): + class FakePool: + processes = 4 + + pool_mod.set_global_pool(FakePool()) + try: + with ( + caplog.at_level(logging.WARNING, logger="lance_ray.pool"), + pool_mod.get_or_create_pool(processes=None, ray_remote_args=None) as pool, + ): + assert pool is pool_mod.get_global_pool() + finally: + pool_mod.clear_global_pool() + + # processes=None must not raise and must not emit a size-mismatch warning. + assert "workers will be ignored" not in caplog.text + + def test_set_global_pool_can_clear_without_closing(): events = [] diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 70bf648f..d7e60c13 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -159,7 +159,12 @@ def create_fragment_handler(): events.append("create_handler") return lambda fragment_ids: {"status": "success", "fragment_ids": fragment_ids} - monkeypatch.setattr(index_mod, "Pool", FakePool) + # _map_async_with_pool now goes through get_or_create_pool, which constructs + # and closes/joins the Pool in lance_ray.pool. Patch it there and make sure + # no global Pool is configured so the local close-and-join path is taken. + pool_mod = sys.modules[index_mod.get_or_create_pool.__module__] + monkeypatch.setattr(pool_mod, "_GLOBAL_POOL", None) + monkeypatch.setattr(pool_mod, "Pool", FakePool) assert index_mod._map_async_with_pool( create_fragment_handler=create_fragment_handler,