[GCP] Fix --disk-size for Custom Machine Images

docs/source/cloud-setup/cloud-permissions/gcp.rst

@@ -134,17 +134,26 @@ User
     compute.firewalls.list
     compute.firewalls.update
 
-8. Click **Create** to create the role.
-9. Go back to the "IAM" tab and click on **GRANT ACCESS**.
-10. Fill in the email address of the user in the “Add principals” section, and select ``minimal-skypilot-role`` in the “Assign roles” section. Click **Save**.
+8. **Optional**: If the user needs to use custom machine images with ``sky launch --image-id``, you can additionally add the following permissions:
+
+.. code-block:: text
+
+    compute.disks.get
+    compute.disks.resize
+    compute.images.get
+    compute.images.useReadOnly
+
+9. Click **Create** to create the role.
+10. Go back to the "IAM" tab and click on **GRANT ACCESS**.
+11. Fill in the email address of the user in the “Add principals” section, and select ``minimal-skypilot-role`` in the “Assign roles” section. Click **Save**.
 
 
 .. image:: ../../images/screenshots/gcp/create-iam.png
     :width: 80%
     :align: center
     :alt: GCP Grant Access
 
-11. The user should receive an invitation to the project and should be able to setup SkyPilot by following the instructions in :ref:`Installation <installation-gcp>`.
+12. The user should receive an invitation to the project and should be able to setup SkyPilot by following the instructions in :ref:`Installation <installation-gcp>`.
 
 .. note::
 

sky/skylet/providers/gcp/node.py

@@ -37,6 +37,7 @@ class inheriting from ``GCPNode``. Those classes are essentially dicts
 
 from googleapiclient.discovery import Resource
 from googleapiclient.errors import HttpError
+from googleapiclient import discovery
 
 from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_NAME
 
@@ -290,6 +291,16 @@ def create_instance(
         """
         return
 
+    @abc.abstractmethod
+    def resize_disk(
+        self, base_config: dict, instance_name: str, wait_for_operation: bool = True
+    ) -> Tuple[dict, str]:
+        """Resize a Google Cloud disk based on the provided configuration.
+
+        Returns the response of resize operation.
+        """
+        return
+
     def create_instances(
         self,
         base_config: dict,
@@ -518,7 +529,6 @@ def _convert_resources_to_urls(
     def create_instance(
         self, base_config: dict, labels: dict, wait_for_operation: bool = True
     ) -> Tuple[dict, str]:
-
         config = self._convert_resources_to_urls(base_config)
         # removing TPU-specific default key set in config.py
         config.pop("networkConfig", None)
@@ -621,6 +631,70 @@ def delete_instance(self, node_id: str, wait_for_operation: bool = True) -> dict
 
         return result
 
+    def resize_disk(
+        self, base_config: dict, instance_name: str, wait_for_operation: bool = True
+    ) -> Tuple[dict, str]:
+        """Resize a Google Cloud disk based on the provided configuration."""
+
+        # Extract the specified disk size from the configuration
+        new_size_gb = base_config["disks"][0]["initializeParams"]["diskSizeGb"]
+
+        # Fetch the instance details to get the disk name and current disk size
+        response = (
+            self.resource.instances()
+            .get(
+                project=self.project_id,
+                zone=self.availability_zone,
+                instance=instance_name,
+            )
+            .execute()
+        )
+        disk_name = response["disks"][0]["source"].split("/")[-1]
+        cur_size_gb = response["disks"][0]["diskSizeGb"]
+
+        # If the new disk size is the same as the existing disk size
+        # resizing is unnecessary
+        if int(new_size_gb) == int(cur_size_gb):
+            # Return empty list instead of raising exception to not break
+            # ray down.
+            return []
+
+        elif int(new_size_gb) < int(cur_size_gb):
+            logger.warning(
+                "Disk resize operation aborted: Specified new disk size (%s GB) is smaller than the current disk size (%s GB)",
+                new_size_gb,
+                cur_size_gb,
+            )
+            return []
+
+        try:
+            # Execute the resize request and return the response
+            operation = (
+                self.resource.disks()
+                .resize(
+                    project=self.project_id,
+                    zone=self.availability_zone,
+                    disk=disk_name,
+                    body={
+                        "sizeGb": str(new_size_gb),
+                    },
+                )
+                .execute()
+            )
+        except HttpError as e:
+            # Catch HttpError when provided with invalid value for new disk size.
+            # Raises error when the new disk size is smaller than the exisiting
+            # disk size
+            logger.warning(f"googleapiclient.errors.HttpError: {e.reason}")
+            return []
+
+        if wait_for_operation:
+            result = self.wait_for_operation(operation)
+        else:
+            result = operation
+
+        return result
+
 
 class GCPTPU(GCPResource):
     """Abstraction around GCP TPU resource"""
@@ -698,7 +772,6 @@ def _list_instances(
         label_filters[TAG_RAY_CLUSTER_NAME] = self.cluster_name
 
         def filter_instance(instance: GCPTPUNode) -> bool:
-
             labels = instance.get_labels()
             if label_filters:
                 for key, value in label_filters.items():

sky/skylet/providers/gcp/node_provider.py

@@ -280,12 +280,17 @@ def get_order_key(node):
                     )
                     for node_id in reuse_node_ids:
                         result = resource.start_instance(node_id)
+                        if "sourceMachineImage" in base_config:
+                            resource.resize_disk(base_config, node_id)
                         result_dict[node_id] = {node_id: result}
                     for node_id in reuse_node_ids:
                         self.set_node_tags(node_id, tags)
                     count -= len(reuse_node_ids)
             if count:
                 results = resource.create_instances(base_config, labels, count)
+                if "sourceMachineImage" in base_config:
+                    for _, instance_id in results:
+                        resource.resize_disk(base_config, instance_id)
                 result_dict.update(
                     {instance_id: result for result, instance_id in results}
                 )

sky/templates/gcp-ray.yml.j2

@@ -18,7 +18,7 @@ docker:
 {%- endif %}
 
 provider:
-  # We use a custom node provider for GCP to support instance stop and reuse.
+  # We use a custom node provider for GCP to create, stop and reuse instances.
   type: external  # type: gcp
   module: sky.skylet.providers.gcp.GCPNodeProvider
   region: {{region}}