wehub-resource-sync
312 lines
12 KiB
Python
312 lines
12 KiB
Python
import logging
|
|
from typing import Iterator, List, Optional
|
|
|
|
import pyarrow as pa
|
|
|
|
from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
|
|
from ray.data._internal.execution.interfaces import PhysicalOperator
|
|
from ray.data._internal.execution.operators.actor_pool_map_operator import (
|
|
ActorPoolMapOperator,
|
|
)
|
|
from ray.data._internal.execution.operators.map_operator import MapOperator
|
|
from ray.data._internal.execution.operators.map_transformer import (
|
|
BlockMapTransformFn,
|
|
MapTransformer,
|
|
)
|
|
from ray.data._internal.logical.operators import Download
|
|
from ray.data._internal.output_buffer import OutputBlockSizeOption
|
|
from ray.data._internal.planner._obstore_download import (
|
|
OBSTORE_AVAILABLE,
|
|
_log_fallback_warning,
|
|
_plan_obstore_routing,
|
|
download_bytes_async,
|
|
)
|
|
from ray.data._internal.planner.download_partition_actor import (
|
|
URI_DOWNLOAD_MAX_WORKERS,
|
|
AsyncPartitionActor,
|
|
PartitionActor,
|
|
)
|
|
from ray.data._internal.util import (
|
|
RetryingPyFileSystem,
|
|
_iter_arrow_table_for_target_max_block_size,
|
|
make_async_gen,
|
|
)
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.context import DataContext
|
|
from ray.data.datasource.path_util import (
|
|
_resolve_paths_and_filesystem,
|
|
_validate_and_wrap_filesystem,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def plan_download_op(
|
|
op: Download,
|
|
physical_children: List[PhysicalOperator],
|
|
data_context: DataContext,
|
|
) -> MapOperator:
|
|
"""Plan the download operation with partitioning and downloading stages."""
|
|
assert len(physical_children) == 1
|
|
input_physical_dag = physical_children[0]
|
|
|
|
upstream_op_is_download = False
|
|
if len(input_physical_dag._logical_operators) == 1 and isinstance(
|
|
input_physical_dag._logical_operators[0], Download
|
|
):
|
|
upstream_op_is_download = True
|
|
|
|
uri_column_names = op.uri_column_names
|
|
uri_column_names_str = ", ".join(uri_column_names)
|
|
output_bytes_column_names = op.output_bytes_column_names
|
|
ray_remote_args = op.ray_remote_args
|
|
filesystem = op.filesystem
|
|
|
|
# Import _get_udf from the main planner file
|
|
from ray.data._internal.planner.plan_udf_map_op import (
|
|
_generate_transform_fn_for_map_batches,
|
|
_get_udf,
|
|
)
|
|
|
|
# If we have multiple download operators in a row, we should only include the partition actor
|
|
# at the start of the chain. This is primarily done to prevent partition actors from bottlenecking
|
|
# the chain becuase the interleaved operators would be a single actor. As a result, the
|
|
# URIDownloader physical operator is responsible for outputting appropriately sized blocks.
|
|
# Decide obstore vs threaded upfront. For fsspec-S3 filesystems backed by
|
|
# a session we can't statically introspect (Okta / STS / profile-based),
|
|
# _plan_obstore_routing emits a warning and returns use_obstore=False so
|
|
# we fall back to the threaded PyArrow path — which uses the user's
|
|
# filesystem directly and resolves credentials correctly.
|
|
use_obstore_path = False
|
|
if OBSTORE_AVAILABLE:
|
|
use_obstore_path, _ = _plan_obstore_routing(filesystem)
|
|
|
|
partition_map_operator = None
|
|
if not upstream_op_is_download:
|
|
partition_cls = AsyncPartitionActor if use_obstore_path else PartitionActor
|
|
# PartitionActor / AsyncPartitionActor are callable classes, so we need
|
|
# ActorPoolStrategy.
|
|
partition_compute = ActorPoolStrategy(
|
|
size=1, enable_true_multi_threading=True
|
|
) # Use single actor for partitioning
|
|
|
|
fn, init_fn = _get_udf(
|
|
partition_cls,
|
|
(),
|
|
{},
|
|
(uri_column_names, data_context, filesystem),
|
|
{},
|
|
compute=partition_compute,
|
|
)
|
|
block_fn = _generate_transform_fn_for_map_batches(fn)
|
|
|
|
partition_transform_fns = [
|
|
BlockMapTransformFn(
|
|
block_fn,
|
|
# NOTE: Disable block-shaping to produce blocks as is
|
|
disable_block_shaping=True,
|
|
),
|
|
]
|
|
partition_map_transformer = MapTransformer(
|
|
partition_transform_fns,
|
|
init_fn=init_fn,
|
|
)
|
|
|
|
partition_map_operator = ActorPoolMapOperator(
|
|
partition_map_transformer,
|
|
input_physical_dag,
|
|
data_context,
|
|
name=f"Partition({uri_column_names_str})",
|
|
# NOTE: Partition actor doesn't use the user-provided `ray_remote_args`
|
|
# since those only apply to the actual download tasks. Partitioning is
|
|
# a lightweight internal operation that doesn't need custom resource
|
|
# requirements.
|
|
ray_remote_args=None,
|
|
compute_strategy=partition_compute, # Use actor-based compute for callable class
|
|
# NOTE: We set `_generator_backpressure_num_objects` to -1 to unblock
|
|
# backpressure since partitioning is extremely fast. Without this, the
|
|
# partition actor gets bottlenecked by the Ray Data scheduler, which
|
|
# can prevent Ray Data from launching enough download tasks.
|
|
ray_actor_task_remote_args={"_generator_backpressure_num_objects": -1},
|
|
)
|
|
|
|
if use_obstore_path:
|
|
download_fn = download_bytes_async
|
|
logger.debug("Using obstore async download path.")
|
|
else:
|
|
download_fn = download_bytes_threaded
|
|
# The "obstore not installed" warning is only relevant when obstore is
|
|
# missing entirely. When obstore is available but the filesystem can't
|
|
# be routed through it, _plan_obstore_routing already logged the reason
|
|
# (a WARNING for fsspec-S3-unextractable, DEBUG otherwise).
|
|
if not OBSTORE_AVAILABLE:
|
|
_log_fallback_warning()
|
|
|
|
fn, init_fn = _get_udf(
|
|
download_fn,
|
|
(uri_column_names, output_bytes_column_names, data_context, filesystem),
|
|
{},
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
|
|
download_transform_fn = _generate_transform_fn_for_map_batches(fn)
|
|
transform_fns = [
|
|
BlockMapTransformFn(
|
|
download_transform_fn,
|
|
output_block_size_option=OutputBlockSizeOption.of(
|
|
target_max_block_size=data_context.target_max_block_size
|
|
),
|
|
),
|
|
]
|
|
|
|
download_compute = TaskPoolStrategy()
|
|
download_map_transformer = MapTransformer(
|
|
transform_fns,
|
|
init_fn=init_fn,
|
|
)
|
|
|
|
download_map_operator = MapOperator.create(
|
|
download_map_transformer,
|
|
partition_map_operator if partition_map_operator else input_physical_dag,
|
|
data_context,
|
|
name=f"Download({uri_column_names_str})",
|
|
compute_strategy=download_compute,
|
|
ray_remote_args=ray_remote_args,
|
|
)
|
|
|
|
return download_map_operator
|
|
|
|
|
|
def download_bytes_threaded(
|
|
block: pa.Table,
|
|
uri_column_names: List[str],
|
|
output_bytes_column_names: List[str],
|
|
data_context: DataContext,
|
|
filesystem: Optional["pa.fs.FileSystem"] = None,
|
|
) -> Iterator[pa.Table]:
|
|
"""Optimized version that uses make_async_gen for concurrent downloads.
|
|
|
|
Supports downloading from multiple URI columns in a single operation.
|
|
|
|
Args:
|
|
block: Input PyArrow table containing URI columns.
|
|
uri_column_names: Names of columns containing URIs to download.
|
|
output_bytes_column_names: Names for the output columns containing downloaded bytes.
|
|
data_context: Ray Data context for configuration.
|
|
filesystem: PyArrow filesystem to use for reading remote files.
|
|
If None, the filesystem is auto-detected from the path scheme.
|
|
|
|
Yields:
|
|
pa.Table: PyArrow table with the downloaded bytes added as new columns.
|
|
"""
|
|
if not isinstance(block, pa.Table):
|
|
block = BlockAccessor.for_block(block).to_arrow()
|
|
|
|
output_block = block
|
|
|
|
# Download each URI column and add it to the output block
|
|
for uri_column_name, output_bytes_column_name in zip(
|
|
uri_column_names, output_bytes_column_names
|
|
):
|
|
# Extract URIs from PyArrow table
|
|
uris = output_block.column(uri_column_name).to_pylist()
|
|
|
|
if len(uris) == 0:
|
|
continue
|
|
|
|
# Resolve the filesystem once before spawning workers; otherwise each
|
|
# worker infers its own S3FileSystem and fires a duplicate IMDS
|
|
# credential fetch. Normalize fsspec inputs so RetryingPyFileSystem.wrap
|
|
# can forward open_input_stream.
|
|
resolved_fs = _validate_and_wrap_filesystem(filesystem)
|
|
if resolved_fs is None:
|
|
for probe_uri in uris:
|
|
if probe_uri is None:
|
|
continue
|
|
try:
|
|
paths, candidate_fs = _resolve_paths_and_filesystem(probe_uri, None)
|
|
except Exception as e:
|
|
logger.debug(f"Could not infer filesystem from '{probe_uri}': {e}")
|
|
continue
|
|
# Skip results that drop the URI (([], ...)) or yield no FS.
|
|
if paths and candidate_fs is not None:
|
|
resolved_fs = candidate_fs
|
|
break
|
|
|
|
if resolved_fs is None:
|
|
# No URI resolved a filesystem; workers would only repeat the same
|
|
# failed inference. Yield None for every row and skip the pool.
|
|
logger.warning(
|
|
"Could not resolve a filesystem from any URI in column "
|
|
f"{uri_column_name!r} ({len(uris)} URIs). Yielding None for "
|
|
"all rows."
|
|
)
|
|
output_block = output_block.add_column(
|
|
len(output_block.column_names),
|
|
output_bytes_column_name,
|
|
pa.array([None] * len(uris), type=pa.binary()),
|
|
)
|
|
continue
|
|
|
|
wrapped_fs = RetryingPyFileSystem.wrap(
|
|
resolved_fs, retryable_errors=data_context.retried_io_errors
|
|
)
|
|
|
|
def load_uri_bytes(
|
|
uri_iterator,
|
|
wrapped_fs=wrapped_fs,
|
|
resolved_fs=resolved_fs,
|
|
uri_column_name=uri_column_name,
|
|
):
|
|
"""Download bytes for each URI using the pre-resolved filesystem."""
|
|
for uri in uri_iterator:
|
|
read_bytes = None
|
|
try:
|
|
if uri is None:
|
|
continue
|
|
# Normalize the path only; FS is supplied so no network I/O.
|
|
resolved_paths, _ = _resolve_paths_and_filesystem(
|
|
uri, filesystem=resolved_fs
|
|
)
|
|
resolved_path = resolved_paths[0] if resolved_paths else None
|
|
if resolved_path is None:
|
|
continue
|
|
with wrapped_fs.open_input_stream(resolved_path) as f:
|
|
read_bytes = f.read()
|
|
except OSError as e:
|
|
logger.debug(
|
|
f"OSError reading uri '{uri}' for column '{uri_column_name}': {e}"
|
|
)
|
|
except Exception as e:
|
|
# Catch unexpected errors like pyarrow.lib.ArrowInvalid caused by an invalid uri like
|
|
# `foo://bar` to avoid failing because of one invalid uri.
|
|
logger.warning(
|
|
f"Unexpected error reading uri '{uri}' for column '{uri_column_name}': {e}"
|
|
)
|
|
finally:
|
|
yield read_bytes
|
|
|
|
# Use make_async_gen to resolve and download URI bytes concurrently
|
|
# preserve_ordering=True ensures results are returned in the same order as input URIs
|
|
uri_bytes = list(
|
|
make_async_gen(
|
|
base_iterator=iter(uris),
|
|
fn=load_uri_bytes,
|
|
preserve_ordering=True,
|
|
num_workers=URI_DOWNLOAD_MAX_WORKERS,
|
|
)
|
|
)
|
|
|
|
# Add the new column to the PyArrow table
|
|
output_block = output_block.add_column(
|
|
len(output_block.column_names),
|
|
output_bytes_column_name,
|
|
pa.array(uri_bytes),
|
|
)
|
|
|
|
yield from _iter_arrow_table_for_target_max_block_size(
|
|
output_block, data_context.target_max_block_size
|
|
)
|