Repository: skypilot-org/skypilot
Pull request: [Storage] Support for storage info in cluster table's metadata
<!-- Describe the changes in this PR -->
This PR resolves #1203.

Currently, when we run `sky stop <cluster-name>` on a cluster that has cloud storage mounted, the mounting will be lost when restarted with `sky start <cluster-name>` as the mounting command does not get ran. 

In this PR, I save the `storage_mounts` dictionary as a metadata in `cluster` table's `storage_mounts_metadata` when `sky launch` is initially ran. Then, this is loaded to be used to re-mount when `sky start <cluster-name>` is ran.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Manual tests on `gcsfuse`, `goofys`, and `rclone`.
- [x] Relevant individual smoke tests: `pytest tests/test_smoke.py::TestStorageWithCredentials` 
- [x] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_file_mounts`
- [x] `state.db` created with `sky launch` in `master` branch, `sky stop`/`start` ran from PR branch
- [x] `state.db`created with `sky launch` in `master` branch, `sky stop` in master branch, `sky start` in PR branch
- [x] Running `sky storage delete` where storage has 3 stores and 1 of those are externally deleted
- [x] Running `sky start` with a cluster that has a bucket mounted and deleted externally
- [x] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_aws_storage_mounts_with_stop --aws`
- [x] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_gcp_storage_mounts_with_stop --gcp`
- [x] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_kubernetes_storage_mounts_with_stop --kubernetes`

Problem statement:
[Storage] VM loses S3 mount after stop/start
To reproduce: make `cluster.yml`:

```
resources:
  cloud: aws
  instance_type: g4dn.xlarge
  disk_size: 500

workdir: .

num_nodes: 4

file_mounts:
  /dataset:
    source: s3://raysort-kl
    store: s3

setup: |
  echo "Hello, SkyPilot!"

run: |
  echo "Hello, SkyPilot!"
```

Then

```
sky launch cluster.yml -c g4
# login, verify content of /dataset is present
sky stop -y g4
sky start -y g4
# login, /dataset becomes empty
```

I _think_ this also happens on a single-node VM but I haven't verified it yet.

Hints from issue discussion:
Yes... current `start` will not mount the s3 bucket again. Please try to run ~~`sky launch --no-setup -y`~~ `sky launch cluster.yml -c g4 --no-setup`.
This problem was recently brought up by @lhqing as well, in #1157. 

TODO: Add the `sky launch --no-setup` in the doc for the storage mount.
Thanks! Just to make sure, I should do `sky launch cluster.yml -c g4 --no-setup` in this case?
> Thanks! Just to make sure, I should do `sky launch cluster.yml -c g4 --no-setup` in this case?

Ahh, yes. Sorry for the incomplete command.
There's no `--no-setup` in 0.1.2 yet. Is this on master?
Yes, it is in the master right now and will be in our next release. Please feel free to switch to the latest master branch with the following commands:
```
git clone https://github.com/skypilot-org/skypilot.git
cd ./skypilot
pip install -e .
```
We are planning to make the `sky start` automatically mount the cloud storage previously mounted. Please stay tuned for the update. ; )
This issue is stale because it has been open 120 days with no activity. Remove stale label or comment or this will be closed in 10 days.
This should be kept open...

Commit message: format

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 941a60ec299..f1e57f52198 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -2844,6 +2844,7 @@ def _sync_file_mounts(
         """Mounts all user files to the remote nodes."""
         self._execute_file_mounts(handle, all_file_mounts)
         self._execute_storage_mounts(handle, storage_mounts)
+        self._set_cluster_metadata_storage_mounts(handle, storage_mounts)
 
     def _setup(self, handle: CloudVmRayResourceHandle, task: task_lib.Task,
                detach_setup: bool) -> None:
@@ -4240,6 +4241,45 @@ def _execute_storage_mounts(self, handle: CloudVmRayResourceHandle,
         end = time.time()
         logger.debug(f'Storage mount sync took {end - start} seconds.')
 
+    def _set_cluster_metadata_storage_mounts(
+            self, handle: CloudVmRayResourceHandle,
+            storage_mounts: Dict[Path, storage_lib.Storage]) -> None:
+        """Sets 'storage_mounts' object in cluster's storage metadata"""
+        if not storage_mounts:
+            return
+
+        cluster_name = handle.cluster_name
+        cluster_metadata = global_user_state.get_cluster_metadata(cluster_name)
+        if cluster_metadata is None:
+            return None
+        if not 'storage' in cluster_metadata.keys():
+            cluster_metadata['storage'] = {}
+        for _, storage_obj in storage_mounts.items():
+            for _, store_obj in storage_obj.stores.items():
+                # Update some of the non-picklabe attributes of store instances
+                store_obj.make_picklable()
+        cluster_metadata['storage']['storage_mounts'] = storage_mounts
+        global_user_state.set_cluster_metadata(cluster_name, cluster_metadata)
+
+    def get_cluster_metadata_storage_mounts(
+        self, handle: CloudVmRayResourceHandle
+    ) -> Optional[Dict[Path, storage_lib.Storage]]:
+        """Gets 'storage_mounts' object from cluster's storage metadata"""
+        cluster_name = handle.cluster_name
+        cluster_metadata = global_user_state.get_cluster_metadata(cluster_name)
+        if cluster_metadata is None:
+            return None
+        if not 'storage' in cluster_metadata.keys():
+            return None
+        if not 'storage_mounts' in cluster_metadata['storage'].keys():
+            return None
+        storage_mounts = cluster_metadata['storage']['storage_mounts']
+        for _, storage_obj in storage_mounts.items():
+            for _, store_obj in storage_obj.stores.items():
+                # Restoring removed attributes due to non-picklable character
+                store_obj.bucket = store_obj.get_handle()
+        return storage_mounts
+
     def _execute_task_one_node(self, handle: CloudVmRayResourceHandle,
                                task: task_lib.Task, job_id: int,
                                detach_run: bool) -> None:
diff --git a/sky/core.py b/sky/core.py
index 64a3161a943..7be868534c1 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -208,6 +208,8 @@ def _start(
                                stream_logs=True,
                                cluster_name=cluster_name,
                                retry_until_up=retry_until_up)
+    storage_mounts = backend.get_cluster_metadata_storage_mounts(handle)
+    backend.sync_file_mounts(handle, None, storage_mounts)
     if idle_minutes_to_autostop is not None:
         backend.set_autostop(handle, idle_minutes_to_autostop, down=down)
     return handle
diff --git a/sky/data/storage.py b/sky/data/storage.py
index 9634ce774bd..34c812bb11f 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -315,6 +315,11 @@ def mount_command(self, mount_path: str) -> str:
         """
         raise NotImplementedError
 
+    def make_picklable(self) -> None:
+        """Sets the attributes that are non-picklable to None
+        """
+        raise NotImplementedError
+
     def __deepcopy__(self, memo):
         # S3 Client and GCS Client cannot be deep copied, hence the
         # original Store object is returned
@@ -1313,6 +1318,12 @@ def _delete_s3_bucket(self, bucket_name: str) -> bool:
             time.sleep(0.1)
         return True
 
+    def make_picklable(self) -> None:
+        """Sets the attributes that are non-picklable to None
+        """
+        self.client = None
+        self.bucket = None
+
 
 class GcsStore(AbstractStore):
     """GcsStore inherits from Storage Object and represents the backend
@@ -1753,6 +1764,12 @@ def _delete_gcs_bucket(self, bucket_name: str) -> bool:
                     raise exceptions.StorageBucketDeleteError(
                         f'Failed to delete GCS bucket {bucket_name}.')
 
+    def make_picklable(self) -> None:
+        """Sets the attributes that are non-picklable to None
+        """
+        self.client = None
+        self.bucket = None
+
 
 class R2Store(AbstractStore):
     """R2Store inherits from S3Store Object and represents the backend
@@ -2113,6 +2130,12 @@ def _delete_r2_bucket(self, bucket_name: str) -> bool:
             time.sleep(0.1)
         return True
 
+    def make_picklable(self) -> None:
+        """Sets the attributes that are non-picklable to None
+        """
+        self.client = None
+        self.bucket = None
+
 
 class IBMCosStore(AbstractStore):
     """IBMCosStore inherits from Storage Object and represents the backend
@@ -2498,3 +2521,9 @@ def _delete_cos_bucket(self):
             if e.__class__.__name__ == 'NoSuchBucket':
                 logger.debug('bucket already removed')
         Rclone.delete_rclone_bucket_profile(self.name, Rclone.RcloneClouds.IBM)
+
+    def make_picklable(self) -> None:
+        """Sets the attributes that are non-picklable to None
+        """
+        self.client = None
+        self.bucket = None
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 37d7c9ba903..db9806c6ff7 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -99,8 +99,7 @@ def create_table(cursor, conn):
     db_utils.add_column_to_table(cursor, conn, 'clusters', 'autostop',
                                  'INTEGER DEFAULT -1')
 
-    db_utils.add_column_to_table(cursor, conn, 'clusters', 'metadata',
-                                 'TEXT DEFAULT "{}"')
+    db_utils.add_column_to_table(cursor, conn, 'clusters', 'metadata', 'BLOB')
 
     db_utils.add_column_to_table(cursor, conn, 'clusters', 'to_down',
                                  'INTEGER DEFAULT 0')
@@ -374,13 +373,18 @@ def get_cluster_metadata(cluster_name: str) -> Optional[Dict[str, Any]]:
     for (metadata,) in rows:
         if metadata is None:
             return None
-        return json.loads(metadata)
+        try:
+            metadata_obj = pickle.loads(metadata)
+        except TypeError as e:
+            if 'bytes-like object is required' in str(e):
+                metadata_obj = json.loads(metadata)
+        return metadata_obj
     return None
 
 
 def set_cluster_metadata(cluster_name: str, metadata: Dict[str, Any]) -> None:
     _DB.cursor.execute('UPDATE clusters SET metadata=(?) WHERE name=(?)', (
-        json.dumps(metadata),
+        pickle.dumps(metadata),
         cluster_name,
     ))
     count = _DB.cursor.rowcount
@@ -518,6 +522,11 @@ def get_cluster_from_name(
         (name, launched_at, handle, last_use, status, autostop, metadata,
          to_down, owner, cluster_hash) = row[:10]
         # TODO: use namedtuple instead of dict
+        try:
+            metadata_obj = pickle.loads(metadata)
+        except TypeError as e:
+            if 'bytes-like object is required' in str(e):
+                metadata_obj = json.loads(metadata)
         record = {
             'name': name,
             'launched_at': launched_at,
@@ -527,7 +536,7 @@ def get_cluster_from_name(
             'autostop': autostop,
             'to_down': bool(to_down),
             'owner': _load_owner(owner),
-            'metadata': json.loads(metadata),
+            'metadata': metadata_obj,
             'cluster_hash': cluster_hash,
         }
         return record
@@ -542,6 +551,11 @@ def get_clusters() -> List[Dict[str, Any]]:
         (name, launched_at, handle, last_use, status, autostop, metadata,
          to_down, owner, cluster_hash) = row[:10]
         # TODO: use namedtuple instead of dict
+        try:
+            metadata_obj = pickle.loads(metadata)
+        except TypeError as e:
+            if 'bytes-like object is required' in str(e):
+                metadata_obj = json.loads(metadata)
 
         record = {
             'name': name,
@@ -552,7 +566,7 @@ def get_clusters() -> List[Dict[str, Any]]:
             'autostop': autostop,
             'to_down': bool(to_down),
             'owner': _load_owner(owner),
-            'metadata': json.loads(metadata),
+            'metadata': metadata_obj,
             'cluster_hash': cluster_hash,
         }
 


Code review comments:
diff_hunk: @@ -4240,6 +4241,45 @@ def _execute_storage_mounts(self, handle: CloudVmRayResourceHandle,
         end = time.time()
         logger.debug(f'Storage mount sync took {end - start} seconds.')
 
+    def _set_cluster_metadata_storage_mounts(
+            self, handle: CloudVmRayResourceHandle,
+            storage_mounts: Dict[Path, storage_lib.Storage]) -> None:
+        """Sets 'storage_mounts' object in cluster's storage metadata"""
+        if not storage_mounts:
+            return
+
+        cluster_name = handle.cluster_name
+        cluster_metadata = global_user_state.get_cluster_metadata(cluster_name)
+        if cluster_metadata is None:
+            return None
+        if not 'storage' in cluster_metadata.keys():
+            cluster_metadata['storage'] = {}
+        for _, storage_obj in storage_mounts.items():
+            for _, store_obj in storage_obj.stores.items():
+                # Update some of the non-picklabe attributes of store instances
+                store_obj.make_picklable()
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
Can we use and store `StorageMetadata` here? It was designed to be a picklable representation of Storage objects. 

@author:
Thanks for pointing this out!

diff_hunk: @@ -4240,6 +4241,45 @@ def _execute_storage_mounts(self, handle: CloudVmRayResourceHandle,
         end = time.time()
         logger.debug(f'Storage mount sync took {end - start} seconds.')
 
+    def _set_cluster_metadata_storage_mounts(
+            self, handle: CloudVmRayResourceHandle,
+            storage_mounts: Dict[Path, storage_lib.Storage]) -> None:
+        """Sets 'storage_mounts' object in cluster's storage metadata"""
+        if not storage_mounts:
+            return
+
+        cluster_name = handle.cluster_name
+        cluster_metadata = global_user_state.get_cluster_metadata(cluster_name)
path: sky/backends/cloud_vm_ray_backend.py
text: @user2:
I'm not sure what counts as metadata for a cluster. In other words, should storage objects be instead stored as a part of the `handle`? Two reasons for this:

1. `get_cluster_metadata` and `set_cluster_metadata` appear to be unused before this... @user1 - was there a specific use case for these methods? 
2. `handle` seems to store all metadata for a cluster (e.g., name, resources etc.)

@author:
To add on, the current implementation I have, it stores `storage_mounts` dictionary, which is `dst[str] : Storage` obj. pair, in `cluster`'s `metadata`. This is useful as I can pass it into `sync_file_mounts` api just like we run it on `sky launch`. We need to consider a way to store `dst` and keep it as a pair with `Storage` obj. if we want to store only `Storage` object in handle.

diff_hunk: @@ -208,6 +208,8 @@ def _start(
                                stream_logs=True,
                                cluster_name=cluster_name,
                                retry_until_up=retry_until_up)
+    storage_mounts = backend.get_cluster_metadata_storage_mounts(handle)
+    backend.sync_file_mounts(handle, None, storage_mounts)
path: sky/core.py
text: @user1:
This will also trigger our file mount syncing code. This can be problematic - e.g., if a user's local path has been deleted, this will likely raise an error and prevent the user from starting a stopped cluster. 

Since rsync file_mounts and `COPY` mode objects would anyway be persistent on the remote disk volume, is there a way we can execute only `MOUNT` mode file_mounts and skip the local->remote upload for them, if any?

@author:
When `file_mounts` parameter is set to `None`, `_execute_file_mounts()` will return immediately without taking any actions. Please correct me if I'm overlooking anything with this appraoch!

@user1:
Ah thanks for clarifying - can we add a short comment above stating that passing `None` will ensure that file_mounts will not be synced. Perhaps also use the kwarg names in calling the method to make it easier to read:

```
backend.sync_file_mounts(handle = ..., all_file_mounts = None, storage_mounts = storage_mounts)
```