deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/utils/ray_utils/concurrency.py

@@ -88,7 +88,7 @@ def round_robin_options_provider(
     **kwargs,
 ) -> Dict[str, Any]:
     """Returns a resource dictionary that can be included with ray remote
-    options to round robin indexed tasks or actors across a list of resource
+    options to round-robin indexed tasks or actors across a list of resource
     keys. For example, the following code round-robins 100 tasks across all
     live cluster nodes:
     ```

deltacat/utils/ray_utils/dataset.py

@@ -2,7 +2,11 @@ import logging
 from typing import Callable, Dict, List, Optional, Union
 
 from fsspec import AbstractFileSystem
+
+import pyarrow as pa
 from pyarrow import csv as pacsv
+import pyarrow.fs as pafs
+
 from ray.data import Dataset
 from ray.data.datasource import FilenameProvider
 
@@ -16,7 +20,7 @@ def write_parquet(
     dataset: Dataset,
     base_path: str,
     *,
-    filesystem: AbstractFileSystem,
+    filesystem: Optional[Union[AbstractFileSystem, pafs.FileSystem]],
     block_path_provider: Union[Callable, FilenameProvider],
     **kwargs,
 ) -> None:
@@ -34,16 +38,36 @@ def write_csv(
     dataset: Dataset,
     base_path: str,
     *,
-    filesystem: AbstractFileSystem,
+    filesystem: Optional[Union[AbstractFileSystem, pafs.FileSystem]],
     block_path_provider: Union[Callable, FilenameProvider],
     **kwargs,
 ) -> None:
+    """
+    Write a Ray Dataset to a CSV file (or other delimited text format).
+    """
+    # Extract CSV-specific options from kwargs
+    delimiter = kwargs.pop("delimiter", ",")
+    quoting_style = kwargs.pop("quoting_style", None)
+    include_header = kwargs.pop("include_header", False)
+
+    # Create a function that will generate WriteOptions inside the worker process
+    def arrow_csv_args_fn():
+        write_options = pacsv.WriteOptions(
+            delimiter=delimiter,
+            include_header=include_header,
+            quoting_style=quoting_style,
+        )
+        return {"write_options": write_options}
+
+    # Check if the block_path_provider will generate .gz files to avoid double compression
+    pa_open_stream_args = {}
+    if not (
+        hasattr(block_path_provider, "content_encoding")
+        and block_path_provider.content_encoding == ContentEncoding.GZIP
+    ):
+        # Block path provider will not generate .gz files, so we need to apply explicit compression
+        pa_open_stream_args["compression"] = ContentEncoding.GZIP.value
 
-    # column names are kept in table metadata, so omit header
-    arrow_csv_args_fn = lambda: {
-        "write_options": pacsv.WriteOptions(include_header=False)
-    }
-    pa_open_stream_args = {"compression": ContentEncoding.GZIP.value}
     dataset.write_csv(
         base_path,
         arrow_open_stream_args=pa_open_stream_args,
@@ -55,12 +79,76 @@ def write_csv(
     )
 
 
+def write_json(
+    dataset: Dataset,
+    base_path: str,
+    *,
+    filesystem: Optional[Union[AbstractFileSystem, pafs.FileSystem]],
+    block_path_provider: Union[Callable, FilenameProvider],
+    **kwargs,
+) -> None:
+    """
+    Write a Ray Dataset to a JSON file using Ray's native JSON writer.
+    """
+    # Check if the block_path_provider will generate .gz files to avoid double compression
+    pa_open_stream_args = {}
+    if not (
+        hasattr(block_path_provider, "content_encoding")
+        and block_path_provider.content_encoding == ContentEncoding.GZIP
+    ):
+        # Block path provider will not generate .gz files, so we need to apply explicit compression
+        pa_open_stream_args["compression"] = ContentEncoding.GZIP.value
+
+    dataset.write_json(
+        base_path,
+        arrow_open_stream_args=pa_open_stream_args,
+        filesystem=filesystem,
+        try_create_dir=False,
+        filename_provider=block_path_provider,
+        **kwargs,
+    )
+
+
 CONTENT_TYPE_TO_DATASET_WRITE_FUNC: Dict[str, Callable] = {
+    ContentType.UNESCAPED_TSV.value: write_csv,
+    ContentType.TSV.value: write_csv,
     ContentType.CSV.value: write_csv,
+    ContentType.PSV.value: write_csv,
     ContentType.PARQUET.value: write_parquet,
+    ContentType.JSON.value: write_json,
 }
 
 
+def content_type_to_writer_kwargs(content_type: str) -> Dict[str, any]:
+    """
+    Returns writer kwargs for the given content type when writing with Ray Dataset.
+    """
+    if content_type == ContentType.UNESCAPED_TSV.value:
+        return {
+            "delimiter": "\t",
+            "include_header": False,
+            "quoting_style": "none",
+        }
+    if content_type == ContentType.TSV.value:
+        return {
+            "delimiter": "\t",
+            "include_header": False,
+        }
+    if content_type == ContentType.CSV.value:
+        return {
+            "delimiter": ",",
+            "include_header": False,
+        }
+    if content_type == ContentType.PSV.value:
+        return {
+            "delimiter": "|",
+            "include_header": False,
+        }
+    if content_type in {ContentType.PARQUET.value, ContentType.JSON.value}:
+        return {}
+    raise ValueError(f"Unsupported content type: {content_type}")
+
+
 def slice_dataset(dataset: Dataset, max_len: Optional[int]) -> List[Dataset]:
     """
     Returns equally-sized dataset slices of up to `max_len` records each.
@@ -88,9 +176,10 @@ def dataset_size(dataset: Dataset) -> int:
 def dataset_to_file(
     table: Dataset,
     base_path: str,
-    file_system: AbstractFileSystem,
+    filesystem: Optional[Union[AbstractFileSystem, pafs.FileSystem]],
     block_path_provider: Union[Callable, FilenameProvider],
     content_type: str = ContentType.PARQUET.value,
+    schema: Optional[pa.Schema] = None,
     **kwargs,
 ) -> None:
     """
@@ -103,10 +192,12 @@ def dataset_to_file(
             f" implemented. Known content types: "
             f"{CONTENT_TYPE_TO_DATASET_WRITE_FUNC.keys}"
         )
+    writer_kwargs = content_type_to_writer_kwargs(content_type)
+    writer_kwargs.update(kwargs)
     writer(
         table,
         base_path,
-        filesystem=file_system,
+        filesystem=filesystem,
         block_path_provider=block_path_provider,
-        **kwargs,
+        **writer_kwargs,
     )

deltacat/utils/ray_utils/runtime.py

@@ -21,7 +21,7 @@ def node_resource_keys(
     keys = []
     node_dict = ray.nodes()
     if node_dict:
-        for node in ray.nodes():
+        for node in node_dict:
             if filter_fn(node):
                 for key in node["Resources"].keys():
                     if key.startswith("node:"):
@@ -37,7 +37,7 @@ def current_node_resource_key() -> str:
     actors on that node via:
     `foo.options(resources={get_current_node_resource_key(): 0.01}).remote()`
     """
-    current_node_id = ray.get_runtime_context().get_node_id().hex()
+    current_node_id = ray.get_runtime_context().get_node_id()
     keys = node_resource_keys(lambda n: n["NodeID"] == current_node_id)
     assert (
         len(keys) <= 1
@@ -45,6 +45,47 @@ def current_node_resource_key() -> str:
     return keys[0] if len(keys) == 1 else None
 
 
+def current_node_resources() -> Dict[str, float]:
+    """Get's Ray's resources for the current node as a dictionary.
+
+    Example Return Value:
+        >>> {
+        >>>    'memory': 17611605607.0,
+        >>>    'node:127.0.0.1': 1.0,
+        >>>    'node:__internal_head__': 1.0,
+        >>>    'object_store_memory': 2147483648.0,
+        >>>    'CPU': 10.0,
+        >>> }
+    """
+    current_node_id = ray.get_runtime_context().get_node_id()
+    node_dict = ray.nodes()
+    if node_dict:
+        for node in node_dict:
+            if node["NodeID"] == current_node_id:
+                return node["Resources"]
+    else:
+        raise ValueError("No node dictionary found on current node.")
+    return {}
+
+
+def find_max_single_node_resource_type(resource_type: str) -> float:
+    """Finds the max resource amount available on any single cluster node
+    for the given resource type. Returns the max resource amount as a float."""
+    node_dict = ray.nodes()
+    max_single_node_resource_amount = 0
+    if node_dict:
+        for node in node_dict:
+            node_resource_amount = node["Resources"].get(resource_type)
+            if node_resource_amount is not None:
+                max_single_node_resource_amount = max(
+                    max_single_node_resource_amount,
+                    node_resource_amount,
+                )
+    else:
+        raise ValueError("No node dictionary found on current node.")
+    return max_single_node_resource_amount
+
+
 def is_node_alive(node: Dict[str, Any]) -> bool:
     """Takes a node from `ray.nodes()` as input. Returns True if the node is
     alive, and False otherwise."""
@@ -67,6 +108,17 @@ def live_node_waiter(min_live_nodes: int, poll_interval_seconds: float = 0.5) ->
         time.sleep(poll_interval_seconds)
 
 
+def live_cpu_waiter(min_live_cpus: int, poll_interval_seconds: float = 0.5) -> None:
+    """Waits until the given minimum number of live CPUs are present in the
+    cluster. Checks the current number of live CPUs every
+    `poll_interval_seconds`."""
+    live_cpus = cluster_cpus()
+    while live_cpus < min_live_cpus:
+        live_cpus = cluster_cpus()
+        logger.info(f"Waiting for Live CPUs: {live_cpus}/{min_live_cpus}")
+        time.sleep(poll_interval_seconds)
+
+
 def live_node_resource_keys() -> List[str]:
     """Get Ray resource keys for all live cluster nodes as a list of strings of
     the form: "node:{node_resource_name}". The returned keys can be used to
@@ -83,7 +135,7 @@ def other_live_node_resource_keys() -> List[str]:
 
     For example, invoking this function from your Ray application driver on the
     head node returns the resource keys of all live worker nodes."""
-    current_node_id = ray.get_runtime_context().get_node_id().hex()
+    current_node_id = ray.get_runtime_context().get_node_id()
     return node_resource_keys(
         lambda n: n["NodeID"] != current_node_id and is_node_alive(n)
     )
@@ -97,7 +149,7 @@ def other_node_resource_keys() -> List[str]:
 
     For example, invoking this function from your Ray application driver on the
     head node returns the resource keys of all worker nodes."""
-    current_node_id = ray.get_runtime_context().get_node_id().hex()
+    current_node_id = ray.get_runtime_context().get_node_id()
     return node_resource_keys(lambda n: n["NodeID"] != current_node_id)