[AIR - Datasets] Hide tensor extension from UDFs. #27019

clarkzinzow · 2022-07-26T17:14:46Z

We previously added automatic tensor extension casting on Datasets transformation outputs to allow the user to not have to worry about tensor column casting; however, this current state creates several issues:

Not all tensors are supported, which means that we’ll need to have an opaque object dtype (i.e. ndarray of ndarray pointers) fallback for the Pandas-only case. Known unsupported tensor use cases:
a. Heterogeneous-shaped (i.e. ragged) tensors
b. Struct arrays
UDFs will expect a NumPy column and won’t know what to do with our TensorArray type. E.g., torchvision transforms don’t respect the array protocol (which they should), and instead only support Torch tensors and NumPy ndarrays; passing a TensorArray column or a TensorArrayElement (a single item in the TensorArray column) fails.
Implicit casting with object dtype fallback on UDF outputs can make the input type to downstream UDFs nondeterministic, where the user won’t know if they’ll get a TensorArray column or an object dtype column.
The tensor extension cast fallback warning spams the logs.

This PR:

Adds automatic casting of tensor extension columns to NumPy ndarray columns for Datasets UDF inputs, meaning the UDFs will never have to see tensor extensions and that the UDF input column types will be consistent and deterministic; this fixes both (2) and (3).
No longer implicitly falls back to an opaque object dtype when TensorArray casting fails (e.g. for ragged tensors), and instead raises an error; this fixes (4) but removes our support for (1).
Adds a global enable_tensor_extension_casting config flag, which is True by default, that controls whether we perform this automatic casting. Turning off the implicit casting provides a path for (1), where the tensor extension can be avoided if working with ragged tensors in Pandas land. Turning off this flag also allows the user to explicitly control their tensor extension casting, if they want to work with it in their UDFs in order to reap the benefits of less data copies, more efficient slicing, stronger column typing, etc.

Checks

I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

amogkam

Thanks @clarkzinzow!

python/ray/air/util/data_batch_conversion.py

python/ray/data/datasource/image_folder_datasource.py

python/ray/data/preprocessors/concatenator.py

python/ray/air/util/data_batch_conversion.py

python/ray/air/tests/test_data_batch_conversion.py

python/ray/data/preprocessors/concatenator.py

python/ray/data/datasource/image_folder_datasource.py

jiaodong · 2022-07-26T18:59:34Z

thanks for quick progress on this ! Can you also ensure to kick off nightly release test on this PR with the new API: https://sourcegraph.com/github.com/ray-project/ray/-/blob/release/air_tests/air_benchmarks/workloads/gpu_batch_prediction.py

for batch prediction ? Training should be ok given the example you changed covered incremental learning.

Depending Balaji's PR #26975 if he merges first the notebooks could change, as fyi.

amogkam · 2022-07-26T20:15:44Z

On a meta point- I'm struggling a bit to reconcile the performance and UX aspects of this.

From the Dataset perspective, preprocessors and predictors are udfs, which we say should not need to use TensorArray (as Datasets will automatically handle it). Yet, we still should be using TensorArray in Predictors and Preprocessors for performance reasons.

As a user or as a developer seems like I still need to be aware of TensorArray for performance reasons, so it's not completely hidden away. Let's keep iterating on this in the future, to see if there is a solution that can improve the UX without sacrificing performance.

jiaodong · 2022-07-26T20:27:47Z

Let's keep iterating on this in the future, to see if there is a solution that can improve the UX without sacrificing performance.

+1 I also think we should go deeper regarding how to use TensorArray in the right way for both end user and perf. So far we have consensus on UX (hide it from user) but still missing on perf implications (for AIR developer) and future plans (support ragged tensor-like with right UX and perf)

clarkzinzow · 2022-07-26T20:29:58Z

Let's keep iterating on this in the future, to see if there is a solution that can improve the UX without sacrificing performance.

@amogkam The ideal solution is pretty simple: Pandas needs to finish making extension types first-class types (they have a few bugs/missing features here), Torch and TensorFlow need to respect the __array__ protocol which has become standard, and we need to support heterogeneous-shaped tensors in our extension type. At that point, we can always use the tensor extension type, and the user will never have to be aware of the extension type since conversion to/from NumPy/Torch/TensorFlow will happen automatically via the __array__ protocol and the Pandas/Arrow extension type registries.

clarkzinzow · 2022-07-26T20:32:41Z

The current UX push is to make users extension type agnostic while providing the ability to opt-in if they're worried about performance or memory usage, while ensuring that we, as internal library developers, use the extension type in the appropriate places if it results in free performance/efficiency wins. I think this is the current best compromise between end-user UX and perf/efficiency by default.

amogkam · 2022-07-26T20:36:19Z

Yes both your points make sense! Agreed this is a good compromise for now.

clarkzinzow · 2022-07-27T01:49:42Z

@amogkam @jiaodong @matthewdeng @ericl @jianoaix ready for another pass!

clarkzinzow · 2022-07-27T17:06:45Z

doc/source/ray-air/examples/torch_incremental_learning.ipynb

-    "    ])\n",
+    "    df.loc[:, \"image\"] = [\n",
+    "        torchvision_transforms(image).numpy() for image in df[\"image\"]\n",
+    "    ]\n",


@jiaodong Note that this is tensor extension type agnostic, but we still need to convert the Torch tensors back to NumPy ndarrays before storing them in the batch.

ack, having UDF purely tensor type agnostic and working with numpy + df is already big dex ux improvement :) I think modern DL frameworks all have good zero-copy numpy iterop so it should be much less concerning.

im curious what would happen to existing user code / example / tests if they still wrapped output / input with TensorArray, or they're just working with Ragged tensor but also didn't explicitly disable auto casting ?

clarkzinzow · 2022-07-27T17:07:31Z

python/ray/air/_internal/tensorflow_utils.py

+    """
+    if ndarray.dtype.type is np.object_:
+        try:
+            # Try to convert the NumPy ndarray to a non-object dtype.


This allows us to support the non-Datasets (and non-tensor extension type) Serve use case.

clarkzinzow · 2022-07-27T17:10:22Z

python/ray/data/_internal/pandas_block.py

-                            "TensorArray but the conversion failed, leaving column "
-                            f"as-is: {e}"
-                        )
+        df.reset_index(drop=True, inplace=True)


Slicing opaque object columns fails silently in user-code if we don't reset the index here (Pandas foot-gun that our tensor extension type avoids).

amogkam

Thanks @clarkzinzow!

python/ray/air/tests/test_data_batch_conversion.py

python/ray/air/util/data_batch_conversion.py

python/ray/train/predictor.py

jianoaix

LG for ray/data

ericl · 2022-07-27T21:58:58Z

Some test failures:

data = array([array([1, 2]), array([3, 4]), array([5, 6])], dtype=object)
  |  
  | def _transform_np_array(data: np.ndarray):
  | if not isinstance(data, np.ndarray) and hasattr(data, '__array__'):
  | data = np.array(data, copy=False)
  | if len(data.shape) != 2:
  | raise ValueError('Expecting 2 dimensional numpy.ndarray, got: ',
  | >                            data.shape)
  | E           ValueError: ('Expecting 2 dimensional numpy.ndarray, got: ', (3,))

clarkzinzow · 2022-07-27T22:53:30Z

@ericl Yep I'm looking into them

clarkzinzow · 2022-07-28T17:37:09Z

@matthewdeng @ericl All tests and examples are passing, merging to bypass the slow MacOS queue.

We previously added automatic tensor extension casting on Datasets transformation outputs to allow the user to not have to worry about tensor column casting; however, this current state creates several issues: 1. Not all tensors are supported, which means that we’ll need to have an opaque object dtype (i.e. ndarray of ndarray pointers) fallback for the Pandas-only case. Known unsupported tensor use cases: a. Heterogeneous-shaped (i.e. ragged) tensors b. Struct arrays 2. UDFs will expect a NumPy column and won’t know what to do with our TensorArray type. E.g., torchvision transforms don’t respect the array protocol (which they should), and instead only support Torch tensors and NumPy ndarrays; passing a TensorArray column or a TensorArrayElement (a single item in the TensorArray column) fails. Implicit casting with object dtype fallback on UDF outputs can make the input type to downstream UDFs nondeterministic, where the user won’t know if they’ll get a TensorArray column or an object dtype column. 3. The tensor extension cast fallback warning spams the logs. This PR: 1. Adds automatic casting of tensor extension columns to NumPy ndarray columns for Datasets UDF inputs, meaning the UDFs will never have to see tensor extensions and that the UDF input column types will be consistent and deterministic; this fixes both (2) and (3). 2. No longer implicitly falls back to an opaque object dtype when TensorArray casting fails (e.g. for ragged tensors), and instead raises an error; this fixes (4) but removes our support for (1). 3. Adds a global enable_tensor_extension_casting config flag, which is True by default, that controls whether we perform this automatic casting. Turning off the implicit casting provides a path for (1), where the tensor extension can be avoided if working with ragged tensors in Pandas land. Turning off this flag also allows the user to explicitly control their tensor extension casting, if they want to work with it in their UDFs in order to reap the benefits of less data copies, more efficient slicing, stronger column typing, etc.

We previously added automatic tensor extension casting on Datasets transformation outputs to allow the user to not have to worry about tensor column casting; however, this current state creates several issues: 1. Not all tensors are supported, which means that we’ll need to have an opaque object dtype (i.e. ndarray of ndarray pointers) fallback for the Pandas-only case. Known unsupported tensor use cases: a. Heterogeneous-shaped (i.e. ragged) tensors b. Struct arrays 2. UDFs will expect a NumPy column and won’t know what to do with our TensorArray type. E.g., torchvision transforms don’t respect the array protocol (which they should), and instead only support Torch tensors and NumPy ndarrays; passing a TensorArray column or a TensorArrayElement (a single item in the TensorArray column) fails. Implicit casting with object dtype fallback on UDF outputs can make the input type to downstream UDFs nondeterministic, where the user won’t know if they’ll get a TensorArray column or an object dtype column. 3. The tensor extension cast fallback warning spams the logs. This PR: 1. Adds automatic casting of tensor extension columns to NumPy ndarray columns for Datasets UDF inputs, meaning the UDFs will never have to see tensor extensions and that the UDF input column types will be consistent and deterministic; this fixes both (2) and (3). 2. No longer implicitly falls back to an opaque object dtype when TensorArray casting fails (e.g. for ragged tensors), and instead raises an error; this fixes (4) but removes our support for (1). 3. Adds a global enable_tensor_extension_casting config flag, which is True by default, that controls whether we perform this automatic casting. Turning off the implicit casting provides a path for (1), where the tensor extension can be avoided if working with ragged tensors in Pandas land. Turning off this flag also allows the user to explicitly control their tensor extension casting, if they want to work with it in their UDFs in order to reap the benefits of less data copies, more efficient slicing, stronger column typing, etc. Signed-off-by: Stefan van der Kleij <s.vanderkleij@viroteq.com>

clarkzinzow requested review from ericl, scv119, jjyao, jianoaix, maxpumperla, pcmoritz, richardliaw, edoakes and simon-mo as code owners July 26, 2022 17:14

clarkzinzow assigned ericl, matthewdeng, jiaodong and amogkam Jul 26, 2022

amogkam reviewed Jul 26, 2022

View reviewed changes

jiaodong reviewed Jul 26, 2022

View reviewed changes

python/ray/data/preprocessors/concatenator.py Outdated Show resolved Hide resolved

python/ray/data/datasource/image_folder_datasource.py Outdated Show resolved Hide resolved

clarkzinzow force-pushed the air/fix/hide-tensor-extension branch 3 times, most recently from 5abe09a to b730545 Compare July 27, 2022 01:41

clarkzinzow added 2 commits July 27, 2022 01:44

Hide tensor extension from UDFs.

77e0ca7

PR feedback

fef615e

clarkzinzow force-pushed the air/fix/hide-tensor-extension branch from b730545 to fef615e Compare July 27, 2022 01:49

clarkzinzow assigned jianoaix Jul 27, 2022

Fixes

a48e34d

clarkzinzow commented Jul 27, 2022

View reviewed changes

amogkam approved these changes Jul 27, 2022

View reviewed changes

python/ray/air/tests/test_data_batch_conversion.py Show resolved Hide resolved

python/ray/air/util/data_batch_conversion.py Show resolved Hide resolved

python/ray/train/predictor.py Outdated Show resolved Hide resolved

python/ray/train/predictor.py Outdated Show resolved Hide resolved

jianoaix approved these changes Jul 27, 2022

View reviewed changes

PR feedback

9b951e3

ericl approved these changes Jul 27, 2022

View reviewed changes

ericl added the v2.0.0-pick label Jul 27, 2022

ericl added the @author-action-required label Jul 27, 2022

clarkzinzow added 2 commits July 27, 2022 23:37

Fixes

f17932b

Fix example.

7ad8c88

clarkzinzow merged commit df124d0 into ray-project:master Jul 28, 2022

clarkzinzow removed the @author-action-required label Jul 28, 2022

clarkzinzow mentioned this pull request Jul 28, 2022

[Cherry-pick] [AIR - Datasets] Hide tensor extension from UDFs. #27196

Merged

2 tasks

scv119 added the v2.0.0-pick-done label Aug 3, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AIR - Datasets] Hide tensor extension from UDFs. #27019

[AIR - Datasets] Hide tensor extension from UDFs. #27019

clarkzinzow commented Jul 26, 2022 •

edited

Loading

amogkam left a comment

jiaodong commented Jul 26, 2022 •

edited

Loading

amogkam commented Jul 26, 2022 •

edited

Loading

jiaodong commented Jul 26, 2022

clarkzinzow commented Jul 26, 2022

clarkzinzow commented Jul 26, 2022

amogkam commented Jul 26, 2022

clarkzinzow commented Jul 27, 2022

clarkzinzow Jul 27, 2022

jiaodong Jul 27, 2022

jiaodong Jul 27, 2022

clarkzinzow Jul 27, 2022

clarkzinzow Jul 27, 2022

amogkam left a comment

jianoaix left a comment

ericl commented Jul 27, 2022

clarkzinzow commented Jul 27, 2022

clarkzinzow commented Jul 28, 2022

[AIR - Datasets] Hide tensor extension from UDFs. #27019

[AIR - Datasets] Hide tensor extension from UDFs. #27019

Conversation

clarkzinzow commented Jul 26, 2022 • edited Loading

Checks

amogkam left a comment

Choose a reason for hiding this comment

jiaodong commented Jul 26, 2022 • edited Loading

amogkam commented Jul 26, 2022 • edited Loading

jiaodong commented Jul 26, 2022

clarkzinzow commented Jul 26, 2022

clarkzinzow commented Jul 26, 2022

amogkam commented Jul 26, 2022

clarkzinzow commented Jul 27, 2022

clarkzinzow Jul 27, 2022

Choose a reason for hiding this comment

jiaodong Jul 27, 2022

Choose a reason for hiding this comment

jiaodong Jul 27, 2022

Choose a reason for hiding this comment

clarkzinzow Jul 27, 2022

Choose a reason for hiding this comment

clarkzinzow Jul 27, 2022

Choose a reason for hiding this comment

amogkam left a comment

Choose a reason for hiding this comment

jianoaix left a comment

Choose a reason for hiding this comment

ericl commented Jul 27, 2022

clarkzinzow commented Jul 27, 2022

clarkzinzow commented Jul 28, 2022

clarkzinzow commented Jul 26, 2022 •

edited

Loading

jiaodong commented Jul 26, 2022 •

edited

Loading

amogkam commented Jul 26, 2022 •

edited

Loading