skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/clouds/runpod.py

@@ -1,11 +1,11 @@
 """ RunPod Cloud. """
 
-import json
 import typing
-from typing import Dict, Iterator, List, Optional, Tuple
+from typing import Dict, Iterator, List, Optional, Tuple, Union
 
 from sky import clouds
 from sky.clouds import service_catalog
+from sky.utils import registry
 from sky.utils import resources_utils
 
 if typing.TYPE_CHECKING:
@@ -16,7 +16,7 @@ _CREDENTIAL_FILES = [
 ]
 
 
-@clouds.CLOUD_REGISTRY.register
+@registry.CLOUD_REGISTRY.register
 class RunPod(clouds.Cloud):
     """ RunPod GPU Cloud
 
@@ -25,32 +25,22 @@ class RunPod(clouds.Cloud):
     _REPR = 'RunPod'
     _CLOUD_UNSUPPORTED_FEATURES = {
         clouds.CloudImplementationFeatures.STOP: 'Stopping not supported.',
-        clouds.CloudImplementationFeatures.SPOT_INSTANCE:
-            ('Spot is not supported, as runpod API does not implement spot.'),
         clouds.CloudImplementationFeatures.MULTI_NODE:
             ('Multi-node not supported yet, as the interconnection among nodes '
              'are non-trivial on RunPod.'),
-        clouds.CloudImplementationFeatures.OPEN_PORTS:
-            ('Opening ports is not '
-             'supported yet on RunPod.'),
-        clouds.CloudImplementationFeatures.IMAGE_ID:
-            ('Specifying image ID is not supported on RunPod.'),
-        clouds.CloudImplementationFeatures.DOCKER_IMAGE:
-            (f'Docker image is currently not supported on {_REPR}.'),
         clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER:
             ('Customizing disk tier is not supported yet on RunPod.'),
         clouds.CloudImplementationFeatures.STORAGE_MOUNTING:
             ('Mounting object stores is not supported on RunPod. To read data '
              'from object stores on RunPod, use `mode: COPY` to copy the data '
              'to local disk.'),
-        clouds.CloudImplementationFeatures.HOST_CONTROLLERS:
-            ('Host controllers are not supported on RunPod.'),
     }
     _MAX_CLUSTER_NAME_LEN_LIMIT = 120
     _regions: List[clouds.Region] = []
 
     PROVISIONER_VERSION = clouds.ProvisionerVersion.SKYPILOT
     STATUS_VERSION = clouds.StatusVersion.SKYPILOT
+    OPEN_PORTS_VERSION = clouds.OpenPortsVersion.LAUNCH_ONLY
 
     @classmethod
     def _unsupported_features_for_resources(
@@ -79,11 +69,8 @@ class RunPod(clouds.Cloud):
                               zone: Optional[str]) -> List[clouds.Region]:
         assert zone is None, 'RunPod does not support zones.'
         del accelerators, zone  # unused
-        if use_spot:
-            return []
-        else:
-            regions = service_catalog.get_region_zones_for_instance_type(
-                instance_type, use_spot, 'runpod')
+        regions = service_catalog.get_region_zones_for_instance_type(
+            instance_type, use_spot, 'runpod')
 
         if region is not None:
             regions = [r for r in regions if r.name == region]
@@ -155,7 +142,7 @@ class RunPod(clouds.Cloud):
 
     @classmethod
     def get_accelerators_from_instance_type(
-            cls, instance_type: str) -> Optional[Dict[str, int]]:
+            cls, instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
         return service_catalog.get_accelerators_from_instance_type(
             instance_type, clouds='runpod')
 
@@ -166,33 +153,54 @@ class RunPod(clouds.Cloud):
     def make_deploy_resources_variables(
             self,
             resources: 'resources_lib.Resources',
-            cluster_name_on_cloud: str,
+            cluster_name: resources_utils.ClusterName,
             region: 'clouds.Region',
             zones: Optional[List['clouds.Zone']],
+            num_nodes: int,
             dryrun: bool = False) -> Dict[str, Optional[str]]:
-        del zones, dryrun  # unused
+        del zones, dryrun, cluster_name  # unused
 
         r = resources
         acc_dict = self.get_accelerators_from_instance_type(r.instance_type)
-        if acc_dict is not None:
-            custom_resources = json.dumps(acc_dict, separators=(',', ':'))
+        custom_resources = resources_utils.make_ray_custom_resources_str(
+            acc_dict)
+
+        if r.image_id is None:
+            image_id = 'runpod/base:0.0.2'
+        elif r.extract_docker_image() is not None:
+            image_id = r.extract_docker_image()
         else:
-            custom_resources = None
+            image_id = r.image_id[r.region]
+
+        instance_type = resources.instance_type
+        use_spot = resources.use_spot
+
+        hourly_cost = self.instance_type_to_hourly_cost(
+            instance_type=instance_type, use_spot=use_spot)
+
+        # default to root
+        docker_username_for_runpod = (resources.docker_username_for_runpod
+                                      if resources.docker_username_for_runpod
+                                      is not None else 'root')
 
         return {
-            'instance_type': resources.instance_type,
+            'instance_type': instance_type,
             'custom_resources': custom_resources,
             'region': region.name,
+            'image_id': image_id,
+            'use_spot': use_spot,
+            'bid_per_gpu': str(hourly_cost),
+            'docker_username_for_runpod': docker_username_for_runpod,
         }
 
     def _get_feasible_launchable_resources(
         self, resources: 'resources_lib.Resources'
-    ) -> Tuple[List['resources_lib.Resources'], List[str]]:
+    ) -> 'resources_utils.FeasibleResources':
         """Returns a list of feasible resources for the given resources."""
         if resources.instance_type is not None:
             assert resources.is_launchable(), resources
             resources = resources.copy(accelerators=None)
-            return ([resources], [])
+            return resources_utils.FeasibleResources([resources], [], None)
 
         def _make(instance_list):
             resource_list = []
@@ -215,9 +223,12 @@ class RunPod(clouds.Cloud):
                 memory=resources.memory,
                 disk_tier=resources.disk_tier)
             if default_instance_type is None:
-                return ([], [])
+                # TODO: Add hints to all return values in this method to help
+                #  users understand why the resources are not launchable.
+                return resources_utils.FeasibleResources([], [], None)
             else:
-                return (_make([default_instance_type]), [])
+                return resources_utils.FeasibleResources(
+                    _make([default_instance_type]), [], None)
 
         assert len(accelerators) == 1, resources
         acc, acc_count = list(accelerators.items())[0]
@@ -231,8 +242,10 @@ class RunPod(clouds.Cloud):
             zone=resources.zone,
             clouds='runpod')
         if instance_list is None:
-            return ([], fuzzy_candidate_list)
-        return (_make(instance_list), fuzzy_candidate_list)
+            return resources_utils.FeasibleResources([], fuzzy_candidate_list,
+                                                     None)
+        return resources_utils.FeasibleResources(_make(instance_list),
+                                                 fuzzy_candidate_list, None)
 
     @classmethod
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
@@ -247,7 +260,7 @@ class RunPod(clouds.Cloud):
                     '    Credentials can be set up by running: \n'
                     f'        $ pip install runpod \n'
                     f'        $ runpod config\n'
-                    '    For more information, see https://skypilot.readthedocs.io/en/latest/getting-started/installation.html#runpod'  # pylint: disable=line-too-long
+                    '    For more information, see https://docs.skypilot.co/en/latest/getting-started/installation.html#runpod'  # pylint: disable=line-too-long
                 )
 
             return True, None
@@ -263,7 +276,7 @@ class RunPod(clouds.Cloud):
         }
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
         # NOTE: used for very advanced SkyPilot functionality
         # Can implement later if desired
         return None
@@ -275,3 +288,9 @@ class RunPod(clouds.Cloud):
         return service_catalog.validate_region_zone(region,
                                                     zone,
                                                     clouds='runpod')
+
+    @classmethod
+    def get_image_size(cls, image_id: str, region: Optional[str]) -> float:
+        # TODO: use 0.0 for now to allow all images. We should change this to
+        # return the docker image size.
+        return 0.0

sky/clouds/scp.py

@@ -4,17 +4,17 @@ This module includes the set of functions
 to access the SCP catalog and check credentials for the SCP access.
 """
 
-import json
 import typing
-from typing import Dict, Iterator, List, Optional, Tuple
+from typing import Dict, Iterator, List, Optional, Tuple, Union
 
 from sky import clouds
 from sky import exceptions
 from sky import sky_logging
-from sky import status_lib
 from sky.clouds import service_catalog
 from sky.clouds.utils import scp_utils
+from sky.utils import registry
 from sky.utils import resources_utils
+from sky.utils import status_lib
 
 if typing.TYPE_CHECKING:
     # Renaming to avoid shadowing variables.
@@ -30,7 +30,7 @@ _SCP_MIN_DISK_SIZE_GB = 100
 _SCP_MAX_DISK_SIZE_GB = 300
 
 
-@clouds.CLOUD_REGISTRY.register
+@registry.CLOUD_REGISTRY.register
 class SCP(clouds.Cloud):
     """SCP Cloud."""
 
@@ -146,10 +146,10 @@ class SCP(clouds.Cloud):
 
     @classmethod
     def get_default_instance_type(
-            cls,
-            cpus: Optional[str] = None,
-            memory: Optional[str] = None,
-            disk_tier: Optional[resources_utils.DiskTier] = None
+        cls,
+        cpus: Optional[str] = None,
+        memory: Optional[str] = None,
+        disk_tier: Optional['resources_utils.DiskTier'] = None
     ) -> Optional[str]:
         return service_catalog.get_default_instance_type(cpus=cpus,
                                                          memory=memory,
@@ -160,7 +160,7 @@ class SCP(clouds.Cloud):
     def get_accelerators_from_instance_type(
         cls,
         instance_type: str,
-    ) -> Optional[Dict[str, int]]:
+    ) -> Optional[Dict[str, Union[int, float]]]:
         return service_catalog.get_accelerators_from_instance_type(
             instance_type, clouds='scp')
 
@@ -179,20 +179,19 @@ class SCP(clouds.Cloud):
     def make_deploy_resources_variables(
             self,
             resources: 'resources_lib.Resources',
-            cluster_name_on_cloud: str,
+            cluster_name: 'resources_utils.ClusterName',
             region: 'clouds.Region',
             zones: Optional[List['clouds.Zone']],
+            num_nodes: int,
             dryrun: bool = False) -> Dict[str, Optional[str]]:
-        del cluster_name_on_cloud, dryrun  # Unused.
+        del cluster_name, dryrun  # Unused.
         assert zones is None, 'SCP does not support zones.'
 
         r = resources
         acc_dict = self.get_accelerators_from_instance_type(r.instance_type)
+        custom_resources = resources_utils.make_ray_custom_resources_str(
+            acc_dict)
 
-        if acc_dict is not None:
-            custom_resources = json.dumps(acc_dict, separators=(',', ':'))
-        else:
-            custom_resources = None
         image_id = self._get_image_id(r.image_id, region.name, r.instance_type)
         return {
             'instance_type': resources.instance_type,
@@ -251,16 +250,18 @@ class SCP(clouds.Cloud):
 
     def _get_feasible_launchable_resources(
         self, resources: 'resources_lib.Resources'
-    ) -> Tuple[List['resources_lib.Resources'], List[str]]:
+    ) -> 'resources_utils.FeasibleResources':
         # Check if the host VM satisfies the min/max disk size limits.
         is_allowed = self._is_disk_size_allowed(resources)
         if not is_allowed:
-            return ([], [])
+            # TODO: Add hints to all return values in this method to help
+            #  users understand why the resources are not launchable.
+            return resources_utils.FeasibleResources([], [], None)
         if resources.instance_type is not None:
             assert resources.is_launchable(), resources
             # Accelerators are part of the instance type in SCP Cloud
             resources = resources.copy(accelerators=None)
-            return ([resources], [])
+            return resources_utils.FeasibleResources([resources], [], None)
 
         def _make(instance_list):
             resource_list = []
@@ -287,9 +288,10 @@ class SCP(clouds.Cloud):
                 memory=resources.memory,
                 disk_tier=resources.disk_tier)
             if default_instance_type is None:
-                return ([], [])
+                return resources_utils.FeasibleResources([], [], None)
             else:
-                return (_make([default_instance_type]), [])
+                return resources_utils.FeasibleResources(
+                    _make([default_instance_type]), [], None)
 
         assert len(accelerators) == 1, resources
         acc, acc_count = list(accelerators.items())[0]
@@ -304,8 +306,10 @@ class SCP(clouds.Cloud):
             zone=resources.zone,
             clouds='scp')
         if instance_list is None:
-            return ([], fuzzy_candidate_list)
-        return (_make(instance_list), fuzzy_candidate_list)
+            return resources_utils.FeasibleResources([], fuzzy_candidate_list,
+                                                     None)
+        return resources_utils.FeasibleResources(_make(instance_list),
+                                                 fuzzy_candidate_list, None)
 
     @classmethod
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
@@ -332,8 +336,8 @@ class SCP(clouds.Cloud):
         }
 
     @classmethod
-    def get_current_user_identity(cls) -> Optional[List[str]]:
-        # TODO(jgoo1): Implement get_current_user_identity for SCP
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
+        # TODO(jgoo1): Implement get_user_identities for SCP
         return None
 
     def instance_type_exists(self, instance_type: str) -> bool:

sky/clouds/service_catalog/__init__.py

@@ -10,6 +10,7 @@ from sky.clouds.service_catalog.constants import CATALOG_DIR
 from sky.clouds.service_catalog.constants import CATALOG_SCHEMA_VERSION
 from sky.clouds.service_catalog.constants import HOSTED_CATALOG_DIR_URL
 from sky.utils import resources_utils
+from sky.utils import subprocess_utils
 
 if typing.TYPE_CHECKING:
     from sky.clouds import cloud
@@ -31,11 +32,10 @@ def _map_clouds_catalog(clouds: CloudFilter, method_name: str, *args, **kwargs):
     if single:
         clouds = [clouds]  # type: ignore
 
-    results = []
-    for cloud in clouds:
+    def _execute_catalog_method(cloud: str):
         try:
             cloud_module = importlib.import_module(
-                f'sky.clouds.service_catalog.{cloud}_catalog')
+                f'sky.clouds.service_catalog.{cloud.lower()}_catalog')
         except ModuleNotFoundError:
             raise ValueError(
                 'Cannot find module "sky.clouds.service_catalog'
@@ -46,7 +46,11 @@ def _map_clouds_catalog(clouds: CloudFilter, method_name: str, *args, **kwargs):
             raise AttributeError(
                 f'Module "{cloud}_catalog" does not '
                 f'implement the "{method_name}" method') from None
-        results.append(method(*args, **kwargs))
+        return method(*args, **kwargs)
+
+    results = subprocess_utils.run_in_parallel(_execute_catalog_method,
+                                               args=list(clouds),
+                                               num_threads=len(clouds))
     if single:
         return results[0]
     return results
@@ -63,7 +67,7 @@ def list_accelerators(
     all_regions: bool = False,
     require_price: bool = True,
 ) -> 'Dict[str, List[common.InstanceTypeInfo]]':
-    """List the names of all accelerators offered by Sky.
+    """Lists the names of all accelerators offered by Sky.
 
     This will include all accelerators offered by Sky, including those
     that may not be available in the user's account.
@@ -91,7 +95,7 @@ def list_accelerator_counts(
     quantity_filter: Optional[int] = None,
     clouds: CloudFilter = None,
 ) -> Dict[str, List[int]]:
-    """List all accelerators offered by Sky and available counts.
+    """Lists all accelerators offered by Sky and available counts.
 
     Returns: A dictionary of canonical accelerator names mapped to a list
     of available counts. See usage in cli.py.
@@ -125,7 +129,7 @@ def list_accelerator_realtime(
     clouds: CloudFilter = None,
     case_sensitive: bool = True,
 ) -> Tuple[Dict[str, List[int]], Dict[str, int], Dict[str, int]]:
-    """List all accelerators offered by Sky with their realtime availability.
+    """Lists all accelerators offered by Sky with their realtime availability.
 
     Realtime availability is the total number of accelerators in the cluster
     and number of accelerators available at the time of the call.
@@ -238,7 +242,7 @@ def get_default_instance_type(cpus: Optional[str] = None,
 
 def get_accelerators_from_instance_type(
         instance_type: str,
-        clouds: CloudFilter = None) -> Optional[Dict[str, int]]:
+        clouds: CloudFilter = None) -> Optional[Dict[str, Union[int, float]]]:
     """Returns the accelerators from a instance type."""
     return _map_clouds_catalog(clouds, 'get_accelerators_from_instance_type',
                                instance_type)
@@ -324,9 +328,8 @@ def get_common_gpus() -> List[str]:
         'A100',
         'A100-80GB',
         'H100',
-        'K80',
         'L4',
-        'M60',
+        'L40S',
         'P100',
         'T4',
         'V100',
@@ -337,10 +340,13 @@ def get_common_gpus() -> List[str]:
 def get_tpus() -> List[str]:
     """Returns a list of TPU names."""
     # TODO(wei-lin): refactor below hard-coded list.
+    # There are many TPU configurations available, we show the some smallest
+    # ones for each generation, and people should find larger ones with
+    # sky show-gpus tpu.
     return [
-        'tpu-v2-8', 'tpu-v2-32', 'tpu-v2-128', 'tpu-v2-256', 'tpu-v2-512',
-        'tpu-v3-8', 'tpu-v3-32', 'tpu-v3-64', 'tpu-v3-128', 'tpu-v3-256',
-        'tpu-v3-512', 'tpu-v3-1024', 'tpu-v3-2048'
+        'tpu-v2-8', 'tpu-v3-8', 'tpu-v4-8', 'tpu-v4-16', 'tpu-v4-32',
+        'tpu-v5litepod-1', 'tpu-v5litepod-4', 'tpu-v5litepod-8', 'tpu-v5p-8',
+        'tpu-v5p-16', 'tpu-v5p-32', 'tpu-v6e-1', 'tpu-v6e-4', 'tpu-v6e-8'
     ]
 
 

sky/clouds/service_catalog/aws_catalog.py

@@ -8,9 +8,7 @@ import hashlib
 import os
 import threading
 import typing
-from typing import Dict, List, Optional, Tuple
-
-import colorama
+from typing import Dict, List, Optional, Tuple, Union
 
 from sky import exceptions
 from sky import sky_logging
@@ -21,6 +19,9 @@ from sky.clouds.service_catalog import config
 from sky.clouds.service_catalog.data_fetchers import fetch_aws
 from sky.utils import common_utils
 from sky.utils import resources_utils
+from sky.utils import rich_utils
+from sky.utils import timeline
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     import pandas as pd
@@ -82,11 +83,10 @@ def _get_az_mappings(aws_user_hash: str) -> Optional['pd.DataFrame']:
         az_mappings = None
         if aws_user_hash != 'default':
             # Fetch az mapping from AWS.
-            print(
-                f'\r{colorama.Style.DIM}AWS: Fetching availability zones '
-                f'mapping...{colorama.Style.RESET_ALL}',
-                end='')
-            az_mappings = fetch_aws.fetch_availability_zone_mappings()
+            with rich_utils.safe_status(
+                    ux_utils.spinner_message('AWS: Fetching availability '
+                                             'zones mapping')):
+                az_mappings = fetch_aws.fetch_availability_zone_mappings()
         else:
             return None
         az_mappings.to_csv(az_mapping_path, index=False)
@@ -101,6 +101,7 @@ def _get_az_mappings(aws_user_hash: str) -> Optional['pd.DataFrame']:
     return az_mappings
 
 
+@timeline.event
 def _fetch_and_apply_az_mapping(df: common.LazyDataFrame) -> 'pd.DataFrame':
     """Maps zone IDs (use1-az1) to zone names (us-east-1x).
 
@@ -123,11 +124,17 @@ def _fetch_and_apply_az_mapping(df: common.LazyDataFrame) -> 'pd.DataFrame':
         with the zone name (e.g. us-east-1a).
     """
     try:
-        user_identity_list = aws.AWS.get_current_user_identity()
+        user_identity_list = aws.AWS.get_active_user_identity()
         assert user_identity_list, user_identity_list
         user_identity = user_identity_list[0]
         aws_user_hash = hashlib.md5(user_identity.encode()).hexdigest()[:8]
-    except exceptions.CloudUserIdentityError:
+    except (exceptions.CloudUserIdentityError, ImportError):
+        # If failed to get user identity, or import aws dependencies, we use the
+        # latest mapping file or the default mapping file.
+        # The import error can happen on the client side when the user does not
+        # have AWS dependencies installed.
+        # TODO(zhwu): we should avoid the dependency of the availability zone
+        # mapping so as to get rid of the import error.
         glob_name = common.get_catalog_path('aws/az_mappings-*.csv')
         # Find the most recent file that matches the glob.
         # We check the existing files because the user could remove the
@@ -244,7 +251,7 @@ def get_default_instance_type(
 
 
 def get_accelerators_from_instance_type(
-        instance_type: str) -> Optional[Dict[str, int]]:
+        instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
     return common.get_accelerators_from_instance_type_impl(
         _get_df(), instance_type)
 
@@ -309,7 +316,17 @@ def list_accelerators(
 
 def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:
     """Returns the image id from the tag."""
-    return common.get_image_id_from_tag_impl(_image_df, tag, region)
+    global _image_df
+
+    image_id = common.get_image_id_from_tag_impl(_image_df, tag, region)
+    if image_id is None:
+        # Refresh the image catalog and try again, if the image tag is not
+        # found.
+        logger.debug('Refreshing the image catalog and trying again.')
+        _image_df = common.read_catalog('aws/images.csv',
+                                        pull_frequency_hours=0)
+        image_id = common.get_image_id_from_tag_impl(_image_df, tag, region)
+    return image_id
 
 
 def is_image_tag_valid(tag: str, region: Optional[str]) -> bool:

sky/clouds/service_catalog/azure_catalog.py

@@ -4,14 +4,32 @@ This module loads the service catalog file and can be used to query
 instance types and pricing information for Azure.
 """
 import re
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from sky import clouds as cloud_lib
+from sky import sky_logging
 from sky.clouds import Azure
 from sky.clouds.service_catalog import common
 from sky.utils import resources_utils
 from sky.utils import ux_utils
 
+logger = sky_logging.init_logger(__name__)
+
+# This list should match the list of regions in
+# skypilot image generation Packer script's replication_regions
+# sky/clouds/service_catalog/images/skypilot-azure-cpu-ubuntu.pkr.hcl
+COMMUNITY_IMAGE_AVAILABLE_REGIONS = {
+    'centralus',
+    'eastus',
+    'eastus2',
+    'northcentralus',
+    'southcentralus',
+    'westcentralus',
+    'westus',
+    'westus2',
+    'westus3',
+}
+
 # The frequency of pulling the latest catalog from the cloud provider.
 # Though the catalog update is manual in our skypilot-catalog repo, we
 # still want to pull the latest catalog periodically to make sure the
@@ -110,7 +128,8 @@ def get_default_instance_type(
         _DEFAULT_INSTANCE_FAMILY)]
 
     def _filter_disk_type(instance_type: str) -> bool:
-        return Azure.check_disk_tier(instance_type, disk_tier)[0]
+        valid, _ = Azure.check_disk_tier(instance_type, disk_tier)
+        return valid
 
     df = df.loc[df['InstanceType'].apply(_filter_disk_type)]
     return common.get_instance_type_for_cpus_mem_impl(df, cpus,
@@ -118,7 +137,7 @@ def get_default_instance_type(
 
 
 def get_accelerators_from_instance_type(
-        instance_type: str) -> Optional[Dict[str, int]]:
+        instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
     return common.get_accelerators_from_instance_type_impl(_df, instance_type)
 
 
@@ -138,6 +157,7 @@ def get_instance_type_for_accelerator(
     if zone is not None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Azure does not support zones.')
+
     return common.get_instance_type_for_accelerator_impl(df=_df,
                                                          acc_name=acc_name,
                                                          acc_count=acc_count,
@@ -175,9 +195,16 @@ def list_accelerators(
 
 def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:
     """Returns the image id from the tag."""
-    # Azure images are not region-specific.
-    del region  # Unused.
-    return common.get_image_id_from_tag_impl(_image_df, tag, None)
+    global _image_df
+    image_id = common.get_image_id_from_tag_impl(_image_df, tag, region)
+    if image_id is None:
+        # Refresh the image catalog and try again, if the image tag is not
+        # found.
+        logger.debug('Refreshing the image catalog and trying again.')
+        _image_df = common.read_catalog('azure/images.csv',
+                                        pull_frequency_hours=0)
+        image_id = common.get_image_id_from_tag_impl(_image_df, tag, region)
+    return image_id
 
 
 def is_image_tag_valid(tag: str, region: Optional[str]) -> bool: