chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,376 @@
|
||||
import warnings
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Dict, List, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.air.data_batch_type import DataBatchType
|
||||
from ray.data.constants import TENSOR_COLUMN_NAME
|
||||
from ray.data.util.expression_utils import _get_setting_with_copy_warning
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pandas as pd
|
||||
|
||||
# TODO: Consolidate data conversion edges for arrow bug workaround.
|
||||
try:
|
||||
import pyarrow
|
||||
except ImportError:
|
||||
pyarrow = None
|
||||
|
||||
# Lazy import to avoid ray init failures without pandas installed and allow
|
||||
# dataset to import modules in this file.
|
||||
_pandas = None
|
||||
|
||||
|
||||
def _lazy_import_pandas():
|
||||
global _pandas
|
||||
if _pandas is None:
|
||||
import pandas
|
||||
|
||||
_pandas = pandas
|
||||
return _pandas
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class BatchFormat(str, Enum):
|
||||
PANDAS = "pandas"
|
||||
# TODO: Remove once Arrow is deprecated as user facing batch format
|
||||
ARROW = "arrow"
|
||||
NUMPY = "numpy" # Either a single numpy array or a Dict of numpy arrays.
|
||||
CUDF = "cudf"
|
||||
|
||||
|
||||
_CUDF_UNSET = object()
|
||||
_cudf = _CUDF_UNSET
|
||||
|
||||
|
||||
def _lazy_import_cudf():
|
||||
"""Lazy import cudf, returning the module or None if not installed."""
|
||||
global _cudf
|
||||
if _cudf is _CUDF_UNSET:
|
||||
try:
|
||||
import cudf
|
||||
|
||||
_cudf = cudf
|
||||
except ImportError:
|
||||
_cudf = None
|
||||
return _cudf
|
||||
|
||||
|
||||
def _convert_batch_type_to_pandas(
|
||||
data: DataBatchType,
|
||||
cast_tensor_columns: bool = False,
|
||||
) -> "pd.DataFrame":
|
||||
"""Convert the provided data to a Pandas DataFrame.
|
||||
|
||||
Args:
|
||||
data: Data of type DataBatchType
|
||||
cast_tensor_columns: Whether tensor columns should be cast to NumPy ndarrays.
|
||||
|
||||
Returns:
|
||||
A pandas Dataframe representation of the input data.
|
||||
|
||||
"""
|
||||
pd = _lazy_import_pandas()
|
||||
|
||||
if isinstance(data, np.ndarray):
|
||||
data = pd.DataFrame({TENSOR_COLUMN_NAME: _ndarray_to_column(data)})
|
||||
elif isinstance(data, dict):
|
||||
tensor_dict = {}
|
||||
for col_name, col in data.items():
|
||||
if not isinstance(col, np.ndarray):
|
||||
raise ValueError(
|
||||
"All values in the provided dict must be of type "
|
||||
f"np.ndarray. Found type {type(col)} for key {col_name} "
|
||||
f"instead."
|
||||
)
|
||||
tensor_dict[col_name] = _ndarray_to_column(col)
|
||||
data = pd.DataFrame(tensor_dict)
|
||||
elif pyarrow is not None and isinstance(data, pyarrow.Table):
|
||||
data = data.to_pandas()
|
||||
else:
|
||||
# Handle cudf.DataFrame (lazy check to avoid import when not used)
|
||||
cudf = _lazy_import_cudf()
|
||||
if cudf is not None and isinstance(data, cudf.DataFrame):
|
||||
data = data.to_pandas()
|
||||
if not isinstance(data, pd.DataFrame):
|
||||
raise ValueError(
|
||||
f"Received data of type: {type(data)}, but expected it to be one "
|
||||
f"of {DataBatchType}"
|
||||
)
|
||||
if cast_tensor_columns:
|
||||
data = _cast_tensor_columns_to_ndarrays(data)
|
||||
return data
|
||||
|
||||
|
||||
def _convert_pandas_to_batch_type(
|
||||
data: "pd.DataFrame",
|
||||
type: BatchFormat,
|
||||
cast_tensor_columns: bool = False,
|
||||
) -> DataBatchType:
|
||||
"""Convert the provided Pandas dataframe to the provided ``type``.
|
||||
|
||||
Args:
|
||||
data: A Pandas DataFrame
|
||||
type: The specific ``BatchFormat`` to convert to.
|
||||
cast_tensor_columns: Whether tensor columns should be cast to our tensor
|
||||
extension type.
|
||||
|
||||
Returns:
|
||||
The input data represented with the provided type.
|
||||
"""
|
||||
if cast_tensor_columns:
|
||||
data = _cast_ndarray_columns_to_tensor_extension(data)
|
||||
|
||||
if type == BatchFormat.PANDAS:
|
||||
return data
|
||||
elif type == BatchFormat.NUMPY:
|
||||
if len(data.columns) == 1:
|
||||
# If just a single column, return as a single numpy array.
|
||||
return data.iloc[:, 0].to_numpy()
|
||||
else:
|
||||
# Else return as a dict of numpy arrays.
|
||||
output_dict = {}
|
||||
for column in data:
|
||||
output_dict[column] = data[column].to_numpy()
|
||||
return output_dict
|
||||
elif type == BatchFormat.ARROW:
|
||||
if not pyarrow:
|
||||
raise ValueError(
|
||||
"Attempted to convert data to Pyarrow Table but Pyarrow "
|
||||
"is not installed. Please do `pip install pyarrow` to "
|
||||
"install Pyarrow."
|
||||
)
|
||||
return pyarrow.Table.from_pandas(data)
|
||||
elif type == BatchFormat.CUDF:
|
||||
cudf = _lazy_import_cudf()
|
||||
if cudf is None:
|
||||
raise ValueError(
|
||||
"Attempted to convert data to cuDF DataFrame but cuDF "
|
||||
"is not installed. Please do `pip install cudf-cu12` to "
|
||||
"install cuDF (GPU required)."
|
||||
)
|
||||
return cudf.from_pandas(data)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Received type {type}, but expected it to be one of {DataBatchType}"
|
||||
)
|
||||
|
||||
|
||||
def _convert_batch_type_to_numpy(
|
||||
data: DataBatchType,
|
||||
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
"""Convert the provided data to a NumPy ndarray or dict of ndarrays.
|
||||
|
||||
Args:
|
||||
data: Data of type DataBatchType
|
||||
|
||||
Returns:
|
||||
A numpy representation of the input data.
|
||||
"""
|
||||
pd = _lazy_import_pandas()
|
||||
|
||||
if isinstance(data, np.ndarray):
|
||||
return data
|
||||
elif isinstance(data, dict):
|
||||
for col_name, col in data.items():
|
||||
if not isinstance(col, np.ndarray):
|
||||
raise ValueError(
|
||||
"All values in the provided dict must be of type "
|
||||
f"np.ndarray. Found type {type(col)} for key {col_name} "
|
||||
f"instead."
|
||||
)
|
||||
return data
|
||||
elif pyarrow is not None and isinstance(data, pyarrow.Table):
|
||||
from ray.data._internal.arrow_ops import transform_pyarrow
|
||||
from ray.data._internal.tensor_extensions.arrow import (
|
||||
get_arrow_extension_fixed_shape_tensor_types,
|
||||
)
|
||||
|
||||
column_values_ndarrays = []
|
||||
|
||||
for col in data.columns:
|
||||
# Combine columnar values arrays to make these contiguous
|
||||
# (making them compatible with numpy format)
|
||||
combined_array = transform_pyarrow.combine_chunked_array(col)
|
||||
|
||||
column_values_ndarrays.append(
|
||||
transform_pyarrow.to_numpy(combined_array, zero_copy_only=False)
|
||||
)
|
||||
|
||||
arrow_fixed_shape_tensor_types = get_arrow_extension_fixed_shape_tensor_types()
|
||||
|
||||
# NOTE: This branch is here for backwards-compatibility
|
||||
if data.column_names == [TENSOR_COLUMN_NAME] and (
|
||||
isinstance(data.schema.types[0], arrow_fixed_shape_tensor_types)
|
||||
):
|
||||
return column_values_ndarrays[0]
|
||||
|
||||
return dict(zip(data.column_names, column_values_ndarrays))
|
||||
elif isinstance(data, pd.DataFrame):
|
||||
return _convert_pandas_to_batch_type(data, BatchFormat.NUMPY)
|
||||
else:
|
||||
# Handle cudf.DataFrame via pandas path
|
||||
cudf = _lazy_import_cudf()
|
||||
if cudf is not None and isinstance(data, cudf.DataFrame):
|
||||
return _convert_pandas_to_batch_type(data.to_pandas(), BatchFormat.NUMPY)
|
||||
raise ValueError(
|
||||
f"Received data of type: {type(data)}, but expected it to be one "
|
||||
f"of {DataBatchType}"
|
||||
)
|
||||
|
||||
|
||||
def _ndarray_to_column(arr: np.ndarray) -> Union["pd.Series", List[np.ndarray]]:
|
||||
"""Convert a NumPy ndarray into an appropriate column format for insertion into a
|
||||
pandas DataFrame.
|
||||
|
||||
If conversion to a pandas Series fails (e.g. if the ndarray is multi-dimensional),
|
||||
fall back to a list of NumPy ndarrays.
|
||||
"""
|
||||
pd = _lazy_import_pandas()
|
||||
try:
|
||||
# Try to convert to Series, falling back to a list conversion if this fails
|
||||
# (e.g. if the ndarray is multi-dimensional).
|
||||
return pd.Series(arr)
|
||||
except ValueError:
|
||||
return list(arr)
|
||||
|
||||
|
||||
def _unwrap_ndarray_object_type_if_needed(arr: np.ndarray) -> np.ndarray:
|
||||
"""Unwrap an object-dtyped NumPy ndarray containing ndarray pointers into a single
|
||||
contiguous ndarray, if needed/possible.
|
||||
"""
|
||||
if arr.dtype.type is np.object_:
|
||||
try:
|
||||
# Try to convert the NumPy ndarray to a non-object dtype.
|
||||
arr = np.array([np.asarray(v) for v in arr])
|
||||
except Exception:
|
||||
# This may fail if the subndarrays are of heterogeneous shape
|
||||
pass
|
||||
return arr
|
||||
|
||||
|
||||
def _cast_ndarray_columns_to_tensor_extension(df: "pd.DataFrame") -> "pd.DataFrame":
|
||||
"""
|
||||
Cast all NumPy ndarray columns in df to our tensor extension type, TensorArray.
|
||||
"""
|
||||
# Get the SettingWithCopyWarning class if available
|
||||
SettingWithCopyWarning = _get_setting_with_copy_warning()
|
||||
|
||||
from ray.data._internal.tensor_extensions.pandas import (
|
||||
TensorArray,
|
||||
column_needs_tensor_extension,
|
||||
)
|
||||
|
||||
# Try to convert any ndarray columns to TensorArray columns.
|
||||
# TODO(Clark): Once Pandas supports registering extension types for type
|
||||
# inference on construction, implement as much for NumPy ndarrays and remove
|
||||
# this. See https://github.com/pandas-dev/pandas/issues/41848
|
||||
# TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of
|
||||
# column names containing tensor columns, to make this an O(# of tensor columns)
|
||||
# check rather than the current O(# of columns) check.
|
||||
|
||||
# Scan dtypes rather than df.items(), which would
|
||||
# materialize a Series for every column just to read its dtype.
|
||||
# The below approach avoids the cost of a Series build for non-tensor columns.
|
||||
#
|
||||
# When column names are unique we select and assign by label.
|
||||
# With duplicate names, ``df[col_name]`` returns a DataFrame
|
||||
# rather than a Series, so we select and assign by position instead.
|
||||
columns_unique = df.columns.is_unique
|
||||
for i, (col_name, dtype) in enumerate(df.dtypes.items()):
|
||||
if (
|
||||
dtype.type is not np.object_
|
||||
): # Short circuit if non-object type before materializing the column
|
||||
continue
|
||||
col = df[col_name] if columns_unique else df.iloc[:, i]
|
||||
if column_needs_tensor_extension(col):
|
||||
try:
|
||||
# Suppress Pandas warnings:
|
||||
# https://github.com/ray-project/ray/issues/29270
|
||||
# We actually want in-place operations so we surpress this warning.
|
||||
# https://stackoverflow.com/a/74193599
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", category=FutureWarning)
|
||||
if SettingWithCopyWarning is not None:
|
||||
warnings.simplefilter("ignore", category=SettingWithCopyWarning)
|
||||
if columns_unique:
|
||||
df[col_name] = TensorArray(col)
|
||||
else:
|
||||
df.isetitem(i, TensorArray(col))
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Tried to cast column {col_name} to the TensorArray tensor "
|
||||
"extension type but the conversion failed. To disable "
|
||||
"automatic casting to this tensor extension, set "
|
||||
"ctx = DataContext.get_current(); "
|
||||
"ctx.enable_tensor_extension_casting = False."
|
||||
) from e
|
||||
return df
|
||||
|
||||
|
||||
def _cast_tensor_columns_to_ndarrays(
|
||||
df: "pd.DataFrame",
|
||||
arrow_schema: "pyarrow.Schema" = None,
|
||||
) -> "pd.DataFrame":
|
||||
"""Cast all tensor extension columns in df to NumPy ndarrays.
|
||||
|
||||
Args:
|
||||
df: The DataFrame whose tensor columns should be converted.
|
||||
arrow_schema: If provided, used to reshape columns that were native
|
||||
``FixedShapeTensorType`` in Arrow. PyArrow's ``to_pandas()``
|
||||
flattens these to 1-D ndarrays; passing the original schema
|
||||
lets us restore the correct shape.
|
||||
|
||||
Returns:
|
||||
The DataFrame with tensor columns converted to NumPy ndarrays.
|
||||
"""
|
||||
# Get the SettingWithCopyWarning class if available
|
||||
SettingWithCopyWarning = _get_setting_with_copy_warning()
|
||||
from ray.data._internal.tensor_extensions.pandas import TensorDtype
|
||||
|
||||
# Try to convert any tensor extension columns to ndarray columns.
|
||||
# TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of
|
||||
# column names containing tensor columns, to make this an O(# of tensor columns)
|
||||
# check rather than the current O(# of columns) check.
|
||||
|
||||
# Reshape native FixedShapeTensorType columns that were flattened by
|
||||
# to_pandas().
|
||||
|
||||
if arrow_schema is not None:
|
||||
from ray.data._internal.utils.transform_pyarrow import (
|
||||
_is_native_tensor_type,
|
||||
)
|
||||
|
||||
for field in arrow_schema:
|
||||
if _is_native_tensor_type(field.type) and field.name in df.columns:
|
||||
shape = tuple(field.type.shape)
|
||||
df[field.name] = [
|
||||
arr.reshape(shape) if arr is not None else None
|
||||
for arr in df[field.name]
|
||||
]
|
||||
|
||||
# Scan dtypes rather than df.items(), which would
|
||||
# materialize a Series for every column just to read its dtype.
|
||||
# The below approach avoids the cost of a Series build for non-tensor columns.
|
||||
#
|
||||
# When column names are unique we select and assign by label (the fast,
|
||||
# cached path). With duplicate names, ``df[col_name]`` returns a DataFrame
|
||||
# rather than a Series, so we select and assign by position instead.
|
||||
columns_unique = df.columns.is_unique
|
||||
for i, (col_name, dtype) in enumerate(df.dtypes.items()):
|
||||
if isinstance(dtype, TensorDtype):
|
||||
# Suppress Pandas warnings:
|
||||
# https://github.com/ray-project/ray/issues/29270
|
||||
# We actually want in-place operations so we surpress this warning.
|
||||
# https://stackoverflow.com/a/74193599
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", category=FutureWarning)
|
||||
if SettingWithCopyWarning is not None:
|
||||
warnings.simplefilter("ignore", category=SettingWithCopyWarning)
|
||||
if columns_unique:
|
||||
df[col_name] = list(df[col_name].to_numpy())
|
||||
else:
|
||||
df.isetitem(i, list(df.iloc[:, i].to_numpy()))
|
||||
return df
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Utility functions for expression-based operations."""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Callable, Hashable, List, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.expressions import (
|
||||
AliasExpr,
|
||||
BinaryExpr,
|
||||
ColumnExpr,
|
||||
DownloadExpr,
|
||||
Expr,
|
||||
LiteralExpr,
|
||||
MonotonicallyIncreasingIdExpr,
|
||||
RandomExpr,
|
||||
UDFExpr,
|
||||
UnaryExpr,
|
||||
UUIDExpr,
|
||||
)
|
||||
|
||||
|
||||
def _get_setting_with_copy_warning() -> Optional[type]:
|
||||
"""Get the SettingWithCopyWarning class from pandas, if available.
|
||||
|
||||
Pandas has moved/renamed this warning across versions, and pandas 3.x may not
|
||||
expose it at all. This function handles the version differences gracefully
|
||||
using hasattr checks instead of try-except blocks.
|
||||
|
||||
Returns:
|
||||
The SettingWithCopyWarning class if found, None otherwise.
|
||||
"""
|
||||
import pandas as pd
|
||||
|
||||
# Use hasattr to avoid try-catch blocks as suggested
|
||||
if hasattr(pd.core.common, "SettingWithCopyWarning"):
|
||||
return pd.core.common.SettingWithCopyWarning
|
||||
elif hasattr(pd.errors, "SettingWithCopyWarning"):
|
||||
return pd.errors.SettingWithCopyWarning
|
||||
else:
|
||||
# Warning not available in this pandas version
|
||||
return None
|
||||
|
||||
|
||||
def _create_callable_class_udf_init_fn(
|
||||
exprs: List["Expr"],
|
||||
) -> Optional[Callable[[], None]]:
|
||||
"""Create an init_fn to initialize all callable class UDFs in expressions.
|
||||
|
||||
This function collects all _CallableClassUDF instances from the given expressions,
|
||||
groups them by their callable_class_spec key, and returns an init_fn that
|
||||
initializes each group at actor startup. UDFs with the same key (same class and
|
||||
constructor args) share a single instance to ensure all are properly initialized.
|
||||
|
||||
Args:
|
||||
exprs: List of expressions to collect callable class UDFs from.
|
||||
|
||||
Returns:
|
||||
An init_fn that initializes all callable class UDFs, or None if there are
|
||||
no callable class UDFs in the expressions.
|
||||
"""
|
||||
from ray.data._internal.planner.plan_expression.expression_visitors import (
|
||||
_CallableClassUDFCollector,
|
||||
)
|
||||
|
||||
callable_class_udfs = []
|
||||
for expr in exprs:
|
||||
collector = _CallableClassUDFCollector()
|
||||
collector.visit(expr)
|
||||
callable_class_udfs.extend(collector.get_callable_class_udfs())
|
||||
|
||||
if not callable_class_udfs:
|
||||
return None
|
||||
|
||||
# Group UDFs by callable_class_spec key.
|
||||
# Multiple _CallableClassUDF objects may have the same key (same class + args).
|
||||
# We need to initialize ALL of them, sharing a single instance per key.
|
||||
udfs_by_key = {}
|
||||
for udf in callable_class_udfs:
|
||||
key = udf.callable_class_spec.make_key()
|
||||
if key not in udfs_by_key:
|
||||
udfs_by_key[key] = []
|
||||
udfs_by_key[key].append(udf)
|
||||
|
||||
def init_fn():
|
||||
for udfs_with_same_key in udfs_by_key.values():
|
||||
# Initialize the first UDF to create the instance
|
||||
first_udf = udfs_with_same_key[0]
|
||||
first_udf.init()
|
||||
# Share the instance with all other UDFs that have the same key
|
||||
for other_udf in udfs_with_same_key[1:]:
|
||||
other_udf._instance = first_udf._instance
|
||||
|
||||
return init_fn
|
||||
|
||||
|
||||
def _call_udf_instance_with_async_bridge(
|
||||
instance: Any,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> Any:
|
||||
"""Call a UDF instance, bridging from sync context to async if needed.
|
||||
|
||||
This handles the complexity of calling callable class UDF instances that may
|
||||
be sync, async coroutine, or async generator functions.
|
||||
|
||||
Args:
|
||||
instance: The callable instance to call
|
||||
*args: Positional arguments
|
||||
**kwargs: Keyword arguments
|
||||
|
||||
Returns:
|
||||
The result of calling the instance
|
||||
"""
|
||||
import asyncio
|
||||
import inspect
|
||||
|
||||
# Check if the instance's __call__ is async
|
||||
if inspect.iscoroutinefunction(instance.__call__):
|
||||
# Async coroutine: bridge from sync to async
|
||||
return asyncio.run(instance(*args, **kwargs))
|
||||
elif inspect.isasyncgenfunction(instance.__call__):
|
||||
# Async generator: collect results
|
||||
async def _collect():
|
||||
results = []
|
||||
async for item in instance(*args, **kwargs):
|
||||
results.append(item)
|
||||
# In expressions, the UDF must return a single array with the same
|
||||
# length as the input (one output element per input row).
|
||||
# If the async generator yields multiple arrays, we take the last one
|
||||
# since expressions don't support multi-batch output semantics.
|
||||
if not results:
|
||||
return None
|
||||
elif len(results) == 1:
|
||||
return results[0]
|
||||
else:
|
||||
import logging
|
||||
|
||||
logging.warning(
|
||||
f"Async generator yielded {len(results)} values in expression context; "
|
||||
"only the last (most recent) is returned. Use map_batches for multi-yield support."
|
||||
)
|
||||
return results[-1]
|
||||
|
||||
return asyncio.run(_collect())
|
||||
else:
|
||||
# Synchronous instance - direct call
|
||||
return instance(*args, **kwargs)
|
||||
|
||||
|
||||
def _make_hashable(value: Any) -> Hashable:
|
||||
try:
|
||||
hash(value)
|
||||
return value
|
||||
except TypeError:
|
||||
pass
|
||||
|
||||
if isinstance(value, list):
|
||||
return tuple(_make_hashable(v) for v in value)
|
||||
if isinstance(value, tuple):
|
||||
return tuple(_make_hashable(v) for v in value)
|
||||
if isinstance(value, dict):
|
||||
return tuple(
|
||||
sorted(
|
||||
((k, _make_hashable(v)) for k, v in value.items()),
|
||||
key=lambda item: repr(item[0]),
|
||||
)
|
||||
)
|
||||
if isinstance(value, set):
|
||||
return frozenset(_make_hashable(v) for v in value)
|
||||
|
||||
return repr(value)
|
||||
|
||||
|
||||
def _data_type_key(expr: "Expr") -> Hashable:
|
||||
return repr(getattr(expr, "data_type", None))
|
||||
|
||||
|
||||
def _udf_function_key(fn: Any) -> Hashable:
|
||||
from ray.data.expressions import _CallableClassUDF
|
||||
|
||||
if isinstance(fn, _CallableClassUDF):
|
||||
return ("callable_class", fn.callable_class_spec.make_key())
|
||||
return ("function", _make_hashable(fn))
|
||||
|
||||
|
||||
def _column_fingerprint_key(expr: "ColumnExpr") -> Hashable:
|
||||
return ("column", expr.name)
|
||||
|
||||
|
||||
def _literal_fingerprint_key(expr: "LiteralExpr") -> Hashable:
|
||||
return ("literal", type(expr.value), _make_hashable(expr.value))
|
||||
|
||||
|
||||
def _binary_fingerprint_key(
|
||||
expr: "BinaryExpr", left_key: Hashable, right_key: Hashable
|
||||
) -> Hashable:
|
||||
return ("binary", expr.op, left_key, right_key)
|
||||
|
||||
|
||||
def _unary_fingerprint_key(expr: "UnaryExpr", operand_key: Hashable) -> Hashable:
|
||||
return ("unary", expr.op, operand_key)
|
||||
|
||||
|
||||
def _udf_fingerprint_key(
|
||||
expr: "UDFExpr",
|
||||
arg_keys: tuple[Hashable, ...],
|
||||
kwarg_keys: tuple[tuple[str, Hashable], ...],
|
||||
) -> Hashable:
|
||||
from ray.data.expressions import PyArrowComputeUDFExpr
|
||||
|
||||
if isinstance(expr, PyArrowComputeUDFExpr):
|
||||
return (
|
||||
"pyarrow_compute_udf",
|
||||
_make_hashable(expr.pc_func),
|
||||
_make_hashable(expr.pc_positional),
|
||||
_make_hashable(expr.pc_kwargs),
|
||||
arg_keys,
|
||||
kwarg_keys,
|
||||
_data_type_key(expr),
|
||||
)
|
||||
|
||||
return (
|
||||
"udf",
|
||||
_udf_function_key(expr.fn),
|
||||
arg_keys,
|
||||
kwarg_keys,
|
||||
_data_type_key(expr),
|
||||
)
|
||||
|
||||
|
||||
def _alias_fingerprint_key(expr: "AliasExpr", child_key: Hashable) -> Hashable:
|
||||
return (
|
||||
"alias",
|
||||
expr.name,
|
||||
expr._is_rename,
|
||||
child_key,
|
||||
_data_type_key(expr),
|
||||
)
|
||||
|
||||
|
||||
def _download_fingerprint_key(expr: "DownloadExpr") -> Hashable:
|
||||
return ("download", expr.uri_column_name)
|
||||
|
||||
|
||||
def _star_fingerprint_key() -> Hashable:
|
||||
return ("star",)
|
||||
|
||||
|
||||
def _monotonically_increasing_id_fingerprint_key(
|
||||
expr: "MonotonicallyIncreasingIdExpr",
|
||||
) -> Hashable:
|
||||
return ("monotonically_increasing_id", expr._instance_id)
|
||||
|
||||
|
||||
def _random_fingerprint_key(expr: "RandomExpr") -> Hashable:
|
||||
return (
|
||||
"random",
|
||||
expr.seed,
|
||||
expr.reseed_after_execution,
|
||||
_data_type_key(expr),
|
||||
)
|
||||
|
||||
|
||||
def _uuid_fingerprint_key(expr: "UUIDExpr") -> Hashable:
|
||||
return ("uuid", _data_type_key(expr))
|
||||
@@ -0,0 +1,416 @@
|
||||
import logging
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
Iterable,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
import numpy as np
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import jax
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_GLOBAL_MESH_1D_AXIS = "data"
|
||||
|
||||
NumpyBatch = Union[np.ndarray, Dict[str, np.ndarray]]
|
||||
JaxBatch = Union["jax.Array", Dict[str, "jax.Array"]]
|
||||
DTypeLikeSpec = Union["jax.typing.DTypeLike", Dict[str, "jax.typing.DTypeLike"]]
|
||||
Scalar = Union[int, float, bool]
|
||||
PaddingsSpec = Union[Scalar, Dict[str, Scalar]]
|
||||
|
||||
|
||||
def _get_column_value(mapping_or_value: Any, key: str) -> Any:
|
||||
"""Get the value for a specific column from a mapping or a single value."""
|
||||
if isinstance(mapping_or_value, dict):
|
||||
return mapping_or_value[key]
|
||||
return mapping_or_value
|
||||
|
||||
|
||||
def _unwrap_single_column_value(mapping_or_value: Any, name: str) -> Any:
|
||||
"""Unwrap a single value from a mapping if it's a dictionary."""
|
||||
if isinstance(mapping_or_value, dict):
|
||||
if len(mapping_or_value) != 1:
|
||||
raise ValueError(
|
||||
f"When constructing a single-tensor batch, only a single {name} "
|
||||
f"should be given, instead got: {mapping_or_value}"
|
||||
)
|
||||
return next(iter(mapping_or_value.values()))
|
||||
return mapping_or_value
|
||||
|
||||
|
||||
def _create_sharding_1d(axis_name: str) -> "jax.sharding.Sharding":
|
||||
"""Create a 1D JAX sharding, preferably using topology-aware mesh_utils."""
|
||||
import jax
|
||||
from jax.sharding import Mesh, NamedSharding, PartitionSpec
|
||||
|
||||
devices = None
|
||||
try:
|
||||
from jax.experimental import mesh_utils
|
||||
|
||||
# Attempt to create a topology-aware mesh (e.g. for TPU/GPU interconnects)
|
||||
devices = mesh_utils.create_device_mesh((jax.device_count(),))
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to use jax.experimental.mesh_utils: {e}")
|
||||
|
||||
if devices is None:
|
||||
# Fallback to default device ordering if mesh_utils fails or is unavailable
|
||||
devices = np.array(jax.devices())
|
||||
|
||||
physical_mesh = Mesh(devices, (axis_name,))
|
||||
return NamedSharding(physical_mesh, PartitionSpec(axis_name))
|
||||
|
||||
|
||||
def _convert_ndarray_to_jax_array(
|
||||
ndarray: np.ndarray,
|
||||
sharding: "jax.sharding.Sharding", # noqa: F821
|
||||
dtype: Optional["jax.typing.DTypeLike"] = None,
|
||||
) -> "jax.Array": # noqa: F821
|
||||
import jax
|
||||
|
||||
local_batch_size = ndarray.shape[0]
|
||||
host_count = jax.process_count()
|
||||
|
||||
# Global shape assumes each host gets the exact same local batch size.
|
||||
global_shape = (local_batch_size * host_count,) + ndarray.shape[1:]
|
||||
|
||||
# Construct the globally aware 1D array from process-local data.
|
||||
# This automatically shards the local ndarray across the local devices
|
||||
# assigned to this process by the provided sharding.
|
||||
physical_array = jax.make_array_from_process_local_data(
|
||||
sharding, ndarray, global_shape
|
||||
)
|
||||
|
||||
if dtype is not None:
|
||||
physical_array = physical_array.astype(dtype)
|
||||
|
||||
return physical_array
|
||||
|
||||
|
||||
def _convert_batch(
|
||||
ndarrays: NumpyBatch,
|
||||
sharding: "jax.sharding.Sharding",
|
||||
dtypes: Optional[DTypeLikeSpec] = None,
|
||||
) -> JaxBatch:
|
||||
"""Convert a NumPy ndarray batch to a globally sharded JAX Array batch.
|
||||
|
||||
Args:
|
||||
ndarrays: A single NumPy ndarray or dictionary of NumPy ndarrays.
|
||||
sharding: The JAX sharding to use for the conversion.
|
||||
dtypes: A single JAX dtype or dictionary of JAX dtypes for the created arrays.
|
||||
|
||||
Returns:
|
||||
A globally sharded JAX Array (or dictionary of arrays) residing
|
||||
in TPU/GPU memory.
|
||||
"""
|
||||
if isinstance(ndarrays, np.ndarray):
|
||||
dtype = _unwrap_single_column_value(dtypes, "dtype")
|
||||
jax_batch = _convert_ndarray_to_jax_array(ndarrays, sharding, dtype=dtype)
|
||||
else:
|
||||
jax_batch = {}
|
||||
for col_name, col_ndarray in ndarrays.items():
|
||||
dtype = _get_column_value(dtypes, col_name)
|
||||
try:
|
||||
jax_batch[col_name] = _convert_ndarray_to_jax_array(
|
||||
col_ndarray, sharding, dtype=dtype
|
||||
)
|
||||
except ValueError as e:
|
||||
raise ValueError(
|
||||
f"JAX Array Conversion Error for column '{col_name}'"
|
||||
) from e
|
||||
|
||||
return jax_batch
|
||||
|
||||
|
||||
def _get_batch_size(batch: NumpyBatch) -> int:
|
||||
"""Get the batch size of a NumPy ndarray or dictionary of NumPy ndarrays."""
|
||||
if isinstance(batch, dict):
|
||||
# Use the first column to determine the batch size
|
||||
try:
|
||||
return len(next(iter(batch.values())))
|
||||
except StopIteration:
|
||||
return 0
|
||||
return len(batch)
|
||||
|
||||
|
||||
def _pad_array(arr: np.ndarray, target_size: int, pad_value: Scalar) -> np.ndarray:
|
||||
"""Pad a single array to target_size using pad_value."""
|
||||
current_size = len(arr)
|
||||
if current_size == target_size:
|
||||
return arr
|
||||
padding_shape = (target_size - current_size,) + arr.shape[1:]
|
||||
padding = np.full(padding_shape, pad_value, dtype=arr.dtype)
|
||||
return np.concatenate([arr, padding], axis=0)
|
||||
|
||||
|
||||
def _dummy_array(arr: np.ndarray, target_size: int, pad_value: Scalar) -> np.ndarray:
|
||||
"""Create a dummy array of target_size filled with pad_value."""
|
||||
shape = (target_size,) + arr.shape[1:]
|
||||
return np.full(shape, pad_value, dtype=arr.dtype)
|
||||
|
||||
|
||||
def _pad_batch(
|
||||
batch: NumpyBatch,
|
||||
target_size: int,
|
||||
paddings: PaddingsSpec,
|
||||
) -> NumpyBatch:
|
||||
"""Pad a batch to target_size using paddings."""
|
||||
|
||||
if isinstance(batch, dict):
|
||||
return {
|
||||
k: _pad_array(v, target_size, _get_column_value(paddings, k))
|
||||
for k, v in batch.items()
|
||||
}
|
||||
return _pad_array(
|
||||
batch,
|
||||
target_size,
|
||||
_unwrap_single_column_value(paddings, "padding"),
|
||||
)
|
||||
|
||||
|
||||
def _create_dummy_batch(
|
||||
template_batch: NumpyBatch,
|
||||
target_size: int,
|
||||
paddings: PaddingsSpec,
|
||||
) -> NumpyBatch:
|
||||
"""Create a dummy batch of target_size filled with paddings."""
|
||||
|
||||
if isinstance(template_batch, dict):
|
||||
return {
|
||||
k: _dummy_array(v, target_size, _get_column_value(paddings, k))
|
||||
for k, v in template_batch.items()
|
||||
}
|
||||
return _dummy_array(
|
||||
template_batch,
|
||||
target_size,
|
||||
_unwrap_single_column_value(paddings, "padding"),
|
||||
)
|
||||
|
||||
|
||||
def _yield_batches_no_sync(
|
||||
iterator: Iterator[NumpyBatch],
|
||||
sharding: "jax.sharding.Sharding",
|
||||
num_local_devices: int,
|
||||
batch_size: int,
|
||||
paddings: Optional[PaddingsSpec],
|
||||
dtypes: Optional[DTypeLikeSpec] = None,
|
||||
) -> Iterator[JaxBatch]:
|
||||
"""Yield batches without multi-host synchronization."""
|
||||
for batch in iterator:
|
||||
local_batch_size = _get_batch_size(batch)
|
||||
|
||||
if local_batch_size == 0:
|
||||
continue
|
||||
|
||||
if paddings is not None:
|
||||
if local_batch_size < batch_size:
|
||||
batch = _pad_batch(batch, batch_size, paddings)
|
||||
elif local_batch_size % num_local_devices != 0:
|
||||
# Without padding, batch size must be divisible by num_local_devices
|
||||
raise ValueError(
|
||||
f"The local batch size ({local_batch_size}) must be evenly "
|
||||
f"divisible by the number of local JAX devices "
|
||||
f"({num_local_devices}) on this host. "
|
||||
f"To safely truncate or pad the batch, "
|
||||
f"set `drop_last=True` or provide a `paddings` in `iter_jax_batches()`."
|
||||
)
|
||||
|
||||
yield _convert_batch(batch, sharding, dtypes=dtypes)
|
||||
|
||||
|
||||
def _fetch_lookahead_batches(
|
||||
iterator: Iterator[NumpyBatch],
|
||||
lookahead: int,
|
||||
) -> Tuple[List[Optional[NumpyBatch]], List[int], Optional[NumpyBatch]]:
|
||||
"""Fetch a window of batches and prepare synchronization info."""
|
||||
local_batches = []
|
||||
local_infos = []
|
||||
template_batch: Optional[NumpyBatch] = None
|
||||
for _ in range(lookahead):
|
||||
try:
|
||||
batch = next(iterator)
|
||||
has_batch = True
|
||||
local_batch_size = _get_batch_size(batch)
|
||||
if template_batch is None:
|
||||
template_batch = batch
|
||||
except StopIteration:
|
||||
batch = None
|
||||
has_batch = False
|
||||
local_batch_size = 0
|
||||
|
||||
local_batches.append(batch)
|
||||
local_infos.extend([int(has_batch), local_batch_size])
|
||||
if not has_batch:
|
||||
break
|
||||
return local_batches, local_infos, template_batch
|
||||
|
||||
|
||||
def _yield_batches_with_sync(
|
||||
iterator: Iterator[NumpyBatch],
|
||||
sharding: "jax.sharding.Sharding",
|
||||
num_local_devices: int,
|
||||
drop_last: bool,
|
||||
batch_size: int,
|
||||
paddings: Optional[PaddingsSpec],
|
||||
synchronize_lookahead: int,
|
||||
dtypes: Optional[DTypeLikeSpec] = None,
|
||||
) -> Iterator[JaxBatch]:
|
||||
"""Yield batches with multi-host synchronization."""
|
||||
import jax.numpy as jnp
|
||||
from jax.experimental.multihost_utils import process_allgather
|
||||
|
||||
template_batch: Optional[NumpyBatch] = None
|
||||
while True:
|
||||
local_batches, local_infos, window_template = _fetch_lookahead_batches(
|
||||
iterator, synchronize_lookahead
|
||||
)
|
||||
if template_batch is None:
|
||||
template_batch = window_template
|
||||
|
||||
# Pad local_infos to 2 * synchronize_lookahead
|
||||
padding_needed = 2 * synchronize_lookahead - len(local_infos)
|
||||
if padding_needed > 0:
|
||||
local_infos.extend([0] * padding_needed)
|
||||
|
||||
gathered = process_allgather(jnp.array(local_infos, dtype=jnp.int32))
|
||||
|
||||
for i in range(synchronize_lookahead):
|
||||
h = gathered[:, 2 * i]
|
||||
s = gathered[:, 2 * i + 1]
|
||||
|
||||
all_have_batch = bool(h.all())
|
||||
any_have_batch = bool(h.any())
|
||||
min_batch_size = int(s.min())
|
||||
max_batch_size = int(s.max())
|
||||
|
||||
if not any_have_batch:
|
||||
return
|
||||
|
||||
if not all_have_batch:
|
||||
# Some workers have exhausted their data while others have more.
|
||||
if drop_last:
|
||||
# If drop_last=True, we stop as soon as any worker is exhausted.
|
||||
return
|
||||
elif paddings is not None:
|
||||
# If paddings is set, we continue until all workers are exhausted.
|
||||
# Workers that are already exhausted will yield dummy batches.
|
||||
pass
|
||||
else:
|
||||
raise ValueError(
|
||||
"Uneven number of batches detected across JAX workers. "
|
||||
"To safely drop orphaned batches without hanging, "
|
||||
"set `drop_last=True` or provide a `paddings` in `iter_jax_batches()`."
|
||||
)
|
||||
|
||||
if paddings is not None:
|
||||
batch = local_batches[i]
|
||||
if batch is None:
|
||||
if template_batch is None:
|
||||
raise ValueError(
|
||||
"Cannot create dummy batches for synchronization because this "
|
||||
"JAX host has not received any data batches to use as a "
|
||||
"template. This usually happens if one JAX host's dataset "
|
||||
"shard is completely empty while others have data. "
|
||||
"Ensure that all JAX hosts have at least one batch of data, "
|
||||
"or use `drop_last=True` to avoid yielding dummy batches."
|
||||
)
|
||||
batch = _create_dummy_batch(template_batch, batch_size, paddings)
|
||||
else:
|
||||
local_batch_size = _get_batch_size(batch)
|
||||
if local_batch_size < batch_size:
|
||||
batch = _pad_batch(batch, batch_size, paddings)
|
||||
assert batch is not None
|
||||
yield _convert_batch(batch, sharding, dtypes=dtypes)
|
||||
else:
|
||||
if max_batch_size > min_batch_size:
|
||||
raise ValueError(
|
||||
"Uneven batch sizes detected across JAX workers. "
|
||||
f"Host batch sizes range from {min_batch_size} to {max_batch_size}. "
|
||||
"To handle uneven batch sizes, provide a `paddings` in `iter_jax_batches()`."
|
||||
)
|
||||
|
||||
if min_batch_size % num_local_devices != 0:
|
||||
raise ValueError(
|
||||
f"The globally minimum batch size ({min_batch_size}) must be evenly "
|
||||
f"divisible by the number of local JAX devices "
|
||||
f"({num_local_devices}) on this host. "
|
||||
f"To safely truncate or pad the batch, "
|
||||
f"set `drop_last=True` or provide a `paddings` in `iter_jax_batches()`."
|
||||
)
|
||||
|
||||
batch = local_batches[i]
|
||||
assert batch is not None
|
||||
yield _convert_batch(batch, sharding, dtypes=dtypes)
|
||||
|
||||
|
||||
def jax_sync_generator(
|
||||
batch_iterable: Iterable[NumpyBatch],
|
||||
drop_last: bool,
|
||||
batch_size: int = 256,
|
||||
paddings: Optional[PaddingsSpec] = None,
|
||||
dtypes: Optional[DTypeLikeSpec] = None,
|
||||
synchronize_batches: bool = False,
|
||||
synchronize_lookahead: int = 10,
|
||||
) -> Iterator[JaxBatch]:
|
||||
"""A generator that synchronizes and shards batches across JAX workers.
|
||||
|
||||
This generator wraps a locally yielded batch iterable and ensures that all JAX
|
||||
workers within a multi-host training setup receive the exact same number of batches
|
||||
and identical batch shapes, which is required for JAX's SPMD execution.
|
||||
|
||||
Args:
|
||||
batch_iterable: An iterable yielding local data batches (either a NumPy ndarray
|
||||
or a dictionary of NumPy ndarrays).
|
||||
drop_last: Whether to drop partial or uneven batches.
|
||||
batch_size: The target batch size for each host.
|
||||
paddings: The value to use for padding uneven batches to `batch_size`.
|
||||
If a dictionary is provided, it must map column names to padding values.
|
||||
If None, padding is disabled.
|
||||
dtypes: A single JAX dtype or dictionary of JAX dtypes for the created arrays.
|
||||
synchronize_batches: Whether to synchronize batch shapes across all hosts.
|
||||
Setting this to False can improve performance if you guarantee that all
|
||||
hosts produce identical batch shapes and counts beforehand.
|
||||
synchronize_lookahead: The number of batches to look ahead and synchronize at
|
||||
once. Increasing this value reduces synchronization overhead but may
|
||||
increase memory usage as more batches are buffered locally.
|
||||
|
||||
Yields:
|
||||
JaxBatch: Globally sharded batches.
|
||||
"""
|
||||
import jax
|
||||
|
||||
# Physical Sharding (1D across the _GLOBAL_MESH_1D_AXIS dimension)
|
||||
# The sharding is created once for the lifetime of this generator and reused
|
||||
# across all batches.
|
||||
sharding = _create_sharding_1d(_GLOBAL_MESH_1D_AXIS)
|
||||
|
||||
num_local_devices = jax.local_device_count()
|
||||
iterator = iter(batch_iterable)
|
||||
|
||||
if not synchronize_batches or jax.process_count() == 1:
|
||||
yield from _yield_batches_no_sync(
|
||||
iterator,
|
||||
sharding,
|
||||
num_local_devices,
|
||||
batch_size,
|
||||
paddings,
|
||||
dtypes=dtypes,
|
||||
)
|
||||
else:
|
||||
yield from _yield_batches_with_sync(
|
||||
iterator,
|
||||
sharding,
|
||||
num_local_devices,
|
||||
drop_last,
|
||||
batch_size,
|
||||
paddings,
|
||||
synchronize_lookahead,
|
||||
dtypes=dtypes,
|
||||
)
|
||||
@@ -0,0 +1,497 @@
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import pyarrow
|
||||
import torch
|
||||
|
||||
from ray._common.utils import env_bool
|
||||
from ray.data.collate_fn import (
|
||||
TensorBatchReturnType,
|
||||
TensorBatchType,
|
||||
_is_nested_tensor_sequence,
|
||||
_is_tensor,
|
||||
_is_tensor_mapping,
|
||||
_is_tensor_sequence,
|
||||
_is_tensor_sequence_mapping,
|
||||
)
|
||||
from ray.data.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
|
||||
|
||||
# Default non-blocking transfer for tensors.
|
||||
DEFAULT_TENSOR_NON_BLOCKING_TRANSFER = env_bool(
|
||||
"RAY_AIR_DEFAULT_TENSOR_NON_BLOCKING_TRANSFER",
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
def convert_table_to_torch_tensor(
|
||||
data_batch: pyarrow.Table,
|
||||
columns: Optional[Union[List[str], List[List[str]]]] = None,
|
||||
column_dtypes: Optional[Union[torch.dtype, List[torch.dtype]]] = None,
|
||||
unsqueeze: bool = True,
|
||||
) -> Union[torch.Tensor, List[torch.Tensor]]:
|
||||
"""Converts a PyArrow table to a torch Tensor or list of torch Tensors.
|
||||
|
||||
The return type matches the format of ``columns``: a flat list of column
|
||||
names produces a single tensor; a list of lists produces a list of tensors
|
||||
(one per group). If ``columns`` is None, all columns are used.
|
||||
|
||||
Args:
|
||||
data_batch: The PyArrow table to convert.
|
||||
columns: Column names to include. A list of lists returns one tensor
|
||||
per group (useful for multi-input models). None uses all columns.
|
||||
column_dtypes: Torch dtype(s) for the output. A single dtype applies
|
||||
to all columns/groups. A list must match the number of groups when
|
||||
``columns`` is a list of lists.
|
||||
unsqueeze: If True, reshape each per-column tensor from (N,) to (N, 1)
|
||||
before concatenating. Defaults to True.
|
||||
|
||||
Returns:
|
||||
A single tensor of shape (N, len(columns)), or a list of tensors when
|
||||
``columns`` is a list of lists.
|
||||
"""
|
||||
multi_input = columns and isinstance(columns[0], (list, tuple))
|
||||
|
||||
if columns is None:
|
||||
columns = data_batch.column_names
|
||||
|
||||
if not multi_input and column_dtypes and not isinstance(column_dtypes, torch.dtype):
|
||||
raise TypeError(
|
||||
"If `columns` is a list of strings, "
|
||||
"`column_dtypes` must be None or a single `torch.dtype`."
|
||||
f"Got {type(column_dtypes)} instead."
|
||||
)
|
||||
|
||||
if multi_input:
|
||||
if not isinstance(column_dtypes, (list, tuple)):
|
||||
column_dtypes = [column_dtypes] * len(columns)
|
||||
|
||||
return [
|
||||
_columns_to_tensor(data_batch, group, dtype, unsqueeze)
|
||||
for group, dtype in zip(columns, column_dtypes)
|
||||
]
|
||||
|
||||
return _columns_to_tensor(data_batch, columns, column_dtypes, unsqueeze)
|
||||
|
||||
|
||||
def _columns_to_tensor(
|
||||
table: pyarrow.Table,
|
||||
column_names: List[str],
|
||||
dtype: Optional[torch.dtype],
|
||||
unsqueeze: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Convert selected columns from a PyArrow table into a single tensor."""
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
|
||||
numpy_batch = ArrowBlockAccessor(table).to_numpy(columns=column_names)
|
||||
|
||||
tensors = []
|
||||
for col in column_names:
|
||||
try:
|
||||
t = convert_ndarray_to_torch_tensor(numpy_batch[col], dtype)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Failed to convert column {col} to a Torch Tensor of dtype "
|
||||
f"{dtype}. See above exception chain for the exact failure."
|
||||
) from e
|
||||
|
||||
if unsqueeze:
|
||||
t = t.unsqueeze(1)
|
||||
|
||||
tensors.append(t)
|
||||
|
||||
if len(tensors) > 1:
|
||||
return torch.cat(tensors, dim=1)
|
||||
|
||||
return tensors[0]
|
||||
|
||||
|
||||
def convert_ndarray_to_torch_tensor(
|
||||
ndarray: np.ndarray,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
pin_memory: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Convert a NumPy ndarray to a Torch Tensor.
|
||||
|
||||
Args:
|
||||
ndarray: A NumPy ndarray that we wish to convert to a Torch Tensor.
|
||||
dtype: A Torch dtype for the created tensor; if None, the dtype will be
|
||||
inferred from the NumPy ndarray data.
|
||||
device: The device on which the tensor(s) should be placed; if None, the Torch
|
||||
tensor(s) will be constructed on the CPU.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
|
||||
Returns:
|
||||
A Torch Tensor.
|
||||
"""
|
||||
ndarray = _unwrap_ndarray_object_type_if_needed(ndarray)
|
||||
|
||||
# Object dtype cannot be converted into PyTorch Tensor.
|
||||
if ndarray.dtype.type is np.object_:
|
||||
raise RuntimeError(
|
||||
"Numpy array of object dtype cannot be converted to a Torch Tensor. This "
|
||||
"may because the numpy array is a ragged tensor--it contains items of "
|
||||
"different sizes. If using `iter_torch_batches()` API, you can pass in a "
|
||||
"`collate_fn` argument to specify custom logic to convert the Numpy array "
|
||||
"batch to a Torch tensor batch."
|
||||
)
|
||||
|
||||
# The numpy array is not always writeable as it can come from the Ray object store.
|
||||
# Numpy will throw a verbose warning here, which we suppress, as we don't write
|
||||
# to the tensors. We also don't want to copy the array to avoid memory overhead.
|
||||
# Original warning: https://github.com/pytorch/pytorch/blob/v1.13.0/
|
||||
# torch/csrc/utils/tensor_numpy.cpp#L198-L206
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
result = torch.as_tensor(ndarray, dtype=dtype, device=device)
|
||||
|
||||
if pin_memory:
|
||||
assert result.device.type == "cpu", (
|
||||
"Pin memory is only supported for CPU tensors. "
|
||||
f"Got device: {result.device} and pin_memory: {pin_memory}."
|
||||
)
|
||||
result = result.pin_memory()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def convert_ndarray_batch_to_torch_tensor_batch(
|
||||
ndarrays: Union[np.ndarray, Dict[str, np.ndarray]],
|
||||
dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
pin_memory: bool = False,
|
||||
) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
"""Convert a NumPy ndarray batch to a Torch Tensor batch.
|
||||
|
||||
Args:
|
||||
ndarrays: A (dict of) NumPy ndarray(s) that we wish to convert to a Torch Tensor.
|
||||
dtypes: A (dict of) Torch dtype(s) for the created tensor; if None, the dtype
|
||||
will be inferred from the NumPy ndarray data.
|
||||
device: The device on which the tensor(s) should be placed; if None, the Torch
|
||||
tensor(s) will be constructed on the CPU.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
|
||||
Returns:
|
||||
A (dict of) Torch Tensor(s).
|
||||
"""
|
||||
if isinstance(ndarrays, np.ndarray):
|
||||
# Single-tensor case.
|
||||
if isinstance(dtypes, dict):
|
||||
if len(dtypes) != 1:
|
||||
raise ValueError(
|
||||
"When constructing a single-tensor batch, only a single dtype "
|
||||
f"should be given, instead got: {dtypes}"
|
||||
)
|
||||
dtypes = next(iter(dtypes.values()))
|
||||
batch = convert_ndarray_to_torch_tensor(
|
||||
ndarrays,
|
||||
dtype=dtypes,
|
||||
device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
else:
|
||||
# Multi-tensor case.
|
||||
batch = {
|
||||
col_name: convert_ndarray_to_torch_tensor(
|
||||
col_ndarray,
|
||||
dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for col_name, col_ndarray in ndarrays.items()
|
||||
}
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
def convert_ndarray_list_to_torch_tensor_list(
|
||||
ndarrays: Dict[str, List[np.ndarray]],
|
||||
dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
pin_memory: bool = False,
|
||||
) -> Dict[str, List[torch.Tensor]]:
|
||||
"""Convert a dict mapping column names to lists of ndarrays to Torch Tensors.
|
||||
|
||||
Args:
|
||||
ndarrays: A dict mapping column names to lists of ndarrays that we wish to convert
|
||||
to Torch Tensors.
|
||||
dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype
|
||||
will be inferred from the NumPy ndarray data.
|
||||
device: The device on which the tensor(s) should be placed; if None, the Torch
|
||||
tensor(s) will be constructed on the CPU.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
|
||||
Returns:
|
||||
A dict mapping column names to lists of Tensors.
|
||||
"""
|
||||
return {
|
||||
col_name: [
|
||||
convert_ndarray_batch_to_torch_tensor_batch(
|
||||
ndarray,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for ndarray in col_ndarrays
|
||||
]
|
||||
for col_name, col_ndarrays in ndarrays.items()
|
||||
}
|
||||
|
||||
|
||||
def arrow_batch_to_tensors(
|
||||
batch: pyarrow.Table,
|
||||
dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
|
||||
combine_chunks: bool = False,
|
||||
pin_memory: bool = False,
|
||||
threadpool: Optional[ThreadPoolExecutor] = None,
|
||||
) -> Union[Dict[str, torch.Tensor], Dict[str, List[torch.Tensor]]]:
|
||||
"""Convert PyArrow batch to PyTorch tensors.
|
||||
|
||||
Args:
|
||||
batch: PyArrow batch to convert
|
||||
dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype
|
||||
will be inferred from the NumPy ndarray data.
|
||||
combine_chunks: If True, combine chunks in Arrow batch before converting to
|
||||
tensors.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
threadpool: Optional ThreadPoolExecutor for parallel processing. If provided,
|
||||
columns/arrays will be processed in parallel. If None, processing is
|
||||
sequential.
|
||||
|
||||
Returns:
|
||||
When combine_chunks=True: A dictionary of column name to single tensor.
|
||||
When combine_chunks=False: A dictionary of column name to list of tensors.
|
||||
"""
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
from ray.data._internal.arrow_ops import transform_pyarrow
|
||||
|
||||
if combine_chunks:
|
||||
numpy_batch = ArrowBlockAccessor(batch).to_batch_format("numpy")
|
||||
num_columns = len(numpy_batch)
|
||||
|
||||
if num_columns > 1 and threadpool is not None:
|
||||
# Process columns in parallel using provided threadpool
|
||||
def process_column(
|
||||
col_name_col_array: Tuple[str, np.ndarray]
|
||||
) -> Tuple[str, torch.Tensor]:
|
||||
col_name, col_array = col_name_col_array
|
||||
return col_name, convert_ndarray_batch_to_torch_tensor_batch(
|
||||
col_array,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
|
||||
# Submit all columns to threadpool and collect results
|
||||
processed_cols = threadpool.map(process_column, numpy_batch.items())
|
||||
return dict(processed_cols)
|
||||
else:
|
||||
# Sequential processing for single column or single worker
|
||||
return {
|
||||
col_name: convert_ndarray_batch_to_torch_tensor_batch(
|
||||
col_array,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for col_name, col_array in numpy_batch.items()
|
||||
}
|
||||
else:
|
||||
numpy_list = transform_pyarrow.table_to_numpy_dict_chunked(
|
||||
batch,
|
||||
)
|
||||
# Count total number of arrays across all columns
|
||||
total_arrays = sum(len(arrays) for arrays in numpy_list.values())
|
||||
num_columns = len(numpy_list)
|
||||
|
||||
if total_arrays > 1 and threadpool is not None:
|
||||
# Process arrays in parallel using provided threadpool
|
||||
def process_array(
|
||||
array_item: Tuple[str, int, np.ndarray]
|
||||
) -> Tuple[str, int, torch.Tensor]:
|
||||
col_name, array_index, array = array_item
|
||||
return (
|
||||
col_name,
|
||||
array_index,
|
||||
convert_ndarray_batch_to_torch_tensor_batch(
|
||||
array,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
pin_memory=pin_memory,
|
||||
),
|
||||
)
|
||||
|
||||
# Flatten arrays with column name and index for parallel processing
|
||||
array_items = [
|
||||
(col_name, idx, array)
|
||||
for col_name, arrays in numpy_list.items()
|
||||
for idx, array in enumerate(arrays)
|
||||
]
|
||||
|
||||
# Submit all arrays to threadpool and collect results
|
||||
processed_arrays = list(threadpool.map(process_array, array_items))
|
||||
|
||||
# Initialize result with all columns from numpy_list, including empty ones
|
||||
# Pre-allocate lists of the correct size for each column
|
||||
result: Dict[str, List[torch.Tensor]] = {
|
||||
col_name: [None] * len(arrays)
|
||||
for col_name, arrays in numpy_list.items()
|
||||
}
|
||||
|
||||
# Populate result with processed tensors
|
||||
for col_name, array_index, tensor in processed_arrays:
|
||||
result[col_name][array_index] = tensor
|
||||
|
||||
return result
|
||||
else:
|
||||
# Sequential processing
|
||||
return convert_ndarray_list_to_torch_tensor_list(
|
||||
numpy_list,
|
||||
dtypes=dtypes,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def concat_tensors_to_device(
|
||||
tensor_sequence: Sequence[torch.Tensor],
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER,
|
||||
) -> torch.Tensor:
|
||||
"""Stack sequence of tensors into a contiguous GPU tensor.
|
||||
|
||||
Args:
|
||||
tensor_sequence: Sequence of tensors to stack
|
||||
device: The device to move tensors to. If None, tensors are not moved.
|
||||
non_blocking: If True, perform device transfer without forcing a
|
||||
synchronization.
|
||||
|
||||
Returns:
|
||||
A contiguous tensor on the target device
|
||||
"""
|
||||
# Assumes tensors have the same shape/dtype
|
||||
assert (
|
||||
tensor_sequence
|
||||
), f"Cannot stack empty sequence of tensors. Received: {tensor_sequence}"
|
||||
|
||||
assert all(
|
||||
isinstance(t, torch.Tensor) for t in tensor_sequence
|
||||
), "All items must be torch.Tensor. Found invalid types: " + str(
|
||||
[type(t) for t in tensor_sequence if not isinstance(t, torch.Tensor)]
|
||||
)
|
||||
|
||||
# If there is only one tensor and its device already matches, return it directly.
|
||||
if len(tensor_sequence) == 1 and (
|
||||
device is None or tensor_sequence[0].device == torch.device(device)
|
||||
):
|
||||
return tensor_sequence[0]
|
||||
|
||||
first_dtype = tensor_sequence[0].dtype
|
||||
assert all(t.dtype == first_dtype for t in tensor_sequence), (
|
||||
"All tensors must have the same dtype. "
|
||||
f"Expected: {first_dtype}, got: {[t.dtype for t in tensor_sequence]}"
|
||||
)
|
||||
|
||||
first_shape = tensor_sequence[0].shape[1:]
|
||||
assert all(t.shape[1:] == first_shape for t in tensor_sequence), (
|
||||
"All tensors must have the same shape[1:]. "
|
||||
f"Expected: {first_shape}, got: {[t.shape[1:] for t in tensor_sequence]}"
|
||||
)
|
||||
|
||||
first = tensor_sequence[0]
|
||||
dtype = first.dtype
|
||||
shape_tail = first.shape[1:]
|
||||
total_rows = sum(t.shape[0] for t in tensor_sequence)
|
||||
|
||||
# Allocate an empty Tensor on device
|
||||
result = torch.empty((total_rows, *shape_tail), dtype=dtype, device=device)
|
||||
|
||||
row_start = 0
|
||||
for t in tensor_sequence:
|
||||
row_end = row_start + t.shape[0]
|
||||
result[row_start:row_end].copy_(t, non_blocking=non_blocking)
|
||||
row_start = row_end
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_type_str(batch: Any) -> str:
|
||||
"""Get a string representation of the possibly nested type of the batch.
|
||||
|
||||
>>> import torch
|
||||
>>> _get_type_str([1, 2, "???"])
|
||||
'list[int | str]'
|
||||
>>> _get_type_str({"a": [1, 2, 3], "b": 4})
|
||||
'dict[str, int | list[int]]'
|
||||
>>> _get_type_str({"a": torch.tensor(1), "b": [torch.tensor(2)]})
|
||||
'dict[str, Tensor | list[Tensor]]'
|
||||
>>> _get_type_str({"a": torch.tensor(1), "b": {"c": torch.tensor(2)}})
|
||||
'dict[str, Tensor | dict[str, Tensor]]'
|
||||
"""
|
||||
curr_type = type(batch).__name__
|
||||
if isinstance(batch, (list, tuple)):
|
||||
val_types = " | ".join(sorted({_get_type_str(v) for v in batch}))
|
||||
invalid_type_str = f"{curr_type}[{val_types}]"
|
||||
elif isinstance(batch, dict):
|
||||
val_types = " | ".join(sorted({_get_type_str(v) for v in batch.values()}))
|
||||
invalid_type_str = f"{curr_type}[str, {val_types}]"
|
||||
else:
|
||||
invalid_type_str = curr_type
|
||||
return invalid_type_str
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def move_tensors_to_device(
|
||||
batch: TensorBatchType,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER,
|
||||
) -> TensorBatchReturnType:
|
||||
"""Move tensors to the specified device.
|
||||
|
||||
Concatenate nested lists/tuples of tensors along the first (batch) dimension.
|
||||
For example, for the input
|
||||
((feature_0_chunk_0,), (feature_1_chunk_0, feature_1_chunk_1))
|
||||
the output will be (feature_0_chunk_0, feature_1_chunk_0+1)
|
||||
where each feature is concatenated along the batch dimension.
|
||||
|
||||
Args:
|
||||
batch: A tensor or collection of tensors to move to device. Can be:
|
||||
- A single tensor
|
||||
- A sequence of tensors
|
||||
- A sequence of sequences of tensors. The inner sequence of tensors is
|
||||
combined during GPU transfer.
|
||||
- A mapping (e.g., dict) of keys to tensors or sequences of tensors. The
|
||||
sequence of tensors is combined during GPU transfer.
|
||||
device: The device to move tensors to. If None, tensors are not moved.
|
||||
non_blocking: If True, perform device transfer without forcing a
|
||||
synchronization.
|
||||
|
||||
Returns:
|
||||
The input tensors moved to the specified device
|
||||
"""
|
||||
if device is None:
|
||||
return batch
|
||||
|
||||
if _is_tensor(batch):
|
||||
return batch.to(device, non_blocking=non_blocking)
|
||||
elif _is_tensor_sequence(batch):
|
||||
return type(batch)([t.to(device, non_blocking=non_blocking) for t in batch])
|
||||
elif _is_nested_tensor_sequence(batch):
|
||||
return type(batch)(
|
||||
[concat_tensors_to_device(t, device, non_blocking) for t in batch]
|
||||
)
|
||||
elif _is_tensor_mapping(batch):
|
||||
return {k: t.to(device, non_blocking=non_blocking) for k, t in batch.items()}
|
||||
elif _is_tensor_sequence_mapping(batch):
|
||||
return {
|
||||
k: concat_tensors_to_device(v, device, non_blocking)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid input type: {_get_type_str(batch)}.\n"
|
||||
"Expected one of the following: "
|
||||
"torch.Tensor, "
|
||||
"List/Tuple[torch.Tensor], "
|
||||
"Dict[str, torch.Tensor], "
|
||||
"Mapping[str, List/Tuple[torch.Tensor]]"
|
||||
)
|
||||
Reference in New Issue
Block a user