84 lines
4.8 KiB
Python
84 lines
4.8 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import importlib
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
PATCH_MARKER = "AIRFRANS_SELECTED_OFFER_ID_PATCH"
|
||
|
|
SDK_API_KEY_MARKER = "AIRFRANS_VAST_SDK_API_KEY_PATCH"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class PatchStatus:
|
||
|
|
path: Path | None
|
||
|
|
installed: bool
|
||
|
|
patched: bool
|
||
|
|
message: str
|
||
|
|
|
||
|
|
|
||
|
|
def locate_vast_utils() -> Path:
|
||
|
|
try:
|
||
|
|
module = importlib.import_module("sky.provision.vast.utils")
|
||
|
|
except ModuleNotFoundError as exc:
|
||
|
|
raise RuntimeError("SkyPilot is not importable in this Python environment") from exc
|
||
|
|
path = getattr(module, "__file__", None)
|
||
|
|
if not path:
|
||
|
|
raise RuntimeError("Could not locate sky.provision.vast.utils source file")
|
||
|
|
return Path(path)
|
||
|
|
|
||
|
|
|
||
|
|
def patch_status() -> PatchStatus:
|
||
|
|
try:
|
||
|
|
path = locate_vast_utils()
|
||
|
|
except RuntimeError as exc:
|
||
|
|
return PatchStatus(path=None, installed=False, patched=False, message=str(exc))
|
||
|
|
text = path.read_text()
|
||
|
|
patched = PATCH_MARKER in text and SDK_API_KEY_MARKER in text
|
||
|
|
if patched:
|
||
|
|
return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch is installed")
|
||
|
|
return PatchStatus(path=path, installed=True, patched=False, message="SkyPilot Vast selected-offer patch is missing")
|
||
|
|
|
||
|
|
|
||
|
|
def apply_patch() -> PatchStatus:
|
||
|
|
path = locate_vast_utils()
|
||
|
|
text = path.read_text()
|
||
|
|
already_selected = PATCH_MARKER in text
|
||
|
|
already_api_key = SDK_API_KEY_MARKER in text
|
||
|
|
if already_selected and already_api_key:
|
||
|
|
return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch already installed")
|
||
|
|
|
||
|
|
if not already_selected:
|
||
|
|
old = ''' instance_list = vast.vast().search_offers(query=query_str)\n\n if isinstance(instance_list, int) or len(instance_list) == 0:\n raise RuntimeError('Failed to create instances, could not find an '\n 'offer that satisfies the requirements '\n f'"{query_str}".')\n\n instance_touse = instance_list[0]\n\n # Start with user-provided kwargs as the base\n launch_params: Dict[str, Any] = dict(create_instance_kwargs or {})\n # Remove None values to avoid overriding defaults\n launch_params = {k: v for k, v in launch_params.items() if v is not None}\n'''
|
||
|
|
new = f''' # {PATCH_MARKER}: allow callers to bypass SkyPilot's Vast offer\n # search after selecting a vetted Vast offer themselves. This keeps\n # SkyPilot's create/wait/bootstrap/run lifecycle while preserving\n # external marketplace-quality filtering.\n launch_params: Dict[str, Any] = dict(create_instance_kwargs or {{}})\n launch_params = {{k: v for k, v in launch_params.items() if v is not None}}\n selected_offer_id = launch_params.pop('selected_offer_id', None)\n\n if selected_offer_id is not None:\n logger.info(f'Using externally selected Vast offer {{selected_offer_id}}.')\n try:\n instance_touse = {{'id': int(selected_offer_id)}}\n except (TypeError, ValueError) as e:\n raise RuntimeError(\n f'Invalid selected_offer_id for Vast: {{selected_offer_id!r}}'\n ) from e\n else:\n instance_list = vast.vast().search_offers(query=query_str)\n\n if isinstance(instance_list, int) or len(instance_list) == 0:\n raise RuntimeError('Failed to create instances, could not find an '\n 'offer that satisfies the requirements '\n f'"{{query_str}}".')\n\n instance_touse = instance_list[0]\n\n'''
|
||
|
|
if old not in text:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Could not apply SkyPilot Vast selected-offer patch; expected source block not found in {path}"
|
||
|
|
)
|
||
|
|
text = text.replace(old, new, 1)
|
||
|
|
|
||
|
|
if not already_api_key:
|
||
|
|
old_key = " f'echo \"{vast.vast().client.api_key}\" > ~/.vast_api_key',"
|
||
|
|
new_key = (
|
||
|
|
f" # {SDK_API_KEY_MARKER}: vastai-sdk exposes api_key directly in current releases.\n"
|
||
|
|
" f'echo \"{vast.vast().api_key}\" > ~/.vast_api_key',"
|
||
|
|
)
|
||
|
|
if old_key not in text:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Could not apply SkyPilot Vast API-key patch; expected source line not found in {path}"
|
||
|
|
)
|
||
|
|
text = text.replace(old_key, new_key, 1)
|
||
|
|
|
||
|
|
path.write_text(text)
|
||
|
|
return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch installed")
|
||
|
|
|
||
|
|
|
||
|
|
def require_patch() -> Path:
|
||
|
|
status = patch_status()
|
||
|
|
if not status.installed:
|
||
|
|
raise RuntimeError(status.message)
|
||
|
|
if not status.patched:
|
||
|
|
raise RuntimeError(f"{status.message}; run `remote-run doctor --apply-skypilot-patch`")
|
||
|
|
assert status.path is not None
|
||
|
|
return status.path
|