diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index d476ea573..98e167c9f 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -574,6 +574,8 @@ "backend": "Backend", "backend_plural": "Backends", "availability_not_available": "Not available", + "availability_available": "Available", + "availability_unknown": "Unknown", "availability_no_quota": "No quota", "availability_no_balance": "No balance", "availability_busy": "Busy", diff --git a/frontend/src/pages/Offers/List/helpers.test.ts b/frontend/src/pages/Offers/List/helpers.test.ts index 4f8b5a2d9..2585e8d9d 100644 --- a/frontend/src/pages/Offers/List/helpers.test.ts +++ b/frontend/src/pages/Offers/List/helpers.test.ts @@ -1,9 +1,13 @@ +import { createElement } from 'react'; +import { renderToStaticMarkup } from 'react-dom/server'; + import { formatGpuVendor, getOfferAvailability, getOfferResources, getOfferSearchParams, rangeToObject, + renderRangeJSX, toCpuCard, toOfferCard, } from './helpers'; @@ -19,6 +23,16 @@ const makeGpu = (name: string, price: { min: number; max: number }, availability }); describe('Offers helpers', () => { + test.each([ + [{ min: 52.7436, max: 217.54 }, '52.74..217.54'], + [{ min: 1.236, max: 1.236 }, '1.24'], + [{ min: 2, max: 2 }, '2'], + ] as const)('displays price %j as %s', (price, expected) => { + const html = renderToStaticMarkup(createElement('span', null, renderRangeJSX(price))); + + expect(html.replace(/<[^>]*>/g, '')).toBe(expected); + }); + test('GPU selection preserves precise memory and explicit fleet scope', () => { const gpu = { ...makeGpu('H100', { min: 1, max: 8 }), memory_mib: 275032 }; const params = getOfferSearchParams(new URLSearchParams('project_name=main&fleet=chosen&group_by=gpu'), gpu); @@ -104,20 +118,42 @@ describe('Offers helpers', () => { }); test.each([ - [['no_quota', 'no_balance', 'not_available'], 'not_available'], - [['no_quota', 'no_balance'], 'no_balance'], - [['no_quota'], 'no_quota'], - [['idle'], 'idle'], - [['busy'], 'busy'], - [['no_quota', 'available'], undefined], - [['not_available', 'idle'], 'idle'], - [['busy', 'idle'], 'idle'], - [['idle', 'available'], undefined], - [['no_balance', 'unknown'], undefined], - [[], undefined], - ] as [TAvailability[], TAvailability | undefined][])('availability %j shows %s', (statuses, expected) => - expect(getOfferAvailability(statuses)).toBe(expected), - ); + [['available'], ['available']], + [['unknown'], ['unknown']], + [ + ['unknown', 'available', 'unknown'], + ['available', 'unknown'], + ], + [ + ['no_quota', 'no_balance', 'not_available', 'no_quota'], + ['no_balance', 'no_quota', 'not_available'], + ], + [ + ['not_available', 'no_quota'], + ['no_quota', 'not_available'], + ], + [ + ['no_quota', 'not_available'], + ['no_quota', 'not_available'], + ], + [ + ['no_quota', 'available'], + ['available', 'no_quota'], + ], + [ + ['idle', 'busy'], + ['busy', 'idle'], + ], + [ + ['no_balance', 'unknown'], + ['no_balance', 'unknown'], + ], + [[], []], + ] as [TAvailability[], TAvailability[]][])('availability %j shows the complete sorted set %j', (statuses, expected) => { + const original = [...statuses]; + expect(getOfferAvailability(statuses)).toEqual(expected); + expect(statuses).toEqual(original); + }); test('rangeToObject parses open and closed ranges', () => { expect(rangeToObject('0')).toEqual({ min: 0, max: 0 }); diff --git a/frontend/src/pages/Offers/List/helpers.tsx b/frontend/src/pages/Offers/List/helpers.tsx index 47f177968..c2c9170ef 100644 --- a/frontend/src/pages/Offers/List/helpers.tsx +++ b/frontend/src/pages/Offers/List/helpers.tsx @@ -43,8 +43,6 @@ export const getPropertyFilterOptions = (gpus: (IGpu | ICpuGroup)[]) => { }; }; -const AVAILABLE_STATUSES = new Set(['unknown', 'available']); - export const formatGpuVendor = (vendor: string) => (vendor === 'google' ? 'TPU' : vendor.toUpperCase()); export type OfferCard = IGpu & { @@ -133,12 +131,7 @@ export const getOfferSearchParams = (searchParams: URLSearchParams, offer?: Offe return params; }; -export const getOfferAvailability = (availability: TAvailability[]) => - availability.some((status) => AVAILABLE_STATUSES.has(status)) - ? undefined - : (['idle', 'not_available', 'no_balance', 'no_quota', 'busy'] as const).find((status) => - availability.includes(status), - ); +export const getOfferAvailability = (availability: TAvailability[]) => Array.from(new Set(availability)).sort(); export const getFleetFilterValue = (fleet: IFleet, selectedProjectName?: string) => { if (selectedProjectName && fleet.project_name === selectedProjectName) { @@ -162,15 +155,15 @@ export const renderRangeJSX = (range: { min?: number; max?: number }) => { if (typeof range.min === 'number' && typeof range.max === 'number' && range.max != range.min) { return ( <> - {Number(range.min.toFixed(4))} + {Number(range.min.toFixed(2))} {rangeSeparator} - {Number(range.max.toFixed(4))} + {Number(range.max.toFixed(2))} ); } const value = range.min ?? range.max; - return value === undefined ? undefined : Number(value.toFixed(4)); + return value === undefined ? undefined : Number(value.toFixed(2)); }; export const rangeToObject = ( diff --git a/frontend/src/pages/Offers/List/index.test.ts b/frontend/src/pages/Offers/List/index.test.ts new file mode 100644 index 000000000..0f2cc271b --- /dev/null +++ b/frontend/src/pages/Offers/List/index.test.ts @@ -0,0 +1,22 @@ +import { getRequestParams } from './index'; + +jest.mock('@cloudscape-design/components/badge', () => ({})); +jest.mock('components', () => ({})); +jest.mock('hooks', () => ({})); +jest.mock('libs', () => ({})); +jest.mock('services/gpu', () => ({})); +jest.mock('./hooks/useFilters', () => ({})); + +describe('Offer request parameters', () => { + test('includes spot and on-demand offers when no spot policy is selected', () => { + const params = getRequestParams({ project_name: 'main' }); + + expect(params.run_spec.configuration.spot_policy).toBe('auto'); + }); + + test.each(['spot', 'on-demand', 'auto'] as const)('preserves an explicit %s policy', (spot_policy) => { + const params = getRequestParams({ project_name: 'main', spot_policy }); + + expect(params.run_spec.configuration.spot_policy).toBe(spot_policy); + }); +}); diff --git a/frontend/src/pages/Offers/List/index.tsx b/frontend/src/pages/Offers/List/index.tsx index 51a682be2..aebc2ac66 100644 --- a/frontend/src/pages/Offers/List/index.tsx +++ b/frontend/src/pages/Offers/List/index.tsx @@ -3,7 +3,7 @@ import { useTranslation } from 'react-i18next'; import Badge from '@cloudscape-design/components/badge'; import { skipToken } from '@reduxjs/toolkit/query'; -import { Alert, Button, Cards, CardsProps, MultiselectCSD, Popover, PropertyFilter, SpaceBetween, Toggle } from 'components'; +import { Alert, Button, Cards, CardsProps, MultiselectCSD, PropertyFilter, SpaceBetween, Toggle } from 'components'; import { useCollection } from 'hooks'; import { getServerError } from 'libs'; @@ -88,7 +88,7 @@ export const getRequestParams = ({ ...(gpuMemoryMinMax ? { memory: gpuMemoryMinMax as unknown as TRange } : {}), }, }, - spot_policy, + spot_policy: spot_policy ?? 'auto', volumes: [], files: [], setup: [], @@ -293,9 +293,6 @@ export const OfferList: React.FC = ({ id: 'resources', content: (gpu) => { const availability = getOfferAvailability(gpu.availability); - const label = availability && ( - {t(`offer.availability_${availability}`)} - ); return (
@@ -313,19 +310,9 @@ export const OfferList: React.FC = ({
{gpu.spot.join(', ') || '-'}
{ungrouped && gpu.region &&
{gpu.region}
} - {label && - (gpu.availability.length === 1 ? ( - label - ) : ( - t(`offer.availability_${a}`)).join(', ')} - > - {label} - - ))} + + {availability.map((a) => t(`offer.availability_${a}`)).join(', ')} +
); }, diff --git a/src/dstack/_internal/cli/utils/gpu.py b/src/dstack/_internal/cli/utils/gpu.py index bf85b4677..d5aa6dd1e 100644 --- a/src/dstack/_internal/cli/utils/gpu.py +++ b/src/dstack/_internal/cli/utils/gpu.py @@ -4,9 +4,8 @@ from rich.table import Table from dstack._internal.cli.models.offers import OfferCommandGroupByGpuOutput, OfferRequirements -from dstack._internal.cli.utils.common import console, format_instance_availability +from dstack._internal.cli.utils.common import console from dstack._internal.core.models.gpus import GpuGroup -from dstack._internal.core.models.instances import InstanceAvailability from dstack._internal.core.models.profiles import SpotPolicy from dstack._internal.core.models.runs import Requirements, RunSpec, get_policy_map @@ -61,7 +60,7 @@ def print_gpu_table(gpus: List[GpuGroup], run_spec: RunSpec, group_by: List[str] table.add_column(backend_column, style="grey58", ratio=2) if region_column: table.add_column(region_column, style="grey58", ratio=2) - table.add_column() + table.add_column("AVAILABILITY") for i, gpu_group in enumerate(gpus, start=1): backend_text = "" @@ -116,15 +115,10 @@ def print_gpu_table(gpus: List[GpuGroup], run_spec: RunSpec, group_by: List[str] max_formatted = f"{gpu_group.price.max:.4f}".rstrip("0").rstrip(".") price_display = f"{min_formatted}..{max_formatted}" - availability = "" - has_available = any( - av in (InstanceAvailability.UNKNOWN, InstanceAvailability.AVAILABLE) - for av in gpu_group.availability + availability = ", ".join( + av.value.replace("_", " ") + for av in sorted(set(gpu_group.availability), key=lambda av: av.value) ) - if not has_available: - availability = ", ".join( - map(format_instance_availability, set(gpu_group.availability)) - ) secondary_style = "grey58" row_data = [ diff --git a/src/dstack/_internal/core/models/gpus.py b/src/dstack/_internal/core/models/gpus.py index bb55b62a7..ffc58cad2 100644 --- a/src/dstack/_internal/core/models/gpus.py +++ b/src/dstack/_internal/core/models/gpus.py @@ -14,7 +14,7 @@ class BackendGpu(CoreModel): name: str memory_mib: int vendor: gpuhunt.AcceleratorVendor - availability: InstanceAvailability + availability: List[InstanceAvailability] spot: bool count: int price: float diff --git a/src/dstack/_internal/server/services/gpus.py b/src/dstack/_internal/server/services/gpus.py index a9bf530df..3020ac1b8 100644 --- a/src/dstack/_internal/server/services/gpus.py +++ b/src/dstack/_internal/server/services/gpus.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Literal, Optional, Tuple +from typing import Dict, Iterable, List, Literal, Optional, Tuple from sqlalchemy.ext.asyncio import AsyncSession @@ -6,7 +6,10 @@ from dstack._internal.core.errors import ServerClientError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.gpus import BackendGpu, BackendGpus, CpuGroup, GpuGroup -from dstack._internal.core.models.instances import InstanceOfferWithAvailability +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceOfferWithAvailability, +) from dstack._internal.core.models.profiles import CreationPolicy from dstack._internal.core.models.resources import Range from dstack._internal.core.models.runs import RunSpec @@ -86,19 +89,35 @@ def _get_gpus_grouped( def _get_cpus_grouped( offers: list[InstanceOfferWithAvailability], group_by: set[str], reverse: bool ) -> list[CpuGroup]: - groups: dict[tuple, list[InstanceOfferWithAvailability]] = {} + configurations: dict[tuple, list[InstanceOfferWithAvailability]] = {} for offer in offers: - if offer.instance.resources.gpus: + resources = offer.instance.resources + if resources.gpus: continue + key = ( + offer.backend, + offer.region, + resources.spot, + resources.cpu_arch, + resources.cpus, + resources.memory_mib, + ) + configurations.setdefault(key, []).append(offer) + + groups: dict[tuple, list[InstanceOfferWithAvailability]] = {} + for configuration in configurations.values(): + offer = configuration[0] + availability = _get_availability_group_key(o.availability for o in configuration) key = ( offer.backend if "backend" in group_by else None, offer.region if "region" in group_by else None, offer.instance.resources.cpus if "count" in group_by else None, + availability, ) - groups.setdefault(key, []).append(offer) + groups.setdefault(key, []).extend(configuration) rows = [] - for (backend, region, _), group in groups.items(): + for (backend, region, _, _), group in groups.items(): resources = [offer.instance.resources for offer in group] rows.append( CpuGroup( @@ -114,7 +133,7 @@ def _get_cpus_grouped( min=min(offer.price for offer in group), max=max(offer.price for offer in group), ), - availability=list(dict.fromkeys(offer.availability for offer in group)), + availability=sorted({o.availability for o in group}, key=lambda av: av.value), spot=sorted({"spot" if r.spot else "on-demand" for r in resources}), backend=backend, backends=sorted({offer.backend for offer in group}) if backend is None else None, @@ -137,6 +156,21 @@ def _get_cpus_grouped( ) +def _get_availability_group_key( + availability: Iterable[InstanceAvailability], +) -> tuple[InstanceAvailability, ...]: + """Treat available and unknown as equivalent only when grouping.""" + return tuple( + sorted( + { + InstanceAvailability.AVAILABLE if av == InstanceAvailability.UNKNOWN else av + for av in availability + }, + key=lambda av: av.value, + ) + ) + + async def _get_gpu_offers( session: AsyncSession, project: ProjectModel, @@ -228,19 +262,24 @@ def _process_offers_into_backend_gpus( name=gpu_name, memory_mib=gpu_memory_mib, vendor=gpu_vendor, - availability=offer.availability, + availability=[offer.availability], spot=offer.instance.resources.spot, count=gpu_count_in_offer, price=offer.price, region=offer.region, ) + else: + gpu = backend_data[backend_type]["gpus"][instance_config_key] + gpu.availability = sorted( + set(gpu.availability) | {offer.availability}, key=lambda av: av.value + ) backend_gpus_list = [] for backend_type, data in backend_data.items(): gpus_list = sorted( list(data["gpus"].values()), key=lambda g: ( - not g.availability.is_available(), + not any(av.is_available() for av in g.availability), g.vendor.value, g.name, g.memory_mib, @@ -260,12 +299,15 @@ def _update_gpu_group(row: GpuGroup, gpu: BackendGpu, backend_type: BackendType) """Updates an existing GpuGroup with new data from another GPU offer.""" spot_type: Literal["spot", "on-demand"] = "spot" if gpu.spot else "on-demand" - if gpu.availability not in row.availability: - row.availability.append(gpu.availability) + row.availability = sorted( + set(row.availability) | set(gpu.availability), key=lambda av: av.value + ) if spot_type not in row.spot: row.spot.append(spot_type) if row.backends and backend_type not in row.backends: row.backends.append(backend_type) + if row.regions is not None and gpu.region not in row.regions: + row.regions = sorted([*row.regions, gpu.region]) # FIXME: Consider using non-optional range assert row.count.min is not None @@ -288,7 +330,12 @@ def _get_gpus_with_no_grouping(backend_gpus: List[BackendGpus]) -> List[GpuGroup gpu_rows: Dict[Tuple, GpuGroup] = {} for backend in backend_gpus: for gpu in backend.gpus: - key = (gpu.name, gpu.memory_mib, gpu.vendor) + key = ( + gpu.name, + gpu.memory_mib, + gpu.vendor, + _get_availability_group_key(gpu.availability), + ) if key not in gpu_rows: per_gpu_price = gpu.price / gpu.count price_range = Range[float](min=per_gpu_price, max=per_gpu_price) @@ -297,7 +344,7 @@ def _get_gpus_with_no_grouping(backend_gpus: List[BackendGpus]) -> List[GpuGroup name=gpu.name, memory_mib=gpu.memory_mib, vendor=gpu.vendor, - availability=[gpu.availability], + availability=gpu.availability.copy(), spot=["spot" if gpu.spot else "on-demand"], count=Range[int](min=gpu.count, max=gpu.count), price=price_range, @@ -326,20 +373,26 @@ def _get_gpus_grouped_by_backend(backend_gpus: List[BackendGpus]) -> List[GpuGro gpu_rows: Dict[Tuple, GpuGroup] = {} for backend in backend_gpus: for gpu in backend.gpus: - key = (gpu.name, gpu.memory_mib, gpu.vendor, backend.backend_type) + key = ( + gpu.name, + gpu.memory_mib, + gpu.vendor, + backend.backend_type, + _get_availability_group_key(gpu.availability), + ) if key not in gpu_rows: per_gpu_price = gpu.price / gpu.count gpu_rows[key] = GpuGroup( name=gpu.name, memory_mib=gpu.memory_mib, vendor=gpu.vendor, - availability=[gpu.availability], + availability=gpu.availability.copy(), spot=["spot" if gpu.spot else "on-demand"], count=Range[int](min=gpu.count, max=gpu.count), price=Range[float](min=per_gpu_price, max=per_gpu_price), offer_price=Range[float](min=gpu.price, max=gpu.price), backend=backend.backend_type, - regions=backend.regions.copy(), + regions=[gpu.region], ) else: _update_gpu_group(gpu_rows[key], gpu, backend.backend_type) @@ -362,14 +415,21 @@ def _get_gpus_grouped_by_backend_and_region(backend_gpus: List[BackendGpus]) -> gpu_rows: Dict[Tuple, GpuGroup] = {} for backend in backend_gpus: for gpu in backend.gpus: - key = (gpu.name, gpu.memory_mib, gpu.vendor, backend.backend_type, gpu.region) + key = ( + gpu.name, + gpu.memory_mib, + gpu.vendor, + backend.backend_type, + gpu.region, + _get_availability_group_key(gpu.availability), + ) if key not in gpu_rows: per_gpu_price = gpu.price / gpu.count gpu_rows[key] = GpuGroup( name=gpu.name, memory_mib=gpu.memory_mib, vendor=gpu.vendor, - availability=[gpu.availability], + availability=gpu.availability.copy(), spot=["spot" if gpu.spot else "on-demand"], count=Range[int](min=gpu.count, max=gpu.count), price=Range[float](min=per_gpu_price, max=per_gpu_price), @@ -399,14 +459,20 @@ def _get_gpus_grouped_by_count(backend_gpus: List[BackendGpus]) -> List[GpuGroup gpu_rows: Dict[Tuple, GpuGroup] = {} for backend in backend_gpus: for gpu in backend.gpus: - key = (gpu.name, gpu.memory_mib, gpu.vendor, gpu.count) + key = ( + gpu.name, + gpu.memory_mib, + gpu.vendor, + gpu.count, + _get_availability_group_key(gpu.availability), + ) if key not in gpu_rows: per_gpu_price = gpu.price / gpu.count gpu_rows[key] = GpuGroup( name=gpu.name, memory_mib=gpu.memory_mib, vendor=gpu.vendor, - availability=[gpu.availability], + availability=gpu.availability.copy(), spot=["spot" if gpu.spot else "on-demand"], count=Range[int](min=gpu.count, max=gpu.count), price=Range[float](min=per_gpu_price, max=per_gpu_price), @@ -434,20 +500,27 @@ def _get_gpus_grouped_by_backend_and_count(backend_gpus: List[BackendGpus]) -> L gpu_rows: Dict[Tuple, GpuGroup] = {} for backend in backend_gpus: for gpu in backend.gpus: - key = (gpu.name, gpu.memory_mib, gpu.vendor, backend.backend_type, gpu.count) + key = ( + gpu.name, + gpu.memory_mib, + gpu.vendor, + backend.backend_type, + gpu.count, + _get_availability_group_key(gpu.availability), + ) if key not in gpu_rows: per_gpu_price = gpu.price / gpu.count gpu_rows[key] = GpuGroup( name=gpu.name, memory_mib=gpu.memory_mib, vendor=gpu.vendor, - availability=[gpu.availability], + availability=gpu.availability.copy(), spot=["spot" if gpu.spot else "on-demand"], count=Range[int](min=gpu.count, max=gpu.count), price=Range[float](min=per_gpu_price, max=per_gpu_price), offer_price=Range[float](min=gpu.price, max=gpu.price), backend=backend.backend_type, - regions=backend.regions.copy(), + regions=[gpu.region], ) else: _update_gpu_group(gpu_rows[key], gpu, backend.backend_type) @@ -480,6 +553,7 @@ def _get_gpus_grouped_by_backend_region_and_count( backend.backend_type, gpu.region, gpu.count, + _get_availability_group_key(gpu.availability), ) if key not in gpu_rows: per_gpu_price = gpu.price / gpu.count @@ -487,7 +561,7 @@ def _get_gpus_grouped_by_backend_region_and_count( name=gpu.name, memory_mib=gpu.memory_mib, vendor=gpu.vendor, - availability=[gpu.availability], + availability=gpu.availability.copy(), spot=["spot" if gpu.spot else "on-demand"], count=Range[int](min=gpu.count, max=gpu.count), price=Range[float](min=per_gpu_price, max=per_gpu_price), diff --git a/src/tests/_internal/cli/utils/test_gpu.py b/src/tests/_internal/cli/utils/test_gpu.py new file mode 100644 index 000000000..f09077802 --- /dev/null +++ b/src/tests/_internal/cli/utils/test_gpu.py @@ -0,0 +1,64 @@ +import json +from unittest.mock import patch + +import gpuhunt +import pytest +from rich.console import Console + +from dstack._internal.cli.utils.gpu import print_gpu_json, print_gpu_table +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.gpus import GpuGroup +from dstack._internal.core.models.instances import InstanceAvailability +from dstack._internal.core.models.resources import Range +from dstack._internal.server.testing.common import get_run_spec + + +def _get_group(availability: list[InstanceAvailability]) -> GpuGroup: + return GpuGroup( + name="H200", + memory_mib=144384, + vendor=gpuhunt.AcceleratorVendor.NVIDIA, + count=Range[int](min=8, max=8), + price=Range[float](min=2, max=2), + availability=availability, + spot=["on-demand"], + backends=[BackendType.NEBIUS], + ) + + +class TestPrintGpuTable: + @pytest.mark.parametrize( + ("statuses", "expected"), + [ + (["available"], "available"), + (["unknown"], "unknown"), + (["unknown", "available", "unknown"], "available, unknown"), + (["not_available", "no_quota", "not_available"], "no quota, not available"), + (["no_quota", "available"], "available, no quota"), + (["no_balance", "unknown"], "no balance, unknown"), + ], + ) + def test_displays_complete_sorted_set(self, statuses, expected): + console = Console(width=200, color_system=None) + with ( + patch("dstack._internal.cli.utils.gpu.console", console), + console.capture() as capture, + ): + print_gpu_table( + [_get_group([InstanceAvailability(status) for status in statuses])], + get_run_spec(repo_id="test-repo"), + ["gpu"], + "main", + ) + + assert expected in capture.get() + + +class TestPrintGpuJson: + def test_keeps_availability_as_a_list(self, capsys): + group = _get_group([InstanceAvailability.NO_QUOTA, InstanceAvailability.NOT_AVAILABLE]) + print_gpu_json([group], get_run_spec(repo_id="test-repo"), ["gpu"], "main") + + output = json.loads(capsys.readouterr().out) + assert output["gpus"][0]["availability"] == ["no_quota", "not_available"] + assert "status" not in output["gpus"][0] diff --git a/src/tests/_internal/server/routers/test_gpus.py b/src/tests/_internal/server/routers/test_gpus.py index 82729f171..12f9e13a9 100644 --- a/src/tests/_internal/server/routers/test_gpus.py +++ b/src/tests/_internal/server/routers/test_gpus.py @@ -1,3 +1,4 @@ +from itertools import product from typing import Any, Dict, List, Optional, Union from unittest.mock import Mock, patch @@ -161,6 +162,275 @@ async def call_gpus_api( ) +class TestAvailabilitySetGrouping: + @pytest.mark.parametrize("dimension", ["backend", "region", "spot", "count"]) + def test_different_sets_do_not_merge_across_dimensions(self, dimension): + available = create_gpu_offer(BackendType.NEBIUS, "H200", 144384, 2) + unavailable = available.model_copy(deep=True) + unavailable.availability = InstanceAvailability.NO_QUOTA + if dimension == "backend": + unavailable.backend = BackendType.AWS + elif dimension == "region": + unavailable.region = "other-region" + elif dimension == "spot": + unavailable.instance.resources.spot = True + else: + unavailable.instance.resources.gpus *= 8 + + rows = gpus_services._get_gpus_grouped([available, unavailable], set(), reverse=False) + + assert [row.availability for row in rows] == [ + [InstanceAvailability.AVAILABLE], + [InstanceAvailability.NO_QUOTA], + ] + + @pytest.mark.parametrize("cpu", [False, True]) + @pytest.mark.parametrize("reverse_input", [False, True]) + @pytest.mark.parametrize( + ("group_by", "expected_count"), + [ + (set(), 1), + ({"backend"}, 2), + ({"backend", "region"}, 4), + ({"count"}, 2), + ({"backend", "count"}, 4), + ({"backend", "region", "count"}, 8), + ], + ) + def test_identical_sets_merge_across_unselected_dimensions( + self, cpu, reverse_input, group_by, expected_count + ): + offers = [] + for backend, region, count, spot in product( + [BackendType.AWS, BackendType.NEBIUS], ["east", "west"], [1, 8], [False, True] + ): + for status in ["not_available", "no_quota", "not_available"]: + offer = create_gpu_offer( + backend, + "H200", + 144384, + count * 2, + region=region, + gpu_count=count, + spot=spot, + availability=InstanceAvailability(status), + ) + if cpu: + offer.instance.resources.gpus = [] + offers.append(offer) + if reverse_input: + offers.reverse() + + group_offers = gpus_services._get_cpus_grouped if cpu else gpus_services._get_gpus_grouped + rows = group_offers(offers, group_by, reverse=False) + + assert len(rows) == expected_count + for row in rows: + assert [av.value for av in row.availability] == ["no_quota", "not_available"] + assert set(row.spot) == {"on-demand", "spot"} + if "count" in group_by: + assert row.count.min == row.count.max + else: + assert row.count == Range[int](min=4 if cpu else 1, max=32 if cpu else 8) + if "backend" in group_by: + assert row.backend in {BackendType.AWS, BackendType.NEBIUS} + else: + assert set(row.backends or []) == {BackendType.AWS, BackendType.NEBIUS} + if "region" in group_by: + assert row.region in {"east", "west"} + elif "backend" in group_by: + assert row.regions == ["east", "west"] + + @pytest.mark.parametrize("cpu", [False, True]) + @pytest.mark.parametrize("reverse_input", [False, True]) + @pytest.mark.parametrize( + ("group_by", "expected_count"), + [ + (set(), 1), + ({"backend"}, 2), + ({"backend", "region"}, 4), + ({"count"}, 2), + ({"backend", "count"}, 4), + ({"backend", "region", "count"}, 8), + ], + ) + def test_available_and_unknown_merge_with_actual_statuses_preserved( + self, cpu, reverse_input, group_by, expected_count + ): + offers = [] + for backend, region, count, spot in product( + [BackendType.AWS, BackendType.NEBIUS], ["east", "west"], [1, 8], [False, True] + ): + statuses = ( + ["unknown", "available", "unknown"] + if spot + else ["available" if count == 1 else "unknown"] + ) + for status in statuses: + offer = create_gpu_offer( + backend, + "H200", + 144384, + count * 2, + region=region, + gpu_count=count, + spot=spot, + availability=InstanceAvailability(status), + ) + if cpu: + offer.instance.resources.gpus = [] + offers.append(offer) + if reverse_input: + offers.reverse() + original_statuses = [offer.availability for offer in offers] + + group_offers = gpus_services._get_cpus_grouped if cpu else gpus_services._get_gpus_grouped + rows = group_offers(offers, group_by, reverse=False) + + assert len(rows) == expected_count + assert all( + row.availability == [InstanceAvailability.AVAILABLE, InstanceAvailability.UNKNOWN] + for row in rows + ) + assert all(set(row.spot) == {"on-demand", "spot"} for row in rows) + assert [offer.availability for offer in offers] == original_statuses + + @pytest.mark.parametrize("cpu", [False, True]) + @pytest.mark.parametrize("status", ["available", "unknown"]) + def test_single_positive_status_is_preserved(self, cpu, status): + offer = create_gpu_offer( + BackendType.NEBIUS, "H200", 144384, 2, availability=InstanceAvailability(status) + ) + if cpu: + offer.instance.resources.gpus = [] + + group_offers = gpus_services._get_cpus_grouped if cpu else gpus_services._get_gpus_grouped + rows = group_offers([offer], set(), reverse=False) + + assert len(rows) == 1 + assert rows[0].availability == [InstanceAvailability(status)] + + @pytest.mark.parametrize("cpu", [False, True]) + @pytest.mark.parametrize("reverse_input", [False, True]) + def test_merging_available_and_unknown_keeps_other_status_distinctions( + self, cpu, reverse_input + ): + offers = [] + for count, statuses in [ + (1, ["available", "no_quota"]), + (2, ["unknown", "no_quota"]), + (3, ["unknown", "available", "no_quota"]), + (4, ["no_quota"]), + (5, ["available", "no_balance"]), + (6, ["unknown", "not_available"]), + (7, ["idle"]), + (8, ["available"]), + ]: + for status in statuses: + offer = create_gpu_offer( + BackendType.NEBIUS, + "H200", + 144384, + count * 2, + gpu_count=count, + availability=InstanceAvailability(status), + ) + if cpu: + offer.instance.resources.gpus = [] + offers.append(offer) + if reverse_input: + offers.reverse() + + group_offers = gpus_services._get_cpus_grouped if cpu else gpus_services._get_gpus_grouped + rows = group_offers(offers, set(), reverse=False) + by_set = {tuple(av.value for av in row.availability): row for row in rows} + + assert set(by_set) == { + ("available", "no_quota", "unknown"), + ("no_quota",), + ("available", "no_balance"), + ("not_available", "unknown"), + ("idle",), + ("available",), + } + merged = by_set[("available", "no_quota", "unknown")] + assert merged.count == Range[int](min=4 if cpu else 1, max=12 if cpu else 3) + quota = by_set[("no_quota",)] + assert quota.count == Range[int](min=16 if cpu else 4, max=16 if cpu else 4) + + @pytest.mark.parametrize("cpu", [False, True]) + @pytest.mark.parametrize("reverse_input", [False, True]) + def test_distinct_sets_keep_their_own_resources_and_prices(self, cpu, reverse_input): + offers = [] + for count, statuses in [ + (1, ["not_available", "no_quota", "no_quota"]), + (2, ["no_quota", "not_available"]), + (4, ["no_quota"]), + (8, ["not_available"]), + ]: + for status in statuses: + offer = create_gpu_offer( + BackendType.NEBIUS, + "H200", + 144384, + count * count, + gpu_count=count, + region=f"region-{count}", + availability=InstanceAvailability(status), + ) + if cpu: + offer.instance.resources.gpus = [] + offers.append(offer) + if reverse_input: + offers.reverse() + + group_offers = gpus_services._get_cpus_grouped if cpu else gpus_services._get_gpus_grouped + rows = group_offers(offers, {"backend"}, reverse=False) + by_set = {tuple(av.value for av in row.availability): row for row in rows} + + assert len(rows) == 3 + both = by_set[("no_quota", "not_available")] + assert both.count == Range[int](min=4 if cpu else 1, max=8 if cpu else 2) + assert both.price == Range[float](min=1, max=4 if cpu else 2) + assert both.regions == ["region-1", "region-2"] + quota = by_set[("no_quota",)] + assert quota.count == Range[int](min=16 if cpu else 4, max=16 if cpu else 4) + assert quota.regions == ["region-4"] + capacity = by_set[("not_available",)] + assert capacity.count == Range[int](min=32 if cpu else 8, max=32 if cpu else 8) + assert capacity.regions == ["region-8"] + + @pytest.mark.parametrize( + "group_by", + [ + set(), + {"backend"}, + {"backend", "region"}, + {"count"}, + {"backend", "count"}, + {"backend", "region", "count"}, + ], + ) + def test_available_single_gpu_is_separate_from_unavailable_eight_gpus(self, group_by): + offers = [ + create_gpu_offer(BackendType.NEBIUS, "H200", 144384, 2), + create_gpu_offer( + BackendType.NEBIUS, + "H200", + 144384, + 16, + gpu_count=8, + availability=InstanceAvailability.NOT_AVAILABLE, + ), + ] + rows = gpus_services._get_gpus_grouped(offers, group_by, reverse=False) + + assert [(row.count.min, row.count.max, row.availability) for row in rows] == [ + (1, 1, [InstanceAvailability.AVAILABLE]), + (8, 8, [InstanceAvailability.NOT_AVAILABLE]), + ] + + class TestListGpus: def test_reverse_sorts_by_whole_offer_price(self): offers = [ @@ -236,7 +506,7 @@ async def test_price_order( @pytest.mark.parametrize( ("group_by", "prices"), [ - ([], [(1, 10)]), + ([], [(1, 4), (10, 10)]), (["backend"], [(1, 4), (10, 10)]), (["backend", "region"], [(1, 4), (2, 2), (10, 10)]), (["count"], [(1, 2), (4, 4), (10, 10)]), @@ -300,11 +570,20 @@ async def test_cpu_groups( "cpu_arch": ["arm", "x86"], "count": {"min": 2, "max": 8}, "memory_mib": {"min": 4096, "max": 16384}, - "price": {"min": 1, "max": 10}, - "availability": ["available", "no_quota"], + "price": {"min": 1, "max": 4}, + "availability": ["available"], "spot": ["on-demand", "spot"], - "backends": ["aws", "azure"], - } + "backends": ["aws"], + }, + { + "cpu_arch": ["x86"], + "count": {"min": 4, "max": 4}, + "memory_mib": {"min": 8192, "max": 8192}, + "price": {"min": 10, "max": 10}, + "availability": ["no_quota"], + "spot": ["on-demand"], + "backends": ["azure"], + }, ] if "count" in group_by: assert all(group["count"]["min"] == group["count"]["max"] for group in cpus) @@ -761,7 +1040,7 @@ async def test_exact_aggregation_values( assert response.status_code == 200 data = response.json() - assert len(data["gpus"]) == 2 + assert len(data["gpus"]) == 3 t4_gpu = next((gpu for gpu in data["gpus"] if gpu["name"] == "T4"), None) rtx_gpu = next((gpu for gpu in data["gpus"] if gpu["name"] == "RTX4090"), None) @@ -785,14 +1064,18 @@ async def test_exact_aggregation_values( assert response_count_grouped.status_code == 200 count_grouped_data = response_count_grouped.json() - assert len(count_grouped_data["gpus"]) == 3 + assert len(count_grouped_data["gpus"]) == 4 t4_single_group = None t4_multi_group = None rtx_single_group = None for gpu in count_grouped_data["gpus"]: - if gpu["name"] == "T4" and gpu["count"]["min"] == 1 and gpu["count"]["max"] == 1: + if ( + gpu["name"] == "T4" + and gpu["count"] == {"min": 1, "max": 1} + and gpu["availability"] == ["available"] + ): t4_single_group = gpu elif gpu["name"] == "T4" and gpu["count"]["min"] == 4 and gpu["count"]["max"] == 4: t4_multi_group = gpu @@ -825,7 +1108,7 @@ async def test_exact_aggregation_values( assert response_backend.status_code == 200 backend_data = response_backend.json() - assert len(backend_data["gpus"]) == 3 + assert len(backend_data["gpus"]) == 4 t4_runpod = next( (