Ray Logging (#1623)

* Adding Raylogger and env var

* pylint

* flake8

* black

* pydocstyle

* reverting test change

* returning Logger only when necessary

* static checks

Author: malachi-constant

awswrangler/_config.py

@@ -57,6 +57,7 @@ class _ConfigArg(NamedTuple):
     "redis_password": _ConfigArg(dtype=str, nullable=True),
     "ignore_reinit_error": _ConfigArg(dtype=bool, nullable=True),
     "include_dashboard": _ConfigArg(dtype=bool, nullable=True),
+    "log_to_driver": _ConfigArg(dtype=bool, nullable=True),
     "object_store_memory": _ConfigArg(dtype=int, nullable=True),
     "cpu_count": _ConfigArg(dtype=int, nullable=True),
     "gpu_count": _ConfigArg(dtype=int, nullable=True),

awswrangler/distributed/__init__.py

@@ -1,10 +1,17 @@
 """Distributed Module."""
 
-from awswrangler.distributed._distributed import initialize_ray, modin_repartition, ray_get, ray_remote  # noqa
+from awswrangler.distributed._distributed import (  # noqa
+    RayLogger,
+    initialize_ray,
+    modin_repartition,
+    ray_get,
+    ray_remote,
+)
 
 __all__ = [
     "initialize_ray",
     "ray_get",
+    "RayLogger",
     "ray_remote",
     "modin_repartition",
 ]

awswrangler/distributed/_distributed.py

@@ -1,12 +1,11 @@
 """Distributed Module (PRIVATE)."""
-
 import logging
 import multiprocessing
 import os
 import sys
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
 
 from awswrangler._config import apply_configs, config
 
@@ -19,6 +18,22 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+class RayLogger:
+    """Create discrete Logger instance for Ray Tasks."""
+
+    def __init__(
+        self,
+        log_level: int = logging.INFO,
+        format: str = "%(asctime)s::%(levelname)-2s::%(name)s::%(message)s",  # pylint: disable=redefined-builtin
+        datefmt: str = "%Y-%m-%d %H:%M:%S",
+    ):
+        logging.basicConfig(level=log_level, format=format, datefmt=datefmt)
+
+    def get_logger(self, name: Union[str, Any] = None) -> Union[logging.Logger, Any]:
+        """Return logger object."""
+        return logging.getLogger(name) if config.distributed else None
+
+
 def ray_get(futures: List[Any]) -> List[Any]:
     """
     Run ray.get on futures if distributed.
@@ -45,7 +60,6 @@ def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]:
     ----------
     function : Callable[..., Any]
         Callable as input to ray.remote
-
     Returns
     -------
     Callable[..., Any]
@@ -94,6 +108,7 @@ def initialize_ray(
     redis_password: Optional[str] = None,
     ignore_reinit_error: Optional[bool] = True,
     include_dashboard: Optional[bool] = False,
+    log_to_driver: Optional[bool] = True,
     object_store_memory: Optional[int] = None,
     cpu_count: Optional[int] = None,
     gpu_count: Optional[int] = 0,
@@ -111,6 +126,8 @@ def initialize_ray(
         If true, Ray suppress errors from calling ray.init() twice, by default True
     include_dashboard : Optional[bool]
         Boolean flag indicating whether or not to start the Ray dashboard, by default False
+    log_to_driver : Optional[bool]
+        Boolean flag to enable routing of all worker logs to the driver, by default True
     object_store_memory : Optional[int]
         The amount of memory (in bytes) to start the object store with, by default None
     cpu_count : Optional[int]
@@ -124,6 +141,7 @@ def initialize_ray(
                 address=address,
                 include_dashboard=include_dashboard,
                 ignore_reinit_error=ignore_reinit_error,
+                log_to_driver=log_to_driver,
             )
         else:
             if not object_store_memory:
@@ -151,6 +169,7 @@ def initialize_ray(
                 "num_gpus": gpu_count,
                 "include_dashboard": include_dashboard,
                 "ignore_reinit_error": ignore_reinit_error,
+                "log_to_driver": log_to_driver,
                 "object_store_memory": object_store_memory,
                 "_redis_password": redis_password,
                 "_memory": object_store_memory,

awswrangler/distributed/_distributed.pyi

@@ -1,13 +1,23 @@
 """Distributed Module (PRIVATE)."""
 
+import logging
 import multiprocessing
 import os
 import sys
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
 
 from awswrangler._config import apply_configs, config
 
+class RayLogger:
+    def __init__(
+        self,
+        log_level: int = logging.INFO,
+        format: str = "%(asctime)s::%(levelname)-2s::%(name)s::%(message)s",
+        datefmt: str = "%Y-%m-%d %H:%M:%S",
+    ): ...
+    def get_logger(self, name: Union[str, Any] = None) -> logging.Logger: ...
+
 def ray_get(futures: List[Any]) -> List[Any]: ...
 def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]: ...
 def modin_repartition(function: Callable[..., Any]) -> Callable[..., Any]: ...
@@ -16,6 +26,7 @@ def initialize_ray(
     redis_password: Optional[str] = None,
     ignore_reinit_error: Optional[bool] = True,
     include_dashboard: Optional[bool] = False,
+    log_to_driver: Optional[bool] = True,
     object_store_memory: Optional[int] = None,
     cpu_count: Optional[int] = None,
     gpu_count: Optional[int] = 0,

awswrangler/lakeformation/_read.py

@@ -11,7 +11,7 @@
 from awswrangler._config import apply_configs
 from awswrangler._threading import _get_executor
 from awswrangler.catalog._utils import _catalog_id, _transaction_id
-from awswrangler.distributed import ray_remote
+from awswrangler.distributed import RayLogger, ray_remote
 from awswrangler.lakeformation._utils import commit_transaction, start_transaction, wait_query
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -23,6 +23,7 @@ def _get_work_unit_results(
     query_id: str,
     token_work_unit: Tuple[str, int],
 ) -> Table:
+    RayLogger().get_logger(name=_get_work_unit_results.__name__)
     _logger.debug("Query id: %s Token work unit: %s", query_id, token_work_unit)
     client_lakeformation: boto3.client = _utils.client(service_name="lakeformation", session=boto3_session)
 

awswrangler/s3/_delete.py

@@ -9,7 +9,7 @@
 
 from awswrangler import _utils, exceptions
 from awswrangler._threading import _get_executor
-from awswrangler.distributed import ray_get, ray_remote
+from awswrangler.distributed import RayLogger, ray_get, ray_remote
 from awswrangler.s3._fs import get_botocore_valid_kwargs
 from awswrangler.s3._list import _path2list
 
@@ -33,6 +33,7 @@ def _delete_objects(
     paths: List[str],
     s3_additional_kwargs: Optional[Dict[str, Any]],
 ) -> None:
+    RayLogger().get_logger(name=_delete_objects.__name__)
     client_s3: boto3.client = _utils.client(
         service_name="s3",
         session=boto3_session,

awswrangler/s3/_read_parquet.py

@@ -17,7 +17,7 @@
 from awswrangler._threading import _get_executor
 from awswrangler.catalog._get import _get_partitions
 from awswrangler.catalog._utils import _catalog_id
-from awswrangler.distributed import ray_get, ray_remote
+from awswrangler.distributed import RayLogger, ray_get, ray_remote
 from awswrangler.s3._fs import open_s3_object
 from awswrangler.s3._list import _path2list
 from awswrangler.s3._read import (
@@ -76,6 +76,7 @@ def _read_parquet_metadata_file(
     version_id: Optional[str] = None,
     coerce_int96_timestamp_unit: Optional[str] = None,
 ) -> pa.schema:
+    RayLogger().get_logger(name=_read_parquet_metadata_file.__name__)
     with open_s3_object(
         path=path,
         mode="rb",

awswrangler/s3/_select.py

@@ -13,7 +13,7 @@
 
 from awswrangler import _data_types, _utils, exceptions
 from awswrangler._threading import _get_executor
-from awswrangler.distributed import ray_get, ray_remote
+from awswrangler.distributed import RayLogger, ray_get, ray_remote
 from awswrangler.s3._describe import size_objects
 from awswrangler.s3._list import _path2list
 from awswrangler.s3._read import _get_path_ignore_suffix
@@ -42,6 +42,7 @@ def _select_object_content(
     args: Dict[str, Any],
     scan_range: Optional[Tuple[int, int]] = None,
 ) -> Union[pa.Table, "ray.ObjectRef[pa.Table]"]:
+    RayLogger().get_logger(name=_select_object_content.__name__)
     client_s3: boto3.client = _utils.client(service_name="s3", session=boto3_session)
 
     if scan_range:
@@ -83,6 +84,7 @@ def _select_query(
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, Any]] = None,
 ) -> List[Union[pa.Table, "ray.ObjectRef[pa.Table]"]]:
+    RayLogger().get_logger(name=_select_query.__name__)
     bucket, key = _utils.parse_path(path)
 
     args: Dict[str, Any] = {

awswrangler/s3/_wait.py

@@ -9,7 +9,7 @@
 from awswrangler import _utils
 from awswrangler._config import config
 from awswrangler._threading import _get_executor
-from awswrangler.distributed import ray_get, ray_remote
+from awswrangler.distributed import RayLogger, ray_get, ray_remote
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -28,6 +28,7 @@ def _wait_object(
 def _wait_object_batch(
     boto3_session: Optional[boto3.Session], paths: List[str], waiter_name: str, delay: int, max_attempts: int
 ) -> None:
+    RayLogger().get_logger(name=_wait_object_batch.__name__)
     for path in paths:
         _wait_object(boto3_session, path, waiter_name, delay, max_attempts)