refactor: reuse the global Ray Pool across distributed operations

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"]

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
 

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

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,14 +103,15 @@ 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."""
     with _GLOBAL_POOL_LOCK:
         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

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 = []
 

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,