Repository: skypilot-org/skypilot
Pull request: [Storage] Use task's region for initializing new stores
Closes #3114. Uses the optimizer's `best_resources` or the the tasks' resource spec to infer the region for the bucket, if available.

Note that currently the bucket is eagerly provisioned in the first region that will be tried by the provisioner. If failover is trigerred, the bucket and compute may lie in a different region. A better fix could have been to make store init lazy and create the bucket in the region where the VM actually gets provisioned, but that requires deeper changes + considerations on wasted VM time while data is uploaded to the object store.

TODO:
- [x] Testing across StoreTypes and regions

Tested:
- [x] YAML with and without region specified in `task.resources` on GCP and AWS. Also tested that the same buckets can be successfully mounted by goofys/gcsfuse on tasks in other regions.
- [x] `pytest tests/test_smoke.py::TestStorageWithCredentials` for gcs and s3
- [x] `pytest tests/test_smoke.py::test_spot_storage --gcp`
- [x] `pytest tests/test_smoke.py::test_spot_storage --aws`
- [x] `pytest tests/test_smoke.py --aws` (Unrelated failures in `spot_recovery_aws` and `test-skyserve-new-autoscaler-update`)
- [x] `pytest tests/test_smoke.py --gcp` (Unrelated failures in `sky-bench`, `test_spot_tpu` )

Problem statement:
[Bug/Feature Request] sky serve doesn't respect regions bucket
It seems sky serve doesn't respect the settings for the region. sky serve created a bucket in us-east-2 and tried to create a controller in the same regions. UPDATE: controller can be configured in config.yaml via serve section.

(Luckily it failed for me, because the specified vpc is only available in my zone)

# service.yaml
```
service:
  readiness_probe: /v1/models
 
resources: 
  use_spot: true
  replicas: 1
  memory: 24+
  ordered: 
    - region: eu-central-1
      accelerators: A10G:1
    - region: eu-central-1
      accelerators: T4:4
    - region: eu-west-1
      accelerators: A10G:1
    - region: eu-west-2
      accelerators: A10G:1
    - region: eu-west-3
      accelerators: A10G:1
    - region: eu-north-1
      accelerators: A10G:1
    - region: eu-west-1
      accelerators: T4:4
    - region: eu-west-2
      accelerators: T4:4
    - region: eu-west-3
      accelerators: T4:4
    - region: eu-north-1
      accelerators: T4:4
      memory: 24+
  
  ports: 
    - 8000

num_nodes: 1
workdir: .

file_mounts:
  ~/sky_workdir/openchat.jinja: openchat.jinja

setup: |
  conda create -n vllm python=3.9 -y
  conda activate vllm
  pip install vllm

run: |
  pwd
  ls -la /
  conda activate vllm
  python -m vllm.entrypoints.openai.api_server \
    --tensor-parallel-size $SKYPILOT_NUM_GPUS_PER_NODE \
    --host 0.0.0.0 --port 8000 \
    --model TheBloke/openchat-3.5-0106-AWQ \
    --served-model-name TheBloke/openchat_3.5-AWQ \
    --quantization awq \
    --chat-template openchat.jinja \
    --dtype auto
```

Hints from issue discussion:
@manzke Thanks for reporting this! Actually the `resources` section under the Service YAML is the resources requirement for service replicas, not the controller. To specify controller resources, pls check here: https://skypilot.readthedocs.io/en/latest/serving/sky-serve.html#customizing-skyserve-controller-resources
You are right. the config.yaml and the service ones are confusing. It is not always clear. For example a sky serve controller for me is something specific to the task / service, but not in general.
Also tags for my instances are task / service specific, not general.

Btw the bucket is still generated in us-east-2. Why?

```
I 02-07 19:36:10 storage.py:1330] Created S3 bucket skypilot-workdir-danielmanzke-d64e0564 in us-east-2
I 02-07 19:36:13 storage.py:1330] Created S3 bucket skypilot-filemounts-files-danielmanzke-d64e0564 in us-east-2
```
> You are right. the config.yaml and the service ones are confusing. It is not always clear. For example a sky serve controller for me is something specific to the task / service, but not in general. Also tags for my instances are task / service specific, not general.
> 
> Btw the bucket is still generated in us-east-2. Why?
> 
> ```
> I 02-07 19:36:10 storage.py:1330] Created S3 bucket skypilot-workdir-danielmanzke-d64e0564 in us-east-2
> I 02-07 19:36:13 storage.py:1330] Created S3 bucket skypilot-filemounts-files-danielmanzke-d64e0564 in us-east-2
> ```

> the config.yaml and the service ones are confusing.

Thanks for sharing your feedback! We are thinking to have an controller resources override in service yaml. Does that sounds good to you? Or do you have other options in mind?

> Btw the bucket is still generated in us-east-2. Why?

Currently we won't limit the region of buckets - the `region` section is only for your cloud VMs. cc @romilbhardwaj @landscapepainter for a feature request here
Hi @manzke thanks for the feedback! Do you mind sharing the use case for specifying the regions for the buckets? Would it be necessary for you to specify multiple different regions for buckets? Or a single region to specify is enough?
The use case is simple. 😁 I have all my resources in Europe. I'm running a German company and have to make sure all resources are in Frankfurt (or wherever I decide to host it).
I think the bucket should be close to the resources which is using it.

Example: a bucket in us-east and mounting it in a machine in Europe doesn't make sense at all.
Latency (regulatory) wise.

Also I'm sitting Europe, which could even to lead to a speed up. 

And if you have an instance with a bucket in the same region, you won't pay for traffic.

What's your use case for having buckets in us east and machines somewhere else? (Maybe I'm missing something) -> but I guess all your stuff is in us east anyway 😁
@landscapepainter Thanks for building such a powerful tool.
I am also running into the exactly the same scenario as @manzke . After digging into the implementation, I figured that the s3 buckets are always created in us-east-2 because it was hard-coded to do so. Please look at the screenshots below.

![image](https://github.com/skypilot-org/skypilot/assets/5001683/6a26f043-93e5-4ec8-818f-6e85567b32ff)
![image](https://github.com/skypilot-org/skypilot/assets/5001683/8163076d-f49a-4615-a0bd-25ebe69526bb)

The class [S3Store](https://github.com/skypilot-org/skypilot/blob/master/sky/data/storage.py#L1011) also defines the default region as us-east-2.

![image](https://github.com/skypilot-org/skypilot/assets/5001683/6b6e1caf-c934-4bb8-be24-6eb1fd2d9b33)

The region value is never passed to [_create_s3_backend ](https://github.com/skypilot-org/skypilot/blob/master/sky/data/storage.py#L1313) function. It appears to be a minor bug.

Is it possible to refer to task resources region while creating the buckets?
For region-aware bucket creation, we need to:

1. Make bucket creation lazy (i.e., initialize bucket after the VM has been provisioned). That way we can know exactly which region to create the bucket in.
2. Add the plumbing to pass region info from upstream methods to the create bucket calls.

Two potential, but longer term optimizations:
1. Translation of regions between clouds (e.g., Illinois on fluidstack should create a bucket on us-east-1 on AWS).
2. With lazy bucket creation, parallelize file upload with setup/later stages of provisioning to maximize performance. 

We're looking for help on this, and would love to have PRs for this! Would you be interested in contributing @vjsrinath @manzke?
@romilbhardwaj I would love to but there are a quite a lot of integrations to test.

Can we think about a short-term solution to read the region from resources first? Would that be a viable option?
Thanks for the report @vjsrinath! Could you share the YAML you used?
@concretevitamin Here is the YAML that I tried,

```yaml
name: llm-spot

resources:
  # accelerators: A100:8
  instance_type: p4d.24xlarge
  cloud: aws # optional
  region: us-east-1
  use_spot: true

workdir: train

file_mounts:
  /training/datasets: ./datasets
  /sky-notebook:
    name: ${BUCKET}
    mode: MOUNT

setup: |
  docker pull winglian/axolotl:main-py3.10-cu118-2.0.1

run: |
  docker run --gpus all \
    -v ~/sky_workdir:/sky_workdir \
    -v /root/.cache:/root/.cache \
    winglian/axolotl:main-py3.10-cu118-2.0.1 
  
  docker run --gpus all \
    -v ~/sky_workdir:/sky_workdir \
    -v /root/.cache:/root/.cache \
    -v /sky-notebook:/sky-notebook \
    winglian/axolotl:main-py3.10-cu118-2.0.1 \
    accelerate launch -m axolotl.cli.train /sky_workdir/models/mistral-aria/mistral_domain-adapt_axolotl_lora.yaml

envs:
  BUCKET: <a-unique-bucket-name-to-use>
```

I tried launching this task using the following command,

```sh
BUCKET=llm-training sky spot launch -n axolotl-spot axolotl-spot-new.yaml --env BUCKET -y
```

Thanks for looking into this.

Commit message: make stores region aware

Patch to review:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index 9356a50f365..081d3b0425a 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -791,7 +791,9 @@ def from_metadata(cls, metadata: StorageMetadata,
 
         return storage_obj
 
-    def add_store(self, store_type: Union[str, StoreType]) -> AbstractStore:
+    def add_store(self,
+                  store_type: Union[str, StoreType],
+                  region: Optional[str] = None) -> AbstractStore:
         """Initializes and adds a new store to the storage.
 
         Invoked by the optimizer after it has selected a store to
@@ -799,6 +801,8 @@ def add_store(self, store_type: Union[str, StoreType]) -> AbstractStore:
 
         Args:
           store_type: StoreType; Type of the storage [S3, GCS, AZURE, R2, IBM]
+          region: str; Region to place the bucket in. Caller must ensure that
+            the region is valid for the chosen store_type.
         """
         if isinstance(store_type, str):
             store_type = StoreType(store_type)
@@ -826,6 +830,7 @@ def add_store(self, store_type: Union[str, StoreType]) -> AbstractStore:
             store = store_cls(
                 name=self.name,
                 source=self.source,
+                region=region,
                 sync_on_reconstruction=self.sync_on_reconstruction)
         except exceptions.StorageBucketCreateError:
             # Creation failed, so this must be sky managed store. Add failure
@@ -1310,7 +1315,7 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
         # Store object is being reconstructed for deletion or re-mount with
         # sky start, and error is raised instead.
         if self.sync_on_reconstruction:
-            bucket = self._create_s3_bucket(self.name)
+            bucket = self._create_s3_bucket(self.name, self.region)
             return bucket, True
         else:
             # Raised when Storage object is reconstructed for sky storage
@@ -1360,9 +1365,13 @@ def _create_s3_bucket(self,
             if region is None:
                 s3_client.create_bucket(Bucket=bucket_name)
             else:
-                location = {'LocationConstraint': region}
-                s3_client.create_bucket(Bucket=bucket_name,
-                                        CreateBucketConfiguration=location)
+                if region == 'us-east-1':
+                    # If default us-east-1 region is used, the LocationConstraint must not be specified. https://stackoverflow.com/a/51912090
+                    s3_client.create_bucket(Bucket=bucket_name)
+                else:
+                    location = {'LocationConstraint': region}
+                    s3_client.create_bucket(Bucket=bucket_name,
+                                            CreateBucketConfiguration=location)
                 logger.info(f'Created S3 bucket {bucket_name} in {region}')
         except aws.botocore_exceptions().ClientError as e:
             with ux_utils.print_exception_no_traceback():
@@ -1736,7 +1745,7 @@ def _get_bucket(self) -> Tuple[StorageHandle, bool]:
                 # is being reconstructed for deletion or re-mount with
                 # sky start, and error is raised instead.
                 if self.sync_on_reconstruction:
-                    bucket = self._create_gcs_bucket(self.name)
+                    bucket = self._create_gcs_bucket(self.name, self.region)
                     return bucket, True
                 else:
                     # This is raised when Storage object is reconstructed for
diff --git a/sky/task.py b/sky/task.py
index acca4ff3df5..f1ef20cf48a 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -857,15 +857,16 @@ def update_storage_mounts(
         task_storage_mounts.update(storage_mounts)
         return self.set_storage_mounts(task_storage_mounts)
 
-    def get_preferred_store_type(self) -> storage_lib.StoreType:
+    def get_preferred_store(self) -> Tuple[storage_lib.StoreType, Optional[str]]:
+        """Returns the preferred store type and region for this task."""
         # TODO(zhwu, romilb): The optimizer should look at the source and
         #  destination to figure out the right stores to use. For now, we
         #  use a heuristic solution to find the store type by the following
         #  order:
-        #  1. cloud decided in best_resources.
-        #  2. cloud specified in the task resources.
+        #  1. cloud/region decided in best_resources.
+        #  2. cloud/region specified in the task resources.
         #  3. if not specified or the task's cloud does not support storage,
-        #     use the first enabled storage cloud.
+        #     use the first enabled storage cloud with default region.
         # This should be refactored and moved to the optimizer.
 
         # This check is not needed to support multiple accelerators;
@@ -880,9 +881,11 @@ def get_preferred_store_type(self) -> storage_lib.StoreType:
 
         if self.best_resources is not None:
             storage_cloud = self.best_resources.cloud
+            storage_region = self.best_resources.region
         else:
             resources = list(self.resources)[0]
             storage_cloud = resources.cloud
+            storage_region = resources.region
         if storage_cloud is not None:
             if str(storage_cloud) not in enabled_storage_clouds:
                 storage_cloud = None
@@ -891,9 +894,10 @@ def get_preferred_store_type(self) -> storage_lib.StoreType:
             storage_cloud = clouds.CLOUD_REGISTRY.from_str(
                 enabled_storage_clouds[0])
             assert storage_cloud is not None, enabled_storage_clouds[0]
+            storage_region = None   # Use default region in the Store class
 
         store_type = storage_lib.get_storetype_from_cloud(storage_cloud)
-        return store_type
+        return store_type, storage_region
 
     def sync_storage_mounts(self) -> None:
         """(INTERNAL) Eagerly syncs storage mounts to cloud storage.
@@ -904,9 +908,9 @@ def sync_storage_mounts(self) -> None:
         """
         for storage in self.storage_mounts.values():
             if len(storage.stores) == 0:
-                store_type = self.get_preferred_store_type()
+                store_type, store_region = self.get_preferred_store()
                 self.storage_plans[storage] = store_type
-                storage.add_store(store_type)
+                storage.add_store(store_type, store_region)
             else:
                 # We will download the first store that is added to remote.
                 self.storage_plans[storage] = list(storage.stores.keys())[0]


Code review comments:
diff_hunk: @@ -857,15 +857,16 @@ def update_storage_mounts(
         task_storage_mounts.update(storage_mounts)
         return self.set_storage_mounts(task_storage_mounts)
 
-    def get_preferred_store_type(self) -> storage_lib.StoreType:
+    def get_preferred_store(self) -> Tuple[storage_lib.StoreType, Optional[str]]:
path: sky/task.py
text: @user1:
Can we make this private? Its only use is within this module.