(enhance) configure scheduling options, remove dependencies on internal ray impl (#1734)

* Update ray remote decorator to pass options

* Remove dependency on ray._internal methods

* Add scheduling parameters

* remove pack scheduling; typing/linting fixes

* linting/comments

* fix write_block remote fn call

* remove default scheduling strategy

* Remove _open_input_source overload as it's not used by ParquetReader

* clarify comments

Author: kukushking

awswrangler/distributed/ray/_core.py

@@ -51,27 +51,46 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
     return wrapper
 
 
-def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]:
+def ray_remote(**options: Any) -> Callable[..., Any]:
     """
-    Decorate callable to wrap within ray.remote.
+    Decorate with @ray.remote providing .options().
 
     Parameters
     ----------
-    function : Callable[..., Any]
-        Callable as input to ray.remote.
+    options : Any
+        Ray remote options
 
     Returns
     -------
     Callable[..., Any]
     """
-    # Access the source function if it exists
-    function = getattr(function, "_source_func", function)
 
-    @wraps(function)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
-        return ray.remote(ray_logger(function)).remote(*args, **kwargs)  # type: ignore
+    def remote_decorator(function: Callable[..., Any]) -> Callable[..., Any]:
+        """
+        Decorate callable to wrap within ray.remote.
 
-    return wrapper
+        Parameters
+        ----------
+        function : Callable[..., Any]
+            Callable as input to ray.remote.
+
+        Returns
+        -------
+        Callable[..., Any]
+        """
+        # Access the source function if it exists
+        function = getattr(function, "_source_func", function)
+
+        @wraps(function)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            remote_fn = ray.remote(ray_logger(function))
+            if options:
+                remote_fn = remote_fn.options(**options)
+            return remote_fn.remote(*args, **kwargs)  # type: ignore
+
+        return wrapper
+
+    return remote_decorator
 
 
 def ray_get(futures: Union["ray.ObjectRef[Any]", List["ray.ObjectRef[Any]"]]) -> Any:

awswrangler/distributed/ray/_core.pyi

@@ -13,7 +13,7 @@ class RayLogger:
     def get_logger(self, name: Union[str, Any] = None) -> logging.Logger: ...
 
 def ray_logger(function: Callable[..., Any]) -> Callable[..., Any]: ...
-def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]: ...
+def ray_remote(**options: Any) -> Callable[..., Any]: ...
 def ray_get(futures: List[Any]) -> Any: ...
 def initialize_ray(
     address: Optional[str] = None,

awswrangler/distributed/ray/_register.py

@@ -27,10 +27,15 @@ 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 func in [
         _write_batch,
         _write_df,
     ]:
-        engine.register_func(func, ray_remote(func))
+        engine.register_func(func, ray_remote()(func))
 
     if memory_format.get() == MemoryFormatEnum.MODIN:
         from awswrangler.distributed.ray.modin._data_types import pyarrow_types_from_pandas_distributed