diff --git a/src/dstack/_internal/core/backends/oci/region.py b/src/dstack/_internal/core/backends/oci/region.py index 806ed76fd1..8967e9ce80 100644 --- a/src/dstack/_internal/core/backends/oci/region.py +++ b/src/dstack/_internal/core/backends/oci/region.py @@ -1,8 +1,10 @@ +import threading from dataclasses import dataclass from functools import cached_property from typing import Any, Dict, Iterable, List, Mapping, Set import oci +from cachetools import Cache, cachedmethod from dstack._internal.core.backends.oci.auth import get_client_config from dstack._internal.core.backends.oci.models import AnyOCICreds @@ -15,6 +17,8 @@ class OCIRegionClient: def __init__(self, client_config: Mapping[str, Any]): self.client_config = client_config + self._availability_domains_cache = Cache(maxsize=10) + self._availability_domains_lock = threading.Lock() @property def name(self) -> str: @@ -44,9 +48,14 @@ def virtual_network_client(self) -> oci.core.VirtualNetworkClient: def work_request_client(self) -> oci.work_requests.WorkRequestClient: return oci.work_requests.WorkRequestClient(self.client_config) - @cached_property - def availability_domains(self) -> List[oci.identity.models.AvailabilityDomain]: - return self.identity_client.list_availability_domains(self.client_config["tenancy"]).data + @cachedmethod( + cache=lambda self: self._availability_domains_cache, + lock=lambda self: self._availability_domains_lock, + ) + def availability_domains_in( + self, compartment_id: str + ) -> List[oci.identity.models.AvailabilityDomain]: + return self.identity_client.list_availability_domains(compartment_id).data def make_region_client(region_name: str, creds: AnyOCICreds) -> OCIRegionClient: diff --git a/src/dstack/_internal/core/backends/oci/resources.py b/src/dstack/_internal/core/backends/oci/resources.py index 91e1d890f0..413395a4ca 100644 --- a/src/dstack/_internal/core/backends/oci/resources.py +++ b/src/dstack/_internal/core/backends/oci/resources.py @@ -140,7 +140,7 @@ def list_shapes_in_region(region: OCIRegionClient, compartment_id: str) -> Dict[ """ result = {} - for availability_domain in region.availability_domains: + for availability_domain in region.availability_domains_in(compartment_id): result[availability_domain.name] = list_shapes_in_domain( availability_domain.name, region.compute_client, compartment_id ) @@ -214,7 +214,7 @@ def check_availability_per_domain( all_shapes = set(shape_names) available_shapes_per_domain = {} - for availability_domain in region.availability_domains: + for availability_domain in region.availability_domains_in(compartment_id): shapes_to_check = { shape for shape in all_shapes @@ -277,7 +277,7 @@ def get_available_domains( `shape_name` is available and within `shapes_quota`. """ domains = [] - for domain in region.availability_domains: + for domain in region.availability_domains_in(compartment_id): if shapes_quota.is_within_domain_quota( shape_name, domain.name ) and check_availability_in_domain(