[Storage]Mount cached mode

[landscapepainter]

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Check if mountpoints remain mounted after restarting by running sky stop and sky start:
  • pytest tests/test_smoke.py::test_aws_storage_mounts_with_stop --aws
  • pytest tests/test_smoke.py::test_gcp_storage_mounts_with_stop --gcp
  • Manually running sky stop and sky start with comp_test.yaml to see if all the MOUNT_CACHED remains storage remains mounted.
• Set of comprehensive manual tests with comp_test.yaml:
  • launch tests:
    • sky launch comp_test.yaml --cloud aws -y
    • sky launch comp_test.yaml --cloud gcp -y
  • sky launch two instances mounting same backend storage. Create new file on mountpoint on instance 1, and check if it appears shortly in instance 2.
  • managed spot jobs tests:
    • sky jobs launch comp_test.yaml --use-spot --cloud aws -y
    • sky jobs launch comp_test.yaml --use-spot --cloud gcp -y
  • launch on docker containers:
    • sky launch comp_test.yaml --cloud aws --image-id docker:continuumio/miniconda3:latest -y
    • sky launch comp_test.yaml --cloud gcp --image-id docker:continuumio/miniconda3:latest -y

comp_test.yaml:

file_mounts:
  #### S3 tests ####
  /s3-sky-managed-copy:
    name: s3-sky-managed-copy
    source: ~/sky_logs
    store: s3
    mode: COPY

  /s3-sky-managed-mount:
    name: s3-sky-managed-mount
    source: ~/sky_logs
    store: s3
    mode: MOUNT

  /s3-sky-managed-mount-cached:
    name: s3-sky-managed-mount-cached
    source: ~/sky_logs
    store: s3
    mode: MOUNT_CACHED

  /s3-external-private-by-user:
    source: s3://s3-external-private-by-user-dy
    mode: MOUNT_CACHED

  /s3-external-public-by-user:
    source: s3://s3-external-public-by-user-dy
    mode: MOUNT_CACHED

  /s3-external-public-not-by-user:
    source: s3://digitalcorpora
    mode: MOUNT_CACHED

  #### GCS tests ####
  /gcs-sky-managed-copy:
    name: gcs-sky-managed-copy
    source: ~/sky_logs
    store: gcs
    mode: COPY

  /gcs-sky-managed-mount:
    name: gcs-sky-managed-mount
    source: ~/sky_logs
    store: gcs
    mode: MOUNT

  /gcs-sky-managed-mount-cached:
    name: gcs-sky-managed-mount-cached
    source: ~/sky_logs
    store: gcs
    mode: MOUNT_CACHED

  /gcs-external-private-by-user:
    source: gs://gcs-external-private-by-user-dy
    mode: MOUNT_CACHED

  /gcs-external-public-by-user:
    source: gs://gcs-external-public-by-user-dy
    mode: MOUNT_CACHED

  /gcs-external-public-not-by-user:
    source: gs://gcp-public-data-sentinel-2
    mode: MOUNT_CACHED


workdir: ~/yaml

setup: |
  echo hi

run: |
  # Show workdir
  ls -l .

  # Show private/public storage copy/mount
  ls -l /s3-sky-managed-copy
  ls -l /s3-sky-managed-mount
  ls -l /s3-sky-managed-mount-cached
  ls -l /s3-external-private-by-user
  ls -l /s3-external-public-by-user
  ls -l /s3-external-public-not-by-user
  ls -l /gcs-sky-managed-copy
  ls -l /gcs-sky-managed-mount
  ls -l /gcs-sky-managed-mount-cached
  ls -l /gcs-external-private-by-user
  ls -l /gcs-external-public-by-user
  ls -l /gcs-external-public-not-by-user

  # Write files on a mounted storage with access
  date > /s3-sky-managed-mount/hellotest.txt
  date > /s3-sky-managed-mount-cached/hellotest.txt
  date > /s3-external-private-by-user/hellotest.txt
  date > /s3-external-public-by-user/hellotest.txt
  date > /gcs-sky-managed-mount/hellotest.txt
  date > /gcs-sky-managed-mount-cached/hellotest.txt
  date > /gcs-external-private-by-user/hellotest.txt
  date > /gcs-external-public-by-user/hellotest.txt

  # Confirm file was written
  cat /s3-sky-managed-mount/hellotest.txt
  cat /s3-sky-managed-mount-cached/hellotest.txt
  cat /s3-external-private-by-user/hellotest.txt
  cat /s3-external-public-by-user/hellotest.txt
  cat /gcs-sky-managed-mount/hellotest.txt
  cat /gcs-sky-managed-mount-cached/hellotest.txt
  cat /gcs-external-private-by-user/hellotest.txt
  cat /gcs-external-public-by-user/hellotest.txt

• Specific tests on training with managed spot jobs:
  • Confirm if MOUNT_CACHED mode can recover from preemption to continue training using our vicuna training script.
  • Confirm files written at mountpoint being uploaded to backed storage in order of creation.
  • Confirm if MOUNT_CACHED mode performs as much as or better than our MOUNT mode when writting checkpoints

[romilbhardwaj]

This is great, thanks @landscapepainter! I noticed this is only S3/GCS. Do we plan to support Azure as well in this PR?

[romilbhardwaj]

This is very cool @landscapepainter! Left some quick comments.

[romilbhardwaj]

Can we add R2 as well? Should be very similar to S3 except with a different endpoint.

[romilbhardwaj]

Should we use --vfs-cache-mode full to also cache reads?

[landscapepainter]

@romilbhardwaj @cblmemo @Michaelvll This is ready for a look :)