Skip to content

Update VastAIAPIClient #4038

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jvstme merged 1 commit into from
Jul 16, 2026
Merged

Update VastAIAPIClient #4038

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions src/dstack/_internal/core/backends/vastai/api_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
from typing import Optional, Union

import requests
import httpx

import dstack._internal.utils.docker as docker
from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
Expand Down Expand Up @@ -31,9 +32,11 @@ def __init__(

class VastAIAPIClient:
def __init__(self, api_key: str):
self.api_url = "https://console.vast.ai/api"
self.api_key = api_key
self.s = requests.Session() # TODO: set adequate timeout everywhere the session is used
self.s = httpx.Client(
base_url="https://console.vast.ai/api/",
headers={"Authorization": f"Bearer {api_key}"},
timeout=30,
)

def create_instance(
self,
Expand Down Expand Up @@ -86,10 +89,10 @@ def create_instance(
"create_from": None,
"force": False,
}
resp = self.s.put(self._url(f"/v0/asks/{bundle_id}/"), json=payload)
resp = self.s.put(f"/v0/asks/{bundle_id}/", json=payload)
try:
data = resp.json()
except requests.exceptions.JSONDecodeError:
except json.JSONDecodeError:
raise VastAICreateInstanceError(resp=resp.text)
if resp.status_code != 200 or not data["success"]:
raise VastAICreateInstanceError(
Expand All @@ -98,19 +101,19 @@ def create_instance(
return data["new_contract"]

def destroy_instance(self, instance_id: Union[str, int]) -> None:
resp = self.s.delete(self._url(f"/v0/instances/{instance_id}/"))
resp = self.s.delete(f"/v0/instances/{instance_id}/")
if resp.status_code == 429:
raise VastAIRateLimitError()
try:
data = resp.json()
except requests.exceptions.JSONDecodeError:
except json.JSONDecodeError:
raise ComputeError(resp.text)
if resp.status_code != 200 or not data["success"]:
if data.get("error") != "no_such_instance":
raise ComputeError(resp.text)

def get_instance(self, instance_id: Union[str, int]) -> Optional[dict]:
resp = self.s.get(self._url(f"/v0/instances/{instance_id}/"))
resp = self.s.get(f"/v0/instances/{instance_id}/")
if resp.status_code == 429:
raise VastAIRateLimitError()
resp.raise_for_status()
Expand All @@ -119,10 +122,7 @@ def get_instance(self, instance_id: Union[str, int]) -> Optional[dict]:

def auth_test(self) -> bool:
try:
self.s.get(self._url("/v1/instances/")).raise_for_status()
self.s.get("/v1/instances/").raise_for_status()
return True
except requests.HTTPError:
except httpx.HTTPError:
return False

def _url(self, path):
return f"{self.api_url}/{path.lstrip('/')}?api_key={self.api_key}"
Loading