gpuhunt 0.0.1.dev1__py3-none-any.whl

gpuhunt/__init__.py

@@ -0,0 +1,14 @@
+from functools import lru_cache
+
+from gpuhunt._catalog import Catalog
+
+
+@lru_cache()
+def default_catalog() -> Catalog:
+    catalog = Catalog()
+    catalog.load()
+    return catalog
+
+
+def query() -> list:  # todo
+    return default_catalog().query()

gpuhunt/__main__.py

@@ -0,0 +1,52 @@
+import argparse
+import logging
+import os
+import sys
+
+import gpuhunt._storage as storage
+from gpuhunt.providers.aws import AWSProvider
+from gpuhunt.providers.azure import AzureProvider
+from gpuhunt.providers.gcp import GCPProvider
+from gpuhunt.providers.lambdalabs import LambdaLabsProvider
+from gpuhunt.providers.tensordock import TensorDockProvider
+from gpuhunt.providers.vastai import VastAIProvider
+
+
+def main():
+    parser = argparse.ArgumentParser(prog="python3 -m gpuhunt")
+    parser.add_argument(
+        "provider",
+        choices=["aws", "azure", "gcp", "lambdalabs", "tensordock", "vastai"],
+    )
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--no-filter", action="store_true")
+    args = parser.parse_args()
+    logging.basicConfig(
+        level=logging.INFO,
+        stream=sys.stdout,
+        format="%(asctime)s %(levelname)s %(message)s",
+    )
+
+    if args.provider == "aws":
+        provider = AWSProvider()
+    elif args.provider == "azure":
+        provider = AzureProvider(os.getenv("AZURE_SUBSCRIPTION_ID"))
+    elif args.provider == "gcp":
+        provider = GCPProvider(os.getenv("GCP_PROJECT_ID"))
+    elif args.provider == "lambdalabs":
+        provider = LambdaLabsProvider(os.getenv("LAMBDALABS_TOKEN"))
+    elif args.provider == "tensordock":
+        provider = TensorDockProvider()
+    elif args.provider == "vastai":
+        provider = VastAIProvider()
+    else:
+        exit(f"Unknown provider {args.provider}")
+
+    offers = provider.get()
+    if not args.no_filter:
+        offers = provider.filter(offers)
+    storage.dump(offers, args.output)
+
+
+if __name__ == "__main__":
+    main()

gpuhunt/_catalog.py

@@ -0,0 +1,71 @@
+import csv
+import io
+import logging
+import urllib.request
+import zipfile
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+logger = logging.getLogger(__name__)
+version_url = "https://dstack-gpu-pricing.s3.eu-west-1.amazonaws.com/v1/version"
+catalog_url = (
+    "https://dstack-gpu-pricing.s3.eu-west-1.amazonaws.com/v1/{version}/catalog.zip"
+)
+
+
+@dataclass(frozen=True)
+class CatalogItem:
+    provider: str
+    instance_name: str
+    location: str
+    price: float
+    cpus: int
+    memory: float
+    gpu_count: int
+    gpu_name: Optional[str]
+    gpu_memory: Optional[float]
+    spot: bool
+
+
+class Catalog:
+    def __init__(self):
+        self.catalog = None
+
+    def query(self) -> list[CatalogItem]:
+        return list(self._read_catalog())
+
+    def load(self, version: str = None):
+        if version is None:
+            version = self.get_latest_version()
+        logger.debug("Downloading catalog %s...", version)
+        with urllib.request.urlopen(catalog_url.format(version=version)) as f:
+            self.catalog = io.BytesIO(f.read())
+
+    @staticmethod
+    def get_latest_version() -> str:
+        with urllib.request.urlopen(version_url) as f:
+            return f.read().decode("utf-8").strip()
+
+    def _read_catalog(self) -> Iterable[CatalogItem]:
+        with zipfile.ZipFile(self.catalog) as zip_file:
+            providers = [f[:-4] for f in zip_file.namelist() if f.endswith(".csv")]
+            for provider in providers:
+                with zip_file.open(f"{provider}.csv", "r") as csv_file:
+                    reader: Iterable[dict[str, str]] = csv.DictReader(
+                        io.TextIOWrapper(csv_file, "utf-8")
+                    )
+                    for row in reader:
+                        yield CatalogItem(
+                            provider=provider,
+                            instance_name=row["instance_name"],
+                            location=row["location"],
+                            price=float(row["price"]),
+                            cpus=int(row["cpu"]),
+                            memory=float(row["memory"]),
+                            gpu_count=int(row["gpu_count"]),
+                            gpu_name=row["gpu_name"] or None,
+                            gpu_memory=float(row["gpu_memory"])
+                            if row["gpu_memory"]
+                            else None,
+                            spot=row["spot"] == "True",
+                        )

gpuhunt/_models.py

@@ -0,0 +1,24 @@
+from typing import Any, Optional
+
+from pydantic import BaseModel, model_validator
+
+
+class InstanceOffer(BaseModel):
+    instance_name: str
+    location: Optional[str] = None  # region or zone
+    price: Optional[float] = None  # $ per hour
+    cpu: Optional[int] = None
+    memory: Optional[float] = None  # in GB
+    gpu_count: Optional[int] = None
+    gpu_name: Optional[str] = None
+    gpu_memory: Optional[float] = None  # in GB
+    spot: Optional[bool] = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def parse_empty_as_none(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            for key, value in data.items():
+                if value == "":
+                    data[key] = None
+        return data

gpuhunt/_storage.py

@@ -0,0 +1,26 @@
+import csv
+from typing import Iterable
+
+from gpuhunt._models import InstanceOffer
+
+
+def dump(offers: list[InstanceOffer], path: str):
+    with open(path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=list(InstanceOffer.model_fields.keys()))
+        writer.writeheader()
+        for offer in offers:
+            writer.writerow(offer.model_dump())
+
+
+def load(path: str) -> list[InstanceOffer]:
+    offers = []
+    with open(path, "r", newline="") as f:
+        reader: Iterable[dict[str, str]] = csv.DictReader(f)
+        for row in reader:
+            offer = InstanceOffer.model_validate(row)
+            offers.append(offer)
+    return offers
+
+
+def sort_key(offer: InstanceOffer):
+    return offer.gpu_count, offer.instance_name, offer.price, offer.location

gpuhunt/providers/__init__.py

@@ -0,0 +1,13 @@
+from abc import ABC, abstractmethod
+
+from gpuhunt._models import InstanceOffer
+
+
+class AbstractProvider(ABC):
+    @abstractmethod
+    def get(self) -> list[InstanceOffer]:
+        pass
+
+    @classmethod
+    def filter(cls, offers: list[InstanceOffer]) -> list[InstanceOffer]:
+        return offers

gpuhunt/providers/aws.py

@@ -0,0 +1,195 @@
+import copy
+import csv
+import datetime
+import logging
+import os
+import re
+import tempfile
+from collections import defaultdict
+from typing import Iterable, Optional
+
+import boto3
+import requests
+from botocore.exceptions import ClientError, EndpointConnectionError
+
+from gpuhunt._models import InstanceOffer
+from gpuhunt.providers import AbstractProvider
+
+logger = logging.getLogger(__name__)
+ec2_pricing_url = "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonEC2/current/index.csv"
+disclaimer_rows_skip = 5
+
+
+class AWSProvider(AbstractProvider):
+    """
+    AWSProvider parses Bulk API index file for AmazonEC2 in all regions and fills missing GPU details
+
+    Required IAM permissions:
+    * `ec2:DescribeInstanceTypes`
+    """
+
+    def __init__(self, cache_path: Optional[str] = None):
+        if cache_path:
+            self.cache_path = cache_path
+        else:
+            self.temp_dir = tempfile.TemporaryDirectory()
+            self.cache_path = self.temp_dir.name + "/index.csv"
+        # todo aws creds
+        self.filters = {
+            "TermType": ["OnDemand"],
+            "Tenancy": ["Shared"],
+            "Operating System": ["Linux"],
+            "CapacityStatus": ["Used"],
+            "Unit": ["Hrs"],
+            "Currency": ["USD"],
+            "Pre Installed S/W": ["", "NA"],
+        }
+        self.preview_gpus = {
+            "p4de.24xlarge": ("A100", 80.0),
+        }
+
+    def get(self) -> list[InstanceOffer]:
+        if not os.path.exists(self.cache_path):
+            logger.info("Downloading EC2 prices to %s", self.cache_path)
+            with requests.get(ec2_pricing_url, stream=True) as r:
+                r.raise_for_status()
+                with open(self.cache_path, "wb") as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        f.write(chunk)
+
+        offers = []
+        with open(self.cache_path, "r", 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):
+                    continue
+                offer = InstanceOffer(
+                    instance_name=row["Instance Type"],
+                    location=row["Region Code"],
+                    price=float(row["PricePerUnit"]),
+                    cpu=int(row["vCPU"]),
+                    memory=parse_memory(row["Memory"]),
+                    gpu_count=parse_optional_count(row["GPU"]),
+                    spot=False,
+                )
+                offers.append(offer)
+        self.fill_gpu_details(offers)
+        return self.add_spots(offers)
+
+    def skip(self, row: dict[str, str]) -> bool:
+        for key, values in self.filters.items():
+            if row[key] not in values:
+                return True
+        return False
+
+    def fill_gpu_details(self, offers: list[InstanceOffer]):
+        regions = defaultdict(list)
+        for offer in offers:
+            if offer.gpu_count > 0 and offer.instance_name not in self.preview_gpus:
+                regions[offer.location].append(offer.instance_name)
+
+        gpus = copy.deepcopy(self.preview_gpus)
+        while regions:
+            region = max(regions, key=lambda r: len(regions[r]))
+            instance_types = regions.pop(region)
+
+            logger.info("Fetching GPU details for %s", region)
+            client = boto3.client("ec2", region_name=region)
+            paginator = client.get_paginator("describe_instance_types")
+            for page in paginator.paginate(InstanceTypes=instance_types):
+                for i in page["InstanceTypes"]:
+                    gpu = i["GpuInfo"]["Gpus"][0]
+                    gpus[i["InstanceType"]] = (
+                        gpu["Name"],
+                        gpu["MemoryInfo"]["SizeInMiB"] / 1024,
+                    )
+
+            regions = {
+                region: left
+                for region, names in regions.items()
+                if (left := [i for i in names if i not in instance_types])
+            }
+
+        for offer in offers:
+            if offer.gpu_count > 0:
+                offer.gpu_name, offer.gpu_memory = gpus[offer.instance_name]
+
+    def add_spots(self, offers: list[InstanceOffer]) -> list[InstanceOffer]:
+        region_instances = defaultdict(set)
+        for offer in offers:
+            region_instances[offer.location].add(offer.instance_name)
+
+        spot_prices = dict()
+        for region, instance_types in region_instances.items():
+            logger.info("Fetching spot prices for %s", region)
+            try:
+                client = boto3.client("ec2", region_name=region)  # todo creds
+                pages = client.get_paginator("describe_spot_price_history").paginate(
+                    Filters=[
+                        {
+                            "Name": "product-description",
+                            "Values": ["Linux/UNIX"],
+                        }
+                    ],
+                    InstanceTypes=list(instance_types),
+                    StartTime=datetime.datetime.utcnow(),
+                )
+
+                instance_prices = defaultdict(list)
+                for page in pages:
+                    for item in page["SpotPriceHistory"]:
+                        instance_prices[item["InstanceType"]].append(
+                            float(item["SpotPrice"])
+                        )
+                for (
+                    instance_type,
+                    zone_prices,
+                ) in instance_prices.items():  # reduce zone prices to a single value
+                    spot_prices[(instance_type, region)] = min(zone_prices)
+            except (ClientError, EndpointConnectionError) as e:
+                pass
+
+        spot_offers = []
+        for offer in offers:
+            if (
+                price := spot_prices.get((offer.instance_name, offer.location))
+            ) is None:
+                continue
+            spot_offer = copy.deepcopy(offer)
+            spot_offer.spot = True
+            spot_offer.price = price
+            spot_offers.append(spot_offer)
+        return offers + spot_offers
+
+    @classmethod
+    def filter(cls, offers: list[InstanceOffer]) -> list[InstanceOffer]:
+        return [
+            i
+            for i in offers
+            if any(
+                i.instance_name.startswith(family)
+                for family in [
+                    "t2.small",
+                    "c5.",
+                    "m5.",
+                    "p3.",
+                    "g5.",
+                    "g4dn.",
+                    "p4d.",
+                    "p4de.",
+                ]
+            )
+        ]
+
+
+def parse_memory(s: str) -> float:
+    r = re.match(r"^([0-9.]+) GiB$", s)
+    return float(r.group(1))
+
+
+def parse_optional_count(s: str) -> int:
+    if not s:
+        return 0
+    return int(s)

gpuhunt/providers/azure.py

@@ -0,0 +1,222 @@
+import json
+import logging
+import os
+import re
+import urllib.parse
+from collections import namedtuple
+from queue import Queue
+from threading import Thread
+from typing import Iterable, Optional, Tuple
+
+import requests
+from azure.core.credentials import TokenCredential
+from azure.identity import DefaultAzureCredential
+from azure.mgmt.compute import ComputeManagementClient
+
+from gpuhunt._models import InstanceOffer
+from gpuhunt.providers import AbstractProvider
+
+logger = logging.getLogger(__name__)
+prices_url = "https://prices.azure.com/api/retail/prices"
+prices_version = "2023-01-01-preview"
+prices_filters = [
+    "serviceName eq 'Virtual Machines'",
+    "priceType eq 'Consumption'",
+    "contains(productName, 'Windows') eq false",
+    "contains(productName, 'Dedicated') eq false",
+    "contains(meterName, 'Low Priority') eq false",  # retires in 2025
+]
+VMSeries = namedtuple("VMSeries", ["pattern", "gpu_name", "gpu_memory"])
+gpu_vm_series = [
+    VMSeries(r"NC(\d+)ads_A100_v4", "A100", 80.0),  # NC A100 v4-series [A100 80GB]
+    VMSeries(r"NC(\d+)ads_A10_v4", "A10", 24.0),  # NC A10 v4-series [A10]
+    VMSeries(r"NC(\d+)as_T4_v3", "T4", 16.0),  # NCasT4_v3-series [T4]
+    VMSeries(r"NC(\d+)r?s_v3", "V100", 16.0),  # NCv3-series [V100 16GB]
+    VMSeries(r"ND(\d+)amsr_A100_v4", "A100", 80.0),  # NDm A100 v4-series [8xA100 80GB]
+    VMSeries(r"ND(\d+)asr_v4", "A100", 40.0),  # ND A100 v4-series [8xA100 40GB]
+    VMSeries(r"ND(\d+)rs_v2", "V100", 32.0),  # NDv2-series [8xV100 32GB]
+    VMSeries(r"NG(\d+)adm?s_V620_v1", "V620", None),  # NGads V620-series [V620]  # todo
+    VMSeries(r"NV(\d+)adm?s_A10_v5", "A10", 24.0),  # NVadsA10 v5-series [A10]
+    VMSeries(r"NV(\d+)as_v4", "MI25", None),  # NVv4-series [MI25]  # todo
+    VMSeries(r"NV(\d+)s_v3", "M60", None),  # NVv3-series [M60]  # todo
+]
+# https://learn.microsoft.com/en-us/azure/virtual-machines/sizes-previous-gen
+retired_vm_series = [
+    r"Basic_A(\d+)",
+    r"Standard_A(\d+)",
+    r"Standard_D(\d+)",
+    r"Standard_DC(\d+)s",
+    r"Standard_DS(\d+)",
+    r"Standard_F(\d+)",
+    r"Standard_F(\d+)s",
+    r"Standard_G(\d+)",
+    r"Standard_GS(\d+)",
+    r"Standard_L(\d+)s",
+    r"Standard_NC(\d+)r?",
+    r"Standard_NC(\d+)r?s_v2",
+    r"Standard_ND(\d+)r?s",
+    r"Standard_NV(\d+)",
+    r"Standard_NV(\d+)s_v2",
+]
+
+
+class AzureProvider(AbstractProvider):
+    def __init__(
+        self,
+        subscription_id: str,
+        credential: Optional[TokenCredential] = None,
+        cache_dir: Optional[str] = None,
+    ):
+        self.cache_dir = cache_dir
+        self.client = ComputeManagementClient(
+            credential=credential or DefaultAzureCredential(),
+            subscription_id=subscription_id,
+        )
+
+    def get_pages(self, threads: int = 8) -> Iterable[list[dict]]:
+        q = Queue()
+        workers = [
+            Thread(target=self._get_pages_worker, args=(q, threads, i), daemon=True)
+            for i in range(threads)
+        ]
+        for worker in workers:
+            worker.start()
+
+        exited = 0
+        while exited < threads:
+            page = q.get()
+            if page is None:
+                exited += 1
+            else:
+                yield page
+            q.task_done()
+
+    def _get_pages_worker(self, q: Queue, stride: int, worker_id: int):
+        page_id = worker_id
+        while True:
+            cached_page = None
+            if self.cache_dir is not None:
+                cached_page = os.path.join(self.cache_dir, f"{page_id:04}.json")
+            if cached_page is not None and os.path.exists(cached_page):
+                with open(cached_page, "r") as f:
+                    data = json.load(f)
+            else:
+                logger.info("Worker %s fetches pricing page %s", worker_id, page_id)
+                data = requests.get(
+                    prices_url,
+                    params={
+                        "api-version": "2023-01-01-preview",
+                        "$filter": " and ".join(prices_filters),
+                        "$skip": page_id * 100,
+                    },
+                ).json()
+                if cached_page is not None:
+                    with open(cached_page, "w") as f:
+                        json.dump(data, f)
+
+            if not data["Items"]:
+                q.put(None)
+                logger.info("Worker %s exited", worker_id)
+                return
+            q.put(data["Items"])
+            page_id += stride
+
+    def get(self) -> list[InstanceOffer]:
+        offers = []
+        for page in self.get_pages():
+            for item in page:
+                if is_retired(item["armSkuName"]):
+                    continue
+                if not item["armSkuName"]:
+                    continue
+                offer = InstanceOffer(
+                    instance_name=item["armSkuName"],
+                    location=item["armRegionName"],
+                    price=item["retailPrice"],
+                    spot="Spot" in item["meterName"],
+                )
+                offers.append(offer)
+        return self.fill_details(offers)
+
+    def fill_details(self, offers: list[InstanceOffer]) -> list[InstanceOffer]:
+        logger.info("Fetching instance details")
+        instances = {}
+        resources = self.client.resource_skus.list()
+        for resource in resources:
+            if resource.resource_type != "virtualMachines":
+                continue
+            if is_retired(resource.name):
+                continue
+            capabilities = {pair.name: pair.value for pair in resource.capabilities}
+            gpu_count, gpu_name, gpu_memory = 0, None, None
+            if "GPUs" in capabilities:
+                gpu_count = int(capabilities["GPUs"])
+                gpu_name, gpu_memory = get_gpu_name_memory(resource.name)
+            instances[resource.name] = InstanceOffer(
+                instance_name=resource.name,
+                cpu=capabilities["vCPUs"],
+                memory=float(capabilities["MemoryGB"]),
+                gpu_count=gpu_count,
+                gpu_name=gpu_name,
+                gpu_memory=gpu_memory,
+            )
+        with_details = []
+        without_details = []
+        for offer in offers:
+            if (resources := instances.get(offer.instance_name)) is None:
+                without_details.append(offer)
+                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
+            with_details.append(offer)
+        return with_details + without_details
+
+    @classmethod
+    def filter(cls, offers: list[InstanceOffer]) -> list[InstanceOffer]:
+        vm_series = [
+            VMSeries(r"D(\d+)s_v3", None, None),  # Dsv3-series
+            VMSeries(r"E(\d+)i?s_v4", None, None),  # Esv4-series
+            VMSeries(r"E(\d+)-(\d+)s_v4", None, None),  # Esv4-series (constrained vCPU)
+            VMSeries(r"NC(\d+)s_v3", "V100", 16 * 1024),  # NCv3-series [V100 16GB]
+            VMSeries(r"NC(\d+)as_T4_v3", "T4", 16 * 1024),  # NCasT4_v3-series [T4]
+            VMSeries(r"ND(\d+)rs_v2", "V100", 32 * 1024),  # NDv2-series [8xV100 32GB]
+            VMSeries(
+                r"NV(\d+)adm?s_A10_v5", "A10", 24 * 1024
+            ),  # NVadsA10 v5-series [A10]
+            VMSeries(
+                r"NC(\d+)ads_A100_v4", "A100", 80 * 1024
+            ),  # NC A100 v4-series [A100 80GB]
+            VMSeries(
+                r"ND(\d+)asr_v4", "A100", 40 * 1024
+            ),  # ND A100 v4-series [8xA100 40GB]
+            VMSeries(
+                r"ND(\d+)amsr_A100_v4", "A100", 80 * 1024
+            ),  # NDm A100 v4-series [8xA100 80GB]
+        ]
+        vm_series_pattern = re.compile(
+            f"^Standard_({'|'.join(series.pattern for series in vm_series)})$"
+        )
+        return [i for i in offers if vm_series_pattern.match(i.instance_name)]
+
+
+def get_gpu_name_memory(vm_name: str) -> Tuple[Optional[str], Optional[float]]:
+    for pattern, gpu_name, gpu_memory in gpu_vm_series:
+        m = re.match(f"^Standard_{pattern}$", vm_name)
+        if m is None:
+            continue
+        if gpu_name == "A10" and vm_name.endswith("_v4"):
+            gpu_memory = gpu_memory * min(1.0, int(m.group(1)) / 16)
+        elif gpu_name == "A10" and vm_name.endswith("_v5"):
+            gpu_memory = gpu_memory * min(1.0, int(m.group(1)) / 36)
+
+        return gpu_name, gpu_memory
+    return None, None
+
+
+def is_retired(name: str) -> bool:
+    if re.match(f"^({'|'.join(retired_vm_series)})$", name):
+        return True
+    return False