7d11dfa688
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
"""Async wrappers around the vast.ai CLI. Cannibalized from dtrain."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
|
|
async def _run(*args: str) -> str | None:
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"vastai", *args,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
return None
|
|
return stdout.decode()
|
|
|
|
|
|
async def search_offers(min_vram: int, limit: int = 10) -> list[dict]:
|
|
query = f"gpu_ram>={min_vram} inet_down>=100"
|
|
output = await _run(
|
|
"search", "offers", query, "-o", "dph_total",
|
|
"--limit", str(limit), "--raw",
|
|
)
|
|
if not output:
|
|
return []
|
|
return json.loads(output)
|
|
|
|
|
|
async def rent_instance(offer_id: int, image: str, disk: int = 20) -> str | None:
|
|
return await _run(
|
|
"create", "instance", str(offer_id),
|
|
"--image", image, "--disk", str(disk),
|
|
)
|
|
|
|
|
|
async def get_instances() -> list[dict]:
|
|
output = await _run("show", "instances", "--raw")
|
|
if not output:
|
|
return []
|
|
return json.loads(output)
|
|
|
|
|
|
async def destroy_instance(instance_id: int) -> None:
|
|
await _run("destroy", "instance", str(instance_id))
|