Add unallocated_resources option to Compute.get_offers()

src/dstack/_internal/cli/commands/offer.py

@@ -64,6 +64,11 @@ def _register(self):
             action="store_true",
             help="Show full offers not adjusted by requirements",
         )
+        self._parser.add_argument(
+            "--unallocated",
+            action="store_true",
+            help="Subtract allocated resources to show only unallocated resources",
+        )
         resources_group = self._parser.add_argument_group("Resources")
         register_resources_args(resources_group)
         # TODO: register only relevant options
@@ -85,6 +90,7 @@ def _list_offers(self, args: argparse.Namespace) -> None:
             run_spec=run_spec,
             max_offers=args.max_offers,
             full_offers=args.full_offers,
+            unallocated_resources=args.unallocated,
         )
         job_plan = run_plan.job_plans[0]
         if args.format == "plain":
@@ -110,6 +116,7 @@ def _list_gpus(self, args: argparse.Namespace, group_by: list[str]) -> None:
             run_spec=run_spec,
             group_by=[g for g in group_by if g != "gpu"],
             full_offers=args.full_offers,
+            unallocated_resources=args.unallocated,
         )
         if args.format == "plain":
             print_gpu_table(gpus, run_spec, group_by, self.api.project)

src/dstack/_internal/cli/services/configurators/run.py

@@ -145,6 +145,7 @@ def get_plan(
                 ssh_identity_file=configurator_args.ssh_identity_file,
                 max_offers=configurator_args.max_offers,
                 full_offers=configurator_args.full_offers,
+                unallocated_resources=configurator_args.unallocated,
             )
         return run_plan, repo
 
@@ -394,6 +395,11 @@ def register_args(cls, parser: argparse.ArgumentParser):
             action="store_true",
             help="Show full offers not adjusted by requirements",
         )
+        configuration_group.add_argument(
+            "--unallocated",
+            action="store_true",
+            help="Subtract allocated resources to show only unallocated resources",
+        )
         cls.register_env_args(configuration_group)
         register_resources_args(configuration_group)
         register_profile_args(parser)

src/dstack/_internal/core/backends/aws/compute.py

@@ -157,7 +157,9 @@ def __init__(
         self._security_group_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600))
         self._image_id_and_username_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600))
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.AWS,
             locations=self.config.regions,
@@ -186,17 +188,21 @@ def get_offers_modifiers(
     ) -> Iterable[OfferModifier]:
         return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)]
 
-    def _get_offers_cached_key(self, requirements: Requirements) -> int:
+    def get_offers_post_filter(
+        self, requirements: Requirements
+    ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]:
+        return self._get_offers_post_filter_cached(requirements)
+
+    def _get_offers_post_filter_cached_key(self, requirements: Requirements) -> int:
         # Requirements is not hashable, so we use a hack to get arguments hash
         return hash(requirements.json())
 
-    # For `pyright: ignore` directive, see: https://github.com/tkem/cachetools/issues/394
     @cachedmethod(
         cache=lambda self: self._offers_post_filter_cache.cache,
-        key=_get_offers_cached_key,
+        key=_get_offers_post_filter_cached_key,
         lock=lambda self: self._offers_post_filter_cache.lock,
     )
-    def get_offers_post_filter(  # pyright: ignore[reportIncompatibleMethodOverride]
+    def _get_offers_post_filter_cached(
         self, requirements: Requirements
     ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]:
         if requirements.reservation:

src/dstack/_internal/core/backends/azure/compute.py

@@ -102,7 +102,9 @@ def __init__(self, config: AzureConfig, credential: TokenCredential):
             credential=credential, subscription_id=config.subscription_id
         )
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.AZURE,
             locations=self.config.regions,

src/dstack/_internal/core/backends/base/compute.py

@@ -10,7 +10,7 @@
 from enum import Enum
 from functools import lru_cache
 from pathlib import Path
-from typing import Callable, ClassVar, Dict, List, Optional
+from typing import Callable, ClassVar, Dict, List, Optional, Union
 
 import git
 import requests
@@ -110,7 +110,7 @@ class Compute(ABC):
 
     @abstractmethod
     def get_offers(
-        self, requirements: Requirements, full_offers: bool
+        self, requirements: Requirements, full_offers: bool, unallocated_resources: bool
     ) -> Iterator[InstanceOfferWithAvailability]:
         """
         Returns offers with availability matching `requirements`.
@@ -126,6 +126,13 @@ def get_offers(
         generate synthetic offers from discovered nodes on the fly; these synthetic offers should
         reflect either resources that would be allocated based on `requirements`
         (`full_offers=False`) or full allocatable node resources (`full_offers=True`).
+
+        if `unallocated_resources` set to `True`, the method should exclude (subtract) already
+        allocated resources from offer's resources. As with `full_offers`, this flag has no meaning
+        for most backends and is intended for backends such as Kubernetes and Slurm where many jobs
+        can coexist on a single node. With such backends, the method should return full allocatable
+        node resources if `unallocated_resources=False` and only unallocated (available) resources
+        if `unallocated_resources=True`.
         """
         pass
 
@@ -188,16 +195,29 @@ class ComputeWithAllOffersCached(ABC):
     It caches all offers with availability and post-filters by requirements.
     """
 
+    unallocated_resources_argument_has_effect: ClassVar[bool] = False
+    """
+    Set to `True` if `get_all_offers_with_availability()` produces different results based on
+    the `unallocated_resources` value. Doubles the amount of cached data.
+    """
+
     def __init__(self) -> None:
         super().__init__()
         self._offers_cache_lock = threading.Lock()
         self._offers_cache_execution_lock = threading.Lock()
-        self._offers_cache = TTLCache(maxsize=1, ttl=180)
+        self._offers_cache = TTLCache(
+            maxsize=(2 if self.unallocated_resources_argument_has_effect else 1),
+            ttl=180,
+        )
 
     @abstractmethod
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         """
         Returns all backend offers with availability.
+
+        See `Compute.get_offers()` for the `unallocated_resources` argument description.
         """
         pass
 
@@ -223,12 +243,12 @@ def get_offers_post_filter(
         return None
 
     def get_offers(
-        self, requirements: Requirements, full_offers: bool
+        self, requirements: Requirements, full_offers: bool, unallocated_resources: bool
     ) -> Iterator[InstanceOfferWithAvailability]:
         with self._offers_cache_execution_lock:
             # Cache lock does not prevent concurrent execution.
             # We use a separate lock to avoid requesting offers in parallel, re-doing the work and hitting rate limits.
-            cached_offers = self._get_all_offers_with_availability_cached()
+            cached_offers = self._get_all_offers_with_availability_cached(unallocated_resources)
         offers = self.__apply_modifiers(
             cached_offers, self.get_offers_modifiers(requirements, full_offers)
         )
@@ -238,12 +258,20 @@ def get_offers(
             offers = (o for o in offers if post_filter(o))
         return offers
 
+    def _get_all_offers_with_availability_cached_key(self, unallocated_resources: bool) -> int:
+        if self.unallocated_resources_argument_has_effect:
+            return hash(unallocated_resources)
+        return hash(None)
+
     @cachedmethod(
         cache=lambda self: self._offers_cache,
+        key=_get_all_offers_with_availability_cached_key,
         lock=lambda self: self._offers_cache_lock,
     )
-    def _get_all_offers_with_availability_cached(self) -> List[InstanceOfferWithAvailability]:
-        return self.get_all_offers_with_availability()
+    def _get_all_offers_with_availability_cached(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
+        return self.get_all_offers_with_availability(unallocated_resources)
 
     @staticmethod
     def __apply_modifiers(
@@ -271,6 +299,12 @@ class ComputeWithFilteredOffersCached(ABC):
     the `full_offers` value. Doubles the amount of cached data.
     """
 
+    unallocated_resources_argument_has_effect: ClassVar[bool] = False
+    """
+    Set to `True` if `get_offers_by_requirements()` produces different results based on
+    the `unallocated_resources` value. Doubles the amount of cached data.
+    """
+
     def __init__(self) -> None:
         super().__init__()
         self._offers_cache_lock = threading.Lock()
@@ -281,37 +315,47 @@ def get_offers_by_requirements(
         self,
         requirements: Requirements,
         full_offers: bool,
+        unallocated_resources: bool,
     ) -> List[InstanceOfferWithAvailability]:
         """
         Returns backend offers with availability matching requirements.
 
         See `Compute.get_offers()` for the `full_offers` argument description.
         Set the class variable `full_offers_argument_has_effect` to `True` if the `full_offers`
         value has an effect on the offers produced by this method.
+
+        See `Compute.get_offers()` for the `unallocated_resources` argument description.
+        Set the class variable `unallocated_resources_argument_has_effect` to `True` if
+        the `unallocated_resources` value has an effect on the offers produced by this method.
         """
         pass
 
     def get_offers(
-        self, requirements: Requirements, full_offers: bool
+        self, requirements: Requirements, full_offers: bool, unallocated_resources: bool
     ) -> Iterator[InstanceOfferWithAvailability]:
-        return iter(self._get_offers_cached(requirements, full_offers))
+        return iter(self._get_offers_cached(requirements, full_offers, unallocated_resources))
 
-    def _get_offers_cached_key(self, requirements: Requirements, full_offers: bool) -> int:
+    def _get_offers_cached_key(
+        self, requirements: Requirements, full_offers: bool, unallocated_resources: bool
+    ) -> int:
+        hash_items: list[Union[str, bool]] = []
         # Requirements is not hashable, so we use a hack to get arguments hash
-        hashable_requirements = requirements.json()
+        hash_items.append(requirements.json())
         if self.full_offers_argument_has_effect:
-            return hash((hashable_requirements, full_offers))
-        return hash(hashable_requirements)
+            hash_items.append(full_offers)
+        if self.unallocated_resources_argument_has_effect:
+            hash_items.append(unallocated_resources)
+        return hash(tuple(hash_items))
 
     @cachedmethod(
         cache=lambda self: self._offers_cache,
         key=_get_offers_cached_key,
         lock=lambda self: self._offers_cache_lock,
     )
     def _get_offers_cached(
-        self, requirements: Requirements, full_offers: bool
+        self, requirements: Requirements, full_offers: bool, unallocated_resources: bool
     ) -> List[InstanceOfferWithAvailability]:
-        return self.get_offers_by_requirements(requirements, full_offers)
+        return self.get_offers_by_requirements(requirements, full_offers, unallocated_resources)
 
 
 class ComputeWithCreateInstanceSupport(ABC):

src/dstack/_internal/core/backends/cloudrift/compute.py

@@ -40,7 +40,9 @@ def __init__(self, config: CloudRiftConfig):
         self.config = config
         self.client = RiftClient(self.config.creds.api_key)
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.CLOUDRIFT,
             locations=self.config.regions or None,

src/dstack/_internal/core/backends/crusoe/compute.py

@@ -143,7 +143,9 @@ def __init__(self, config: CrusoeConfig):
             )
         )
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.CRUSOE,
             locations=self.config.regions or None,

src/dstack/_internal/core/backends/digitalocean_base/compute.py

@@ -56,7 +56,9 @@ def __init__(self, config: BaseDigitalOceanConfig, api_url: str, type: BackendTy
             DigitalOceanProvider(api_key=config.creds.api_key, api_url=api_url)
         )
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=self.BACKEND_TYPE,
             locations=self.config.regions,

src/dstack/_internal/core/backends/gcp/compute.py

@@ -134,7 +134,9 @@ def __init__(self, config: GCPConfig):
         # Smaller TTL since we check the reservation's in_use_count, which can change often
         self._reservation_cache = ComputeTTLCache(cache=TTLCache(maxsize=8, ttl=20))
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         regions = get_or_error(self.config.regions)
         zones_by_key: Dict[Tuple, List[str]] = {}
         catalog_item_filter = _make_catalog_item_filter(regions, zones_by_key)

src/dstack/_internal/core/backends/hotaisle/compute.py

@@ -52,7 +52,9 @@ def __init__(self, config: HotAisleConfig):
             HotAisleProvider(api_key=config.creds.api_key, team_handle=config.team_handle)
         )
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.HOTAISLE,
             locations=self.config.regions or None,

src/dstack/_internal/core/backends/jarvislabs/compute.py

@@ -79,7 +79,9 @@ def __init__(self, config: JarvisLabsConfig):
         self._catalog = gpuhunt.Catalog(balance_resources=False, auto_reload=False)
         self._catalog.add_provider(JarvisLabsProvider(api_key=self.config.creds.api_key))
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.JARVISLABS,
             locations=self.config.regions or None,

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -140,20 +140,24 @@ class KubernetesCompute(
     ComputeWithMultinodeSupport,
     Compute,
 ):
+    unallocated_resources_argument_has_effect = True
+
     def __init__(self, config: KubernetesConfig):
         super().__init__()
         self.region_cluster_map = {c.region: c for c in get_clusters_from_backend_config(config)}
         self.skip_offer_cache = RegionalSkipOfferCache(ttl=60)
 
-    def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> list[InstanceOfferWithAvailability]:
         offers: list[InstanceOfferWithAvailability] = []
         with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
             future_cluster_map: dict[
                 concurrent.futures.Future[list[InstanceOfferWithAvailability]], Cluster
             ] = {}
             for region, cluster in self.region_cluster_map.items():
                 api = client.CoreV1Api(cluster.api_client)
-                future = executor.submit(get_instance_offers, api, region)
+                future = executor.submit(get_instance_offers, api, region, unallocated_resources)
                 future_cluster_map[future] = cluster
             for future in concurrent.futures.as_completed(future_cluster_map):
                 try:

src/dstack/_internal/core/backends/kubernetes/resources.py

@@ -385,16 +385,24 @@ def is_taint_tolerated(taint: V1Taint) -> bool:
     return taint.key in (NVIDIA_GPU_NODE_TAINT, AMD_GPU_NODE_TAINT)
 
 
-def get_instance_offers(api: CoreV1Api, region: str) -> list[InstanceOfferWithAvailability]:
-    nodes_allocated_resources = _get_nodes_allocated_resources(api)
+def get_instance_offers(
+    api: CoreV1Api, region: str, unallocated_resources: bool
+) -> list[InstanceOfferWithAvailability]:
+    nodes_allocated_resources: Optional[dict[str, KubernetesResources]] = None
+    if unallocated_resources:
+        nodes_allocated_resources = _get_nodes_allocated_resources(api)
     offers: list[InstanceOfferWithAvailability] = []
     for node in api.list_node().items:
         if (node_name := get_node_name(node)) is None:
             continue
         offer = _get_instance_offer_from_node(
             node=node,
             node_name=node_name,
-            node_allocated_resources=nodes_allocated_resources.get(node_name),
+            node_allocated_resources=(
+                None
+                if nodes_allocated_resources is None
+                else nodes_allocated_resources.get(node_name)
+            ),
             region=region,
         )
         if offer is not None:

src/dstack/_internal/core/backends/lambdalabs/compute.py

@@ -42,7 +42,9 @@ def __init__(self, config: LambdaConfig):
         self.config = config
         self.api_client = LambdaAPIClient(config.creds.api_key)
 
-    def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
+    def get_all_offers_with_availability(
+        self, unallocated_resources: bool
+    ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.LAMBDA,
             locations=self.config.regions or None,