[Image] Support passing AMIs for different regions

[Michaelvll]

Closes #1345

Main changes:

1. Support a dict for image_id
2. Added 3 smoke tests for the image_id dict.

Notes:

1. The failover will only go through the regions listed in the image_id.
2. If a region or zone is specified and not in the image_ids, it will throw an error.
3. If a task to be run on a existing cluster has a different image_id (even if it is the same in the same region the cluster is in), we will just error out.

Tested:

• sky launch -c test-image-dict test.yaml

# test.yaml
resources:
  cloud: aws
  image_id:
    us-east-1: ami-0fb70a87623506e3e
    us-west-2: ami-03e40ee6f58eea284

• sky launch -c test-image-dict-failover test.yaml

# test.yaml
resources:
  cloud: aws
  image_id:
    us-east-1: ami-0fb70a87623506e3e
    us-west-2: invalid

• tests/run_smoke_tests.sh test_image_id_dict
• tests/run_smoke_tests.sh test_image_id_dict_with_region
• tests/run_smoke_tests.sh test_image_id_dict_with_zone

[concretevitamin]

May not be able to get to this today. Tagging @WoosukKwon as well.

[romilbhardwaj]

Taking a look at this now

[romilbhardwaj]

Thanks for the awesome work @Michaelvll! Left some comments and faced a bug, mentioned below.

[romilbhardwaj]

If there are many image ids, sky status becomes a little crowded:

NAME  LAUNCHED        RESOURCES                                                                  STATUS  AUTOSTOP  COMMAND
clus  a few secs ago  1x AWS(m6i.2xlarge, image_id={'us-east-1': 'ami-0fb70a87623506e3e', 'u...  UP      -         sky launch -c clus test.y...

I'm trying to think of a better way to do this but not sure what's the best:

1. List only regions/zones: image_id={us-east-1,us-west-1}. Pros - gives a hint which images are specified. Cons - can be confusing, no details about the specific image used.
2. Show the count of images specified: image_id={5 images}. Pros - concise. Cons - no details about region/images
3. Do not show image_id in sky status. Pros - concise. Cons - No details.
4. Keep as it is. Pros - detailed. Cons - too long

[romilbhardwaj]

On thinking more, this field is shown only if image_id is specified. And if a user is specifying an image_id, they would want to know which image is being used for a task. So choosing option 4 (keep as is) might be okay.

[Michaelvll]

Good call! I narrow down the image_id to only showing the region that the cluster is in, as the following (as only 1 region will be selected after the cluster is launched). Wdyt?

Clusters
NAME                                        LAUNCHED        RESOURCES                                                                  STATUS   AUTOSTOP  COMMAND                       
test-image                                  a few secs ago  1x AWS(m6i.2xlarge, image_id={'us-east-1': 'ami-0729d913a335efca7'})       INIT     -         sky launch -c test-image ...

[romilbhardwaj]

^that's an excellent idea! I like it

[romilbhardwaj]

Bug:
I created a cluster named clus with:

resources:
  cloud: aws
  image_id:
    us-east-1: ami-0fb70a87623506e3e
    us-west-2: ami-03e40ee6f58eea284

My cluster got provisioned on us-east-1. Then I removed us-west-2 from image_id dict:

resources:
  cloud: aws
  image_id:
    us-east-1: ami-0fb70a87623506e3e

However I get this error, even though I feel it should work:

$ sky launch -c clus test.yaml
Task from YAML spec: test.yaml
W 11-11 10:14:45 resources.py:638] image_id in resources is experimental. It only supports AWS/GCP.
sky.exceptions.ResourcesMismatchError: Requested resources do not match the existing cluster.
  Requested:    1x AWS(None, image_id={'us-east-1': 'ami-0fb70a87623506e3e'})
  Existing:     1x AWS(m6i.2xlarge, image_id={'us-east-1': 'ami-0fb70a87623506e3e', 'us-west-2': 'ami-03e40ee6f58eea284'})
To fix: specify a new cluster name, or down the existing cluster first: sky down clus

We might need to update the less_demanding_than method to check if the requested resources image_id dict is a subset of the existing resources image_id dict/str:

skypilot/sky/resources.py

Lines 519 to 520 in 40ab453

	if (self.image_id is not None and self.image_id != other.image_id):
	return False

[Michaelvll]

Thanks for the review @romilbhardwaj! Previously, I intentionally requires the image dict in a task to be the same as launched resources, but after reading your comment and thinking more, I think it makes sense to allow the task to have a less demanding image_id than the cluster. I updated the code.
The current behavior is that:
The task image_id is not specified or contains the same image in the region (even if it has some different images request for other regions) the existing cluster is in: succeed.

[Michaelvll]

Tested:

• tests/run_smoke_tests.sh

[romilbhardwaj]

Nice! Once we have #1384 (comment) also added to this PR, we should be good to go!

[Michaelvll]

Nice! Once we have #1384 (comment) also added to this PR, we should be good to go!

This should have been added. PTAL. : )

[romilbhardwaj]

Nice! Once we have #1384 (comment) also added to this PR, we should be good to go!

This should have been added. PTAL. : )

Oops - I hadn't updated my branch. Looks great to me!

NAME    LAUNCHED    RESOURCES                                                             STATUS  AUTOSTOP  COMMAND
image   1 min ago   1x AWS(m6i.2xlarge, image_id={'us-east-1': 'ami-0fb70a87623506e3e'})  UP      -         sky launch -c image test....