[air] Use filesystem wrapper to exclude files from upload #34102

krfricke · 2023-04-05T17:16:34Z

Why are these changes needed?

Ray Tune uploads experiment state using pyarrow. When cloud checkpointing is configured, the driver will exclude any trial-level checkpoints. Pyarrow does not natively support file exclusion, though - instead, we repeatedly call pyarrow.fs.copy_files on single non-excluded files.

This seems to be inefficient as the connection to the remote filesystem is opened and closed repeatedly. It also means we can never leverage multi-threaded upload. This PR implements a custom fsspec-based local filesystem that excludes files on the selector level. Thus, we can call pyarrow.fs.copy_files exactly once, with a selector that does not see the excluded files.

Edit: See here for benchmark results

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
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 added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
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 :(

Signed-off-by: Kai Fricke <[email protected]>

krfricke · 2023-04-05T17:25:17Z

I'll run the restore example with these wheels to check if it improves performance

krfricke · 2023-04-05T17:45:46Z

It does not seem to significantly impact the upload time. I still think this is generally a cleaner version to go for, but I'll look more into this.
Converting to draft for now pending further investigation.

krfricke · 2023-04-21T13:35:32Z

I ran a few benchmarks to compare different versions of our code.

Data: 5 files a 2MB in one folder

Compared methods:

awscli_sync: Direct usage of aws s3 sync
pyarrow_threads: Direct usage of pyarrow.fs.copy_files(source, target, use_threads=False)
pyarrow_no_threads: Direct usage of pyarrow.fs.copy_files(source, target, use_threads=True)
air_default: Usage of AIR's upload_to_uri(source, target) utility. Defaults to no-threads for S3
air_exclude_threads: upload_to_uri(source, target, exclude=something) and with threading enabled
air_exclude_no_threads: upload_to_uri(source, target, exclude=something) and with threading disabled (default)
air_patched_threads: Version of this PR with threading enabled
air_patched_no_threads: Version of this PR with threading disabled.

Results:

Results:
    awscli_sync:         0.76 (0.01) seconds
    pyarrow_threads:     0.69 (0.04) seconds
    pyarrow_no_threads:  1.19 (0.07) seconds
    air_default:         1.19 (0.08) seconds
    air_exclude_threads:         1.43 (0.02) seconds
    air_exclude_no_threads:      1.32 (0.03) seconds
    air_patched_threads:         0.53 (0.01) seconds
    air_patched_no_threads:      1.23 (0.08) seconds

The TLDR of these results is that threading gives an (obvious) advantage.

However, due to apache/arrow#32372, we currently can't enable threading per default.

But, assuming apache/arrow#32372 is fixed at some point, our current exclude logic can't leverage threading, as we call pyarrow.fs.copy_files multiple times, sequentially. Thus we see that air_exclude_threads has the same bad performance as pyarrow_no_threads or air_exclude_no_threads

With the patch from this PR, we can regain the threading benefit, once the pyarrow issue is resolved, as shown in the good performance of air_patched_threads.

Thus, we should aim to land this PR.

I'll work on a workaround for the pyarrow issue separately.

…e-exclude

justinvyu

Looks good! Could we also add a unit test that tests upload_to_uri with and without fsspec? Maybe you can monkeypatch fsspec to None to switch to default.

justinvyu · 2023-04-21T17:09:25Z

python/ray/tune/execution/experiment_state.py

+                fsspec_msg = (
+                    "If your data is small, try installing fsspec "
+                    "(`pip install fsspec`) for more efficient local file parsing. "
+                )


Should ffspec always be recommended, or is it only beneficial for many, small files? Does it perform worse for few, large files?

Initial benchmarks suggested that it always performed better. After #34663 any implementation efficiencies will be hugely shadowed by being able to leverage multi threaded uploads, so we should always recommend fsspec (we'll be printing warnings in #34663 if it's not installed).

justinvyu · 2023-04-21T17:25:29Z

python/ray/tune/execution/experiment_state.py

+
+        self._slow_sync_threshold = float(
+            os.environ.get(
+                "TUNE_WARN_SLOW_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S", "30"


I am wondering if we should always warn users if they don't have fsspec, rather than just if their exp syncing is taking a while? The threshold might mask the problems and the user would never know to do something.

Will there be a separate warning checking if using any cloud storage + no fsspec?

Yes! I have added the warnings in the follow-up PR: #34663

…e-exclude

Signed-off-by: Kai Fricke <[email protected]>

krfricke · 2023-04-24T10:56:28Z

Added a test, ptal!

justinvyu

Thanks for clarifying! One nit on another test case. Also, what's the main difference between the 2 unit tests that were added?

justinvyu · 2023-04-24T17:59:27Z

python/ray/air/tests/test_remote_storage.py

+
+    tmp_source, tmp_target = temp_data_dirs
+
+    upload_to_uri(tmp_source, tmp_target, exclude=["*_exclude.txt", "*_exclude/*"])


Nit: Should we include a test case for an exclude pattern like "exclude/"? I remember there is some special logic for the trailing /.

The tests were copied from existing tests that test the same behavior with an in-memory filesystem. The main difference is that in the second case, multiple patterns are tested (not just one).

The exclude/ pattern won't work as it will not match the files contained within the directory - it always has to be exclude/*

…34663) Following up from #34102, our benchmarks showed how inferior unthreaded syncing is over threaded syncing. However, due to pyarrow, we currently can't use threading. Since the bug affects all pyarrow versions <=11, which will likely be used by some users in the future, we have to look into workarounds. One such workaround is to use the fsspec-provided filesystem and prefer it over the native pyarrow fs. Signed-off-by: Kai Fricke <[email protected]>

…y-project#34102)" This reverts commit 8728c77.

…t#34102) Ray Tune uploads experiment state using pyarrow. When cloud checkpointing is configured, the driver will exclude any trial-level checkpoints. Pyarrow does not natively support file exclusion, though - instead, we repeatedly call `pyarrow.fs.copy_files` on single non-excluded files. This seems to be inefficient as the connection to the remote filesystem is opened and closed repeatedly. It also means we can never leverage multi-threaded upload. This PR implements a custom fsspec-based local filesystem that excludes files on the selector level. Thus, we can call pyarrow.fs.copy_files exactly once, with a selector that does not see the excluded files. Edit: [See here for benchmark results](ray-project#34102 (comment)) Signed-off-by: Kai Fricke <[email protected]> Signed-off-by: Jack He <[email protected]>

…ay-project#34663) Following up from ray-project#34102, our benchmarks showed how inferior unthreaded syncing is over threaded syncing. However, due to pyarrow, we currently can't use threading. Since the bug affects all pyarrow versions <=11, which will likely be used by some users in the future, we have to look into workarounds. One such workaround is to use the fsspec-provided filesystem and prefer it over the native pyarrow fs. Signed-off-by: Kai Fricke <[email protected]> Signed-off-by: Jack He <[email protected]>

…t#34102) Ray Tune uploads experiment state using pyarrow. When cloud checkpointing is configured, the driver will exclude any trial-level checkpoints. Pyarrow does not natively support file exclusion, though - instead, we repeatedly call `pyarrow.fs.copy_files` on single non-excluded files. This seems to be inefficient as the connection to the remote filesystem is opened and closed repeatedly. It also means we can never leverage multi-threaded upload. This PR implements a custom fsspec-based local filesystem that excludes files on the selector level. Thus, we can call pyarrow.fs.copy_files exactly once, with a selector that does not see the excluded files. Edit: [See here for benchmark results](ray-project#34102 (comment)) Signed-off-by: Kai Fricke <[email protected]>

…ay-project#34663) Following up from ray-project#34102, our benchmarks showed how inferior unthreaded syncing is over threaded syncing. However, due to pyarrow, we currently can't use threading. Since the bug affects all pyarrow versions <=11, which will likely be used by some users in the future, we have to look into workarounds. One such workaround is to use the fsspec-provided filesystem and prefer it over the native pyarrow fs. Signed-off-by: Kai Fricke <[email protected]>

Kai Fricke added 2 commits April 5, 2023 11:34

Exclude local filesystem

a5cc36a

Signed-off-by: Kai Fricke <[email protected]>

[air] Use filesystem wrapper to exclude files from upload

7120a65

Signed-off-by: Kai Fricke <[email protected]>

krfricke requested review from richardliaw, xwjiang2010, amogkam, matthewdeng, Yard1, maxpumperla and a team as code owners April 5, 2023 17:16

krfricke marked this pull request as draft April 5, 2023 17:46

krfricke marked this pull request as ready for review April 21, 2023 13:35

Merge remote-tracking branch 'upstream/master' into air/remote-storag…

bbda0cc

…e-exclude

krfricke assigned richardliaw and xwjiang2010 Apr 21, 2023

krfricke mentioned this pull request Apr 21, 2023

[air] remote_storage: Prefer fsspec filesystems over native pyarrow #34663

Merged

8 tasks

justinvyu reviewed Apr 21, 2023

View reviewed changes

Kai Fricke added 2 commits April 24, 2023 11:24

Merge remote-tracking branch 'upstream/master' into air/remote-storag…

f085b98

…e-exclude

Add test

03c278c

Signed-off-by: Kai Fricke <[email protected]>

krfricke requested a review from justinvyu April 24, 2023 10:56

justinvyu approved these changes Apr 24, 2023

View reviewed changes

maxpumperla approved these changes Apr 25, 2023

View reviewed changes

krfricke merged commit 8728c77 into ray-project:master Apr 25, 2023

krfricke deleted the air/remote-storage-exclude branch April 25, 2023 12:08

krfricke pushed a commit to krfricke/ray that referenced this pull request Apr 28, 2023

Revert "[air] Use filesystem wrapper to exclude files from upload (ra…

318ff6a

…y-project#34102)" This reverts commit 8728c77.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[air] Use filesystem wrapper to exclude files from upload #34102

[air] Use filesystem wrapper to exclude files from upload #34102

krfricke commented Apr 5, 2023 •

edited

Loading

krfricke commented Apr 5, 2023

krfricke commented Apr 5, 2023

krfricke commented Apr 21, 2023

justinvyu left a comment

justinvyu Apr 21, 2023

krfricke Apr 24, 2023

justinvyu Apr 21, 2023

krfricke Apr 24, 2023

krfricke commented Apr 24, 2023

justinvyu left a comment

justinvyu Apr 24, 2023 •

edited

Loading

krfricke Apr 25, 2023


		tmp_source, tmp_target = temp_data_dirs

		upload_to_uri(tmp_source, tmp_target, exclude=["_exclude.txt", "_exclude/*"])

[air] Use filesystem wrapper to exclude files from upload #34102

[air] Use filesystem wrapper to exclude files from upload #34102

Conversation

krfricke commented Apr 5, 2023 • edited Loading

Why are these changes needed?

Related issue number

Checks

krfricke commented Apr 5, 2023

krfricke commented Apr 5, 2023

krfricke commented Apr 21, 2023

justinvyu left a comment

Choose a reason for hiding this comment

justinvyu Apr 21, 2023

Choose a reason for hiding this comment

krfricke Apr 24, 2023

Choose a reason for hiding this comment

justinvyu Apr 21, 2023

Choose a reason for hiding this comment

krfricke Apr 24, 2023

Choose a reason for hiding this comment

krfricke commented Apr 24, 2023

justinvyu left a comment

Choose a reason for hiding this comment

justinvyu Apr 24, 2023 • edited Loading

Choose a reason for hiding this comment

krfricke Apr 25, 2023

Choose a reason for hiding this comment

krfricke commented Apr 5, 2023 •

edited

Loading

justinvyu Apr 24, 2023 •

edited

Loading