Add scheduling parameters

Author: kukushking

awswrangler/distributed/ray/_register.py

@@ -25,10 +25,16 @@ def register_ray() -> None:
         _select_query,
         _select_object_content,
         _wait_object_batch,
+    ]:
+        # Schedule for maximum concurrency
+        engine.register_func(func, ray_remote(scheduling_strategy="SPREAD")(func))
+
+    for pack_func in [
         _write_batch,
         _write_df,
     ]:
-        engine.register_func(func, ray_remote()(func))
+        # Schedule for data locality
+        engine.register_func(pack_func, ray_remote(scheduling_strategy="PACK")(pack_func))
 
     if memory_format.get() == MemoryFormatEnum.MODIN:
         from awswrangler.distributed.ray.modin._data_types import pyarrow_types_from_pandas_distributed

awswrangler/distributed/ray/modin/_data_types.py

@@ -13,7 +13,7 @@ def pyarrow_types_from_pandas_distributed(
     df: pd.DataFrame, index: bool, ignore_cols: Optional[List[str]] = None, index_left: bool = False
 ) -> Dict[str, pa.DataType]:
     """Extract the related Pyarrow data types from a pandas DataFrame."""
-    func = ray_remote()(pyarrow_types_from_pandas)
+    func = ray_remote(scheduling_strategy="PACK")(pyarrow_types_from_pandas)
     first_block_object_ref = ray.data.from_modin(df).get_internal_block_refs()[0]
     return ray_get(  # type: ignore
         func(

awswrangler/distributed/ray/modin/_utils.py

@@ -14,7 +14,7 @@
 from awswrangler._arrow import _table_to_df
 
 
-@ray_remote()
+@ray_remote(scheduling_strategy="PACK")
 def _block_to_df(
     block: Any,
     to_pandas_kwargs: Dict[str, Any],

awswrangler/distributed/ray/modin/s3/_write_dataset.py

@@ -140,7 +140,7 @@ def _to_partitions_distributed(  # pylint: disable=unused-argument
         _to_partitions_func = engine.dispatch_func(_to_partitions, PandasDataFrame)
         func = engine.dispatch_func(func, PandasDataFrame)
 
-        @ray_remote()
+        @ray_remote(scheduling_strategy="PACK")
         def write_partitions(df: pd.DataFrame, block_index: int) -> Tuple[List[str], Dict[str, List[str]]]:
             paths, partitions_values = _to_partitions_func(
                 # Passing a copy of the data frame because data in ray object store is immutable