dstack 0.18.43__py3-none-any.whl → 0.19.0rc1__py3-none-any.whl

dstack/_internal/core/{models/backends/oci.py → backends/oci/models.py}

@@ -1,18 +1,10 @@
-from typing import Dict
+from typing import Annotated, Dict, List, Literal, Optional, Union
 
 from pydantic import Field, root_validator
-from typing_extensions import Annotated, List, Literal, Optional, Union
 
-from dstack._internal.core.models.backends.base import ConfigMultiElement
 from dstack._internal.core.models.common import CoreModel
 
 
-class OCIConfigInfo(CoreModel):
-    type: Literal["oci"] = "oci"
-    regions: Optional[List[str]] = None
-    compartment_id: Optional[str] = None
-
-
 class OCIClientCreds(CoreModel):
     type: Annotated[Literal["client"], Field(description="The type of credentials")] = "client"
     user: Annotated[str, Field(description="User OCID")]
@@ -62,27 +54,34 @@ class OCICreds(CoreModel):
     __root__: AnyOCICreds = Field(..., discriminator="type")
 
 
-class OCIConfigInfoWithCreds(OCIConfigInfo):
-    creds: AnyOCICreds
-
-
-AnyOCIConfigInfo = Union[OCIConfigInfo, OCIConfigInfoWithCreds]
+class OCIBackendConfig(CoreModel):
+    type: Annotated[Literal["oci"], Field(description="The type of backend")] = "oci"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of OCI regions. Omit to use all regions"),
+    ] = None
+    compartment_id: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "Compartment where `dstack` will create all resources."
+                " Omit to instruct `dstack` to create a new compartment"
+            )
+        ),
+    ] = None
 
 
-class OCIConfigInfoWithCredsPartial(CoreModel):
-    type: Literal["oci"] = "oci"
-    creds: Optional[AnyOCICreds]
-    regions: Optional[List[str]]
-    compartment_id: Optional[str]
+class OCIBackendConfigWithCreds(OCIBackendConfig):
+    creds: Annotated[AnyOCICreds, Field(description="The credentials", discriminator="type")]
 
 
-class OCIConfigValues(CoreModel):
-    type: Literal["oci"] = "oci"
-    default_creds: bool = False
-    regions: Optional[ConfigMultiElement]
-    compartment_id: Optional[str] = None
+AnyOCIBackendConfig = Union[OCIBackendConfig, OCIBackendConfigWithCreds]
 
 
-class OCIStoredConfig(OCIConfigInfo):
+class OCIStoredConfig(OCIBackendConfig):
     compartment_id: str
     subnet_ids_per_region: Dict[str, str]
+
+
+class OCIConfig(OCIStoredConfig):
+    creds: AnyOCICreds

dstack/_internal/core/backends/oci/region.py

@@ -5,7 +5,7 @@ from typing import Any, Dict, Iterable, List, Mapping, Set
 import oci
 
 from dstack._internal.core.backends.oci.auth import get_client_config
-from dstack._internal.core.models.backends.oci import AnyOCICreds
+from dstack._internal.core.backends.oci.models import AnyOCICreds
 
 
 class OCIRegionClient:

dstack/_internal/core/backends/runpod/__init__.py

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.runpod.compute import RunpodCompute
-from dstack._internal.core.backends.runpod.config import RunpodConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class RunpodBackend(Backend):
-    TYPE: BackendType = BackendType.RUNPOD
-
-    def __init__(self, config: RunpodConfig):
-        self.config = config
-        self._compute = RunpodCompute(self.config)
-
-    def compute(self) -> RunpodCompute:
-        return self._compute

dstack/_internal/core/backends/runpod/backend.py

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.runpod.compute import RunpodCompute
+from dstack._internal.core.backends.runpod.models import RunpodConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class RunpodBackend(Backend):
+    TYPE = BackendType.RUNPOD
+    COMPUTE_CLASS = RunpodCompute
+
+    def __init__(self, config: RunpodConfig):
+        self.config = config
+        self._compute = RunpodCompute(self.config)
+
+    def compute(self) -> RunpodCompute:
+        return self._compute

dstack/_internal/core/backends/runpod/compute.py

@@ -3,8 +3,9 @@ import uuid
 from datetime import timedelta
 from typing import List, Optional
 
-from dstack._internal.core.backends.base import Compute
+from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithVolumeSupport,
     generate_unique_instance_name,
     generate_unique_volume_name,
     get_docker_commands,
@@ -12,7 +13,7 @@ from dstack._internal.core.backends.base.compute import (
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.runpod.api_client import RunpodApiClient
-from dstack._internal.core.backends.runpod.config import RunpodConfig
+from dstack._internal.core.backends.runpod.models import RunpodConfig
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.errors import (
     BackendError,
@@ -39,7 +40,10 @@ MAX_RESOURCE_NAME_LEN = 60
 CONTAINER_REGISTRY_AUTH_CLEANUP_INTERVAL = 60 * 60 * 24  # 24 hour
 
 
-class RunpodCompute(Compute):
+class RunpodCompute(
+    ComputeWithVolumeSupport,
+    Compute,
+):
     _last_cleanup_time = None
 
     def __init__(self, config: RunpodConfig):
@@ -52,8 +56,9 @@ class RunpodCompute(Compute):
     ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.RUNPOD,
-            locations=self.config.regions,
+            locations=self.config.regions or None,
             requirements=requirements,
+            extra_filter=lambda o: _is_secure_cloud(o.region) or self.config.allow_community_cloud,
         )
         offers = [
             InstanceOfferWithAvailability(
@@ -102,13 +107,22 @@ class RunpodCompute(Compute):
         bid_per_gpu = None
         if instance_offer.instance.resources.spot and gpu_count:
             bid_per_gpu = instance_offer.price / gpu_count
+        if _is_secure_cloud(instance_offer.region):
+            cloud_type = "SECURE"
+            data_center_id = instance_offer.region
+            country_code = None
+        else:
+            cloud_type = "COMMUNITY"
+            data_center_id = None
+            country_code = instance_offer.region
 
         resp = self.api_client.create_pod(
             name=pod_name,
             image_name=job.job_spec.image_name,
             gpu_type_id=instance_offer.instance.name,
-            cloud_type="SECURE",  # ["ALL", "COMMUNITY", "SECURE"]:
-            data_center_id=instance_offer.region,
+            cloud_type=cloud_type,
+            data_center_id=data_center_id,
+            country_code=country_code,
             gpu_count=gpu_count,
             container_disk_in_gb=disk_size,
             min_vcpu_count=instance_offer.instance.resources.cpus,
@@ -257,3 +271,11 @@ def _get_volume_price(size: int) -> float:
     if size < 1000:
         return 0.07 * size
     return 0.05 * size
+
+
+def _is_secure_cloud(region: str) -> str:
+    """
+    Secure cloud regions are datacenter IDs: CA-MTL-1, EU-NL-1, etc.
+    Community cloud regions are country codes: CA, NL, etc.
+    """
+    return "-" in region

dstack/_internal/core/backends/runpod/configurator.py

@@ -0,0 +1,59 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.runpod import api_client
+from dstack._internal.core.backends.runpod.backend import RunpodBackend
+from dstack._internal.core.backends.runpod.models import (
+    AnyRunpodBackendConfig,
+    RunpodBackendConfig,
+    RunpodBackendConfigWithCreds,
+    RunpodConfig,
+    RunpodCreds,
+    RunpodStoredConfig,
+)
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class RunpodConfigurator(Configurator):
+    TYPE = BackendType.RUNPOD
+    BACKEND_CLASS = RunpodBackend
+
+    def validate_config(self, config: RunpodBackendConfigWithCreds, default_creds_enabled: bool):
+        self._validate_runpod_api_key(config.creds.api_key)
+
+    def create_backend(
+        self, project_name: str, config: RunpodBackendConfigWithCreds
+    ) -> BackendRecord:
+        return BackendRecord(
+            config=RunpodStoredConfig(
+                **RunpodBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=RunpodCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyRunpodBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return RunpodBackendConfigWithCreds.__response__.parse_obj(config)
+        return RunpodBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> RunpodBackend:
+        config = self._get_config(record)
+        return RunpodBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> RunpodConfig:
+        return RunpodConfig(
+            **json.loads(record.config),
+            creds=RunpodCreds.parse_raw(record.auth),
+        )
+
+    def _validate_runpod_api_key(self, api_key: str):
+        client = api_client.RunpodApiClient(api_key=api_key)
+        if not client.validate_api_key():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"]])

dstack/_internal/core/backends/runpod/models.py

@@ -0,0 +1,54 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+RUNPOD_COMMUNITY_CLOUD_DEFAULT = True
+
+
+class RunpodAPIKeyCreds(CoreModel):
+    type: Literal["api_key"] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyRunpodCreds = RunpodAPIKeyCreds
+RunpodCreds = AnyRunpodCreds
+
+
+class RunpodBackendConfig(CoreModel):
+    type: Literal["runpod"] = "runpod"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of RunPod regions. Omit to use all regions"),
+    ] = None
+    community_cloud: Annotated[
+        Optional[bool],
+        Field(
+            description=(
+                "Whether Community Cloud offers can be suggested in addition to Secure Cloud."
+                f" Defaults to `{str(RUNPOD_COMMUNITY_CLOUD_DEFAULT).lower()}`"
+            )
+        ),
+    ] = None
+
+
+class RunpodBackendConfigWithCreds(RunpodBackendConfig):
+    creds: Annotated[AnyRunpodCreds, Field(description="The credentials")]
+
+
+AnyRunpodBackendConfig = Union[RunpodBackendConfig, RunpodBackendConfigWithCreds]
+
+
+class RunpodStoredConfig(RunpodBackendConfig):
+    pass
+
+
+class RunpodConfig(RunpodStoredConfig):
+    creds: AnyRunpodCreds
+
+    @property
+    def allow_community_cloud(self) -> bool:
+        if self.community_cloud is not None:
+            return self.community_cloud
+        return RUNPOD_COMMUNITY_CLOUD_DEFAULT

dstack/_internal/core/backends/tensordock/__init__.py

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.tensordock.compute import TensorDockCompute
-from dstack._internal.core.backends.tensordock.config import TensorDockConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class TensorDockBackend(Backend):
-    TYPE: BackendType = BackendType.TENSORDOCK
-
-    def __init__(self, config: TensorDockConfig):
-        self.config = config
-        self._compute = TensorDockCompute(self.config)
-
-    def compute(self) -> TensorDockCompute:
-        return self._compute

dstack/_internal/core/backends/tensordock/backend.py

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.tensordock.compute import TensorDockCompute
+from dstack._internal.core.backends.tensordock.models import TensorDockConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class TensorDockBackend(Backend):
+    TYPE = BackendType.TENSORDOCK
+    COMPUTE_CLASS = TensorDockCompute
+
+    def __init__(self, config: TensorDockConfig):
+        self.config = config
+        self._compute = TensorDockCompute(self.config)
+
+    def compute(self) -> TensorDockCompute:
+        return self._compute

dstack/_internal/core/backends/tensordock/compute.py

@@ -3,25 +3,23 @@ from typing import List, Optional
 
 import requests
 
-from dstack._internal.core.backends.base import Compute
+from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
+    ComputeWithCreateInstanceSupport,
     generate_unique_instance_name,
-    get_job_instance_name,
     get_shim_commands,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.tensordock.api_client import TensorDockAPIClient
-from dstack._internal.core.backends.tensordock.config import TensorDockConfig
+from dstack._internal.core.backends.tensordock.models import TensorDockConfig
 from dstack._internal.core.errors import BackendError, NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
     InstanceConfiguration,
     InstanceOfferWithAvailability,
-    SSHKey,
 )
-from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
-from dstack._internal.core.models.volumes import Volume
+from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -31,7 +29,10 @@ logger = get_logger(__name__)
 MAX_INSTANCE_NAME_LEN = 60
 
 
-class TensorDockCompute(Compute):
+class TensorDockCompute(
+    ComputeWithCreateInstanceSupport,
+    Compute,
+):
     def __init__(self, config: TensorDockConfig):
         super().__init__()
         self.config = config
@@ -113,26 +114,6 @@ class TensorDockCompute(Compute):
             backend_data=None,
         )
 
-    def run_job(
-        self,
-        run: Run,
-        job: Job,
-        instance_offer: InstanceOfferWithAvailability,
-        project_ssh_public_key: str,
-        project_ssh_private_key: str,
-        volumes: List[Volume],
-    ) -> JobProvisioningData:
-        instance_config = InstanceConfiguration(
-            project_name=run.project_name,
-            instance_name=get_job_instance_name(run, job),  # TODO: generate name
-            ssh_keys=[
-                SSHKey(public=run.run_spec.ssh_key_pub.strip()),
-                SSHKey(public=project_ssh_public_key.strip()),
-            ],
-            user=run.user,
-        )
-        return self.create_instance(instance_offer, instance_config)
-
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
     ):

dstack/_internal/core/backends/tensordock/configurator.py

@@ -0,0 +1,68 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.tensordock import api_client
+from dstack._internal.core.backends.tensordock.backend import TensorDockBackend
+from dstack._internal.core.backends.tensordock.models import (
+    AnyTensorDockBackendConfig,
+    TensorDockBackendConfig,
+    TensorDockBackendConfigWithCreds,
+    TensorDockConfig,
+    TensorDockCreds,
+    TensorDockStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+# TensorDock regions are dynamic, currently we don't offer any filtering
+REGIONS = []
+
+
+class TensorDockConfigurator(Configurator):
+    TYPE = BackendType.TENSORDOCK
+    BACKEND_CLASS = TensorDockBackend
+
+    def validate_config(
+        self, config: TensorDockBackendConfigWithCreds, default_creds_enabled: bool
+    ):
+        self._validate_tensordock_creds(config.creds.api_key, config.creds.api_token)
+
+    def create_backend(
+        self, project_name: str, config: TensorDockBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = REGIONS
+        return BackendRecord(
+            config=TensorDockStoredConfig(
+                **TensorDockBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=TensorDockCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyTensorDockBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return TensorDockBackendConfigWithCreds.__response__.parse_obj(config)
+        return TensorDockBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> TensorDockBackend:
+        config = self._get_config(record)
+        return TensorDockBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> TensorDockConfig:
+        return TensorDockConfig.__response__(
+            **json.loads(record.config),
+            creds=TensorDockCreds.parse_raw(record.auth),
+        )
+
+    def _validate_tensordock_creds(self, api_key: str, api_token: str):
+        client = api_client.TensorDockAPIClient(api_key=api_key, api_token=api_token)
+        if not client.auth_test():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"], ["creds", "api_token"]])

dstack/_internal/core/backends/tensordock/models.py

@@ -0,0 +1,38 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class TensorDockAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+    api_token: Annotated[str, Field(description="The API token")]
+
+
+AnyTensorDockCreds = TensorDockAPIKeyCreds
+TensorDockCreds = AnyTensorDockCreds
+
+
+class TensorDockBackendConfig(CoreModel):
+    type: Annotated[Literal["tensordock"], Field(description="The type of backend")] = "tensordock"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of TensorDock regions. Omit to use all regions"),
+    ] = None
+
+
+class TensorDockBackendConfigWithCreds(TensorDockBackendConfig):
+    creds: Annotated[AnyTensorDockCreds, Field(description="The credentials")]
+
+
+AnyTensorDockBackendConfig = Union[TensorDockBackendConfig, TensorDockBackendConfigWithCreds]
+
+
+class TensorDockStoredConfig(TensorDockBackendConfig):
+    pass
+
+
+class TensorDockConfig(TensorDockStoredConfig):
+    creds: AnyTensorDockCreds

dstack/_internal/core/backends/vastai/__init__.py

@@ -1,15 +0,0 @@
-from dstack._internal.core.backends.base import Backend
-from dstack._internal.core.backends.vastai.compute import VastAICompute
-from dstack._internal.core.backends.vastai.config import VastAIConfig
-from dstack._internal.core.models.backends.base import BackendType
-
-
-class VastAIBackend(Backend):
-    TYPE: BackendType = BackendType.VASTAI
-
-    def __init__(self, config: VastAIConfig):
-        self.config = config
-        self._compute = VastAICompute(self.config)
-
-    def compute(self) -> VastAICompute:
-        return self._compute

dstack/_internal/core/backends/vastai/backend.py

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.vastai.compute import VastAICompute
+from dstack._internal.core.backends.vastai.models import VastAIConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class VastAIBackend(Backend):
+    TYPE = BackendType.VASTAI
+    COMPUTE_CLASS = VastAICompute
+
+    def __init__(self, config: VastAIConfig):
+        self.config = config
+        self._compute = VastAICompute(self.config)
+
+    def compute(self) -> VastAICompute:
+        return self._compute

dstack/_internal/core/backends/vastai/compute.py

@@ -3,14 +3,14 @@ from typing import List, Optional
 import gpuhunt
 from gpuhunt.providers.vastai import VastAIProvider
 
-from dstack._internal.core.backends.base import Compute
+from dstack._internal.core.backends.base.backend import Compute
 from dstack._internal.core.backends.base.compute import (
     generate_unique_instance_name_for_job,
     get_docker_commands,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.vastai.api_client import VastAIAPIClient
-from dstack._internal.core.backends.vastai.config import VastAIConfig
+from dstack._internal.core.backends.vastai.models import VastAIConfig
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.errors import ProvisioningError
 from dstack._internal.core.models.backends.base import BackendType

dstack/_internal/core/backends/vastai/configurator.py

@@ -0,0 +1,66 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.vastai import api_client
+from dstack._internal.core.backends.vastai.backend import VastAIBackend
+from dstack._internal.core.backends.vastai.models import (
+    AnyVastAIBackendConfig,
+    VastAIBackendConfig,
+    VastAIBackendConfigWithCreds,
+    VastAIConfig,
+    VastAICreds,
+    VastAIStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+# VastAI regions are dynamic, currently we don't offer any filtering
+REGIONS = []
+
+
+class VastAIConfigurator(Configurator):
+    TYPE = BackendType.VASTAI
+    BACKEND_CLASS = VastAIBackend
+
+    def validate_config(self, config: VastAIBackendConfigWithCreds, default_creds_enabled: bool):
+        self._validate_vastai_creds(config.creds.api_key)
+
+    def create_backend(
+        self, project_name: str, config: VastAIBackendConfigWithCreds
+    ) -> BackendRecord:
+        if config.regions is None:
+            config.regions = REGIONS
+        return BackendRecord(
+            config=VastAIStoredConfig(
+                **VastAIBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=VastAICreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyVastAIBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return VastAIBackendConfigWithCreds.__response__.parse_obj(config)
+        return VastAIBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> VastAIBackend:
+        config = self._get_config(record)
+        return VastAIBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> VastAIConfig:
+        return VastAIConfig.__response__(
+            **json.loads(record.config),
+            creds=VastAICreds.parse_raw(record.auth),
+        )
+
+    def _validate_vastai_creds(self, api_key: str):
+        client = api_client.VastAIAPIClient(api_key=api_key)
+        if not client.auth_test():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"]])

dstack/_internal/core/backends/vastai/models.py

@@ -0,0 +1,37 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class VastAIAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyVastAICreds = VastAIAPIKeyCreds
+VastAICreds = AnyVastAICreds
+
+
+class VastAIBackendConfig(CoreModel):
+    type: Annotated[Literal["vastai"], Field(description="The type of backend")] = "vastai"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of VastAI regions. Omit to use all regions"),
+    ] = None
+
+
+class VastAIBackendConfigWithCreds(VastAIBackendConfig):
+    creds: Annotated[AnyVastAICreds, Field(description="The credentials")]
+
+
+AnyVastAIBackendConfig = Union[VastAIBackendConfig, VastAIBackendConfigWithCreds]
+
+
+class VastAIStoredConfig(VastAIBackendConfig):
+    pass
+
+
+class VastAIConfig(VastAIStoredConfig):
+    creds: AnyVastAICreds