wehub-resource-sync
126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
import pickle
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Optional, Union
|
|
|
|
import ray
|
|
from ray.data.block import Block, BlockAccessor, CallableClass
|
|
|
|
if TYPE_CHECKING:
|
|
from ray._raylet import StreamingGeneratorStats
|
|
from ray.data._internal.execution.interfaces import RefBundle
|
|
from ray.data.block import BlockMetadataWithSchema
|
|
|
|
|
|
def merge_label_selector(
|
|
ray_remote_args: Dict[str, Any],
|
|
ctx_label_selector: Optional[Dict[str, str]],
|
|
) -> Dict[str, Any]:
|
|
"""Merge a DataContext-level label_selector into ``ray_remote_args``.
|
|
|
|
Operator-level keys (already in ``ray_remote_args["label_selector"]``) win on
|
|
conflict so existing node-pin selectors are preserved. Returns a new dict;
|
|
the input is not mutated. If ``ctx_label_selector`` is falsy, returns the
|
|
input unchanged.
|
|
"""
|
|
if not ctx_label_selector:
|
|
return ray_remote_args
|
|
op_selector = ray_remote_args.get("label_selector") or {}
|
|
merged = {**ctx_label_selector, **op_selector}
|
|
out = dict(ray_remote_args)
|
|
out["label_selector"] = merged
|
|
return out
|
|
|
|
|
|
def make_ref_bundles(simple_data: List[List[Any]]) -> List["RefBundle"]:
|
|
"""Create ref bundles from a list of block data.
|
|
|
|
One bundle is created for each input block.
|
|
"""
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
|
|
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
|
|
|
|
output = []
|
|
for block in simple_data:
|
|
block = pd.DataFrame({"id": block})
|
|
output.append(
|
|
RefBundle(
|
|
[
|
|
BlockEntry(
|
|
ray.put(block),
|
|
BlockAccessor.for_block(block).get_metadata(),
|
|
)
|
|
],
|
|
owns_blocks=True,
|
|
schema=pa.lib.Schema.from_pandas(block, preserve_index=False),
|
|
)
|
|
)
|
|
return output
|
|
|
|
|
|
memory_units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
|
|
|
|
|
|
def memory_string(num_bytes: float) -> str:
|
|
"""Return a human-readable memory string for the given amount of bytes."""
|
|
k = 0
|
|
while num_bytes >= 1024 and k < len(memory_units) - 1:
|
|
num_bytes /= 1024
|
|
k += 1
|
|
return f"{num_bytes:.1f}{memory_units[k]}"
|
|
|
|
|
|
def locality_string(locality_hits: int, locality_misses) -> str:
|
|
"""Return a human-readable string for object locality stats."""
|
|
if not locality_misses:
|
|
return "[all objects local]"
|
|
return f"[{locality_hits}/{locality_hits + locality_misses} objects local]"
|
|
|
|
|
|
def yield_block_with_stats(
|
|
block: Block,
|
|
build_metadata: "Callable[[Optional[float]], BlockMetadataWithSchema]",
|
|
) -> Generator[Union[Block, bytes], "StreamingGeneratorStats", None]:
|
|
"""Yield a block then its pickled metadata, per the streaming-gen protocol.
|
|
|
|
Args:
|
|
block: The block to emit.
|
|
build_metadata: Given the block serialization time in seconds (or ``None``
|
|
if Ray didn't report it), returns the block's metadata to pickle.
|
|
|
|
Yields:
|
|
Union[Block, bytes]: The block, followed by its pickled
|
|
``BlockMetadataWithSchema``.
|
|
"""
|
|
gen_stats: "StreamingGeneratorStats" = yield block
|
|
block_ser_time_s = gen_stats.object_creation_dur_s if gen_stats else None
|
|
yield pickle.dumps(build_metadata(block_ser_time_s))
|
|
|
|
|
|
def make_callable_class_single_threaded(callable_cls: CallableClass) -> CallableClass:
|
|
"""Returns a thread-safe CallableClass with the same logic as the provided
|
|
`callable_cls`.
|
|
|
|
This function allows the usage of concurrent actors by safeguarding user logic
|
|
behind a separate thread.
|
|
|
|
This allows batch slicing and formatting to occur concurrently, to overlap with the
|
|
user provided UDF.
|
|
"""
|
|
|
|
class _SingleThreadedWrapper(callable_cls):
|
|
def __init__(self, *args, **kwargs):
|
|
self.thread_pool_executor = ThreadPoolExecutor(max_workers=1)
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def __repr__(self):
|
|
return super().__repr__()
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
# ThreadPoolExecutor will reuse the same thread for every submit call.
|
|
future = self.thread_pool_executor.submit(super().__call__, *args, **kwargs)
|
|
return future.result()
|
|
|
|
return _SingleThreadedWrapper
|