[Storage] Fix spill/restore error when using Arrow S3FS

[frank-lsf]

Why are these changes needed?

This PR fixes two issues.

1. Occasionally when restoring objects from S3 we get ValueError: Obtained data has a size of 100102955, although it is supposed to have the size of 100115855. This is usually because the offset in the object URI is wrong.

The fix is twofold: 1. manually call f.flush() because Arrow's S3FS is not guaranteed to flush the output when closing (https://github.com/apache/arrow/blob/master/python/pyarrow/io.pxi#L99), and 2. replace f.tell() with manually adding the written bytes to the offset, because f.write() is asynchronous in S3FS and as such the tell() position is not advanced at this point in the code.

2. When restoring objects from S3, the restore worker makes a lot of small reads which are currently not buffered in ExternalStorageRayStorageImpl. This causes unnecessary extra S3 requests, costing time and money.

By opening the input stream with a buffer size set (https://github.com/apache/arrow/blob/master/python/pyarrow/_fs.pyx#L618), Arrow will buffer the reads and increase I/O efficiency.

The solution is not very pretty right now because we use a private Arrow API. But this is not avoidable until Arrow provides a seek() method for BufferedInputStream. I opened an issue here https://issues.apache.org/jira/projects/ARROW/issues/ARROW-16351.

UPDATE: Reverted (2) for now since it is not running stably.

Related issue number

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 :(

[stephanie-wang]

Nice find! I think you might be able to test this in CI with the ray_storage_object_spilling_config fixture. Can you try that out on one of the object spilling tests? If not, it's okay for now.

[frank-lsf]

@stephanie-wang I think this is already covered by existing tests? The issue occurs quite rarely even when using S3. I can add a test for ray_storage + S3 since it looks like that path is not tested yet.

[stephanie-wang]

@stephanie-wang I think this is already covered by existing tests? The issue occurs quite rarely even when using S3. I can add a test for ray_storage + S3 since it looks like that path is not tested yet.

Hmm what do you mean that it was already covered? Do you mean we have existing CI tests that failed with the same error?

Probably we don't want to depend on S3 in CI tests, but I think I saw this error before using ray_storage and the local filesystem.

[frank-lsf]

@stephanie-wang I see. I can add a test using the local filesystem then. I originally thought the tell() and flush() issues only happen with Arrow S3FS since it is asynchronous; but looks like Arrow's local filesystem might see this issue too.

[ericl]

Triggering rebuild.

[ericl]

Test_object_spilling seems to be failing on the windows build. Maybe either try to fix or disable it there?