diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 113618b..8707949 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -32,7 +32,7 @@ jobs: strategy: matrix: os: [ubuntu-latest] - python-version: ["3.10", "3.11", "3.12", "3.13"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} @@ -43,6 +43,10 @@ jobs: run: | python -m pip install --upgrade pip pip install '.[all,dev]' + - name: Run pyright + uses: jakebailey/pyright-action@v3 + with: + pylance-version: latest-release - name: Run doctest run: pytest --doctest-modules src/gpuhunt - name: Run pytest diff --git a/pyproject.toml b/pyproject.toml index 739e9b9..6a7bf85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ dev = [ "pre-commit", "pytest~=7.0", "pytest-mock", + "pyright==1.1.403", # Should match the pinned version in CI "ruff==0.5.3", # Should match .pre-commit-config.yaml "requests-mock", ] @@ -74,3 +75,9 @@ ignore = [ [tool.ruff.lint.isort] known-first-party = ["gpuhunt"] combine-as-imports = true + +[tool.pyright] +typeCheckingMode = "standard" +include = [ + "src/" +] diff --git a/src/gpuhunt/__init__.py b/src/gpuhunt/__init__.py index b12c82d..32911f4 100644 --- a/src/gpuhunt/__init__.py +++ b/src/gpuhunt/__init__.py @@ -15,6 +15,11 @@ default_catalog as default_catalog, query as query, ) +from gpuhunt._internal.errors import ( + GPUHuntError as GPUHuntError, + MissingCredsError as MissingCredsError, + ProviderError as ProviderError, +) from gpuhunt._internal.models import ( AcceleratorInfo as AcceleratorInfo, AcceleratorVendor as AcceleratorVendor, @@ -24,7 +29,6 @@ IntelAcceleratorInfo as IntelAcceleratorInfo, NvidiaGPUInfo as NvidiaGPUInfo, QueryFilter as QueryFilter, - RawCatalogItem as RawCatalogItem, TenstorrentAcceleratorInfo as TenstorrentAcceleratorInfo, TPUInfo as TPUInfo, ) diff --git a/src/gpuhunt/__main__.py b/src/gpuhunt/__main__.py index 28c4f2a..d99f3df 100644 --- a/src/gpuhunt/__main__.py +++ b/src/gpuhunt/__main__.py @@ -4,6 +4,7 @@ import gpuhunt._internal.storage as storage from gpuhunt._internal.utils import configure_logging +from gpuhunt.providers.base import OfflineProvider def main(): @@ -41,15 +42,11 @@ def main(): elif args.provider == "azure": from gpuhunt.providers.azure import AzureProvider - provider = AzureProvider(os.getenv("AZURE_SUBSCRIPTION_ID")) + provider = AzureProvider(os.environ["AZURE_SUBSCRIPTION_ID"]) elif args.provider == "crusoe": from gpuhunt.providers.crusoe import CrusoeProvider - provider = CrusoeProvider( - access_key=os.getenv("CRUSOE_ACCESS_KEY"), - secret_key=os.getenv("CRUSOE_SECRET_KEY"), - project_id=os.getenv("CRUSOE_PROJECT_ID"), - ) + provider = CrusoeProvider.from_env() elif args.provider == "cloudrift": from gpuhunt.providers.cloudrift import CloudRiftProvider @@ -57,33 +54,30 @@ def main(): elif args.provider == "verda": from gpuhunt.providers.verda import VerdaProvider - provider = VerdaProvider(os.getenv("VERDA_CLIENT_ID"), os.getenv("VERDA_CLIENT_SECRET")) + provider = VerdaProvider( + client_id=os.environ["VERDA_CLIENT_ID"], + client_secret=os.environ["VERDA_CLIENT_SECRET"], + ) elif args.provider == "digitalocean": from gpuhunt.providers.digitalocean import DigitalOceanProvider - provider = DigitalOceanProvider( - api_key=os.getenv("DIGITAL_OCEAN_API_KEY"), api_url=os.getenv("DIGITAL_OCEAN_API_URL") - ) + provider = DigitalOceanProvider.from_env() elif args.provider == "gcp": from gpuhunt.providers.gcp import GCPProvider - provider = GCPProvider(os.getenv("GCP_PROJECT_ID")) + provider = GCPProvider(project=os.environ["GCP_PROJECT_ID"]) elif args.provider == "hotaisle": from gpuhunt.providers.hotaisle import HotAisleProvider - provider = HotAisleProvider( - api_key=os.getenv("HOTAISLE_API_KEY"), team_handle=os.getenv("HOTAISLE_TEAM_HANDLE") - ) + provider = HotAisleProvider.from_env() elif args.provider == "jarvislabs": from gpuhunt.providers.jarvislabs import JarvisLabsProvider - provider = JarvisLabsProvider( - api_key=os.getenv("JL_API_KEY"), api_url=os.getenv("JARVISLABS_API_URL") - ) + provider = JarvisLabsProvider.from_env() elif args.provider == "lambdalabs": from gpuhunt.providers.lambdalabs import LambdaLabsProvider - provider = LambdaLabsProvider(os.getenv("LAMBDALABS_TOKEN")) + provider = LambdaLabsProvider(token=os.environ["LAMBDALABS_TOKEN"]) elif args.provider == "nebius": from nebius.base.service_account.pk_file import Reader as PKReader @@ -95,9 +89,9 @@ def main(): os.getenv("NEBIUS_ACCESS_TOKEN") # or service account credentials or PKReader( - filename=os.getenv("NEBIUS_PRIVATE_KEY_FILE"), - public_key_id=os.getenv("NEBIUS_PUBLIC_KEY_ID"), - service_account_id=os.getenv("NEBIUS_SERVICE_ACCOUNT_ID"), + filename=os.environ["NEBIUS_PRIVATE_KEY_FILE"], + public_key_id=os.environ["NEBIUS_PUBLIC_KEY_ID"], + service_account_id=os.environ["NEBIUS_SERVICE_ACCOUNT_ID"], ) ) ) @@ -120,21 +114,21 @@ def main(): elif args.provider == "seeweb": from gpuhunt.providers.seeweb import SeewebProvider - provider = SeewebProvider(os.getenv("SEEWEB_API_TOKEN")) + provider = SeewebProvider.from_env() elif args.provider == "vastai": from gpuhunt.providers.vastai import VastAIProvider - provider = VastAIProvider() + provider = VastAIProvider.from_env() elif args.provider == "vultr": from gpuhunt.providers.vultr import VultrProvider - provider = VultrProvider() + provider = VultrProvider.from_env() else: exit(f"Unknown provider {args.provider}") logging.info("Fetching offers for %s", args.provider) offers = provider.get() - if not args.no_filter: + if not args.no_filter and isinstance(provider, OfflineProvider): offers = provider.filter(offers) storage.dump(offers, args.output) diff --git a/src/gpuhunt/_internal/catalog.py b/src/gpuhunt/_internal/catalog.py index 864bcd2..8d15118 100644 --- a/src/gpuhunt/_internal/catalog.py +++ b/src/gpuhunt/_internal/catalog.py @@ -1,5 +1,3 @@ -import csv -import dataclasses import heapq import io import logging @@ -13,9 +11,10 @@ from pathlib import Path import gpuhunt._internal.constraints as constraints +import gpuhunt._internal.storage as storage from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, CPUArchitecture, QueryFilter from gpuhunt._internal.utils import parse_compute_capability -from gpuhunt.providers import AbstractProvider +from gpuhunt.providers.base import AbstractProvider logger = logging.getLogger(__name__) @@ -212,9 +211,10 @@ def _load(self, version: str | None = None): for provider in OFFLINE_PROVIDERS: try: with zip_file.open(f"{provider}.csv", "r") as csv_file: - reader = csv.DictReader(io.TextIOWrapper(csv_file, "utf-8")) - for row in reader: - item = CatalogItem.from_dict(row, provider=provider) + items = storage.load( + io.TextIOWrapper(csv_file, "utf-8"), provider=provider + ) + for item in items: catalog.setdefault(provider, []).append(item) except KeyError: logger.error( @@ -253,9 +253,9 @@ def _get_offline_provider_items( catalog_dir = os.getenv("GPUHUNT_CATALOG_DIR") if catalog_dir is not None: with open(Path(catalog_dir) / f"{provider_name}.csv", "rb") as csv_file: - reader = csv.DictReader(io.TextIOWrapper(csv_file, "utf-8")) - for row in reader: - item = CatalogItem.from_dict(row, provider=provider_name) + for item in storage.load( + io.TextIOWrapper(csv_file, "utf-8"), provider=provider_name + ): if constraints.matches(item, query_filter): items.append(item) return items @@ -282,10 +282,9 @@ def _get_online_provider_items( if provider.NAME != provider_name: continue found = True - for i in provider.get( + for item in provider.get( query_filter=query_filter, balance_resources=self.balance_resources ): - item = CatalogItem(provider=provider_name, **dataclasses.asdict(i)) if constraints.matches(item, query_filter): items.append(item) if not found: diff --git a/src/gpuhunt/_internal/constraints.py b/src/gpuhunt/_internal/constraints.py index c33e758..6810910 100644 --- a/src/gpuhunt/_internal/constraints.py +++ b/src/gpuhunt/_internal/constraints.py @@ -15,30 +15,6 @@ TPUInfo, ) -# v5litepod = v5e -_TPU_VERSIONS = ["v2", "v3", "v4", "v5p", "v5litepod", "v6e"] - - -Comparable = TypeVar("Comparable", bound=int | float | tuple[int, int]) - - -def is_between(value: Comparable, left: Comparable | None, right: Comparable | None) -> bool: - if is_below(value, left) or is_above(value, right): - return False - return True - - -def is_below(value: Comparable, limit: Comparable | None) -> bool: - if limit is not None and value < limit: - return True - return False - - -def is_above(value: Comparable, limit: Comparable | None) -> bool: - if limit is not None and value > limit: - return True - return False - def matches(i: CatalogItem, q: QueryFilter) -> bool: """ @@ -53,21 +29,21 @@ def matches(i: CatalogItem, q: QueryFilter) -> bool: """ if q.provider is not None and i.provider.lower() not in map(str.lower, q.provider): return False - if not is_between(i.price, q.min_price, q.max_price): + if not _is_between(i.price, q.min_price, q.max_price): return False if q.spot is not None and i.spot != q.spot: return False if q.cpu_arch and q.cpu_arch != i.cpu_arch: return False - if not is_between(i.cpu, q.min_cpu, q.max_cpu): + if not _is_between(i.cpu, q.min_cpu, q.max_cpu): return False - if not is_between(i.memory, q.min_memory, q.max_memory): + if not _is_between(i.memory, q.min_memory, q.max_memory): return False if not (q.min_gpu_count == 0 and i.gpu_count == 0): # GPU filters should not be applied to non-gpu offers if `q.min_gpu_count == 0`. if q.gpu_vendor and q.gpu_vendor != i.gpu_vendor: return False - if not is_between(i.gpu_count, q.min_gpu_count, q.max_gpu_count): + if not _is_between(i.gpu_count, q.min_gpu_count, q.max_gpu_count): return False if q.gpu_name is not None: if i.gpu_name is None: @@ -80,20 +56,22 @@ def matches(i: CatalogItem, q: QueryFilter) -> bool: if not i.gpu_name: return False cc = get_compute_capability(i.gpu_name) - if not cc or not is_between(cc, q.min_compute_capability, q.max_compute_capability): + if not cc or not _is_between(cc, q.min_compute_capability, q.max_compute_capability): return False - if not is_between( - i.gpu_memory if i.gpu_count > 0 else 0, q.min_gpu_memory, q.max_gpu_memory + if not _is_between( + i.gpu_memory if i.gpu_count > 0 and i.gpu_memory is not None else 0, + q.min_gpu_memory, + q.max_gpu_memory, ): return False - if not is_between( - (i.gpu_count * i.gpu_memory) if i.gpu_count > 0 else 0, + if not _is_between( + (i.gpu_count * i.gpu_memory) if i.gpu_count > 0 and i.gpu_memory is not None else 0, q.min_total_gpu_memory, q.max_total_gpu_memory, ): return False if i.disk_size is not None: - if not is_between(i.disk_size, q.min_disk_size, q.max_disk_size): + if not _is_between(i.disk_size, q.min_disk_size, q.max_disk_size): return False if q.allowed_flags is not None: if any(flag not in q.allowed_flags for flag in i.flags): @@ -116,7 +94,7 @@ def find_accelerators( def get_compute_capability(gpu_name: str) -> tuple[int, int] | None: - if accelerators := find_accelerators(names=[gpu_name], vendors=AcceleratorVendor.NVIDIA): + if accelerators := find_accelerators(names=[gpu_name], vendors=[AcceleratorVendor.NVIDIA]): assert isinstance(accelerators[0], NvidiaGPUInfo) return accelerators[0].compute_capability return None @@ -298,6 +276,10 @@ def is_nvidia_superchip(gpu_name: str) -> bool: ), ] + +# v5litepod = v5e +_TPU_VERSIONS = ["v2", "v3", "v4", "v5p", "v5litepod", "v6e"] + KNOWN_TPUS: list[TPUInfo] = [TPUInfo(name=version, memory=0) for version in _TPU_VERSIONS] KNOWN_INTEL_ACCELERATORS: list[IntelAcceleratorInfo] = [ @@ -326,3 +308,24 @@ def is_nvidia_superchip(gpu_name: str) -> bool: + KNOWN_INTEL_ACCELERATORS + KNOWN_TENSTORRENT_ACCELERATORS ) + + +Comparable = TypeVar("Comparable", int, float, tuple[int, int]) + + +def _is_between(value: Comparable, left: Comparable | None, right: Comparable | None) -> bool: + if _is_below(value, left) or _is_above(value, right): + return False + return True + + +def _is_below(value: Comparable, limit: Comparable | None) -> bool: + if limit is not None and value < limit: + return True + return False + + +def _is_above(value: Comparable, limit: Comparable | None) -> bool: + if limit is not None and value > limit: + return True + return False diff --git a/src/gpuhunt/_internal/default.py b/src/gpuhunt/_internal/default.py index ca0e2e3..a9dbea4 100644 --- a/src/gpuhunt/_internal/default.py +++ b/src/gpuhunt/_internal/default.py @@ -7,9 +7,22 @@ from typing_extensions import ParamSpec from gpuhunt._internal.catalog import Catalog +from gpuhunt._internal.errors import MissingCredsError +from gpuhunt.providers.base import OnlineProvider logger = logging.getLogger(__name__) +# Every provider in `ONLINE_PROVIDERS` must be listed here to be queried by `default_catalog`. +ONLINE_PROVIDER_MODULES = [ + ("gpuhunt.providers.crusoe", "CrusoeProvider"), + ("gpuhunt.providers.digitalocean", "DigitalOceanProvider"), + ("gpuhunt.providers.hotaisle", "HotAisleProvider"), + ("gpuhunt.providers.jarvislabs", "JarvisLabsProvider"), + ("gpuhunt.providers.seeweb", "SeewebProvider"), + ("gpuhunt.providers.vastai", "VastAIProvider"), + ("gpuhunt.providers.vultr", "VultrProvider"), +] + @functools.lru_cache def default_catalog() -> Catalog: @@ -19,23 +32,15 @@ def default_catalog() -> Catalog: """ catalog = Catalog() catalog.load() - for module, provider in [ - ("gpuhunt.providers.vastai", "VastAIProvider"), - ("gpuhunt.providers.crusoe", "CrusoeProvider"), - ("gpuhunt.providers.vultr", "VultrProvider"), - ("gpuhunt.providers.hotaisle", "HotAisleProvider"), - ("gpuhunt.providers.jarvislabs", "JarvisLabsProvider"), - ("gpuhunt.providers.digitalocean", "DigitalOceanProvider"), - ]: + for module_name, class_name in ONLINE_PROVIDER_MODULES: try: - module = importlib.import_module(module) - provider = getattr(module, provider)() - catalog.add_provider(provider) + module = importlib.import_module(module_name) + provider_class: type[OnlineProvider] = getattr(module, class_name) + catalog.add_provider(provider_class.from_env()) except ImportError: - logger.warning("Failed to import provider %s", provider) - except ValueError as e: - # Skip providers that require missing environment variables. Eg: HotAisleProvider - logger.warning("Skipping provider %s: %s", provider, e) + logger.warning("Failed to import provider %s", class_name) + except MissingCredsError as e: + logger.warning("Skipping provider %s: %s", class_name, e) return catalog @@ -45,7 +50,7 @@ def default_catalog() -> Catalog: CatalogMethod = Callable[Concatenate[Catalog, P], R] -def with_signature(method: CatalogMethod) -> Callable[[Method], Method]: +def with_signature(method: CatalogMethod[P, R]) -> Callable[[Method[P, R]], Method[P, R]]: """ Returns: decorator to add the signature of the Catalog method to the decorated method @@ -62,7 +67,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: @with_signature(Catalog.query) -def query(*args: P.args, **kwargs: P.kwargs) -> R: +def query(*args, **kwargs): """ Query the `default_catalog`. See `Catalog.query` for more details on parameters diff --git a/src/gpuhunt/_internal/errors.py b/src/gpuhunt/_internal/errors.py new file mode 100644 index 0000000..c9e3144 --- /dev/null +++ b/src/gpuhunt/_internal/errors.py @@ -0,0 +1,13 @@ +class GPUHuntError(Exception): + pass + + +class ProviderError(GPUHuntError): + pass + + +class MissingCredsError(ProviderError): + """ + Raised by `OnlineProvider.from_env` when a required environment variable is not set. + Online providers are skipped rather than failing the whole catalog when this is raised. + """ diff --git a/src/gpuhunt/_internal/models.py b/src/gpuhunt/_internal/models.py index 2ed3936..4a6599e 100644 --- a/src/gpuhunt/_internal/models.py +++ b/src/gpuhunt/_internal/models.py @@ -1,14 +1,11 @@ import enum -import json from collections.abc import Container -from dataclasses import asdict, dataclass, field, fields +from dataclasses import dataclass, field, fields from typing import ( ClassVar, Union, ) -from gpuhunt._internal.utils import empty_as_none - JSONType = Union[ None, bool, @@ -21,12 +18,6 @@ JSONObject = dict[str, JSONType] -def bool_loader(x: bool | str) -> bool: - if isinstance(x, bool): - return x - return x.lower() == "true" - - class AMDArchitecture(enum.Enum): CDNA = "CDNA" CDNA2 = "CDNA2" @@ -65,86 +56,6 @@ def cast(cls, value: Union["CPUArchitecture", str]) -> "CPUArchitecture": return cls(value.lower()) -@dataclass -class RawCatalogItem: - """ - An item stored in the catalog. - See `CatalogItem` for field descriptions. - """ - - instance_name: str | None - location: str | None - price: float | None - cpu: int | None - memory: float | None - gpu_count: int | None - gpu_name: str | None - gpu_memory: float | None - spot: bool | None - disk_size: float | None - gpu_vendor: str | None = None - flags: list[str] = field(default_factory=list) - cpu_arch: str | None = None - provider_data: JSONObject = field(default_factory=dict) - - def __post_init__(self) -> None: - self._process_gpu_vendor() - self._process_cpu_arch() - - def _process_gpu_vendor(self) -> None: - # This heuristic will be required indefinitely since we support historical catalogs. - is_tpu = False - gpu_name = self.gpu_name - if gpu_name and gpu_name.startswith("tpu-"): - is_tpu = True - self.gpu_name = gpu_name[4:] - gpu_vendor = self.gpu_vendor - if gpu_vendor is None: - if not self.gpu_count: - # None or 0 - return - if is_tpu: - self.gpu_vendor = AcceleratorVendor.GOOGLE.value - else: - self.gpu_vendor = AcceleratorVendor.NVIDIA.value - elif isinstance(gpu_vendor, AcceleratorVendor): - self.gpu_vendor = gpu_vendor.value - - def _process_cpu_arch(self) -> None: - # This heuristic will be required indefinitely since we support historical catalogs. - cpu_arch = self.cpu_arch - if cpu_arch is None: - self.cpu_arch = CPUArchitecture.X86.value - elif isinstance(cpu_arch, CPUArchitecture): - self.cpu_arch = cpu_arch.value - - @staticmethod - def from_dict(v: dict) -> "RawCatalogItem": - return RawCatalogItem( - instance_name=empty_as_none(v.get("instance_name")), - location=empty_as_none(v.get("location")), - price=empty_as_none(v.get("price"), loader=float), - cpu_arch=empty_as_none(v.get("cpu_arch")), - cpu=empty_as_none(v.get("cpu"), loader=int), - memory=empty_as_none(v.get("memory"), loader=float), - gpu_vendor=empty_as_none(v.get("gpu_vendor")), - gpu_count=empty_as_none(v.get("gpu_count"), loader=int), - gpu_name=empty_as_none(v.get("gpu_name")), - gpu_memory=empty_as_none(v.get("gpu_memory"), loader=float), - spot=empty_as_none(v.get("spot"), loader=bool_loader), - disk_size=empty_as_none(v.get("disk_size"), loader=float), - flags=v.get("flags", "").split(), - provider_data=json.loads(v.get("provider_data", "{}")), - ) - - def dict(self) -> dict[str, str | int | float | bool | None]: - return { - **asdict(self), - "flags": " ".join(self.flags), - "provider_data": json.dumps(self.provider_data), - } - - @dataclass class CatalogItem: """ @@ -172,6 +83,7 @@ class CatalogItem: Prefer defining a TypedDict within provider implementation. """ + provider: str instance_name: str location: str price: float @@ -182,41 +94,14 @@ class CatalogItem: gpu_memory: float | None spot: bool disk_size: float | None - provider: str gpu_vendor: AcceleratorVendor | None = None flags: list[str] = field(default_factory=list) - cpu_arch: CPUArchitecture | None = None + cpu_arch: CPUArchitecture = CPUArchitecture.X86 provider_data: JSONObject = field(default_factory=dict) def __post_init__(self) -> None: - self._process_gpu_vendor() - self._process_cpu_arch() - - def _process_gpu_vendor(self) -> None: - # This heuristic is only required until we update all providers to always set the vendor. - gpu_vendor = self.gpu_vendor - if gpu_vendor is None: - if not self.gpu_count: - # None or 0 - return - # GCPProvider already sets gpu_vendor, and all other providers only support Nvidia - self.gpu_vendor = AcceleratorVendor.NVIDIA - else: - # This cast to the enum is always required since RawCatalogItem.gpu_vendor - # is a string field (for (de)serialization purposes). - self.gpu_vendor = AcceleratorVendor.cast(gpu_vendor) - - def _process_cpu_arch(self) -> None: - # This heuristic is only required until we update all providers to always set the arch. - cpu_arch = self.cpu_arch - if cpu_arch is None: - self.cpu_arch = CPUArchitecture.X86 - else: - self.cpu_arch = CPUArchitecture.cast(cpu_arch) - - @staticmethod - def from_dict(v: dict, *, provider: str | None = None) -> "CatalogItem": - return CatalogItem(provider=provider, **asdict(RawCatalogItem.from_dict(v))) + if self.gpu_count and self.gpu_vendor is None: + raise ValueError(f"gpu_vendor is required when gpu_count is non-zero: {self}") @dataclass diff --git a/src/gpuhunt/_internal/storage.py b/src/gpuhunt/_internal/storage.py index 81e3f80..027d278 100644 --- a/src/gpuhunt/_internal/storage.py +++ b/src/gpuhunt/_internal/storage.py @@ -1,8 +1,14 @@ import csv -import dataclasses -from typing import TypeVar +import json +import logging +from collections.abc import Callable, Iterable, Iterator, Mapping +from typing import IO, TypeVar, overload -from gpuhunt._internal.models import RawCatalogItem +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, CPUArchitecture + +R = TypeVar("R") + +logger = logging.getLogger(__name__) CATALOG_V1_FIELDS = [ "instance_name", @@ -17,15 +23,99 @@ "disk_size", "gpu_vendor", ] -T = TypeVar("T", bound=RawCatalogItem) +# The columns of a v2 catalog file, in order. `provider` is not stored: the file name +# carries it. Listed explicitly rather than derived from `CatalogItem` so that adding a +# field to the model cannot change the published format. +CATALOG_V2_FIELDS = [ + "instance_name", + "location", + "price", + "cpu", + "memory", + "gpu_count", + "gpu_name", + "gpu_memory", + "spot", + "disk_size", + "gpu_vendor", + "flags", + "cpu_arch", + "provider_data", +] + +def item_to_row(item: CatalogItem) -> dict[str, str]: + return { + "instance_name": item.instance_name, + "location": item.location, + "price": str(item.price), + "cpu": str(item.cpu), + "memory": str(item.memory), + "gpu_count": str(item.gpu_count), + "gpu_name": _dump_optional(item.gpu_name), + "gpu_memory": _dump_optional(item.gpu_memory), + "spot": str(item.spot), + "disk_size": _dump_optional(item.disk_size), + "gpu_vendor": _dump_optional(item.gpu_vendor.value if item.gpu_vendor else None), + "flags": " ".join(item.flags), + "cpu_arch": item.cpu_arch.value, + "provider_data": json.dumps(item.provider_data), + } -def dump(items: list[T], path: str, *, cls: type[T] = RawCatalogItem): + +def item_from_row(row: Mapping[str, str], *, provider: str) -> CatalogItem: + gpu_count = _load_required(row, "gpu_count", int) + gpu_name = _load_optional(row, "gpu_name") + raw_gpu_vendor = _load_optional(row, "gpu_vendor") + gpu_vendor = AcceleratorVendor.cast(raw_gpu_vendor) if raw_gpu_vendor else None + # Catalogs published before the `gpu_vendor` column existed encode TPUs by prefixing + # the accelerator name, and imply Nvidia otherwise. Required as long as we support + # historical catalogs. + if gpu_name and gpu_name.startswith("tpu-"): + gpu_name = gpu_name[4:] + if gpu_vendor is None: + gpu_vendor = AcceleratorVendor.GOOGLE + elif gpu_vendor is None and gpu_count: + gpu_vendor = AcceleratorVendor.NVIDIA + # `cpu_arch` predates its column too, and x86 is what those catalogs contain. + raw_cpu_arch = _load_optional(row, "cpu_arch") + cpu_arch = CPUArchitecture.cast(raw_cpu_arch) if raw_cpu_arch else CPUArchitecture.X86 + return CatalogItem( + provider=provider, + instance_name=_load_required(row, "instance_name"), + location=_load_required(row, "location"), + price=_load_required(row, "price", float), + cpu=_load_required(row, "cpu", int), + memory=_load_required(row, "memory", float), + gpu_count=gpu_count, + gpu_name=gpu_name, + gpu_memory=_load_optional(row, "gpu_memory", float), + spot=_load_required(row, "spot", _load_bool), + disk_size=_load_optional(row, "disk_size", float), + gpu_vendor=gpu_vendor, + flags=(row.get("flags") or "").split(), + cpu_arch=cpu_arch, + provider_data=json.loads(row.get("provider_data") or "{}"), + ) + + +def dump(items: Iterable[CatalogItem], path: str) -> None: with open(path, "w", newline="") as f: - writer = csv.DictWriter(f, fieldnames=[field.name for field in dataclasses.fields(cls)]) + writer = csv.DictWriter(f, fieldnames=CATALOG_V2_FIELDS) writer.writeheader() for item in items: - writer.writerow(item.dict()) + writer.writerow(item_to_row(item)) + + +def load(f: IO[str], *, provider: str) -> Iterator[CatalogItem]: + reader = csv.DictReader(f) + for row in reader: + try: + yield item_from_row(row, provider=provider) + except ValueError as e: + logger.warning( + "Skipping malformed row in %s catalog at line %s: %s", provider, reader.line_num, e + ) def convert_catalog_v2_to_v1(path_v2: str, path_v1: str) -> None: @@ -36,3 +126,58 @@ def convert_catalog_v2_to_v1(path_v2: str, path_v1: str) -> None: for row in reader: if not row.get("flags"): writer.writerow(row) + + +def _dump_optional(value: str | float | None) -> str: + return "" if value is None else str(value) + + +@overload +def _load_required(row: Mapping[str, str], field: str, loader: Callable[[str], R]) -> R: ... + + +@overload +def _load_required(row: Mapping[str, str], field: str, loader: None = None) -> str: ... + + +def _load_required( + row: Mapping[str, str], field: str, loader: Callable[[str], R] | None = None +) -> str | R: + value = row.get(field) + if value is None: + raise ValueError(f"Required field {field!r} is missing") + if value == "": + raise ValueError(f"Required field {field!r} is empty") + return _apply_loader(field, value, loader) + + +@overload +def _load_optional(row: Mapping[str, str], field: str, loader: Callable[[str], R]) -> R | None: ... + + +@overload +def _load_optional(row: Mapping[str, str], field: str, loader: None = None) -> str | None: ... + + +def _load_optional( + row: Mapping[str, str], field: str, loader: Callable[[str], R] | None = None +) -> str | R | None: + value = row.get(field) + if not value: + return None + return _apply_loader(field, value, loader) + + +def _apply_loader(field: str, value: str, loader: Callable[[str], R] | None) -> str | R: + if loader is None: + return value + try: + return loader(value) + except ValueError as e: + raise ValueError(f"Cannot parse field {field!r}: {e}") from e + + +def _load_bool(value: str) -> bool: + if value.lower() not in ("true", "false"): + raise ValueError(f"Not a boolean: {value!r}") + return value.lower() == "true" diff --git a/src/gpuhunt/_internal/utils.py b/src/gpuhunt/_internal/utils.py index 4464505..99cac00 100644 --- a/src/gpuhunt/_internal/utils.py +++ b/src/gpuhunt/_internal/utils.py @@ -1,6 +1,6 @@ import logging import sys -from collections.abc import Callable +from typing import TypeVar def configure_logging() -> None: @@ -11,12 +11,16 @@ def configure_logging() -> None: ) -def empty_as_none(value: str | None, loader: Callable | None = None): - if value is None or value == "": - return None - if loader is not None: - return loader(value) - return value +T = TypeVar("T") + + +def get_or_error(v: T | None, name: str = "value") -> T: + """ + Unpacks an optional value. Used to denote that None is not possible in the current context. + """ + if v is None: + raise ValueError(f"Expected {name} to be set") + return v def parse_compute_capability( diff --git a/src/gpuhunt/providers/__init__.py b/src/gpuhunt/providers/__init__.py index b939bbb..e69de29 100644 --- a/src/gpuhunt/providers/__init__.py +++ b/src/gpuhunt/providers/__init__.py @@ -1,45 +0,0 @@ -from abc import ABC, abstractmethod - -from gpuhunt._internal.models import QueryFilter, RawCatalogItem - - -class AbstractProvider(ABC): - """ - Abstract class for cloud provider implementations. - - Attributes: - NAME: (class variable) The name of the provider. - """ - - NAME: str = "abstract" # Override in subclasses - - @abstractmethod - def get( - self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: - """ - Return a list of available instance offers. Offers should be ordered by priority. In most - cases - by price, ascending. - - Args: - query_filter: Set of filters requested by the user. Only used with online providers. - Filters are safe to ignore, as they are also enforced by `gpuhunt` after calling - `get`. However, they can be used to reduce the number or size of API requests if - the provider's API supports filtering by GPU, RAM, region, etc. - balance_resources: Whether the instance resources (CPU, RAM, disk) should be - adjusted to better match the GPU. Only used with online providers. Only relevant - to cloud providers that allow configuring instance CPU, RAM, and disk. - """ - - pass - - @classmethod - def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: - """ - Return a subset of offers that should be stored in the catalog. - - Only used with offline providers. Only implement this method if there are reasons to omit - some offers from the catalog. - """ - - return offers diff --git a/src/gpuhunt/providers/aws.py b/src/gpuhunt/providers/aws.py index aa2b444..835d41f 100644 --- a/src/gpuhunt/providers/aws.py +++ b/src/gpuhunt/providers/aws.py @@ -13,8 +13,8 @@ import requests from botocore.exceptions import ClientError, ConnectTimeoutError, EndpointConnectionError -from gpuhunt._internal.models import QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) ec2_pricing_url = ( @@ -91,7 +91,7 @@ } -class AWSProvider(AbstractProvider): +class AWSProvider(OfflineProvider): """ AWSProvider parses Bulk API index file for AmazonEC2 in all regions and fills missing GPU details @@ -115,28 +115,29 @@ def __init__(self, cache_path: str | None = None): def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: if not os.path.exists(self.cache_path): self._download_pricing_file() - offers = [] + offers: list[CatalogItem] = [] with open(self.cache_path, newline="") as f: for _ in range(disclaimer_rows_skip): f.readline() reader: Iterable[dict[str, str]] = csv.DictReader(f) for row in reader: - if self.skip(row): + if self._skip(row): continue gpu_count = _parse_gpu_count(row["GPU"]) if gpu_count is None: continue - offer = RawCatalogItem( + offer = CatalogItem( + provider=AWSProvider.NAME, instance_name=row["Instance Type"], location=row["Region Code"], price=float(row["PricePerUnit"]), cpu=int(row["vCPU"]), memory=_parse_memory(row["Memory"]), - gpu_vendor=None, + gpu_vendor=AcceleratorVendor.NVIDIA if gpu_count else None, gpu_count=gpu_count, spot=False, gpu_name=None, @@ -144,11 +145,11 @@ def get( disk_size=None, ) offers.append(offer) - self.fill_gpu_details(offers) - offers = self.add_spots(offers) + self._fill_gpu_details(offers) + offers = self._with_spot_offers(offers) return sorted(offers, key=lambda i: i.price) - def skip(self, row: dict[str, str]) -> bool: + def _skip(self, row: dict[str, str]) -> bool: if any(row["Instance Type"].startswith(family) for family in previous_generation_families): return True for key, values in pricing_filters.items(): @@ -156,7 +157,7 @@ def skip(self, row: dict[str, str]) -> bool: return True return False - def fill_gpu_details(self, offers: list[RawCatalogItem]): + def _fill_gpu_details(self, offers: list[CatalogItem]) -> None: regions = defaultdict(list) non_ec2_api_regions = set() for offer in offers: @@ -189,7 +190,7 @@ def fill_gpu_details(self, offers: list[RawCatalogItem]): if "GpuInfo" in i: gpu = i["GpuInfo"]["Gpus"][0] gpus[i["InstanceType"]] = ( - GPU_NAME_MAPPING.get(gpu["Name"], gpu["Name"]), + GPU_NAME_MAPPING.get(gpu["Name"]) or gpu["Name"], _get_gpu_memory_gib( gpu["Name"], gpu["MemoryInfo"]["SizeInMiB"] ), @@ -258,7 +259,7 @@ def _add_spots_worker( } ], InstanceTypes=list(instance_types), - StartTime=datetime.datetime.utcnow(), + StartTime=datetime.datetime.now(tz=datetime.timezone.utc), ) instance_prices = defaultdict(list) @@ -331,7 +332,7 @@ def _download_pricing_file(self) -> None: e, ) - def add_spots(self, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: + def _with_spot_offers(self, offers: list[CatalogItem]) -> list[CatalogItem]: region_instances = defaultdict(set) non_ec2_api_regions = set() for offer in offers: @@ -354,7 +355,7 @@ def add_spots(self, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: for future in as_completed(future_to_region): spot_prices.update(future.result()) - spot_offers = [] + spot_offers: list[CatalogItem] = [] for offer in offers: if (price := spot_prices.get((offer.instance_name, offer.location))) is None: continue @@ -365,7 +366,7 @@ def add_spots(self, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: return offers + spot_offers @classmethod - def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: + def filter(cls, offers: list[CatalogItem]) -> list[CatalogItem]: return [ i for i in offers @@ -416,6 +417,8 @@ def _get_gpu_memory_gib(gpu_name: str, reported_memory_mib: int) -> float: def _parse_memory(s: str) -> float: r = re.match(r"^([0-9.]+) GiB$", s) + if r is None: + raise ValueError(f"Cannot parse memory: {s!r}") return float(r.group(1)) @@ -430,7 +433,7 @@ def _parse_gpu_count(s: str) -> int | None: def _get_ec2_api_regions() -> set[str]: - session = boto3.session.Session() + session = boto3.session.Session() # pyright: ignore[reportAttributeAccessIssue] return { region for partition in session.get_available_partitions() diff --git a/src/gpuhunt/providers/azure.py b/src/gpuhunt/providers/azure.py index 44dc452..2f6ffcd 100644 --- a/src/gpuhunt/providers/azure.py +++ b/src/gpuhunt/providers/azure.py @@ -6,6 +6,7 @@ import time from collections import namedtuple from collections.abc import Iterable +from dataclasses import dataclass from queue import Queue from threading import Thread @@ -15,8 +16,10 @@ from azure.identity import DefaultAzureCredential from azure.mgmt.compute import ComputeManagementClient -from gpuhunt._internal.models import QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt import AcceleratorVendor +from gpuhunt._internal.models import CatalogItem, QueryFilter +from gpuhunt._internal.utils import get_or_error +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) prices_url = "https://prices.azure.com/api/retail/prices" @@ -65,7 +68,34 @@ ] -class AzureProvider(AbstractProvider): +@dataclass +class _InstanceSpec: + instance_name: str + cpu: int + memory: float + gpu_count: int + gpu_name: str | None + gpu_memory: float | None + gpu_vendor: AcceleratorVendor | None + + def to_catalog_item(self, *, location: str, price: float, spot: bool) -> CatalogItem: + return CatalogItem( + provider=AzureProvider.NAME, + instance_name=self.instance_name, + location=location, + price=price, + cpu=self.cpu, + memory=self.memory, + gpu_count=self.gpu_count, + gpu_name=self.gpu_name, + gpu_memory=self.gpu_memory, + spot=spot, + disk_size=None, + gpu_vendor=self.gpu_vendor, + ) + + +class AzureProvider(OfflineProvider): NAME = "azure" def __init__( @@ -140,44 +170,43 @@ def _get_pages_worker(self, q: Queue, stride: int, worker_id: int): def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: - offers = [] + ) -> list[CatalogItem]: + offers: list[CatalogItem] = [] + instance_name_to_spec_map = self.get_instance_specs() for page in self.get_pages(): - for item in page: - if is_retired(item["armSkuName"]): + for sku_item in page: + if is_retired(sku_item["armSkuName"]): continue - if not item["armSkuName"]: + if not sku_item["armSkuName"]: continue - price = float(item["retailPrice"]) + price = float(sku_item["retailPrice"]) if math.isclose(price, 0): continue - offer = RawCatalogItem( - instance_name=item["armSkuName"], - location=item["armRegionName"], + spec = instance_name_to_spec_map.get(sku_item["armSkuName"]) + if spec is None: + continue + offer = spec.to_catalog_item( + location=sku_item["armRegionName"], price=price, - spot="Spot" in item["meterName"], - cpu=None, - memory=None, - gpu_vendor=None, - gpu_count=None, - gpu_name=None, - gpu_memory=None, - disk_size=None, + spot="Spot" in sku_item["meterName"], ) offers.append(offer) - offers = self.fill_details(offers) return sorted(offers, key=lambda i: i.price) - def fill_details(self, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: + def get_instance_specs(self) -> dict[str, _InstanceSpec]: logger.info("Fetching instance details") - instances = {} + instance_name_to_spec_map = {} resources = self.client.resource_skus.list() for resource in resources: + assert resource.name is not None if resource.resource_type != "virtualMachines": continue if is_retired(resource.name): continue - capabilities = {pair.name: pair.value for pair in resource.capabilities} + capabilities = { + pair.name: pair.value + for pair in get_or_error(resource.capabilities, "resource capabilities") + } cpu = capabilities.get("vCPUs") memory = capabilities.get("MemoryGB") if not cpu: @@ -188,39 +217,24 @@ def fill_details(self, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: continue gpu_count, gpu_name, gpu_memory = 0, None, None if "GPUs" in capabilities: - gpu_count = int(capabilities["GPUs"]) + gpu_count = int(get_or_error(capabilities["GPUs"], "GPUs capability")) gpu_name, gpu_memory = get_gpu_name_memory(resource.name) if gpu_name is None and gpu_count: logger.warning("Can't parse VM name: %s", resource.name) continue - instances[resource.name] = RawCatalogItem( + instance_name_to_spec_map[resource.name] = _InstanceSpec( instance_name=resource.name, - cpu=cpu, + cpu=int(cpu), memory=float(memory), - gpu_vendor=None, + gpu_vendor=AcceleratorVendor.NVIDIA if gpu_count else None, gpu_count=gpu_count, gpu_name=gpu_name, gpu_memory=gpu_memory, - location=None, - price=None, - spot=None, - disk_size=None, ) - with_details = [] - for offer in offers: - if (resources := instances.get(offer.instance_name)) is None: - continue - offer.cpu = resources.cpu - offer.memory = resources.memory - offer.gpu_count = resources.gpu_count - offer.gpu_name = resources.gpu_name - offer.gpu_memory = resources.gpu_memory - offer.gpu_vendor = resources.gpu_vendor - with_details.append(offer) - return with_details + return instance_name_to_spec_map @classmethod - def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: + def filter(cls, offers: list[CatalogItem]) -> list[CatalogItem]: vm_series = [ VMSeries(r"D(\d+)s_v6", None, None), # Dsv6-series VMSeries( diff --git a/src/gpuhunt/providers/base.py b/src/gpuhunt/providers/base.py new file mode 100644 index 0000000..ea3f6fc --- /dev/null +++ b/src/gpuhunt/providers/base.py @@ -0,0 +1,92 @@ +import os +from abc import ABC, abstractmethod + +from typing_extensions import Self + +from gpuhunt._internal.errors import MissingCredsError +from gpuhunt._internal.models import CatalogItem, QueryFilter + + +class AbstractProvider(ABC): + """ + Abstract class for cloud provider implementations. + + Implement `OnlineProvider` or `OfflineProvider` rather than subclassing this directly. + + Attributes: + NAME: (class variable) The name of the provider. + """ + + NAME: str = "abstract" # Override in subclasses + + @abstractmethod + def get( + self, query_filter: QueryFilter | None = None, balance_resources: bool = True + ) -> list[CatalogItem]: + """ + Return a list of available instance offers. Offers should be ordered by priority. In most + cases - by price, ascending. + + Args: + query_filter: Set of filters requested by the user. Only used with online providers. + Filters are safe to ignore, as they are also enforced by `gpuhunt` after calling + `get`. However, they can be used to reduce the number or size of API requests if + the provider's API supports filtering by GPU, RAM, region, etc. + balance_resources: Whether the instance resources (CPU, RAM, disk) should be + adjusted to better match the GPU. Only used with online providers. Only relevant + to cloud providers that allow configuring instance CPU, RAM, and disk. + """ + + pass + + +class OnlineProvider(AbstractProvider, ABC): + """ + A provider queried at request time, listed in `ONLINE_PROVIDERS`. + + Online providers are constructed by `default_catalog()` in the user's process, so they must + be constructible from the environment alone. + """ + + @classmethod + @abstractmethod + def from_env(cls) -> Self: + """ + Construct the provider from environment variables. + + Raises: + MissingCredsError: If a required environment variable is not set. The provider is + then skipped instead of failing the whole catalog. + """ + + pass + + +class OfflineProvider(AbstractProvider, ABC): + """ + A provider collected into the published catalog, listed in `OFFLINE_PROVIDERS`. + Credentials are supplied by the caller, so a missing one is an error. + """ + + @classmethod + def filter(cls, offers: list[CatalogItem]) -> list[CatalogItem]: + """ + Return a subset of offers that should be stored in the catalog. + Implement this method if there are reasons to omit some offers from the catalog. + """ + + return offers + + +def get_creds_env(name: str) -> str: + """ + Reads an environment variable required to construct a provider. + + Raises: + MissingCredsError: If the variable is not set or empty. + """ + + value = os.getenv(name) + if not value: + raise MissingCredsError(f"Set the {name} environment variable") + return value diff --git a/src/gpuhunt/providers/cloudrift.py b/src/gpuhunt/providers/cloudrift.py index 38b6965..28b7e38 100644 --- a/src/gpuhunt/providers/cloudrift.py +++ b/src/gpuhunt/providers/cloudrift.py @@ -3,9 +3,9 @@ import requests -from gpuhunt import QueryFilter, RawCatalogItem +from gpuhunt import CatalogItem, QueryFilter from gpuhunt._internal.models import AcceleratorVendor -from gpuhunt.providers import AbstractProvider +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) @@ -13,25 +13,25 @@ CLOUDRIFT_API_VERSION = "2025-03-21" -class CloudRiftProvider(AbstractProvider): +class CloudRiftProvider(OfflineProvider): NAME = "cloudrift" def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: instance_types = self._get_instance_types() - instance_types = [ - inst for instance in instance_types for inst in generate_instances(instance) - ] - return sorted(instance_types, key=lambda x: x.price) + offers = [offer for instance in instance_types for offer in _make_offers(instance)] + return sorted(offers, key=lambda i: i.price) - def _get_instance_types(self): + def _get_instance_types(self) -> list[dict]: request_data = {"selector": {"ByServiceAndLocation": {"services": ["vm"]}}} response_data = _make_request("instance-types/list", request_data) + if not isinstance(response_data, dict): + raise ValueError(f"Unexpected instance-types/list response: {response_data!r}") return response_data["instance_types"] -def generate_instances(instance) -> list[RawCatalogItem]: +def _make_offers(instance: dict) -> list[CatalogItem]: instance_gpu_brand = instance["brand_short"] gpu_info = next( (gpu_record for gpu_record in GPU_MAP if gpu_record[0] in instance_gpu_brand), None @@ -44,7 +44,8 @@ def generate_instances(instance) -> list[RawCatalogItem]: instance_types = [] for variant in instance["variants"]: for location, _count in variant["nodes_per_dc"].items(): - raw = RawCatalogItem( + raw = CatalogItem( + provider=CloudRiftProvider.NAME, instance_name=variant["name"], location=location, spot=False, @@ -70,7 +71,7 @@ def generate_instances(instance) -> list[RawCatalogItem]: ] -def _make_request(endpoint: str, request_data: dict) -> dict | str | None: +def _make_request(endpoint: str, request_data: dict) -> dict | str: server = os.environ.get("CLOUDRIFT_SERVER_ADDRESS", CLOUDRIFT_SERVER_ADDRESS) response = requests.request( "POST", @@ -80,10 +81,7 @@ def _make_request(endpoint: str, request_data: dict) -> dict | str | None: ) if not response.ok: response.raise_for_status() - try: - response_json = response.json() - if isinstance(response_json, str): - return response_json - return response_json["data"] - except requests.exceptions.JSONDecodeError: - return None + response_json = response.json() + if isinstance(response_json, str): + return response_json + return response_json["data"] diff --git a/src/gpuhunt/providers/crusoe.py b/src/gpuhunt/providers/crusoe.py index 4d9b13b..d0e5b47 100644 --- a/src/gpuhunt/providers/crusoe.py +++ b/src/gpuhunt/providers/crusoe.py @@ -1,21 +1,19 @@ import base64 -import copy import datetime import hashlib import hmac import logging -import os from collections import defaultdict import requests from gpuhunt._internal.models import ( AcceleratorVendor, + CatalogItem, CPUArchitecture, QueryFilter, - RawCatalogItem, ) -from gpuhunt.providers import AbstractProvider +from gpuhunt.providers.base import OnlineProvider, get_creds_env logger = logging.getLogger(__name__) @@ -63,29 +61,30 @@ } -class CrusoeProvider(AbstractProvider): +class CrusoeProvider(OnlineProvider): NAME = "crusoe" def __init__( self, - access_key: str | None = None, - secret_key: str | None = None, - project_id: str | None = None, + access_key: str, + secret_key: str, + project_id: str, ): - self.access_key = access_key or os.getenv("CRUSOE_ACCESS_KEY") - self.secret_key = secret_key or os.getenv("CRUSOE_SECRET_KEY") - self.project_id = project_id or os.getenv("CRUSOE_PROJECT_ID") - - if not self.access_key: - raise ValueError("Set the CRUSOE_ACCESS_KEY environment variable.") - if not self.secret_key: - raise ValueError("Set the CRUSOE_SECRET_KEY environment variable.") - if not self.project_id: - raise ValueError("Set the CRUSOE_PROJECT_ID environment variable.") + self.access_key = access_key + self.secret_key = secret_key + self.project_id = project_id + + @classmethod + def from_env(cls) -> "CrusoeProvider": + return cls( + access_key=get_creds_env("CRUSOE_ACCESS_KEY"), + secret_key=get_creds_env("CRUSOE_SECRET_KEY"), + project_id=get_creds_env("CRUSOE_PROJECT_ID"), + ) def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: instance_types = self._get_instance_types() type_specs = {t["product_name"]: t for t in instance_types} @@ -96,15 +95,15 @@ def get( capacities = self._get_capacities() available = _get_available_type_locations(capacities) - offers = [] + offers: list[CatalogItem] = [] for product_name, locations in available.items(): spec = type_specs.get(product_name) if spec is None: logger.warning("Capacity for unknown instance type %s, skipping", product_name) continue - items = _make_catalog_items(product_name, spec, locations) - offers.extend(items) + product_offers = _make_offers(product_name, spec, locations) + offers.extend(product_offers) return sorted(offers, key=lambda i: i.price) @@ -143,11 +142,11 @@ def _request(self, method: str, path: str, params: dict | None = None) -> reques return requests.request(method, url, headers=headers, params=params, timeout=TIMEOUT) -def _get_cpu_arch(spec: dict) -> str: +def _get_cpu_arch(spec: dict) -> CPUArchitecture: cpu_type = spec.get("cpu_type", "") if cpu_type == "arm64": - return CPUArchitecture.ARM.value - return CPUArchitecture.X86.value + return CPUArchitecture.ARM + return CPUArchitecture.X86 def _get_available_type_locations(capacities: list[dict]) -> dict[str, list[str]]: @@ -163,21 +162,19 @@ def _get_available_type_locations(capacities: list[dict]) -> dict[str, list[str] return dict(result) -def _make_catalog_items( - product_name: str, spec: dict, locations: list[str] -) -> list[RawCatalogItem]: +def _make_offers(product_name: str, spec: dict, locations: list[str]) -> list[CatalogItem]: gpu_type = spec.get("gpu_type", "") num_gpu = spec.get("num_gpu", 0) if num_gpu > 0 and gpu_type: - return _make_gpu_items(product_name, spec, gpu_type, locations) + return _make_gpu_offers(product_name, spec, gpu_type, locations) else: - return _make_cpu_items(product_name, spec, locations) + return _make_cpu_offers(product_name, spec, locations) -def _make_gpu_items( +def _make_gpu_offers( product_name: str, spec: dict, gpu_type: str, locations: list[str] -) -> list[RawCatalogItem]: +) -> list[CatalogItem]: gpu_info = GPU_TYPE_MAP.get(gpu_type) if gpu_info is None: logger.warning("Unknown GPU type %s for %s, skipping", gpu_type, product_name) @@ -191,69 +188,58 @@ def _make_gpu_items( gpu_name, gpu_vendor, gpu_memory = gpu_info on_demand_per_gpu, spot_per_gpu = pricing num_gpu = spec["num_gpu"] - - template = RawCatalogItem( - instance_name=product_name, - location=None, - price=None, - cpu=spec["cpu_cores"], - memory=float(spec["memory_gb"]), - gpu_vendor=gpu_vendor.value, - gpu_count=num_gpu, - gpu_name=gpu_name, - gpu_memory=gpu_memory, - spot=None, - disk_size=float(spec["disk_gb"]) if spec.get("disk_gb") else None, - cpu_arch=_get_cpu_arch(spec), - # disk_gb: ephemeral NVMe size in GB (0 = no ephemeral disk). - # Used by dstack to decide whether to create a persistent data disk. - provider_data={"disk_gb": spec.get("disk_gb", 0)}, - ) - - items = [] + offers: list[CatalogItem] = [] for location in locations: - on_demand = copy.deepcopy(template) - on_demand.location = location - on_demand.spot = False - on_demand.price = round(num_gpu * on_demand_per_gpu, 2) - items.append(on_demand) - + on_demand_item = CatalogItem( + provider=CrusoeProvider.NAME, + instance_name=product_name, + location=location, + price=round(num_gpu * on_demand_per_gpu, 2), + cpu=spec["cpu_cores"], + memory=float(spec["memory_gb"]), + gpu_vendor=gpu_vendor, + gpu_count=num_gpu, + gpu_name=gpu_name, + gpu_memory=gpu_memory, + spot=False, + disk_size=float(spec["disk_gb"]) if spec.get("disk_gb") else None, + cpu_arch=_get_cpu_arch(spec), + # disk_gb: ephemeral NVMe size in GB (0 = no ephemeral disk). + # Used by dstack to decide whether to create a persistent data disk. + provider_data={"disk_gb": spec.get("disk_gb", 0)}, + ) + offers.append(on_demand_item) # TODO: Enable spot offers once we confirm how to request spot billing # via the VM create API (POST /v1alpha5/projects/{pid}/compute/vms/instances). # The API schema doesn't have an obvious spot/billing_type field. - return items + return offers -def _make_cpu_items(product_name: str, spec: dict, locations: list[str]) -> list[RawCatalogItem]: +def _make_cpu_offers(product_name: str, spec: dict, locations: list[str]) -> list[CatalogItem]: prefix = product_name.split(".")[0] per_vcpu = CPU_PRICING.get(prefix) if per_vcpu is None: logger.warning("No pricing for CPU prefix %s (%s), skipping", prefix, product_name) return [] - cpu_cores = spec["cpu_cores"] - template = RawCatalogItem( - instance_name=product_name, - location=None, - price=None, - cpu=cpu_cores, - memory=float(spec["memory_gb"]), - gpu_vendor=None, - gpu_count=0, - gpu_name=None, - gpu_memory=None, - spot=False, - disk_size=float(spec["disk_gb"]) if spec.get("disk_gb") else None, - cpu_arch=_get_cpu_arch(spec), - provider_data={"disk_gb": spec.get("disk_gb", 0)}, - ) - - items = [] + offers: list[CatalogItem] = [] for location in locations: - item = copy.deepcopy(template) - item.location = location - item.price = round(cpu_cores * per_vcpu, 2) - items.append(item) - - return items + item = CatalogItem( + provider=CrusoeProvider.NAME, + instance_name=product_name, + location=location, + price=round(cpu_cores * per_vcpu, 2), + cpu=cpu_cores, + memory=float(spec["memory_gb"]), + gpu_vendor=None, + gpu_count=0, + gpu_name=None, + gpu_memory=None, + spot=False, + disk_size=float(spec["disk_gb"]) if spec.get("disk_gb") else None, + cpu_arch=_get_cpu_arch(spec), + provider_data={"disk_gb": spec.get("disk_gb", 0)}, + ) + offers.append(item) + return offers diff --git a/src/gpuhunt/providers/digitalocean.py b/src/gpuhunt/providers/digitalocean.py index b6ef23c..7235235 100644 --- a/src/gpuhunt/providers/digitalocean.py +++ b/src/gpuhunt/providers/digitalocean.py @@ -4,8 +4,8 @@ import requests from gpuhunt._internal.constraints import get_gpu_vendor -from gpuhunt._internal.models import QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import CatalogItem, QueryFilter +from gpuhunt.providers.base import OnlineProvider, get_creds_env logger = logging.getLogger(__name__) @@ -13,26 +13,30 @@ STANDARD_CLOUD_API_URL = "https://api.digitalocean.com" -class DigitalOceanProvider(AbstractProvider): +class DigitalOceanProvider(OnlineProvider): NAME = "digitalocean" - def __init__(self, api_key: str | None = None, api_url: str | None = None): - self.api_key = api_key or os.getenv("DIGITAL_OCEAN_API_KEY") - if not self.api_key: - raise ValueError("Set the DIGITAL_OCEAN_API_KEY environment variable.") + def __init__(self, api_key: str, api_url: str = STANDARD_CLOUD_API_URL): + self.api_key = api_key + self.api_url = api_url - self.api_url = api_url or os.getenv("DIGITAL_OCEAN_API_URL", STANDARD_CLOUD_API_URL) + @classmethod + def from_env(cls) -> "DigitalOceanProvider": + return cls( + api_key=get_creds_env("DIGITAL_OCEAN_API_KEY"), + api_url=os.getenv("DIGITAL_OCEAN_API_URL", STANDARD_CLOUD_API_URL), + ) def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: offers = self.fetch_offers() return sorted(offers, key=lambda i: i.price) - def fetch_offers(self) -> list[RawCatalogItem]: + def fetch_offers(self) -> list[CatalogItem]: url = "/v2/sizes" response = self._make_request("GET", url) - return convert_response_to_raw_catalog_items(response) + return _make_offers(response) def _make_request(self, method: str, url: str): full_url = f"{self.api_url}{url}" @@ -48,9 +52,9 @@ def _make_request(self, method: str, url: str): return response -def convert_response_to_raw_catalog_items(response) -> list[RawCatalogItem]: +def _make_offers(response) -> list[CatalogItem]: data = response.json() - offers = [] + offers: list[CatalogItem] = [] for size in data["sizes"]: gpu_info = size.get("gpu_info") @@ -83,7 +87,8 @@ def convert_response_to_raw_catalog_items(response) -> list[RawCatalogItem]: # Creates an offer for each available region. # If regions list is empty, instance type is not available. for region in size["regions"]: - offer = RawCatalogItem( + offer = CatalogItem( + provider=DigitalOceanProvider.NAME, instance_name=size["slug"], location=region, price=size["price_hourly"], diff --git a/src/gpuhunt/providers/gcp.py b/src/gpuhunt/providers/gcp.py index 7066ab3..1d619b6 100644 --- a/src/gpuhunt/providers/gcp.py +++ b/src/gpuhunt/providers/gcp.py @@ -1,4 +1,5 @@ import copy +import dataclasses import enum import importlib.resources import json @@ -18,14 +19,13 @@ from google.cloud.location import locations_pb2 from typing_extensions import NotRequired, TypedDict -from gpuhunt._internal.models import AcceleratorVendor, QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) compute_service = "services/6F81-5844-456A" AcceleratorDetails = namedtuple("AcceleratorDetails", ["name", "memory"]) -# As of 2024-14-08, this mapping contains only Nvidia accelerators; update gpu_vendor -# inferring code in fill_gpu_vendors_and_names() if a non-Nvidia accelerator is added +# Update gpu_vendor detection if adding non-nvidia GPUs. accelerator_details = { "nvidia-b200": AcceleratorDetails("B200", 180.0), "nvidia-a100-80gb": AcceleratorDetails("A100", 80.0), @@ -146,7 +146,20 @@ def load_tpu_pricing() -> dict: TPU_PRICING_TABLE = load_tpu_pricing() -class GCPProvider(AbstractProvider): +@dataclass +class _MachineType: + """A machine type in a zone, with accelerators attached but no price yet.""" + + instance_name: str + location: str # zone + cpu: int + memory: float + gpu_count: int + gpu_name: str | None + gpu_memory: float | None + + +class GCPProvider(OfflineProvider): NAME = "gcp" def __init__(self, project: str): @@ -157,10 +170,65 @@ def __init__(self, project: str): self.regions_client = compute_v1.RegionsClient() self.cloud_catalog_client = billing_v1.CloudCatalogClient() - def list_preconfigured_instances(self) -> list[RawCatalogItem]: - def _list_zone_instances(zone: str) -> list[RawCatalogItem]: - zone_instances = [] - logger.info("Fetching instances for zone %s", zone) + def get( + self, query_filter: QueryFilter | None = None, balance_resources: bool = True + ) -> list[CatalogItem]: + machine_types = self._list_machine_types() + machine_types += self._make_gpu_machine_types(machine_types) + offers = self._make_offers(machine_types) + offers.extend(_make_tpu_offers(self.project)) + _set_flags(offers) + offers = _with_legacy_g4_preview(offers) + return sorted(offers, key=lambda i: i.price) + + @classmethod + def filter(cls, offers: list[CatalogItem]) -> list[CatalogItem]: + return [ + i + for i in offers + if ( + any( + i.instance_name.startswith(family) + for family in [ + "m4-", + "c4-", + "n4-", + "h3-", + "n2-", + "e2-medium", + "e2-standard-", + "e2-highmem-", + "e2-highcpu-", + "m1-", + "a2-", + "g2-", + ] + ) + or (i.gpu_name and i.gpu_name not in ["K80", "P4"]) + ) + and not ( + # Filter out on-demand offers that are not actually available on demand. + # https://cloud.google.com/compute/docs/accelerator-optimized-machines#consumption_option_availability_by_machine_type + i.spot == False + and not cast(GCPCatalogItemProviderData, i.provider_data).get( + "is_dws_calendar_mode" + ) + and ( + i.instance_name.startswith("a4x-") + or i.instance_name.startswith("a4-") + or i.instance_name.startswith("a3-ultragpu-") + or ( + i.instance_name.startswith("a3-highgpu-") + and (i.gpu_count is None or i.gpu_count < 8) + ) + ) + ) + ] + + def _list_machine_types(self) -> list[_MachineType]: + def _list_zone_machine_types(zone: str) -> list[_MachineType]: + zone_machine_types: list[_MachineType] = [] + logger.info("Fetching machine types for zone %s", zone) for machine_type in self.machine_types_client.list(project=self.project, zone=zone): if machine_type.deprecated.state == compute_v1.DeprecationStatus.State.DEPRECATED: continue @@ -172,162 +240,113 @@ def _list_zone_instances(zone: str) -> list[RawCatalogItem]: logger.warning("Unknown accelerator type: %s", accelerator) continue - instance = RawCatalogItem( + machine_type = _MachineType( instance_name=machine_type.name, location=zone, cpu=machine_type.guest_cpus, memory=round(machine_type.memory_mb / 1024, 1), gpu_count=(machine_type.accelerators[0].guest_accelerator_count if gpu else 0), # gpu_name is canonicalized and gpu_vendor is set later - # in fill_gpu_vendors_and_names(), for now we use AcceleratorType.name + # in _make_offers(), for now we use AcceleratorType.name # as a name (it contains a vendor prefix like "nvidia-") gpu_name=( machine_type.accelerators[0].guest_accelerator_type if gpu else None ), - gpu_vendor=None, gpu_memory=gpu.memory if gpu else None, - price=None, - spot=None, - disk_size=None, ) - zone_instances.append(instance) - return zone_instances + zone_machine_types.append(machine_type) + return zone_machine_types - instances = [] + machine_types: list[_MachineType] = [] futures = [] with ThreadPoolExecutor(max_workers=8) as ex: for region in self.regions_client.list(project=self.project): for zone_url in region.zones: zone = zone_url.split("/")[-1] - futures.append(ex.submit(_list_zone_instances, zone)) + futures.append(ex.submit(_list_zone_machine_types, zone)) for future in as_completed(futures): - instances.extend(future.result()) - return instances + machine_types.extend(future.result()) + return machine_types - def add_gpus(self, instances: list[RawCatalogItem]): - def _list_zone_instances( - zone: str, zone_n1_instances: list[RawCatalogItem] - ) -> list[RawCatalogItem]: + def _make_gpu_machine_types(self, machine_types: list[_MachineType]) -> list[_MachineType]: + def _make_zone_gpu_machine_types( + zone: str, zone_n1_machine_types: list[_MachineType] + ) -> list[_MachineType]: logger.info("Fetching GPUs for zone %s", zone) - zone_instances = [] + zone_machine_types: list[_MachineType] = [] for accelerator in self.accelerator_types_client.list(project=self.project, zone=zone): if accelerator.name not in accelerator_limits: continue for n, limit in zip(accelerator_counts, accelerator_limits[accelerator.name]): - for instance in zone_n1_instances: - if instance.cpu > limit.cpu or instance.memory > limit.memory: + for machine_type in zone_n1_machine_types: + if machine_type.cpu > limit.cpu or machine_type.memory > limit.memory: continue - i = copy.deepcopy(instance) - i.gpu_count = n - i.gpu_name = accelerator.name - i.gpu_memory = accelerator_details[accelerator.name].memory - zone_instances.append(i) - return zone_instances - - n1_instances = defaultdict(list) - for instance in instances: - if instance.instance_name.startswith("n1-"): - n1_instances[instance.location].append(instance) - - instances_with_gpus = [] + machine_type_with_gpu = dataclasses.replace( + machine_type, + gpu_count=n, + gpu_name=accelerator.name, + gpu_memory=accelerator_details[accelerator.name].memory, + ) + zone_machine_types.append(machine_type_with_gpu) + return zone_machine_types + + n1_machine_types = defaultdict(list) + for mt in machine_types: + if mt.instance_name.startswith("n1-"): + n1_machine_types[mt.location].append(mt) + + machine_types_with_gpus: list[_MachineType] = [] futures = [] with ThreadPoolExecutor(max_workers=8) as ex: - for zone, zone_n1_instances in n1_instances.items(): - futures.append(ex.submit(_list_zone_instances, zone, zone_n1_instances)) + for zone, zone_n1_machine_types in n1_machine_types.items(): + futures.append( + ex.submit(_make_zone_gpu_machine_types, zone, zone_n1_machine_types) + ) for future in as_completed(futures): - instances_with_gpus.extend(future.result()) - instances += instances_with_gpus + machine_types_with_gpus.extend(future.result()) + return machine_types_with_gpus - def fill_prices(self, instances: list[RawCatalogItem]) -> list[RawCatalogItem]: + def _make_offers(self, machine_types: list[_MachineType]) -> list[CatalogItem]: logger.info("Fetching prices") skus = self.cloud_catalog_client.list_skus(parent=compute_service) prices = Prices() prices.add_skus(skus) - - offers = [] - for instance in instances: + offers: list[CatalogItem] = [] + for machine_type in machine_types: + gpu_vendor = None + gpu_name = None + if machine_type.gpu_name: + if machine_type.gpu_name.startswith("nvidia-"): + gpu_vendor = AcceleratorVendor.NVIDIA + if acc_details := accelerator_details.get(machine_type.gpu_name): + gpu_name = acc_details.name + else: + logger.warning("No accelerator details for %s", machine_type.gpu_name) + continue for capacity_type in CapacityType: - price = prices.get_instance_price(instance, capacity_type) + price = prices.get_instance_price(machine_type, capacity_type) if price is None: continue - - offer = copy.deepcopy(instance) - offer.price = round(price, 6) - offer.spot = capacity_type is CapacityType.SPOT - cast(GCPCatalogItemProviderData, offer.provider_data)["is_dws_calendar_mode"] = ( - capacity_type is CapacityType.DWS_CALENDAR_MODE + item = CatalogItem( + provider=GCPProvider.NAME, + instance_name=machine_type.instance_name, + location=machine_type.location, + price=round(price, 6), + cpu=machine_type.cpu, + memory=machine_type.memory, + gpu_vendor=gpu_vendor, + gpu_count=machine_type.gpu_count, + gpu_name=gpu_name, + gpu_memory=machine_type.gpu_memory, + spot=capacity_type is CapacityType.SPOT, + disk_size=None, + provider_data={ + "is_dws_calendar_mode": capacity_type is CapacityType.DWS_CALENDAR_MODE + }, ) - offers.append(offer) + offers.append(item) return offers - def fill_gpu_vendors_and_names(self, offers: list[RawCatalogItem]) -> None: - # Modifies offers in the list in-place - for offer in offers: - accelerator_type = offer.gpu_name - if not accelerator_type: - continue - offer.gpu_name = accelerator_details[accelerator_type].name - if accelerator_type.startswith("nvidia-"): - offer.gpu_vendor = AcceleratorVendor.NVIDIA.value - else: - logger.warning("Unknown accelerator vendor: %s", accelerator_type) - - def get( - self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: - instances = self.list_preconfigured_instances() - self.add_gpus(instances) - offers = self.fill_prices(instances) - self.fill_gpu_vendors_and_names(offers) - offers.extend(get_tpu_offers(self.project)) - set_flags(offers) - offers = add_legacy_g4_preview(offers) - return sorted(offers, key=lambda i: i.price) - - @classmethod - def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: - return [ - i - for i in offers - if ( - any( - i.instance_name.startswith(family) - for family in [ - "m4-", - "c4-", - "n4-", - "h3-", - "n2-", - "e2-medium", - "e2-standard-", - "e2-highmem-", - "e2-highcpu-", - "m1-", - "a2-", - "g2-", - ] - ) - or (i.gpu_name and i.gpu_name not in ["K80", "P4"]) - ) - and not ( - # Filter out on-demand offers that are not actually available on demand. - # https://cloud.google.com/compute/docs/accelerator-optimized-machines#consumption_option_availability_by_machine_type - i.spot == False - and not cast(GCPCatalogItemProviderData, i.provider_data).get( - "is_dws_calendar_mode" - ) - and ( - i.instance_name.startswith("a4x-") - or i.instance_name.startswith("a4-") - or i.instance_name.startswith("a3-ultragpu-") - or ( - i.instance_name.startswith("a3-highgpu-") - and (i.gpu_count is None or i.gpu_count < 8) - ) - ) - ) - ] - class GCPCatalogItemProviderData(TypedDict): is_dws_calendar_mode: NotRequired[bool] @@ -449,17 +468,20 @@ def _add_price(sku: Sku, family_prices: PricePerRegionCapacityType, price: float family_prices[(region, capacity_type)] = price def get_instance_price( - self, instance: RawCatalogItem, capacity_type: CapacityType + self, machine_type: _MachineType, capacity_type: CapacityType ) -> float | None: - vm_family = self.get_vm_family(instance.instance_name) + vm_family = self.get_vm_family(machine_type.instance_name) if vm_family in ["g1", "f1", "m2"]: # shared-core and reservation-only return None - region_capacity_type = (instance.location[:-2], capacity_type) + region_capacity_type = (machine_type.location[:-2], capacity_type) # For some instances, the price is proportional to the number of GPUs - if instance.gpu_name and region_capacity_type in self.gpu_slice[instance.gpu_name]: - return instance.gpu_count * self.gpu_slice[instance.gpu_name][region_capacity_type] + if machine_type.gpu_name and region_capacity_type in self.gpu_slice[machine_type.gpu_name]: + return ( + machine_type.gpu_count + * self.gpu_slice[machine_type.gpu_name][region_capacity_type] + ) # For others, the price consists of several components price = 0 @@ -468,15 +490,16 @@ def get_instance_price( or region_capacity_type not in self.ram[vm_family] ): return None - price += instance.cpu * self.cpu[vm_family][region_capacity_type] - price += instance.memory * self.ram[vm_family][region_capacity_type] - if instance.gpu_name: - if region_capacity_type not in self.gpu[instance.gpu_name]: + price += machine_type.cpu * self.cpu[vm_family][region_capacity_type] + price += machine_type.memory * self.ram[vm_family][region_capacity_type] + if machine_type.gpu_name: + if region_capacity_type not in self.gpu[machine_type.gpu_name]: return None - price += instance.gpu_count * self.gpu[instance.gpu_name][region_capacity_type] - if instance.instance_name in local_ssd_sizes_gib: + price += machine_type.gpu_count * self.gpu[machine_type.gpu_name][region_capacity_type] + if machine_type.instance_name in local_ssd_sizes_gib: price += ( - local_ssd_sizes_gib[instance.instance_name] * self.local_ssd[region_capacity_type] + local_ssd_sizes_gib[machine_type.instance_name] + * self.local_ssd[region_capacity_type] ) return price @@ -489,8 +512,8 @@ def get_vm_family(instance_name: str) -> str: return instance_name.split("-")[0] -def set_flags(catalog_items: list[RawCatalogItem]) -> None: - for item in catalog_items: +def _set_flags(offers: list[CatalogItem]) -> None: + for item in offers: if cast(GCPCatalogItemProviderData, item.provider_data).get("is_dws_calendar_mode"): item.flags.append("gcp-dws-calendar-mode") if item.instance_name.startswith("a4-"): @@ -500,58 +523,59 @@ def set_flags(catalog_items: list[RawCatalogItem]) -> None: # TODO: drop when dstack 0.19.33 is no longer relevant -def add_legacy_g4_preview(catalog_items: list[RawCatalogItem]) -> list[RawCatalogItem]: +def _with_legacy_g4_preview(offers: list[CatalogItem]) -> list[CatalogItem]: """ For each g4-standard-* instance, add a duplicate item with the "gcp-g4-preview" flag. This is only needed for dstack 0.19.33, where the flag "gcp-g4-preview" is used instead of "gcp-g4". """ - new_items = [] - for item in catalog_items: - new_items.append(item) + new_offers: list[CatalogItem] = [] + for item in offers: + new_offers.append(item) if item.instance_name.startswith("g4-standard-"): preview_item = copy.deepcopy(item) preview_item.flags.remove("gcp-g4") preview_item.flags.append("gcp-g4-preview") - new_items.append(preview_item) - return new_items + new_offers.append(preview_item) + return new_offers -def get_tpu_offers(project_id: str) -> list[RawCatalogItem]: +def _make_tpu_offers(project_id: str) -> list[CatalogItem]: logger.info("Fetching TPU offers") - raw_catalog_items: list[RawCatalogItem] = [] - catalog_items: list[dict] = get_catalog_items(project_id) + offers: list[CatalogItem] = [] + tpu_configs: list[dict] = _get_priced_tpu_configs(project_id) # For some TPU offers in some regions, GCP does not list prices at all. Skip such offers. - filtered_catalog_items = [item for item in catalog_items if item["price"] is not None] - for item in filtered_catalog_items: + priced_tpu_configs = [config for config in tpu_configs if config["price"] is not None] + for item in priced_tpu_configs: hardware_spec = get_tpu_hardware_spec(item["instance_name"]) if hardware_spec is None: logger.debug("No TPU hardware spec for %s", item["instance_name"]) continue - on_demand_item = RawCatalogItem( + on_demand_item = CatalogItem( + provider=GCPProvider.NAME, instance_name=item["instance_name"], location=item["location"], price=item["price"], cpu=hardware_spec.cpu, memory=hardware_spec.memory_gb, - gpu_vendor=AcceleratorVendor.GOOGLE.value, + gpu_vendor=AcceleratorVendor.GOOGLE, gpu_count=1, gpu_name=item["instance_name"], gpu_memory=hardware_spec.hbm_gb, spot=False, disk_size=None, ) - raw_catalog_items.append(on_demand_item) + offers.append(on_demand_item) if item["spot"]: spot_item = copy.deepcopy(on_demand_item) spot_item.price = item["spot"] spot_item.spot = True - raw_catalog_items.append(spot_item) - return raw_catalog_items + offers.append(spot_item) + return offers -def get_catalog_items(project_id: str) -> list[dict]: +def _get_priced_tpu_configs(project_id: str) -> list[dict]: """ Returns TPU configurations with pricing info. Each configuration contains on-demand price and spot price but any price can be missing. @@ -725,7 +749,7 @@ def find_base_price_v5( def get_tpu_configs(project_id: str) -> list[dict]: def _list_zone_configs(zone: str) -> list[dict]: - zone_instances = [] + zone_instances: list[dict] = [] if zone in ["us-east1-b"]: # These zones return # google.api_core.exceptions.ServiceUnavailable: 503 502:Bad Gateway diff --git a/src/gpuhunt/providers/hotaisle.py b/src/gpuhunt/providers/hotaisle.py index ca86aee..f504e0f 100644 --- a/src/gpuhunt/providers/hotaisle.py +++ b/src/gpuhunt/providers/hotaisle.py @@ -1,45 +1,46 @@ import logging -import os from typing import TypedDict, cast import requests from requests import Response from gpuhunt._internal.constraints import find_accelerators -from gpuhunt._internal.models import AcceleratorVendor, JSONObject, QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, JSONObject, QueryFilter +from gpuhunt.providers.base import OnlineProvider, get_creds_env logger = logging.getLogger(__name__) API_URL = "https://admin.hotaisle.app/api" -class HotAisleProvider(AbstractProvider): +class HotAisleProvider(OnlineProvider): NAME = "hotaisle" - def __init__(self, api_key: str | None = None, team_handle: str | None = None): + def __init__(self, api_key: str, team_handle: str): """Hotaisle requries an API key and team handle to access the API.""" - self.api_key = api_key or os.getenv("HOTAISLE_API_KEY") - self.team_handle = team_handle or os.getenv("HOTAISLE_TEAM_HANDLE") - - if not self.api_key: - raise ValueError("Set the HOTAISLE_API_KEY environment variable.") - if not self.team_handle: - raise ValueError("Set the HOTAISLE_TEAM_HANDLE environment variable.") + self.api_key = api_key + self.team_handle = team_handle + + @classmethod + def from_env(cls) -> "HotAisleProvider": + return cls( + api_key=get_creds_env("HOTAISLE_API_KEY"), + team_handle=get_creds_env("HOTAISLE_TEAM_HANDLE"), + ) def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: offers = self.fetch_offers() return sorted(offers, key=lambda i: i.price) - def fetch_offers(self) -> list[RawCatalogItem]: + def fetch_offers(self) -> list[CatalogItem]: """Fetch available virtual machines from HotAisle API. See API documentation(https://admin.hotaisle.app/api/docs) for details.""" url = f"/teams/{self.team_handle}/virtual_machines/available/" response = self._make_request("GET", url) - return convert_response_to_raw_catalog_items(response) + return _make_offers(response) def _make_request(self, method: str, url: str) -> Response: full_url = f"{API_URL}{url}" @@ -64,9 +65,9 @@ def get_gpu_memory(gpu_name: str) -> float | None: return None -def convert_response_to_raw_catalog_items(response: Response) -> list[RawCatalogItem]: +def _make_offers(response: Response) -> list[CatalogItem]: data = response.json() - offers = [] + offers: list[CatalogItem] = [] for item in data: price_in_cents = item["OnDemandPrice"] price = float(price_in_cents) / 100 @@ -82,13 +83,14 @@ def convert_response_to_raw_catalog_items(response: Response) -> list[RawCatalog gpu = gpus[0] gpu_count = gpu["count"] gpu_name = gpu["model"] - gpu_vendor = AcceleratorVendor.AMD.value # All GPUs are AMD with HotAisle. + gpu_vendor = AcceleratorVendor.AMD # All GPUs are AMD with HotAisle. gpu_memory = get_gpu_memory(gpu_name) # Create instance name: cpu_model-cores-ram-gpucount-gpu instance_name = f"{gpu_count}x {gpu_name} {cpu_cores}x {cpu_model}" - offer = RawCatalogItem( + offer = CatalogItem( + provider=HotAisleProvider.NAME, instance_name=instance_name, location="us-michigan-1", # Hardcoded for now, as HotAisle only has one location. price=price, @@ -103,7 +105,7 @@ def convert_response_to_raw_catalog_items(response: Response) -> list[RawCatalog provider_data=cast( JSONObject, HotAisleCatalogItemProviderData( - # The specs object may duplicate some RawCatalogItem fields, but we store it in + # The specs object may duplicate some CatalogItem fields, but we store it in # full because we need to pass it back to the API when creating VMs. vm_specs=specs, ), diff --git a/src/gpuhunt/providers/jarvislabs.py b/src/gpuhunt/providers/jarvislabs.py index 4651baf..f184726 100644 --- a/src/gpuhunt/providers/jarvislabs.py +++ b/src/gpuhunt/providers/jarvislabs.py @@ -6,8 +6,8 @@ from requests import Response from typing_extensions import NotRequired, TypedDict -from gpuhunt._internal.models import AcceleratorVendor, JSONObject, QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, JSONObject, QueryFilter +from gpuhunt.providers.base import OnlineProvider, get_creds_env logger = logging.getLogger(__name__) @@ -41,25 +41,29 @@ class JarvisLabsCatalogItemProviderData(TypedDict): gpu_type: NotRequired[str] -class JarvisLabsProvider(AbstractProvider): +class JarvisLabsProvider(OnlineProvider): NAME = "jarvislabs" - def __init__(self, api_key: str | None = None, api_url: str | None = None): - self.api_key = api_key or os.getenv("JL_API_KEY") - if not self.api_key: - raise ValueError("Set the JL_API_KEY environment variable.") + def __init__(self, api_key: str, api_url: str = API_URL): + self.api_key = api_key + self.api_url = api_url.rstrip("/") - self.api_url = (api_url or os.getenv("JARVISLABS_API_URL", API_URL)).rstrip("/") + @classmethod + def from_env(cls) -> "JarvisLabsProvider": + return cls( + api_key=get_creds_env("JL_API_KEY"), + api_url=os.getenv("JARVISLABS_API_URL", API_URL), + ) def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: offers = self.fetch_offers(query_filter=query_filter) return sorted(offers, key=lambda i: i.price) - def fetch_offers(self, query_filter: QueryFilter | None = None) -> list[RawCatalogItem]: + def fetch_offers(self, query_filter: QueryFilter | None = None) -> list[CatalogItem]: response = self._make_request("GET", SERVER_META_PATH) - return convert_response_to_raw_catalog_items(response.json()) + return _make_offers(response.json()) def _make_request(self, method: str, path: str) -> Response: response = requests.request( @@ -72,15 +76,15 @@ def _make_request(self, method: str, path: str) -> Response: return response -def convert_response_to_raw_catalog_items(data: dict) -> list[RawCatalogItem]: - offers = [] +def _make_offers(data: dict) -> list[CatalogItem]: + offers: list[CatalogItem] = [] for gpu in data.get("server_meta") or []: - offers.extend(_make_gpu_catalog_items(gpu)) - offers.extend(_make_cpu_catalog_items(data.get("cpu_meta") or {})) + offers.extend(_make_gpu_offers(gpu)) + offers.extend(_make_cpu_offers(data.get("cpu_meta") or {})) return offers -def _make_gpu_catalog_items(gpu: dict) -> list[RawCatalogItem]: +def _make_gpu_offers(gpu: dict) -> list[CatalogItem]: region = gpu.get("region") if not region: return [] @@ -122,7 +126,7 @@ def _make_gpu_catalog_items(gpu: dict) -> list[RawCatalogItem]: logger.warning("Skipping JarvisLabs GPU offer without CPU/RAM: %s", gpu_type) return [] - items = _make_gpu_catalog_items_for_price( + offers = _make_gpu_offers_for_price( region=region, gpu_name=gpu_name, gpu_memory=gpu_memory, @@ -137,10 +141,10 @@ def _make_gpu_catalog_items(gpu: dict) -> list[RawCatalogItem]: # JarvisLabs supports spot for containers/templates, not VMs. This provider # only publishes VM-capable offers because dstack provisions JarvisLabs VMs. - return items + return offers -def _make_gpu_catalog_items_for_price( +def _make_gpu_offers_for_price( *, region: str, gpu_name: str, @@ -152,20 +156,21 @@ def _make_gpu_catalog_items_for_price( max_gpus_per_instance: int, provider_data: JSONObject, spot: bool, -) -> list[RawCatalogItem]: - items = [] +) -> list[CatalogItem]: + offers: list[CatalogItem] = [] for gpu_count in _supported_gpu_counts( available_devices=available_devices, max_gpus_per_instance=max_gpus_per_instance, ): - items.append( - RawCatalogItem( + offers.append( + CatalogItem( + provider=JarvisLabsProvider.NAME, instance_name=_gpu_instance_name(gpu_name, gpu_count), location=region, price=round(price * gpu_count, 5), cpu=cpu_per_gpu * gpu_count, memory=ram_per_gpu * gpu_count, - gpu_vendor=AcceleratorVendor.NVIDIA.value, + gpu_vendor=AcceleratorVendor.NVIDIA, gpu_count=gpu_count, gpu_name=gpu_name, gpu_memory=gpu_memory, @@ -174,11 +179,11 @@ def _make_gpu_catalog_items_for_price( provider_data=provider_data, ) ) - return items + return offers -def _make_cpu_catalog_items(cpu_meta: dict) -> list[RawCatalogItem]: - offers = [] +def _make_cpu_offers(cpu_meta: dict) -> list[CatalogItem]: + offers: list[CatalogItem] = [] # The JarvisLabs SDK resolves CPU VMs from cpu_meta.combinations and creates them via # templates/vm/cpu/create; cpu_meta.workload_type is not the GPU workload selector. for combo in cpu_meta.get("combinations") or []: @@ -201,7 +206,8 @@ def _make_cpu_catalog_items(cpu_meta: dict) -> list[RawCatalogItem]: ) continue offers.append( - RawCatalogItem( + CatalogItem( + provider=JarvisLabsProvider.NAME, instance_name=f"cpu-{vcpus}x{int(ram_gb)}", location=region, price=price, @@ -252,18 +258,18 @@ def _gpu_instance_name(gpu_name: str, gpu_count: int) -> str: def _as_int(value: object) -> int | None: - if value is None or value == "": + if not isinstance(value, str | int | float) or value == "": return None try: return int(value) - except (TypeError, ValueError): + except ValueError: return None def _as_float(value: object) -> float | None: - if value is None or value == "": + if not isinstance(value, str | int | float) or value == "": return None try: return float(value) - except (TypeError, ValueError): + except ValueError: return None diff --git a/src/gpuhunt/providers/lambdalabs.py b/src/gpuhunt/providers/lambdalabs.py index 067cc0d..1ccbe29 100644 --- a/src/gpuhunt/providers/lambdalabs.py +++ b/src/gpuhunt/providers/lambdalabs.py @@ -1,4 +1,3 @@ -import copy import logging import re @@ -6,11 +5,12 @@ from gpuhunt._internal.constraints import is_nvidia_superchip from gpuhunt._internal.models import ( + AcceleratorVendor, + CatalogItem, CPUArchitecture, QueryFilter, - RawCatalogItem, ) -from gpuhunt.providers import AbstractProvider +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) INSTANCE_TYPES_URL = "https://cloud.lambdalabs.com/api/v1/instance-types" @@ -20,7 +20,7 @@ FLAG_ARM = "lambda-arm" -class LambdaLabsProvider(AbstractProvider): +class LambdaLabsProvider(OfflineProvider): NAME = "lambdalabs" def __init__(self, token: str): @@ -29,54 +29,45 @@ def __init__(self, token: str): def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: - offers = [] + ) -> list[CatalogItem]: + offers: list[CatalogItem] = [] + regions = self.list_regions() resp = self.session.get(INSTANCE_TYPES_URL, timeout=TIMEOUT) resp.raise_for_status() data = resp.json()["data"] for instance in data.values(): instance = instance["instance_type"] - logger.info(instance["name"]) description = instance["description"] - result = parse_description(description) + result = _parse_description(description) if result is None: logger.warning("Can't parse GPU info from description: %s", description) continue gpu_count, gpu_name, gpu_memory = result - flags: list[str] = [] - cpu_arch = CPUArchitecture.X86 - if is_nvidia_superchip(gpu_name): - cpu_arch = CPUArchitecture.ARM - flags.append(FLAG_ARM) - offer = RawCatalogItem( - instance_name=instance["name"], - price=instance["price_cents_per_hour"] / 100, - cpu_arch=cpu_arch.value, - cpu=instance["specs"]["vcpus"], - memory=float(instance["specs"]["memory_gib"]) * 1.074, - gpu_vendor=None, - gpu_count=gpu_count, - gpu_name=gpu_name, - gpu_memory=gpu_memory, - spot=False, - location=None, - disk_size=float(instance["specs"]["storage_gib"]) * 1.074, - flags=flags, - ) - offers.append(offer) - offers = self.add_regions(offers) + for region in regions: + flags: list[str] = [] + cpu_arch = CPUArchitecture.X86 + if is_nvidia_superchip(gpu_name): + cpu_arch = CPUArchitecture.ARM + flags.append(FLAG_ARM) + offer = CatalogItem( + provider=LambdaLabsProvider.NAME, + instance_name=instance["name"], + price=instance["price_cents_per_hour"] / 100, + cpu_arch=cpu_arch, + cpu=instance["specs"]["vcpus"], + memory=float(instance["specs"]["memory_gib"]) * 1.074, + gpu_vendor=AcceleratorVendor.NVIDIA if gpu_count else None, + gpu_count=gpu_count, + gpu_name=gpu_name, + gpu_memory=gpu_memory, + spot=False, + location=region, + disk_size=float(instance["specs"]["storage_gib"]) * 1.074, + flags=flags, + ) + offers.append(offer) return sorted(offers, key=lambda i: i.price) - def add_regions(self, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: - # TODO: we don't know which regions are actually available for each instance type - region_offers = [] - for region in self.list_regions(): - for offer in offers: - offer = copy.deepcopy(offer) - offer.location = region - region_offers.append(offer) - return region_offers - def list_regions(self) -> list[str]: resp = self.session.get(IMAGES_URL, timeout=TIMEOUT) resp.raise_for_status() @@ -86,7 +77,7 @@ def list_regions(self) -> list[str]: return sorted(regions) -def parse_description(v: str) -> tuple[int, str, float] | None: +def _parse_description(v: str) -> tuple[int, str, float] | None: """Returns gpus count, gpu name, and GPU memory""" r = re.match(r"^(\d)x (?:Tesla )?(.+) \((\d+) GB", v) if r is None: diff --git a/src/gpuhunt/providers/nebius.py b/src/gpuhunt/providers/nebius.py index c0b9eab..910624d 100644 --- a/src/gpuhunt/providers/nebius.py +++ b/src/gpuhunt/providers/nebius.py @@ -34,11 +34,12 @@ from gpuhunt._internal.models import ( AcceleratorInfo, AcceleratorVendor, + CatalogItem, JSONObject, QueryFilter, - RawCatalogItem, ) -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.utils import get_or_error +from gpuhunt.providers.base import OfflineProvider from gpuhunt.version import __version__ logger = logging.getLogger(__name__) @@ -73,7 +74,7 @@ class InfinibandFabric: ] -class NebiusProvider(AbstractProvider): +class NebiusProvider(OfflineProvider): NAME = "nebius" def __init__(self, credentials: Credentials) -> None: @@ -81,8 +82,8 @@ def __init__(self, credentials: Credentials) -> None: def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: - items: list[RawCatalogItem] = [] + ) -> list[CatalogItem]: + offers: list[CatalogItem] = [] sdk = SDK( credentials=self.credentials, user_agent_prefix=f"gpuhunt/{__version__}", @@ -102,15 +103,14 @@ def get( price = get_price( calculator, project_id, platform.metadata.name, preset.name, spot ) - item = make_item( + offer = _make_offer( platform.metadata.name, preset, gpu, region, spot, price ) - if item is not None: - items.append(item) + if offer is not None: + offers.append(offer) finally: sdk.sync_close(timeout=TIMEOUT) - items.sort(key=lambda i: i.price) - return items + return sorted(offers, key=lambda i: i.price) class NebiusCatalogItemProviderData(TypedDict): @@ -151,7 +151,7 @@ def get_price( estimate = calculator.estimate( request=EstimateRequest(resource_spec=ResourceSpec(compute_instance_spec=spec)) ).wait() - return float(estimate.hourly_cost.general.total.cost) + return float(get_or_error(estimate.hourly_cost.general, "general hourly cost").total.cost) def list_platforms(sdk: SDK, project_id: str) -> ListPlatformsResponse: @@ -177,21 +177,22 @@ def get_gpu_info(platform: Platform) -> AcceleratorInfo | None: return accelerator_info[0] -def make_item( +def _make_offer( platform: str, preset: Preset, gpu: AcceleratorInfo | None, region: str, spot: bool, price: float, -) -> RawCatalogItem | None: +) -> CatalogItem | None: fabrics = [] if preset.allow_gpu_clustering: fabrics = [ f.name for f in INFINIBAND_FABRICS if f.platform == platform and f.region == region ] - item = RawCatalogItem( + item = CatalogItem( + provider=NebiusProvider.NAME, instance_name=f"{platform} {preset.name}", location=region, price=price, diff --git a/src/gpuhunt/providers/oci.py b/src/gpuhunt/providers/oci.py index 1ab5396..55fc081 100644 --- a/src/gpuhunt/providers/oci.py +++ b/src/gpuhunt/providers/oci.py @@ -1,7 +1,6 @@ import copy import logging import re -from collections.abc import Iterable from dataclasses import asdict, dataclass from typing import Annotated, TypeVar @@ -12,9 +11,10 @@ from typing_extensions import TypedDict from gpuhunt._internal.constraints import find_accelerators -from gpuhunt._internal.models import AcceleratorVendor, QueryFilter, RawCatalogItem -from gpuhunt._internal.utils import to_camel_case -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.errors import ProviderError +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt._internal.utils import get_or_error, to_camel_case +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) COST_ESTIMATOR_URL_TEMPLATE = "https://www.oracle.com/a/ocom/docs/cloudestimator2/data/{resource}" @@ -37,7 +37,7 @@ class OCICredentials(TypedDict): region: str | None -class OCIProvider(AbstractProvider): +class OCIProvider(OfflineProvider): NAME = "oci" def __init__(self, credentials: OCICredentials): @@ -48,12 +48,15 @@ def __init__(self, credentials: OCICredentials): def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: shapes = self.cost_estimator.get_shapes() products = self.cost_estimator.get_products() - regions: list[Region] = self.api_client.list_regions().data + regions: list[Region] = get_or_error( + self.api_client.list_regions(), "list_regions response" + ).data + region_names = [get_or_error(region.name, "region name") for region in regions] - result = [] + offers: list[CatalogItem] = [] for shape in shapes.items: if ( @@ -72,30 +75,30 @@ def get( "Skipping shape %s due to unexpected Cost Estimator data: %s", shape.name, e ) continue + for region_name in region_names: + on_demand_item = CatalogItem( + provider=OCIProvider.NAME, + instance_name=shape.name, + location=region_name, + price=resources.total_price(), + cpu=resources.cpu.vcpus, + memory=resources.memory.gbs, + gpu_vendor=(AcceleratorVendor.NVIDIA if resources.gpu.units_count else None), + gpu_count=resources.gpu.units_count, + gpu_name=resources.gpu.name, + gpu_memory=resources.gpu.unit_memory_gb, + spot=False, + disk_size=None, + ) + item_variations = [on_demand_item] + if shape.allow_preemptible: + item_variations.append(self._make_spot_offer(on_demand_item)) + offers.extend(item_variations) - on_demand_item = RawCatalogItem( - instance_name=shape.name, - location=None, - price=resources.total_price(), - cpu=resources.cpu.vcpus, - memory=resources.memory.gbs, - gpu_vendor=None, - gpu_count=resources.gpu.units_count, - gpu_name=resources.gpu.name, - gpu_memory=resources.gpu.unit_memory_gb, - spot=False, - disk_size=None, - ) - item_variations = [on_demand_item] - if shape.allow_preemptible: - item_variations.append(self._make_spot_item(on_demand_item)) - for item in item_variations: - result.extend(self._duplicate_item_in_regions(item, regions)) - - return sorted(result, key=lambda i: i.price) + return sorted(offers, key=lambda i: i.price) @staticmethod - def _make_spot_item(item: RawCatalogItem) -> RawCatalogItem: + def _make_spot_offer(item: CatalogItem) -> CatalogItem: item = copy.deepcopy(item) item.spot = True # > Preemptible capacity costs 50% less than on-demand capacity @@ -104,17 +107,6 @@ def _make_spot_item(item: RawCatalogItem) -> RawCatalogItem: item.flags.append("oci-spot") return item - @staticmethod - def _duplicate_item_in_regions( - item: RawCatalogItem, regions: Iterable[Region] - ) -> list[RawCatalogItem]: - result = [] - for region in regions: - regional_item = copy.deepcopy(item) - regional_item.location = region.name - result.append(regional_item) - return result - class CostEstimatorTypeField(BaseModel): value: str @@ -215,7 +207,7 @@ def _get(self, resource: str, ResponseModel: type[ResponseModelT]) -> ResponseMo return ResponseModel.model_validate_json(resp.content) -class CostEstimatorDataError(Exception): +class CostEstimatorDataError(ProviderError): pass diff --git a/src/gpuhunt/providers/runpod.py b/src/gpuhunt/providers/runpod.py index 55a52a6..203fbfd 100644 --- a/src/gpuhunt/providers/runpod.py +++ b/src/gpuhunt/providers/runpod.py @@ -1,4 +1,3 @@ -import copy import logging from concurrent.futures import ThreadPoolExecutor from typing import cast @@ -8,8 +7,8 @@ from typing_extensions import NotRequired, TypedDict from gpuhunt._internal.constraints import find_accelerators -from gpuhunt._internal.models import AcceleratorVendor, QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) API_URL = "https://api.runpod.io/graphql" @@ -21,7 +20,7 @@ class RunpodCatalogItemProviderData(TypedDict): pod_counts: NotRequired[list[int]] -class RunpodProvider(AbstractProvider): +class RunpodProvider(OfflineProvider): NAME = "runpod" # Minimum CUDA version on the host. Used to filter available offers # and should also be used when provisioning pods. @@ -32,12 +31,12 @@ def __init__(self) -> None: def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: offers = self._fetch_offers() return sorted(offers, key=lambda i: i.price or 0) @classmethod - def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: + def filter(cls, offers: list[CatalogItem]) -> list[CatalogItem]: return [ o for o in offers @@ -47,7 +46,7 @@ def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: ] ] - def _fetch_offers(self) -> list[RawCatalogItem]: + def _fetch_offers(self) -> list[CatalogItem]: query_variables = self._build_query_variables() with ThreadPoolExecutor(max_workers=10) as executor: futures = [ @@ -61,16 +60,16 @@ def _fetch_offers(self) -> list[RawCatalogItem]: except RequestException as e: logger.exception("Failed to get pods data: %s", e) - catalog_items = [] + offers: list[CatalogItem] = [] for query_variable, pods in zip(query_variables, pods_by_query): for pod in pods: - catalog_items.extend(self._make_catalog_items(query_variable, pod)) + offers.extend(self._make_offers(query_variable, pod)) - cluster_catalog_items = self._fetch_cluster_offers() - catalog_items.extend(cluster_catalog_items) - cpu_catalog_items = self._fetch_cpu_offers() - catalog_items.extend(cpu_catalog_items) - return catalog_items + cluster_offers = self._fetch_cluster_offers() + offers.extend(cluster_offers) + cpu_offers = self._fetch_cpu_offers() + offers.extend(cpu_offers) + return offers def _build_query_variables(self) -> list[dict]: """Prepare different combinations of API query filters to cover all available GPUs.""" @@ -130,7 +129,7 @@ def _get_pods(self, query_variables: dict) -> list[dict]: ) return resp["data"]["gpuTypes"] - def _make_catalog_items(self, query_variables: dict, pod: dict) -> list[RawCatalogItem]: + def _make_offers(self, query_variables: dict, pod: dict) -> list[CatalogItem]: lowest_price_input_variables = query_variables["lowestPriceInput"] if pod["lowestPrice"]["stockStatus"] is None: return [] @@ -141,40 +140,31 @@ def _make_catalog_items(self, query_variables: dict, pod: dict) -> list[RawCatal if lowest_price_input_variables["secureCloud"]: location = lowest_price_input_variables["dataCenterId"] on_demand_gpu_price = pod["securePrice"] - spot_gpu_price = pod["secureSpotPrice"] else: location = lowest_price_input_variables["countryCode"] on_demand_gpu_price = pod["communityPrice"] - spot_gpu_price = pod["communitySpotPrice"] - item_template = RawCatalogItem( - instance_name=pod["id"], - location=location, - price=None, # set below - cpu=pod["lowestPrice"]["minVcpu"], - memory=pod["lowestPrice"]["minMemory"], - gpu_vendor=listed_gpu_vendor_and_name[0], - gpu_count=lowest_price_input_variables["gpuCount"], - gpu_name=listed_gpu_vendor_and_name[1], - gpu_memory=pod["memoryInGb"], - spot=None, # set below - disk_size=None, - provider_data={}, - ) - items = [] + offers: list[CatalogItem] = [] if on_demand_gpu_price: - item = copy.deepcopy(item_template) - item.spot = False - item.price = item.gpu_count * on_demand_gpu_price - items.append(item) - if spot_gpu_price: - item = copy.deepcopy(item_template) - item.spot = True - item.price = item.gpu_count * spot_gpu_price - items.append(item) - return items - - def _fetch_cluster_offers(self) -> list[RawCatalogItem]: - cluster_catalog_items = [] + offer = CatalogItem( + provider=RunpodProvider.NAME, + instance_name=pod["id"], + location=location, + price=lowest_price_input_variables["gpuCount"] * on_demand_gpu_price, + cpu=pod["lowestPrice"]["minVcpu"], + memory=pod["lowestPrice"]["minMemory"], + gpu_vendor=listed_gpu_vendor_and_name[0], + gpu_count=lowest_price_input_variables["gpuCount"], + gpu_name=listed_gpu_vendor_and_name[1], + gpu_memory=pod["memoryInGb"], + spot=False, + disk_size=None, + provider_data={}, + ) + offers.append(offer) + return offers + + def _fetch_cluster_offers(self) -> list[CatalogItem]: + cluster_offers: list[CatalogItem] = [] query_variables = { "gpuTypesInput": { "cluster": True, @@ -202,7 +192,8 @@ def _fetch_cluster_offers(self) -> list[RawCatalogItem]: logger.warning(f"{pod_type['id']} cluster offer missing minMemory") continue for location in pod_type["nodeGroupDatacenters"]: - catalog_item = RawCatalogItem( + offer = CatalogItem( + provider=RunpodProvider.NAME, instance_name=pod_type["id"], location=location["id"], price=pod_type["clusterPrice"] * pod_type["maxGpuCount"], @@ -220,10 +211,10 @@ def _fetch_cluster_offers(self) -> list[RawCatalogItem]: dict, RunpodCatalogItemProviderData(pod_counts=list(range(2, 9))) ), ) - cluster_catalog_items.append(catalog_item) - return cluster_catalog_items + cluster_offers.append(offer) + return cluster_offers - def _fetch_cpu_offers(self) -> list[RawCatalogItem]: + def _fetch_cpu_offers(self) -> list[CatalogItem]: response = _make_request({"query": cpu_data_centers_query, "variables": {}}) data_centers = [dc["id"] for dc in response["data"]["dataCenters"] if dc["listed"]] if len(data_centers) == 0: @@ -240,13 +231,13 @@ def _fetch_cpu_offers(self) -> list[RawCatalogItem]: except RequestException as e: logger.exception("Failed to get cpuFlavors data for %s: %s", dc_id, e) - catalog_items = [] + offers: list[CatalogItem] = [] for dc_id in data_centers: cpu_flavors = cpu_flavors_by_data_center.get(dc_id) if cpu_flavors is None: continue - catalog_items.extend(self._make_cpu_catalog_items(dc_id, cpu_flavors)) - return catalog_items + offers.extend(self._make_cpu_offers(dc_id, cpu_flavors)) + return offers def _get_cpu_flavors(self, data_center_id: str) -> list[dict]: response = _make_request( @@ -254,10 +245,8 @@ def _get_cpu_flavors(self, data_center_id: str) -> list[dict]: ) return response["data"]["cpuFlavors"] - def _make_cpu_catalog_items( - self, data_center_id: str, cpu_flavors: list[dict] - ) -> list[RawCatalogItem]: - items: list[RawCatalogItem] = [] + def _make_cpu_offers(self, data_center_id: str, cpu_flavors: list[dict]) -> list[CatalogItem]: + offers: list[CatalogItem] = [] for flavor in cpu_flavors: specifics = flavor.get("specifics") or {} if specifics.get("stockStatus") is None: @@ -288,8 +277,9 @@ def _make_cpu_catalog_items( disk_size = float(vcpu * int(disk_limit_per_vcpu)) scale = vcpu / min_vcpu price = base_secure_price * scale - items.append( - RawCatalogItem( + offers.append( + CatalogItem( + provider=RunpodProvider.NAME, instance_name=f"{flavor['id']}-{vcpu}-{memory}", location=data_center_id, price=price, @@ -304,7 +294,7 @@ def _make_cpu_catalog_items( provider_data={}, ) ) - return items + return offers def _get_gpu_vendor_and_name( self, @@ -457,13 +447,10 @@ def _cpu_size_ladder(min_vcpu: int, max_vcpu: int) -> list[int]: displayName memoryInGb securePrice - secureSpotPrice communityPrice - communitySpotPrice oneMonthPrice threeMonthPrice sixMonthPrice - secureSpotPrice __typename } } diff --git a/src/gpuhunt/providers/seeweb.py b/src/gpuhunt/providers/seeweb.py index 0817445..15b55f4 100644 --- a/src/gpuhunt/providers/seeweb.py +++ b/src/gpuhunt/providers/seeweb.py @@ -1,12 +1,11 @@ import logging -import os import re import requests from gpuhunt._internal.constraints import find_accelerators -from gpuhunt._internal.models import AcceleratorVendor, QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt.providers.base import OnlineProvider, get_creds_env logger = logging.getLogger(__name__) @@ -33,7 +32,7 @@ _NON_NVIDIA_MARKERS = ("MI300", "TENSTORRENT", "GRAYSKULL", "WORMHOLE") -class SeewebProvider(AbstractProvider): +class SeewebProvider(OnlineProvider): """Online provider for Seeweb Cloud Server GPU. Seeweb's plan/pricing endpoints require authentication, so this is an online provider queried @@ -44,15 +43,16 @@ class SeewebProvider(AbstractProvider): NAME = "seeweb" - def __init__(self, token: str | None = None): - token = token or os.getenv("SEEWEB_API_TOKEN") - if not token: - raise ValueError("Set the SEEWEB_API_TOKEN environment variable.") + def __init__(self, token: str): self.token = token + @classmethod + def from_env(cls) -> "SeewebProvider": + return cls(token=get_creds_env("SEEWEB_API_TOKEN")) + def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: response = requests.get( f"{API_URL}/plans", headers={"X-APITOKEN": self.token}, @@ -63,16 +63,16 @@ def get( if not isinstance(data, dict) or not isinstance(data.get("plans"), list): raise ValueError("Unexpected response from Seeweb /plans endpoint") - offers: list[RawCatalogItem] = [] + offers: list[CatalogItem] = [] for plan in data["plans"]: if not isinstance(plan, dict): logger.warning("Skipping malformed Seeweb plan: %r", plan) continue - offers.extend(_convert_plan(plan)) + offers.extend(_make_offers(plan)) return sorted(offers, key=lambda item: item.price if item.price is not None else 0.0) -def _convert_plan(plan: dict) -> list[RawCatalogItem]: +def _make_offers(plan: dict) -> list[CatalogItem]: plan_name = plan.get("name") gpu_label = plan.get("gpu_label") # In the /plans response, "available" means that the plan is active. It does not indicate @@ -113,7 +113,7 @@ def _convert_plan(plan: dict) -> list[RawCatalogItem]: ) return [] - offers = [] + offers: list[CatalogItem] = [] seen_regions = set() # These are regions where the plan is active/compatible, not a real-time capacity signal. # Consumers that need current capacity should query /plans/availables separately. @@ -130,7 +130,8 @@ def _convert_plan(plan: dict) -> list[RawCatalogItem]: continue seen_regions.add(location) offers.append( - RawCatalogItem( + CatalogItem( + provider=SeewebProvider.NAME, instance_name=plan_name, location=location, price=price, @@ -140,7 +141,7 @@ def _convert_plan(plan: dict) -> list[RawCatalogItem]: gpu_count=gpu_count, gpu_name=gpu_name, gpu_memory=gpu_memory, - gpu_vendor=AcceleratorVendor.NVIDIA.value, + gpu_vendor=AcceleratorVendor.NVIDIA, spot=False, disk_size=disk_size, ) diff --git a/src/gpuhunt/providers/vastai.py b/src/gpuhunt/providers/vastai.py index f0b34b4..769c323 100644 --- a/src/gpuhunt/providers/vastai.py +++ b/src/gpuhunt/providers/vastai.py @@ -2,15 +2,15 @@ import logging import re from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from typing import Any, Literal, cast import requests from typing_extensions import NotRequired, TypedDict from gpuhunt._internal.constraints import correct_gpu_memory_gib -from gpuhunt._internal.models import QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt.providers.base import OnlineProvider logger = logging.getLogger(__name__) bundles_url = "https://console.vast.ai/api/v0/bundles/" @@ -20,9 +20,13 @@ FilterValue = int | float | str | bool -class VastAIProvider(AbstractProvider): +class VastAIProvider(OnlineProvider): NAME = "vastai" + @classmethod + def from_env(cls) -> "VastAIProvider": + return cls() + def __init__( self, extra_filters: dict[str, dict[Operators, FilterValue]] | None = None, @@ -35,7 +39,7 @@ def __init__( def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: filters: dict[str, Any] = self.make_filters(query_filter or QueryFilter()) if self.extra_filters: for key, constraints in self.extra_filters.items(): @@ -45,7 +49,7 @@ def get( resp.raise_for_status() data = resp.json() - instance_offers = [] + offers: list[CatalogItem] = [] for offer in data["offers"]: cpu_cores = offer["cpu_cores"] # although this is not stated in the docs, the value can be None @@ -62,14 +66,15 @@ def get( gpu_name = get_dstack_gpu_name(offer["gpu_name"]) gpu_memory = correct_gpu_memory_gib(gpu_name, offer["gpu_ram"]) disk_cost = disk_size * offer["storage_cost"] / 30 / 24 - ondemand_offer = RawCatalogItem( + ondemand_offer = CatalogItem( + provider=VastAIProvider.NAME, instance_name=str(offer["id"]), location=get_location(offer["geolocation"]), # storage_cost is $/gb/month price=round(offer["dph_base"] + disk_cost, 5), cpu=int(offer["cpu_cores_effective"]), memory=memory, - gpu_vendor=None, + gpu_vendor=AcceleratorVendor.NVIDIA if offer["num_gpus"] else None, gpu_count=offer["num_gpus"], gpu_name=gpu_name, gpu_memory=float(gpu_memory), @@ -88,11 +93,16 @@ def get( offer_variants.append(spot_offer) offer_variants.sort(key=lambda i: i.price) - instance_offers.extend(offer_variants) - return instance_offers + offers.extend(offer_variants) + return offers + + def make_filters(self, q: QueryFilter) -> dict[str, Any]: + """ + Build the bundles request body: per-field operator constraints, plus the + `limit` and `order` query params. + """ - def make_filters(self, q: QueryFilter) -> dict[str, dict[Operators, FilterValue]]: - filters = defaultdict(dict) + filters: dict[str, Any] = defaultdict(dict) if q.min_cpu is not None: filters["cpu_cores"]["gte"] = q.min_cpu if q.max_cpu is not None: @@ -143,15 +153,18 @@ def make_filters(self, q: QueryFilter) -> dict[str, dict[Operators, FilterValue] return filters @staticmethod - def satisfies_filters(offer: dict, filters: dict[str, dict[Operators, FilterValue]]) -> bool: - for key in filters: + def satisfies_filters(offer: dict, filters: Mapping[str, Any]) -> bool: + for key, constraints in filters.items(): # `datacenter`/`external` are query scope controls. # They don't map to offer fields with strict eq semantics. if key in {"datacenter", "external"}: continue if key not in offer: continue - for op, value in filters[key].items(): + if not isinstance(constraints, dict): + # `limit`/`order` are query params, not per-field constraints + continue + for op, value in constraints.items(): if op == "lt" and offer[key] >= value: return False if op == "lte" and offer[key] > value: diff --git a/src/gpuhunt/providers/verda.py b/src/gpuhunt/providers/verda.py index 0aad521..859e87f 100644 --- a/src/gpuhunt/providers/verda.py +++ b/src/gpuhunt/providers/verda.py @@ -7,8 +7,8 @@ from verda import VerdaClient from verda.instance_types import InstanceType -from gpuhunt import QueryFilter, RawCatalogItem -from gpuhunt.providers import AbstractProvider +from gpuhunt import AcceleratorVendor, CatalogItem, QueryFilter +from gpuhunt.providers.base import OfflineProvider logger = logging.getLogger(__name__) @@ -18,7 +18,7 @@ ] -class VerdaProvider(AbstractProvider): +class VerdaProvider(OfflineProvider): NAME = "verda" def __init__(self, client_id: str, client_secret: str) -> None: @@ -26,15 +26,13 @@ def __init__(self, client_id: str, client_secret: str) -> None: def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: instance_types = self._get_instance_types() locations = self._get_locations() - spots = (True, False) location_codes = [loc["code"] for loc in locations] - instances = generate_instances(spots, location_codes, instance_types) - - return sorted(instances, key=lambda x: x.price) + offers = _make_offers(spots, location_codes, instance_types) + return sorted(offers, key=lambda x: x.price) def _get_instance_types(self) -> list[InstanceType]: return self.verda_client.instance_types.get() @@ -43,30 +41,32 @@ def _get_locations(self) -> list[dict]: return self.verda_client.locations.get() @classmethod - def filter(cls, offers: list[RawCatalogItem]) -> list[RawCatalogItem]: + def filter(cls, offers: list[CatalogItem]) -> list[CatalogItem]: return [o for o in offers if o.gpu_name not in ALL_AMD_GPUS] # skip AMD GPU -def generate_instances( +def _make_offers( spots: Iterable[bool], location_codes: Iterable[str], instance_types: Iterable[InstanceType] -) -> list[RawCatalogItem]: - instances = [] +) -> list[CatalogItem]: + offers: list[CatalogItem] = [] for spot, location, instance in itertools.product(spots, location_codes, instance_types): - item = transform_instance(copy.copy(instance), spot, location) - if item is None: + offer = _make_offer(copy.copy(instance), spot, location) + if offer is None: continue - instances.append(RawCatalogItem.from_dict(item)) - return instances + offers.append(offer) + return offers -def transform_instance(instance: InstanceType, spot: bool, location: str) -> dict | None: +def _make_offer(instance: InstanceType, spot: bool, location: str) -> CatalogItem | None: gpu_memory = None gpu_count = instance.gpu["number_of_gpus"] gpu_name = None + gpu_vendor = None if instance.gpu["number_of_gpus"]: gpu_memory = instance.gpu_memory["size_in_gigabytes"] / instance.gpu["number_of_gpus"] gpu_name = get_gpu_name(instance.gpu["description"]) + gpu_vendor = AcceleratorVendor.NVIDIA if gpu_count and gpu_name is None: logger.warning( @@ -74,18 +74,21 @@ def transform_instance(instance: InstanceType, spot: bool, location: str) -> dic ) return None - raw = dict( + return CatalogItem( + provider=VerdaProvider.NAME, instance_name=instance.instance_type, location=location, spot=spot, - price=instance.spot_price_per_hour if spot else instance.price_per_hour, - cpu=instance.cpu["number_of_cores"], - memory=instance.memory["size_in_gigabytes"], - gpu_count=gpu_count, + # The API reports prices as strings. + price=float(instance.spot_price_per_hour if spot else instance.price_per_hour), + cpu=int(instance.cpu["number_of_cores"]), + memory=float(instance.memory["size_in_gigabytes"]), + gpu_count=int(gpu_count), gpu_name=gpu_name, gpu_memory=gpu_memory, + gpu_vendor=gpu_vendor, + disk_size=None, ) - return raw GPU_MAP = { diff --git a/src/gpuhunt/providers/vultr.py b/src/gpuhunt/providers/vultr.py index 6d37cb5..45ddda1 100644 --- a/src/gpuhunt/providers/vultr.py +++ b/src/gpuhunt/providers/vultr.py @@ -4,31 +4,35 @@ import requests from requests import Response -from gpuhunt import QueryFilter, RawCatalogItem +from gpuhunt import CatalogItem, QueryFilter from gpuhunt._internal.constraints import ( find_accelerators, get_gpu_vendor, is_nvidia_superchip, ) from gpuhunt._internal.models import AcceleratorVendor, CPUArchitecture -from gpuhunt.providers import AbstractProvider +from gpuhunt.providers.base import OnlineProvider logger = logging.getLogger(__name__) API_URL = "https://api.vultr.com/v2" -class VultrProvider(AbstractProvider): +class VultrProvider(OnlineProvider): NAME = "vultr" + @classmethod + def from_env(cls) -> "VultrProvider": + return cls() + def get( self, query_filter: QueryFilter | None = None, balance_resources: bool = True - ) -> list[RawCatalogItem]: + ) -> list[CatalogItem]: offers = fetch_offers() return sorted(offers, key=lambda i: i.price) -def fetch_offers() -> list[RawCatalogItem]: +def fetch_offers() -> list[CatalogItem]: """Fetch plans with types: 1. Cloud GPU (vcg), 2. Bare Metal (vbm), @@ -39,13 +43,13 @@ def fetch_offers() -> list[RawCatalogItem]: All optimized Cloud Types (voc)""" bare_metal_plans_response = _make_request("GET", "/plans-metal?per_page=500") other_plans_response = _make_request("GET", "/plans?type=all&per_page=500") - return convert_response_to_raw_catalog_items(bare_metal_plans_response, other_plans_response) + return _make_offers(bare_metal_plans_response, other_plans_response) -def convert_response_to_raw_catalog_items( +def _make_offers( bare_metal_plans_response: Response, other_plans_response: Response -) -> list[RawCatalogItem]: - catalog_items = [] +) -> list[CatalogItem]: + offers: list[CatalogItem] = [] bare_metal_plans = bare_metal_plans_response.json()["plans_metal"] other_plans = other_plans_response.json()["plans"] @@ -54,18 +58,18 @@ def convert_response_to_raw_catalog_items( for location in plan["locations"]: catalog_item = get_bare_metal_plans(plan, location) if catalog_item: - catalog_items.append(catalog_item) + offers.append(catalog_item) for plan in other_plans: for location in plan["locations"]: catalog_item = get_instance_plans(plan, location) if catalog_item: - catalog_items.append(catalog_item) + offers.append(catalog_item) - return catalog_items + return offers -def get_bare_metal_plans(plan: dict, location: str) -> RawCatalogItem | None: +def get_bare_metal_plans(plan: dict, location: str) -> CatalogItem | None: cpu_arch = CPUArchitecture.X86 gpu_count, gpu_name, gpu_memory, gpu_vendor = 0, None, None, None if "gpu" in plan["id"]: @@ -79,11 +83,12 @@ def get_bare_metal_plans(plan: dict, location: str) -> RawCatalogItem | None: if gpu_vendor is None: logger.warning("Unknown GPU vendor for plan %s, skipping", plan["id"]) return None - return RawCatalogItem( + return CatalogItem( + provider=VultrProvider.NAME, instance_name=plan["id"], location=location, price=plan["hourly_cost"], - cpu_arch=cpu_arch.value, + cpu_arch=cpu_arch, cpu=plan["cpu_threads"], memory=plan["ram"] / 1024, gpu_count=gpu_count, @@ -95,15 +100,16 @@ def get_bare_metal_plans(plan: dict, location: str) -> RawCatalogItem | None: ) -def get_instance_plans(plan: dict, location: str) -> RawCatalogItem | None: +def get_instance_plans(plan: dict, location: str) -> CatalogItem | None: cpu_arch = CPUArchitecture.X86 plan_type = plan["type"] if plan_type in ["vc2", "vhf", "vhp", "voc"]: - return RawCatalogItem( + return CatalogItem( + provider=VultrProvider.NAME, instance_name=plan["id"], location=location, price=plan["hourly_cost"], - cpu_arch=cpu_arch.value, + cpu_arch=cpu_arch, cpu=plan["vcpu_count"], memory=plan["ram"] / 1024, gpu_count=0, @@ -141,11 +147,12 @@ def get_instance_plans(plan: dict, location: str) -> RawCatalogItem | None: gpu_count = max(1, gpu_memory_total // gpu_memory) if is_nvidia_superchip(gpu_name): cpu_arch = CPUArchitecture.ARM - return RawCatalogItem( + return CatalogItem( + provider=VultrProvider.NAME, instance_name=plan["id"], location=location, price=plan["hourly_cost"], - cpu_arch=cpu_arch.value, + cpu_arch=cpu_arch, cpu=plan["vcpu_count"], memory=plan["ram"] / 1024, gpu_count=gpu_count, diff --git a/src/gpuhunt/scripts/catalog_v1/__main__.py b/src/gpuhunt/scripts/catalog_v1/__main__.py index f8b37ea..a8b75e4 100644 --- a/src/gpuhunt/scripts/catalog_v1/__main__.py +++ b/src/gpuhunt/scripts/catalog_v1/__main__.py @@ -21,9 +21,9 @@ def main(args: Sequence[str] | None = None): ) parser.add_argument("--input", type=Path, required=True, help="The v2 catalog file to read") parser.add_argument("--output", type=Path, required=True, help="The v1 catalog file to write") - args = parser.parse_args(args) - storage.convert_catalog_v2_to_v1(path_v2=args.input, path_v1=args.output) - logging.info("Converted %s -> %s", args.input, args.output) + parsed_args = parser.parse_args(args) + storage.convert_catalog_v2_to_v1(path_v2=parsed_args.input, path_v1=parsed_args.output) + logging.info("Converted %s -> %s", parsed_args.input, parsed_args.output) if __name__ == "__main__": diff --git a/src/integrity_tests/test_all.py b/src/integrity_tests/test_all.py index 286ce77..2c1eb00 100644 --- a/src/integrity_tests/test_all.py +++ b/src/integrity_tests/test_all.py @@ -1,11 +1,10 @@ import csv import os -from dataclasses import fields from pathlib import Path import pytest -from gpuhunt._internal.models import RawCatalogItem +from gpuhunt._internal.storage import CATALOG_V2_FIELDS # Fields that are allowed to be empty, including empty strings or empty lists OPTIONAL_CATALOG_ITEM_FIELDS = ["gpu_name", "gpu_memory", "gpu_vendor", "disk_size", "flags"] @@ -24,7 +23,7 @@ def catalog(self, request): @pytest.mark.parametrize( "field", - [f.name for f in fields(RawCatalogItem) if f.name not in OPTIONAL_CATALOG_ITEM_FIELDS], + [f for f in CATALOG_V2_FIELDS if f not in OPTIONAL_CATALOG_ITEM_FIELDS], ) def test_field_present(self, catalog: csv.DictReader, field: str) -> None: for row in catalog: diff --git a/src/integrity_tests/test_hotaisle.py b/src/integrity_tests/test_hotaisle.py index f3bd3f6..9d1db81 100644 --- a/src/integrity_tests/test_hotaisle.py +++ b/src/integrity_tests/test_hotaisle.py @@ -7,9 +7,10 @@ @pytest.fixture def provider(): - api_key = os.environ.get("HOTAISLE_API_KEY") - team_handle = os.environ.get("HOTAISLE_TEAM_HANDLE") - return HotAisleProvider(api_key=api_key, team_handle=team_handle) + return HotAisleProvider( + api_key=os.environ["HOTAISLE_API_KEY"], + team_handle=os.environ["HOTAISLE_TEAM_HANDLE"], + ) @pytest.fixture diff --git a/src/integrity_tests/test_runpod.py b/src/integrity_tests/test_runpod.py index a045bf3..003dd25 100644 --- a/src/integrity_tests/test_runpod.py +++ b/src/integrity_tests/test_runpod.py @@ -1,5 +1,4 @@ import csv -from collections import Counter from pathlib import Path import pytest @@ -43,17 +42,6 @@ def test_locations(data_rows): assert len(expected - locations) <= 4 -def test_spot(data_rows): - spots = select_row(data_rows, "spot") - - expected = set(("True", "False")) - assert set(spots) == expected - - count = Counter(spots) - for spot_key in ("True", "False"): - assert count[spot_key] > 1 - - def test_gpu_present(data_rows): refs = set(name for _, name in get_gpu_map().values()) gpus = set(select_row(data_rows, "gpu_name")) diff --git a/src/tests/_internal/test_catalog.py b/src/tests/_internal/test_catalog.py index 9a04cc8..7867c68 100644 --- a/src/tests/_internal/test_catalog.py +++ b/src/tests/_internal/test_catalog.py @@ -1,7 +1,7 @@ from unittest.mock import Mock import gpuhunt._internal.catalog as internal_catalog -from gpuhunt import Catalog, CatalogItem, RawCatalogItem +from gpuhunt import AcceleratorVendor, Catalog, CatalogItem from gpuhunt.providers.vastai import VastAIProvider from gpuhunt.providers.vultr import VultrProvider @@ -15,7 +15,12 @@ def test_query_merge(self): catalog.add_provider(vultr) vastai = VastAIProvider() - vastai.get = Mock(return_value=[catalog_item(price=2), catalog_item(price=1)]) + vastai.get = Mock( + return_value=[ + catalog_item(provider="vastai", price=2), + catalog_item(provider="vastai", price=1), + ] + ) catalog.add_provider(vastai) assert catalog.query(provider=["vultr", "vastai"]) == [ @@ -43,7 +48,10 @@ def test_provider_filter(self): catalog.add_provider(vastai := VastAIProvider()) vultr_offers = [catalog_item(price=1)] - vastai_offers = [catalog_item(price=2), catalog_item(price=3)] + vastai_offers = [ + catalog_item(provider="vastai", price=2), + catalog_item(provider="vastai", price=3), + ] vultr.get = Mock(return_value=vultr_offers) vastai.get = Mock(return_value=vastai_offers) @@ -74,21 +82,20 @@ def test_gpu_name_filter(self): assert len(catalog.query(gpu_name=["a10", "A100"])) == 3 -def catalog_item(**kwargs) -> CatalogItem | RawCatalogItem: - values = dict( +def catalog_item( + provider: str = "vultr", price: float = 1, gpu_name: str | None = "gpu" +) -> CatalogItem: + return CatalogItem( + provider=provider, instance_name="instance", cpu=1, memory=1, - gpu_vendor="nvidia", + gpu_vendor=AcceleratorVendor.NVIDIA, gpu_count=1, - gpu_name="gpu", + gpu_name=gpu_name, gpu_memory=1, location="location", - price=1, + price=price, spot=False, disk_size=None, ) - values.update(kwargs) - if "provider" in values: - return CatalogItem(**values) - return RawCatalogItem(**values) diff --git a/src/tests/_internal/test_default.py b/src/tests/_internal/test_default.py new file mode 100644 index 0000000..cdb9db3 --- /dev/null +++ b/src/tests/_internal/test_default.py @@ -0,0 +1,37 @@ +import pytest + +from gpuhunt._internal.catalog import Catalog +from gpuhunt._internal.default import default_catalog + +CREDS_ENV_VARS = [ + "CRUSOE_ACCESS_KEY", + "CRUSOE_SECRET_KEY", + "CRUSOE_PROJECT_ID", + "DIGITAL_OCEAN_API_KEY", + "HOTAISLE_API_KEY", + "HOTAISLE_TEAM_HANDLE", + "JL_API_KEY", + "SEEWEB_API_TOKEN", +] + + +@pytest.fixture +def offline_catalog(monkeypatch): + monkeypatch.setattr(Catalog, "load", lambda self, version=None: None) + for var in CREDS_ENV_VARS: + monkeypatch.delenv(var, raising=False) + default_catalog.cache_clear() + yield + default_catalog.cache_clear() + + +class TestDefaultCatalog: + def test_skips_providers_with_missing_creds(self, offline_catalog) -> None: + catalog = default_catalog() + assert sorted(p.NAME for p in catalog.providers) == ["vastai", "vultr"] + + def test_loads_providers_with_creds(self, offline_catalog, monkeypatch) -> None: + monkeypatch.setenv("HOTAISLE_API_KEY", "key") + monkeypatch.setenv("HOTAISLE_TEAM_HANDLE", "team") + catalog = default_catalog() + assert "hotaisle" in [p.NAME for p in catalog.providers] diff --git a/src/tests/_internal/test_models.py b/src/tests/_internal/test_models.py index aff0307..521beeb 100644 --- a/src/tests/_internal/test_models.py +++ b/src/tests/_internal/test_models.py @@ -1,113 +1,7 @@ import pytest from gpuhunt._internal.constraints import KNOWN_AMD_GPUS -from gpuhunt._internal.models import ( - AcceleratorVendor, - AMDArchitecture, - CatalogItem, - CPUArchitecture, - RawCatalogItem, -) - -NVIDIA = AcceleratorVendor.NVIDIA -GOOGLE = AcceleratorVendor.GOOGLE -AMD = AcceleratorVendor.AMD - - -@pytest.mark.parametrize( - ["gpu_count", "gpu_vendor", "gpu_name", "expected_gpu_vendor", "expected_gpu_name"], - [ - pytest.param(None, None, None, None, None, id="none-gpu"), - pytest.param(0, None, None, None, None, id="zero-gpu"), - pytest.param(1, None, "A100", "nvidia", "A100", id="one-gpu"), - pytest.param(1, None, "tpu-v3", "google", "v3", id="one-tpu-vendor-not-set"), - pytest.param(1, "google", "tpu-v5p", "google", "v5p", id="one-tpu-vendor-is-set"), - pytest.param(1, AMD, "MI300X", "amd", "MI300X", id="cast-enum-to-string"), - ], -) -def test_raw_catalog_item_gpu_vendor_heuristic( - gpu_count: int | None, - gpu_vendor: AcceleratorVendor | str | None, - gpu_name: str | None, - expected_gpu_vendor: str | None, - expected_gpu_name: str | None, -): - dct = {} - if gpu_vendor is not None: - dct["gpu_vendor"] = gpu_vendor - if gpu_count is not None: - dct["gpu_count"] = gpu_count - if gpu_name is not None: - dct["gpu_name"] = gpu_name - - item = RawCatalogItem.from_dict(dct) - - assert item.gpu_vendor == expected_gpu_vendor - assert item.gpu_name == expected_gpu_name - - -@pytest.mark.parametrize( - ["gpu_count", "gpu_vendor", "gpu_name", "expected_gpu_vendor"], - [ - pytest.param(None, None, None, None, id="none-gpu"), - pytest.param(0, None, None, None, id="zero-gpu"), - pytest.param(1, None, None, NVIDIA, id="one-gpu-no-name"), - pytest.param(1, None, "v3", NVIDIA, id="one-gpu-with-any-name"), - pytest.param(1, "amd", "MI300X", AMD, id="cast-string-to-enum"), - ], -) -def test_catalog_item_gpu_vendor_heuristic( - gpu_count: int | None, - gpu_vendor: AcceleratorVendor | str | None, - gpu_name: str | None, - expected_gpu_vendor: AcceleratorVendor | None, -): - item = CatalogItem( - instance_name="test-instance", - location="eu-west-1", - price=1.0, - cpu=1, - memory=32.0, - gpu_vendor=gpu_vendor, - gpu_count=gpu_count, - gpu_name=gpu_name, - gpu_memory=8.0, - spot=False, - disk_size=100.0, - provider="test", - ) - - assert item.gpu_vendor == expected_gpu_vendor - - -@pytest.mark.parametrize( - ["cpu_arch", "expected_cpu_arch"], - [ - pytest.param(None, CPUArchitecture.X86, id="non-set"), - pytest.param(CPUArchitecture.X86, CPUArchitecture.X86, id="enum"), - pytest.param("ARM", CPUArchitecture.ARM, id="cast-string-to-enum"), - ], -) -def test_catalog_item_cpu_arch_heuristic( - cpu_arch: CPUArchitecture | str | None, - expected_cpu_arch: CPUArchitecture, -): - item = CatalogItem( - instance_name="test-instance", - location="eu-west-1", - price=1.0, - cpu_arch=cpu_arch, - cpu=1, - memory=32.0, - gpu_count=0, - gpu_name=None, - gpu_memory=8.0, - spot=False, - disk_size=100.0, - provider="test", - ) - - assert item.cpu_arch == expected_cpu_arch +from gpuhunt._internal.models import AMDArchitecture @pytest.mark.parametrize( @@ -131,40 +25,3 @@ def test_amd_gpu_architecture(model: str, architecture: AMDArchitecture, expecte return # If we get here, the test should fail since we could not find the GPU in our known list. assert False - - -def test_raw_catalog_item_to_from_dict() -> None: - item = RawCatalogItem( - instance_name="test-instance", - location="eu-west-1", - price=1.0, - cpu_arch=CPUArchitecture.ARM, - cpu=1, - memory=32.0, - gpu_vendor=AcceleratorVendor.NVIDIA, - gpu_count=1, - gpu_name="A10", - gpu_memory=24.0, - spot=False, - disk_size=100.0, - flags=["f1", "f2", "f3"], - provider_data={"custom_prop": 42}, - ) - item_dict = item.dict() - assert item_dict == { - "instance_name": "test-instance", - "location": "eu-west-1", - "price": 1.0, - "cpu_arch": "arm", - "cpu": 1, - "memory": 32.0, - "gpu_vendor": "nvidia", - "gpu_count": 1, - "gpu_name": "A10", - "gpu_memory": 24.0, - "spot": False, - "disk_size": 100.0, - "flags": "f1 f2 f3", - "provider_data": '{"custom_prop": 42}', - } - assert RawCatalogItem.from_dict(item_dict) == item diff --git a/src/tests/_internal/test_storage.py b/src/tests/_internal/test_storage.py new file mode 100644 index 0000000..17ec1f1 --- /dev/null +++ b/src/tests/_internal/test_storage.py @@ -0,0 +1,131 @@ +import io + +import pytest + +from gpuhunt._internal import storage +from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, CPUArchitecture + + +def catalog_item(**kwargs) -> CatalogItem: + defaults = dict( + provider="test", + instance_name="test-instance", + location="eu-west-1", + price=1.0, + cpu=1, + memory=32.0, + gpu_count=0, + gpu_name=None, + gpu_memory=None, + spot=False, + disk_size=None, + ) + return CatalogItem(**{**defaults, **kwargs}) + + +def row(**kwargs) -> dict[str, str]: + defaults = { + "instance_name": "test-instance", + "location": "eu-west-1", + "price": "1.0", + "cpu": "1", + "memory": "32.0", + "gpu_count": "0", + "gpu_name": "", + "gpu_memory": "", + "spot": "False", + "disk_size": "", + } + return {**defaults, **kwargs} + + +class TestItemToRow: + def test_all_fields(self) -> None: + item = catalog_item( + cpu_arch=CPUArchitecture.ARM, + gpu_vendor=AcceleratorVendor.NVIDIA, + gpu_count=1, + gpu_name="A10", + gpu_memory=24.0, + disk_size=100.0, + flags=["f1", "f2"], + provider_data={"custom_prop": 42}, + ) + assert storage.item_to_row(item) == { + "instance_name": "test-instance", + "location": "eu-west-1", + "price": "1.0", + "cpu": "1", + "memory": "32.0", + "gpu_count": "1", + "gpu_name": "A10", + "gpu_memory": "24.0", + "spot": "False", + "disk_size": "100.0", + "gpu_vendor": "nvidia", + "flags": "f1 f2", + "cpu_arch": "arm", + "provider_data": '{"custom_prop": 42}', + } + + def test_unset_optionals_are_empty(self) -> None: + item_row = storage.item_to_row(catalog_item()) + assert [item_row[f] for f in ("gpu_name", "gpu_memory", "disk_size", "gpu_vendor")] == [ + "", + "", + "", + "", + ] + + +class TestItemFromRow: + def test_missing_required_field_is_rejected(self) -> None: + with pytest.raises(ValueError): + storage.item_from_row(row(price=""), provider="test") + + def test_defaults_for_columns_missing_in_historical_catalogs(self) -> None: + # No cpu_arch, gpu_vendor, flags, or provider_data columns. + item = storage.item_from_row(row(), provider="test") + assert item.cpu_arch == CPUArchitecture.X86 + assert item.gpu_vendor is None + assert item.flags == [] + assert item.provider_data == {} + + def test_gpu_without_vendor_is_nvidia(self) -> None: + item = storage.item_from_row(row(gpu_count="1", gpu_name="A100"), provider="test") + assert item.gpu_vendor == AcceleratorVendor.NVIDIA + + def test_tpu_name_prefix_implies_google(self) -> None: + item = storage.item_from_row(row(gpu_count="1", gpu_name="tpu-v3"), provider="test") + assert item.gpu_name == "v3" + assert item.gpu_vendor == AcceleratorVendor.GOOGLE + + +class TestLoad: + def test_round_trip(self, tmp_path) -> None: + items = [ + catalog_item(), + catalog_item( + price=12.0, + gpu_count=8, + gpu_name="H100", + gpu_memory=80.0, + gpu_vendor=AcceleratorVendor.NVIDIA, + spot=True, + disk_size=500.0, + cpu_arch=CPUArchitecture.ARM, + flags=["f1"], + provider_data={"k": 1}, + ), + ] + path = str(tmp_path / "test.csv") + storage.dump(items, path) + with open(path) as f: + assert list(storage.load(f, provider="test")) == items + + def test_malformed_row_is_skipped(self) -> None: + header = ",".join(storage.CATALOG_V2_FIELDS) + good = "i,loc,1.0,1,1.0,0,,,False,,,,x86,{}" + bad = ",loc,1.0,1,1.0,0,,,False,,,,x86,{}" + f = io.StringIO(f"{header}\n{bad}\n{good}\n") + assert [i.instance_name for i in storage.load(f, provider="test")] == ["i"] diff --git a/src/tests/providers/test_jarvislabs.py b/src/tests/providers/test_jarvislabs.py index e0d47a4..4c89d10 100644 --- a/src/tests/providers/test_jarvislabs.py +++ b/src/tests/providers/test_jarvislabs.py @@ -3,7 +3,7 @@ from gpuhunt._internal.models import QueryFilter from gpuhunt.providers.jarvislabs import ( JarvisLabsProvider, - convert_response_to_raw_catalog_items, + _make_offers, ) SERVER_META_RESPONSE = { @@ -124,8 +124,8 @@ } -def test_convert_response_to_raw_catalog_items(): - offers = convert_response_to_raw_catalog_items(SERVER_META_RESPONSE) +def test_make_offers(): + offers = _make_offers(SERVER_META_RESPONSE) assert not any(o.spot for o in offers) l4_vm = [o for o in offers if o.gpu_name == "L4" and not o.spot] @@ -174,7 +174,7 @@ def test_convert_response_to_raw_catalog_items(): def test_convert_response_warns_and_skips_unsupported_regions(caplog): - convert_response_to_raw_catalog_items(SERVER_META_RESPONSE) + _make_offers(SERVER_META_RESPONSE) assert "Skipping JarvisLabs GPU VM offer in unsupported region unknown-region" in caplog.text assert "Skipping JarvisLabs CPU VM offer in unsupported region unknown-region" in caplog.text @@ -196,7 +196,7 @@ def test_convert_response_skips_unmapped_gpu_types_with_spaces(caplog): ], } - assert convert_response_to_raw_catalog_items(response) == [] + assert _make_offers(response) == [] assert "Skipping JarvisLabs GPU offer with unmapped gpu_type: RTX A6000" in caplog.text @@ -237,7 +237,7 @@ def test_convert_response_skips_malformed_specs(caplog): }, } - offers = convert_response_to_raw_catalog_items(response) + offers = _make_offers(response) assert offers == [] assert "Skipping JarvisLabs GPU offer without price: L4" in caplog.text diff --git a/src/tests/providers/test_providers.py b/src/tests/providers/test_providers.py index 18da7c3..4c4a61c 100644 --- a/src/tests/providers/test_providers.py +++ b/src/tests/providers/test_providers.py @@ -5,7 +5,10 @@ import pytest import gpuhunt.providers +import gpuhunt.providers.base from gpuhunt._internal.catalog import OFFLINE_PROVIDERS, ONLINE_PROVIDERS +from gpuhunt._internal.default import ONLINE_PROVIDER_MODULES +from gpuhunt.providers.base import OfflineProvider, OnlineProvider @pytest.fixture() @@ -22,10 +25,10 @@ def providers(): continue if member.__name__.islower(): continue # skip builtins to avoid CPython bug #89489 in `issubclass` below - if not issubclass(member, gpuhunt.providers.AbstractProvider): - continue - if member.__name__ == "AbstractProvider": + if not issubclass(member, gpuhunt.providers.base.AbstractProvider): continue + if inspect.isabstract(member): + continue # skip AbstractProvider, OnlineProvider, OfflineProvider members.append(member) assert members return members @@ -38,7 +41,7 @@ def test_catalog_providers_is_unique(): def test_all_providers_have_a_names(providers): names = [p.NAME for p in providers] - assert gpuhunt.providers.AbstractProvider.NAME not in names + assert gpuhunt.providers.base.AbstractProvider.NAME not in names assert len(set(names)) == len(names) @@ -47,3 +50,23 @@ def test_catalog_providers(providers): names = [p.NAME for p in providers] assert set(CATALOG_PROVIDERS) == set(names) assert len(CATALOG_PROVIDERS) == len(names) + + +def test_online_providers_subclass_online_provider(providers): + online = [p for p in providers if p.NAME in ONLINE_PROVIDERS] + assert online + for provider in online: + assert issubclass(provider, OnlineProvider), provider + + +def test_offline_providers_subclass_offline_provider(providers): + offline = [p for p in providers if p.NAME in OFFLINE_PROVIDERS] + assert offline + for provider in offline: + assert issubclass(provider, OfflineProvider), provider + + +def test_default_catalog_loads_every_online_provider(providers): + classes_by_name = {p.__name__: p for p in providers} + names = {classes_by_name[class_name].NAME for _, class_name in ONLINE_PROVIDER_MODULES} + assert names == set(ONLINE_PROVIDERS) diff --git a/src/tests/providers/test_runpod.py b/src/tests/providers/test_runpod.py index aa1b953..5d79cf3 100644 --- a/src/tests/providers/test_runpod.py +++ b/src/tests/providers/test_runpod.py @@ -1,7 +1,7 @@ import pytest from requests import RequestException -from gpuhunt._internal.models import RawCatalogItem +from gpuhunt._internal.models import CatalogItem from gpuhunt.providers import runpod as runpod_module from gpuhunt.providers.runpod import RunpodProvider, _cpu_size_ladder @@ -11,7 +11,7 @@ def test_cpu_size_ladder(): assert _cpu_size_ladder(3, 20) == [3, 6, 12, 20] -def test_make_cpu_catalog_items(): +def test_make_cpu_offers(): provider = object.__new__(RunpodProvider) cpu_flavors = [ { @@ -28,7 +28,7 @@ def test_make_cpu_catalog_items(): } ] - items = provider._make_cpu_catalog_items("AP-JP-1", cpu_flavors) + items = provider._make_cpu_offers("AP-JP-1", cpu_flavors) assert [item.instance_name for item in items] == [ "cpu3g-2-8", @@ -50,7 +50,7 @@ def test_make_cpu_catalog_items(): assert items[-1].provider_data == {} -def test_make_cpu_catalog_items_skips_invalid_flavors(): +def test_make_cpu_offers_skips_invalid_flavors(): provider = object.__new__(RunpodProvider) cpu_flavors = [ { @@ -127,7 +127,7 @@ def test_make_cpu_catalog_items_skips_invalid_flavors(): }, ] - assert provider._make_cpu_catalog_items("AP-JP-1", cpu_flavors) == [] + assert provider._make_cpu_offers("AP-JP-1", cpu_flavors) == [] def test_fetch_cpu_offers_handles_partial_datacenter_failures(monkeypatch): @@ -185,7 +185,8 @@ def fake_get_cpu_flavors(dc_id: str): def test_fetch_offers_appends_cpu_items(monkeypatch): provider = object.__new__(RunpodProvider) - cpu_item = RawCatalogItem( + cpu_item = CatalogItem( + provider="runpod", instance_name="cpu3g-2-8", location="AP-JP-1", price=0.08, diff --git a/src/tests/providers/test_seeweb.py b/src/tests/providers/test_seeweb.py index fd1d936..e87ae85 100644 --- a/src/tests/providers/test_seeweb.py +++ b/src/tests/providers/test_seeweb.py @@ -1,6 +1,7 @@ import pytest import requests +from gpuhunt import MissingCredsError from gpuhunt.providers import seeweb as seeweb_module from gpuhunt.providers.seeweb import SeewebProvider, _normalize_gpu @@ -128,15 +129,15 @@ def fake_get(*args, **kwargs): assert request["kwargs"]["timeout"] == 30 -def test_token_defaults_to_environment(monkeypatch): +def test_from_env_reads_token(monkeypatch): monkeypatch.setenv("SEEWEB_API_TOKEN", "from-env") - assert SeewebProvider().token == "from-env" + assert SeewebProvider.from_env().token == "from-env" -def test_missing_token_is_rejected(monkeypatch): +def test_from_env_without_token_is_rejected(monkeypatch): monkeypatch.delenv("SEEWEB_API_TOKEN", raising=False) - with pytest.raises(ValueError, match="SEEWEB_API_TOKEN"): - SeewebProvider() + with pytest.raises(MissingCredsError, match="SEEWEB_API_TOKEN"): + SeewebProvider.from_env() def test_http_error_is_propagated(monkeypatch): diff --git a/src/tests/providers/test_vastai.py b/src/tests/providers/test_vastai.py index 8eb5473..359e039 100644 --- a/src/tests/providers/test_vastai.py +++ b/src/tests/providers/test_vastai.py @@ -4,7 +4,7 @@ def test_make_filters_defaults_to_datacenter_only(): filters = VastAIProvider(community_cloud=False).make_filters(QueryFilter()) - assert filters["datacenter"]["eq"] is True + assert filters["datacenter"] == {"eq": True} assert "external" not in filters diff --git a/src/tests/providers/test_verda.py b/src/tests/providers/test_verda.py index 1203a0f..44aad19 100644 --- a/src/tests/providers/test_verda.py +++ b/src/tests/providers/test_verda.py @@ -1,15 +1,13 @@ -import dataclasses - import pytest import gpuhunt._internal.catalog as internal_catalog -from gpuhunt import AcceleratorVendor, Catalog, CatalogItem, RawCatalogItem +from gpuhunt import AcceleratorVendor, Catalog, CatalogItem from gpuhunt.providers.verda import ( InstanceType, VerdaProvider, - generate_instances, + _make_offer, + _make_offers, get_gpu_name, - transform_instance, ) @@ -185,7 +183,7 @@ def list_available_instances(raw_instance_types, locations): spots = (True, False) locations = [loc["loc"] for loc in locations] instances = [instance_types(raw_instance_types[0])] - list_instances = generate_instances(spots, locations, instances) + list_instances = _make_offers(spots, locations, instances) assert len(list_instances) == 4 assert [i.price for i in list_instances if i.spot] == [1, 70] * 2 @@ -197,14 +195,6 @@ def test_gpu_name(caplog): assert get_gpu_name("") is None -def transform(raw_catalog_items: list[RawCatalogItem]) -> list[CatalogItem]: - items = [] - for raw in raw_catalog_items: - item = CatalogItem(provider="verda", **dataclasses.asdict(raw)) - items.append(item) - return items - - def test_available_query(mocker, raw_instance_types): catalog = Catalog(balance_resources=False, auto_reload=False) @@ -312,15 +302,16 @@ def test_available_query_with_instance(mocker, raw_instance_types): def test_transform_instance(raw_instance_types): location = "ICE-01" is_spot = True - item = transform_instance(instance_types(raw_instance_types[1]), is_spot, location) + item = _make_offer(instance_types(raw_instance_types[1]), is_spot, location) - expected = RawCatalogItem( + expected = CatalogItem( + provider="verda", instance_name="2A6000.20V", location="ICE-01", price=0.7, cpu=20, memory=120, - gpu_vendor=AcceleratorVendor.NVIDIA.value, + gpu_vendor=AcceleratorVendor.NVIDIA, gpu_count=2, gpu_name="A6000", gpu_memory=96 / 2, @@ -328,15 +319,16 @@ def test_transform_instance(raw_instance_types): disk_size=None, ) - assert RawCatalogItem.from_dict(item) == expected + assert item == expected def test_cpu_instance(raw_instance_types): location = "ICE-01" is_spot = False - item = transform_instance(instance_types(raw_instance_types[2]), is_spot, location) + item = _make_offer(instance_types(raw_instance_types[2]), is_spot, location) - expected = RawCatalogItem( + expected = CatalogItem( + provider="verda", instance_name="CPU.120V.480G", location="ICE-01", price=3, @@ -350,7 +342,7 @@ def test_cpu_instance(raw_instance_types): disk_size=None, ) - assert RawCatalogItem.from_dict(item) == expected + assert item == expected def test_order(mocker, raw_instance_types):