From 459c5ae9d56239d90658bf782c56cc6ad36f5dfe Mon Sep 17 00:00:00 2001
From: Hysun He <hysunhe@foxmail.com>
Date: Fri, 3 Jan 2025 21:54:04 +0800
Subject: [PATCH] [OCI] 1. Support specify OS with custom image id. 2. Corner
 case fix (#4524)

* Support specify os type with custom image id.

* trim space

* nit

* comment
---
 sky/clouds/oci.py                | 29 ++++++++++++++++++++---------
 sky/provision/oci/query_utils.py |  7 +++++--
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/sky/clouds/oci.py b/sky/clouds/oci.py
index d4ae6f298d2..b0234e2802c 100644
--- a/sky/clouds/oci.py
+++ b/sky/clouds/oci.py
@@ -232,6 +232,14 @@ def make_deploy_resources_variables(
             listing_id = None
             res_ver = None
 
+        os_type = None
+        if ':' in image_id:
+            # OS type provided in the --image-id. This is the case where
+            # custom image's ocid provided in the --image-id parameter.
+            #  - ocid1.image...aaa:oraclelinux (os type is oraclelinux)
+            #  - ocid1.image...aaa (OS not provided)
+            image_id, os_type = image_id.replace(' ', '').split(':')
+
         cpus = resources.cpus
         instance_type_arr = resources.instance_type.split(
             oci_utils.oci_config.INSTANCE_TYPE_RES_SPERATOR)
@@ -297,15 +305,18 @@ def make_deploy_resources_variables(
             cpus=None if cpus is None else float(cpus),
             disk_tier=resources.disk_tier)
 
-        image_str = self._get_image_str(image_id=resources.image_id,
-                                        instance_type=resources.instance_type,
-                                        region=region.name)
-
-        # pylint: disable=import-outside-toplevel
-        from sky.clouds.service_catalog import oci_catalog
-        os_type = oci_catalog.get_image_os_from_tag(tag=image_str,
-                                                    region=region.name)
-        logger.debug(f'OS type for the image {image_str} is {os_type}')
+        if os_type is None:
+            # OS type is not determined yet. So try to get it from vms.csv
+            image_str = self._get_image_str(
+                image_id=resources.image_id,
+                instance_type=resources.instance_type,
+                region=region.name)
+
+            # pylint: disable=import-outside-toplevel
+            from sky.clouds.service_catalog import oci_catalog
+            os_type = oci_catalog.get_image_os_from_tag(tag=image_str,
+                                                        region=region.name)
+        logger.debug(f'OS type for the image {image_id} is {os_type}')
 
         return {
             'instance_type': instance_type,
diff --git a/sky/provision/oci/query_utils.py b/sky/provision/oci/query_utils.py
index 8cca0629305..3037fcc2703 100644
--- a/sky/provision/oci/query_utils.py
+++ b/sky/provision/oci/query_utils.py
@@ -506,8 +506,11 @@ def find_nsg(cls, region: str, nsg_name: str,
             raise exceptions.ResourcesUnavailableError(
                 'The VCN is not available')
 
-        # Get the primary vnic.
-        assert len(list_vcns_resp.data) > 0
+        # Get the primary vnic. The vnic might be an empty list for the
+        # corner case when the cluster was exited during provision.
+        if not list_vcns_resp.data:
+            return None
+
         vcn = list_vcns_resp.data[0]
 
         list_nsg_resp = net_client.list_network_security_groups(