"""Stateful local Vast.ai API simulator for tests and offline development.""" from __future__ import annotations import argparse import copy import logging from dataclasses import dataclass, field from typing import Any, Mapping, Sequence from aiohttp import web logger = logging.getLogger(__name__) def default_vast_simulator_offers() -> list[dict[str, Any]]: """Return deterministic offers spanning price and capacity tradeoffs.""" common = { "rentable": True, "rented": False, "gpu_arch": "nvidia", "cpu_arch": "amd64", "gpu_frac": 1.0, "gpu_display_active": False, "direct_port_count": 4, "duration": 30 * 86400, "cuda_max_good": 13.0, "verification": "verified", "verified": True, "storage_cost": 0.12, } return [ { **common, "id": 1001, "gpu_name": "RTX 4090", "num_gpus": 1, "gpu_ram": 24 * 1024, "gpu_total_ram": 24 * 1024, "total_flops": 82.6, "cpu_ram": 64 * 1024, "cpu_cores": 16, "cpu_cores_effective": 12.0, "disk_space": 1000.0, "reliability": 0.995, "dlperf": 45.0, "inet_down": 800.0, "inet_up": 500.0, "dph_total": 0.42, "geolocation": "Vancouver, CA", }, { **common, "id": 1002, "gpu_name": "RTX 6000 Ada", "num_gpus": 1, "gpu_ram": 48 * 1024, "gpu_total_ram": 48 * 1024, "total_flops": 91.1, "cpu_ram": 128 * 1024, "cpu_cores": 32, "cpu_cores_effective": 24.0, "disk_space": 2000.0, "reliability": 0.999, "dlperf": 62.0, "inet_down": 1200.0, "inet_up": 800.0, "dph_total": 0.74, "geolocation": "Seattle, US", }, { **common, "id": 1003, "gpu_name": "H100 SXM", "num_gpus": 1, "gpu_ram": 80 * 1024, "gpu_total_ram": 80 * 1024, "total_flops": 989.0, "cpu_ram": 256 * 1024, "cpu_cores": 64, "cpu_cores_effective": 48.0, "disk_space": 4000.0, "reliability": 0.9998, "dlperf": 180.0, "inet_down": 2200.0, "inet_up": 1200.0, "dph_total": 1.96, "geolocation": "Dallas, US", }, ] @dataclass class VastSimulatorState: """Hold mutable marketplace and instance state for the simulator.""" api_key: str = "vast-test-key" offers: list[dict[str, Any]] = field(default_factory=default_vast_simulator_offers) instances: dict[int, dict[str, Any]] = field(default_factory=dict) destroyed_instance_ids: list[int] = field(default_factory=list) request_log: list[dict[str, Any]] = field(default_factory=list) create_failures_remaining: dict[int, int] = field(default_factory=dict) create_failure_status: int = 404 create_failure_message: str = "Offer is no longer available" polls_until_running: int = 2 next_instance_id: int = 9001 def record(self, request: web.Request, body: object | None = None) -> None: """Record one request without retaining its authorization header.""" self.request_log.append( { "method": request.method, "path": request.path, "query": dict(request.query), "body": copy.deepcopy(body), } ) class VastApiSimulator: """Serve the Vast endpoints used by the extension on a local aiohttp app.""" def __init__(self, state: VastSimulatorState | None = None) -> None: """Create one simulator around supplied or default mutable state.""" self.state = state or VastSimulatorState() self.app = web.Application(middlewares=[self._authenticate]) self._register_routes() @web.middleware async def _authenticate( self, request: web.Request, handler: Any, ) -> web.StreamResponse: """Require the configured bearer token on API routes.""" expected = f"Bearer {self.state.api_key}" if request.headers.get("Authorization") != expected: self.state.record(request) return web.json_response({"error": "Unauthorized"}, status=401) return await handler(request) def _register_routes(self) -> None: """Register the supported Vast API surface.""" self.app.router.add_get("/api/v0/users/current/", self.current_user) self.app.router.add_post("/api/v0/bundles/", self.search_offers) self.app.router.add_put("/api/v0/asks/{offer_id}/", self.create_instance) self.app.router.add_get( "/api/v0/instances/{instance_id}/", self.show_instance ) self.app.router.add_put( "/api/v0/instances/{instance_id}/", self.manage_instance ) self.app.router.add_delete( "/api/v0/instances/{instance_id}/", self.destroy_instance ) self.app.router.add_get("/api/v1/instances/", self.list_instances) async def current_user(self, request: web.Request) -> web.Response: """Return a deterministic account record.""" self.state.record(request) return web.json_response( { "id": 42, "email": "simulator@example.invalid", "credit": 100.0, "ssh_key": "ssh-ed25519 simulator-key", } ) async def search_offers(self, request: web.Request) -> web.Response: """Apply documented filter operators and deterministic ordering.""" body = await _json_object(request) self.state.record(request, body) offers = [ copy.deepcopy(offer) for offer in self.state.offers if _offer_matches(offer, body) ] order = body.get("order") if isinstance(order, list): for order_item in reversed(order): if not isinstance(order_item, list) or len(order_item) != 2: continue field_name, direction = str(order_item[0]), str(order_item[1]) offers.sort( key=lambda offer: _sortable_value(offer.get(field_name)), reverse=direction.casefold() == "desc", ) limit = max(1, min(1000, int(body.get("limit", 100)))) return web.json_response({"offers": offers[:limit]}) async def create_instance(self, request: web.Request) -> web.Response: """Rent one available offer or simulate an availability race.""" body = await _json_object(request) self.state.record(request, body) offer_id = int(request.match_info["offer_id"]) offer = next( (candidate for candidate in self.state.offers if candidate["id"] == offer_id), None, ) if offer is None or not offer.get("rentable") or offer.get("rented"): return web.json_response({"error": "Offer is no longer available"}, status=404) failures_remaining = self.state.create_failures_remaining.get(offer_id, 0) if failures_remaining > 0: self.state.create_failures_remaining[offer_id] = failures_remaining - 1 return web.json_response( {"error": self.state.create_failure_message}, status=self.state.create_failure_status, ) if not body.get("cancel_unavail", False): return web.json_response( {"error": "Simulator requires cancel_unavail for deterministic tests."}, status=400, ) instance_id = self.state.next_instance_id self.state.next_instance_id += 1 offer["rented"] = True instance = { **copy.deepcopy(offer), "id": instance_id, "ask_contract_id": offer_id, "actual_status": "loading", "intended_status": "running", "cur_state": "loading", "status_msg": "worker-layer: Download complete", "next_state": "running", "label": body.get("label"), "image_uuid": body.get("image"), "image_runtype": body.get("runtype", "ssh_direct"), "onstart": body.get("onstart"), "extra_env": body.get("env", {}), "disk_space": body.get("disk", offer.get("disk_space", 0)), "ssh_host": f"127.0.0.{(instance_id % 200) + 1}", "ssh_port": 22000 + instance_id % 1000, "_poll_count": 0, } self.state.instances[instance_id] = instance return web.json_response( { "success": True, "new_contract": instance_id, "instance_api_key": f"instance-secret-{instance_id}", } ) async def show_instance(self, request: web.Request) -> web.Response: """Advance loading instances after the configured number of polls.""" self.state.record(request) instance = self._instance(request) if instance is None: return web.json_response({"instances": None}) instance["_poll_count"] = int(instance.get("_poll_count", 0)) + 1 if ( instance["actual_status"] == "loading" and instance["_poll_count"] >= self.state.polls_until_running ): instance["actual_status"] = "running" instance["cur_state"] = "running" instance["status_msg"] = "Worker ready" return web.json_response({"instances": _public_instance(instance)}) async def list_instances(self, request: web.Request) -> web.Response: """Return all non-destroyed simulated instances in one page.""" self.state.record(request) instances = [ _public_instance(instance) for _instance_id, instance in sorted(self.state.instances.items()) ] return web.json_response( { "success": True, "instances_found": len(instances), "total_instances": len(instances), "instances": instances, "next_token": None, } ) async def manage_instance(self, request: web.Request) -> web.Response: """Set instance state or label using Vast's shared PUT route.""" body = await _json_object(request) self.state.record(request, body) instance = self._instance(request) if instance is None: return web.json_response({"error": "Instance not found"}, status=404) state = body.get("state") if state is not None: if state not in {"running", "stopped"}: return web.json_response({"error": "Invalid state"}, status=400) instance["actual_status"] = state instance["intended_status"] = state instance["cur_state"] = state instance["next_state"] = state if "label" in body: instance["label"] = str(body["label"]) return web.json_response({"success": True}) async def destroy_instance(self, request: web.Request) -> web.Response: """Permanently remove one instance and release its offer.""" self.state.record(request) instance_id = int(request.match_info["instance_id"]) instance = self.state.instances.pop(instance_id, None) if instance is None: return web.json_response({"error": "Instance not found"}, status=404) offer_id = instance.get("ask_contract_id") for offer in self.state.offers: if offer.get("id") == offer_id: offer["rented"] = False self.state.destroyed_instance_ids.append(instance_id) return web.json_response( {"success": True, "msg": "Instance destroyed successfully"} ) def _instance(self, request: web.Request) -> dict[str, Any] | None: """Return the instance named by one route.""" return self.state.instances.get(int(request.match_info["instance_id"])) def create_vast_simulator_app( state: VastSimulatorState | None = None, ) -> web.Application: """Return an aiohttp application implementing the local Vast simulator.""" return VastApiSimulator(state).app async def _json_object(request: web.Request) -> dict[str, Any]: """Read one JSON request body and require an object.""" try: payload = await request.json() except ValueError as exc: raise web.HTTPBadRequest(text="Request body must be valid JSON.") from exc if not isinstance(payload, dict): raise web.HTTPBadRequest(text="Request body must be a JSON object.") return payload def _offer_matches(offer: Mapping[str, Any], query: Mapping[str, Any]) -> bool: """Apply the filter subset used by the extension's marketplace queries.""" ignored_fields = {"limit", "order", "type", "allocated_storage"} for field_name, condition in query.items(): if field_name in ignored_fields: continue if not isinstance(condition, Mapping): continue value = offer.get(field_name) if not _matches_condition(value, condition): return False query_type = str(query.get("type") or "on-demand") offer_type = str(offer.get("type") or "on-demand") return query_type == offer_type def _matches_condition(value: object, condition: Mapping[str, Any]) -> bool: """Apply documented eq/neq/range/membership operators.""" for operator, expected in condition.items(): if operator == "eq" and value != expected: return False if operator == "neq" and value == expected: return False if operator == "gte" and not (_numeric(value) >= _numeric(expected)): return False if operator == "lte" and not (_numeric(value) <= _numeric(expected)): return False if operator == "gt" and not (_numeric(value) > _numeric(expected)): return False if operator == "lt" and not (_numeric(value) < _numeric(expected)): return False if operator in {"in", "notin"}: expected_values = expected if isinstance(expected, list) else [] normalized_value = str(value or "").casefold() contained = any( normalized_value == str(candidate).casefold() or normalized_value.endswith(f", {str(candidate).casefold()}") for candidate in expected_values ) if operator == "in" and not contained: return False if operator == "notin" and contained: return False return True def _numeric(value: object) -> float: """Return a number used by simulator filter comparisons.""" try: return float(value or 0.0) except (TypeError, ValueError): return float("-inf") def _sortable_value(value: object) -> tuple[int, object]: """Return a stable key across missing, numeric, and textual fields.""" if value is None: return (2, "") if isinstance(value, (int, float)) and not isinstance(value, bool): return (0, float(value)) return (1, str(value)) def _public_instance(instance: Mapping[str, Any]) -> dict[str, Any]: """Return one instance record without simulator-only fields.""" return { key: copy.deepcopy(value) for key, value in instance.items() if not str(key).startswith("_") } def _parser() -> argparse.ArgumentParser: """Return the standalone simulator command-line parser.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=8099) parser.add_argument("--api-key", default="vast-test-key") return parser def main(argv: Sequence[str] | None = None) -> None: """Run the Vast simulator until interrupted.""" arguments = _parser().parse_args(argv) logging.basicConfig(level=logging.INFO) state = VastSimulatorState(api_key=arguments.api_key) logger.info( "Starting Vast API simulator on http://%s:%d with %d offers.", arguments.host, arguments.port, len(state.offers), ) web.run_app( create_vast_simulator_app(state), host=arguments.host, port=arguments.port, ) if __name__ == "__main__": # pragma: no cover - exercised manually. main() __all__ = [ "VastApiSimulator", "VastSimulatorState", "create_vast_simulator_app", "default_vast_simulator_offers", "main", ]