897 lines
33 KiB
Python
897 lines
33 KiB
Python
"""Direct SSH command and filesystem storage adapters for Vast containers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import shlex
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import zipfile
|
|
from dataclasses import dataclass, field
|
|
from io import BytesIO
|
|
from pathlib import Path, PurePosixPath
|
|
from typing import Any, BinaryIO, Callable, Mapping, Sequence
|
|
|
|
if __package__:
|
|
from .huggingface_assets import HuggingFaceAssetSource
|
|
from .r2_cache import (
|
|
R2DownloadRequest,
|
|
R2UploadPlan,
|
|
R2UploadResult,
|
|
R2WorkerPreflightRequest,
|
|
)
|
|
else: # pragma: no cover - direct debugging imports.
|
|
from huggingface_assets import HuggingFaceAssetSource
|
|
from r2_cache import (
|
|
R2DownloadRequest,
|
|
R2UploadPlan,
|
|
R2UploadResult,
|
|
R2WorkerPreflightRequest,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_VAST_SSH_CONNECT_TIMEOUT_SECONDS = 15
|
|
DEFAULT_VAST_SSH_COMMAND_TIMEOUT_SECONDS = 120.0
|
|
DEFAULT_VAST_SSH_RETRY_ATTEMPTS = 4
|
|
DEFAULT_VAST_SSH_RETRY_BASE_DELAY_SECONDS = 0.5
|
|
DEFAULT_VAST_SSH_RETRY_MAX_DELAY_SECONDS = 4.0
|
|
DEFAULT_VAST_STORAGE_ROOT = PurePosixPath("/storage")
|
|
_TRANSIENT_SSH_TRANSPORT_DIAGNOSTICS = (
|
|
"connection closed by",
|
|
"connection reset by",
|
|
"connection refused",
|
|
"connection timed out",
|
|
"kex_exchange_identification:",
|
|
"network is unreachable",
|
|
"no route to host",
|
|
"operation timed out",
|
|
"remote host has disconnected",
|
|
"ssh_exchange_identification:",
|
|
)
|
|
_ATOMIC_STDIN_WRITER = """\
|
|
import os
|
|
import pathlib
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
|
|
expected_size = int(sys.argv[1])
|
|
target = pathlib.Path(sys.argv[2])
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, name = tempfile.mkstemp(
|
|
prefix=f".{target.name}.",
|
|
suffix=".tmp",
|
|
dir=target.parent,
|
|
)
|
|
try:
|
|
with os.fdopen(fd, "wb", buffering=0) as output:
|
|
shutil.copyfileobj(sys.stdin.buffer, output, length=1048576)
|
|
actual_size = output.tell()
|
|
if actual_size != expected_size:
|
|
raise ValueError(
|
|
f"Incomplete SSH upload: expected {expected_size} bytes, "
|
|
f"received {actual_size}."
|
|
)
|
|
output.flush()
|
|
os.fsync(output.fileno())
|
|
os.chmod(name, 0o600)
|
|
os.replace(name, target)
|
|
except (OSError, ValueError):
|
|
pathlib.Path(name).unlink(missing_ok=True)
|
|
raise
|
|
"""
|
|
|
|
|
|
class VastSshError(RuntimeError):
|
|
"""Raised when a direct Vast SSH or filesystem operation fails."""
|
|
|
|
|
|
class VastSshCancelledError(VastSshError):
|
|
"""Raised when the caller cancels an active direct SSH command."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastSshConnection:
|
|
"""Describe one Vast instance SSH endpoint without private credentials."""
|
|
|
|
host: str
|
|
port: int
|
|
known_hosts_path: Path
|
|
user: str = "root"
|
|
identity_file: Path | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Validate endpoint and local trust-store paths."""
|
|
if not self.host.strip() or any(
|
|
character in self.host for character in ("\x00", "\n", "\r")
|
|
):
|
|
raise ValueError("Vast SSH host must be a non-empty single-line value.")
|
|
if self.host.startswith("-"):
|
|
raise ValueError("Vast SSH host must not begin with an option prefix.")
|
|
if isinstance(self.port, bool) or not (1 <= self.port <= 65535):
|
|
raise ValueError("Vast SSH port must be between 1 and 65535.")
|
|
if not self.user.strip() or not all(
|
|
character.isalnum() or character in "_.-" for character in self.user
|
|
):
|
|
raise ValueError("Vast SSH user contains unsafe characters.")
|
|
if not self.known_hosts_path.is_absolute():
|
|
raise ValueError("Vast known_hosts_path must be absolute.")
|
|
if self.identity_file is not None and not self.identity_file.is_absolute():
|
|
raise ValueError("Vast SSH identity_file must be absolute when configured.")
|
|
|
|
@property
|
|
def target(self) -> str:
|
|
"""Return the OpenSSH user and host destination."""
|
|
return f"{self.user}@{self.host}"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VastSshCommandResult:
|
|
"""Capture a completed direct SSH command."""
|
|
|
|
stdout: bytes
|
|
stderr: bytes
|
|
returncode: int
|
|
|
|
@property
|
|
def stdout_text(self) -> str:
|
|
"""Decode standard output as replacement-safe UTF-8."""
|
|
return self.stdout.decode("utf-8", errors="replace")
|
|
|
|
@property
|
|
def stderr_text(self) -> str:
|
|
"""Decode standard error as replacement-safe UTF-8."""
|
|
return self.stderr.decode("utf-8", errors="replace")
|
|
|
|
|
|
def _is_transient_ssh_transport_failure(result: VastSshCommandResult) -> bool:
|
|
"""Return whether OpenSSH reports a retryable connection-level failure."""
|
|
if result.returncode != 255:
|
|
return False
|
|
diagnostic = f"{result.stderr_text}\n{result.stdout_text}".casefold()
|
|
return any(
|
|
marker in diagnostic for marker in _TRANSIENT_SSH_TRANSPORT_DIAGNOSTICS
|
|
)
|
|
|
|
|
|
def _resolved_input_file(
|
|
input_payload: bytes | None,
|
|
input_file: Path | None,
|
|
) -> Path | None:
|
|
"""Validate mutually exclusive SSH input sources and resolve a file source."""
|
|
if input_payload is not None and input_file is not None:
|
|
raise ValueError("Vast SSH input_payload and input_file are mutually exclusive.")
|
|
resolved = input_file.expanduser().resolve() if input_file is not None else None
|
|
if resolved is not None and not resolved.is_file():
|
|
raise FileNotFoundError(f"Vast SSH input file not found: {resolved}")
|
|
return resolved
|
|
|
|
|
|
@dataclass
|
|
class VastSshRunner:
|
|
"""Run argv-safe commands inside one Vast instance container."""
|
|
|
|
connection: VastSshConnection
|
|
connect_timeout_seconds: int = DEFAULT_VAST_SSH_CONNECT_TIMEOUT_SECONDS
|
|
command_timeout_seconds: float = DEFAULT_VAST_SSH_COMMAND_TIMEOUT_SECONDS
|
|
retry_attempts: int = DEFAULT_VAST_SSH_RETRY_ATTEMPTS
|
|
retry_base_delay_seconds: float = DEFAULT_VAST_SSH_RETRY_BASE_DELAY_SECONDS
|
|
retry_max_delay_seconds: float = DEFAULT_VAST_SSH_RETRY_MAX_DELAY_SECONDS
|
|
ssh_binary: str = "ssh"
|
|
sleep: Callable[[float], None] = field(default=time.sleep, repr=False)
|
|
random_unit: Callable[[], float] = field(default=random.random, repr=False)
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Create the dedicated trust store without weakening global SSH policy."""
|
|
if self.connect_timeout_seconds <= 0 or self.command_timeout_seconds <= 0:
|
|
raise ValueError("Vast SSH timeouts must be positive.")
|
|
if isinstance(self.retry_attempts, bool) or self.retry_attempts <= 0:
|
|
raise ValueError("Vast SSH retry attempts must be positive.")
|
|
if (
|
|
self.retry_base_delay_seconds <= 0
|
|
or self.retry_max_delay_seconds < self.retry_base_delay_seconds
|
|
):
|
|
raise ValueError("Vast SSH retry delays must be positive and ordered.")
|
|
try:
|
|
self.connection.known_hosts_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.connection.known_hosts_path.touch(mode=0o600, exist_ok=True)
|
|
os.chmod(self.connection.known_hosts_path, 0o600)
|
|
except OSError as exc:
|
|
raise VastSshError(
|
|
f"Unable to initialize Vast SSH trust store at "
|
|
f"{self.connection.known_hosts_path}."
|
|
) from exc
|
|
|
|
def command(self, remote_argv: Sequence[str]) -> list[str]:
|
|
"""Return the local OpenSSH argv for one remote argv sequence."""
|
|
if not remote_argv:
|
|
raise ValueError("remote_argv must contain at least one command value.")
|
|
normalized_argv = [str(value) for value in remote_argv]
|
|
if any("\x00" in value for value in normalized_argv):
|
|
raise ValueError("remote_argv must not contain null bytes.")
|
|
command = [
|
|
self.ssh_binary,
|
|
"-p",
|
|
str(self.connection.port),
|
|
"-o",
|
|
"BatchMode=yes",
|
|
"-o",
|
|
"StrictHostKeyChecking=accept-new",
|
|
"-o",
|
|
f"UserKnownHostsFile={self.connection.known_hosts_path}",
|
|
"-o",
|
|
"ClearAllForwardings=yes",
|
|
"-o",
|
|
f"ConnectTimeout={self.connect_timeout_seconds}",
|
|
]
|
|
if self.connection.identity_file is not None:
|
|
command.extend(("-i", str(self.connection.identity_file)))
|
|
command.extend(("--", self.connection.target, shlex.join(normalized_argv)))
|
|
return command
|
|
|
|
def run(
|
|
self,
|
|
remote_argv: Sequence[str],
|
|
*,
|
|
input_payload: bytes | None = None,
|
|
input_file: Path | None = None,
|
|
timeout_seconds: float | None = None,
|
|
check: bool = True,
|
|
cancellation_check: Callable[[], bool] | None = None,
|
|
transport_attempts: int | None = None,
|
|
) -> VastSshCommandResult:
|
|
"""Execute one command with bounded transient transport retries."""
|
|
resolved_input_file = _resolved_input_file(input_payload, input_file)
|
|
resolved_transport_attempts = (
|
|
self.retry_attempts
|
|
if transport_attempts is None
|
|
else transport_attempts
|
|
)
|
|
if (
|
|
isinstance(resolved_transport_attempts, bool)
|
|
or not isinstance(resolved_transport_attempts, int)
|
|
or resolved_transport_attempts <= 0
|
|
):
|
|
raise ValueError("Vast SSH transport attempts must be a positive integer.")
|
|
started_at = time.monotonic()
|
|
for attempt in range(1, resolved_transport_attempts + 1):
|
|
result = self._run_attempt(
|
|
remote_argv,
|
|
input_payload=input_payload,
|
|
input_file=resolved_input_file,
|
|
timeout_seconds=timeout_seconds,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
if _is_transient_ssh_transport_failure(result):
|
|
if attempt >= resolved_transport_attempts:
|
|
raise self._command_error(
|
|
remote_argv,
|
|
result,
|
|
suffix=(
|
|
f" after {attempt} transport attempts"
|
|
if resolved_transport_attempts > 1
|
|
else ""
|
|
),
|
|
)
|
|
delay_seconds = self._retry_delay_seconds(attempt)
|
|
logger.warning(
|
|
"Retrying transient Vast SSH transport failure host=%s "
|
|
"command=%s attempt=%d/%d delay_seconds=%.2f diagnostic=%s",
|
|
self.connection.host,
|
|
remote_argv[0],
|
|
attempt,
|
|
resolved_transport_attempts,
|
|
delay_seconds,
|
|
result.stderr_text.strip() or result.stdout_text.strip(),
|
|
)
|
|
self.sleep(delay_seconds)
|
|
if cancellation_check is not None and cancellation_check():
|
|
raise VastSshCancelledError("Vast SSH command was cancelled.")
|
|
continue
|
|
logger.debug(
|
|
"Finished Vast SSH command host=%s command=%s returncode=%d "
|
|
"attempts=%d elapsed_seconds=%.3f.",
|
|
self.connection.host,
|
|
remote_argv[0],
|
|
result.returncode,
|
|
attempt,
|
|
time.monotonic() - started_at,
|
|
)
|
|
if check and result.returncode != 0:
|
|
raise self._command_error(remote_argv, result)
|
|
return result
|
|
raise RuntimeError("Vast SSH retry loop ended without a result.")
|
|
|
|
def _run_attempt(
|
|
self,
|
|
remote_argv: Sequence[str],
|
|
*,
|
|
input_payload: bytes | None,
|
|
input_file: Path | None,
|
|
timeout_seconds: float | None,
|
|
cancellation_check: Callable[[], bool] | None,
|
|
) -> VastSshCommandResult:
|
|
"""Execute one OpenSSH subprocess attempt, reopening streamed input files."""
|
|
try:
|
|
if input_file is None:
|
|
completed = self._run_subprocess(
|
|
remote_argv,
|
|
input_payload=input_payload,
|
|
input_handle=None,
|
|
timeout_seconds=timeout_seconds,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
else:
|
|
with input_file.open("rb") as input_handle:
|
|
completed = self._run_subprocess(
|
|
remote_argv,
|
|
input_payload=None,
|
|
input_handle=input_handle,
|
|
timeout_seconds=timeout_seconds,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
except FileNotFoundError as exc:
|
|
raise VastSshError(
|
|
f"The SSH client {self.ssh_binary!r} is not installed."
|
|
) from exc
|
|
except subprocess.TimeoutExpired as exc:
|
|
raise VastSshError(
|
|
f"SSH command to Vast instance {self.connection.host!r} timed out "
|
|
f"after {exc.timeout} seconds."
|
|
) from exc
|
|
return VastSshCommandResult(
|
|
stdout=completed.stdout,
|
|
stderr=completed.stderr,
|
|
returncode=completed.returncode,
|
|
)
|
|
|
|
def _retry_delay_seconds(self, failed_attempt: int) -> float:
|
|
"""Return exponential retry delay with bounded per-attempt jitter."""
|
|
base_delay = min(
|
|
self.retry_max_delay_seconds,
|
|
self.retry_base_delay_seconds * (2 ** (failed_attempt - 1)),
|
|
)
|
|
jitter_unit = min(1.0, max(0.0, float(self.random_unit())))
|
|
return base_delay * (0.75 + (0.5 * jitter_unit))
|
|
|
|
@staticmethod
|
|
def _command_error(
|
|
remote_argv: Sequence[str],
|
|
result: VastSshCommandResult,
|
|
*,
|
|
suffix: str = "",
|
|
) -> VastSshError:
|
|
"""Build one bounded, actionable SSH command failure."""
|
|
detail = result.stderr_text.strip() or result.stdout_text.strip()
|
|
return VastSshError(
|
|
f"Vast SSH command {remote_argv[0]!r} failed with exit status "
|
|
f"{result.returncode}{suffix}: {detail or 'no diagnostic output'}"
|
|
)
|
|
|
|
def _run_subprocess(
|
|
self,
|
|
remote_argv: Sequence[str],
|
|
*,
|
|
input_payload: bytes | None,
|
|
input_handle: BinaryIO | None,
|
|
timeout_seconds: float | None,
|
|
cancellation_check: Callable[[], bool] | None,
|
|
) -> subprocess.CompletedProcess[bytes]:
|
|
"""Run OpenSSH with either a small byte payload or a streamed file handle."""
|
|
timeout = (
|
|
self.command_timeout_seconds
|
|
if timeout_seconds is None
|
|
else timeout_seconds
|
|
)
|
|
if cancellation_check is not None:
|
|
return self._run_cancellable_subprocess(
|
|
remote_argv,
|
|
input_payload=input_payload,
|
|
input_handle=input_handle,
|
|
timeout_seconds=timeout,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
if input_handle is not None:
|
|
return subprocess.run(
|
|
self.command(remote_argv),
|
|
stdin=input_handle,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
check=False,
|
|
timeout=timeout,
|
|
)
|
|
return subprocess.run(
|
|
self.command(remote_argv),
|
|
input=input_payload,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
check=False,
|
|
timeout=timeout,
|
|
)
|
|
|
|
def _run_cancellable_subprocess(
|
|
self,
|
|
remote_argv: Sequence[str],
|
|
*,
|
|
input_payload: bytes | None,
|
|
input_handle: BinaryIO | None,
|
|
timeout_seconds: float,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> subprocess.CompletedProcess[bytes]:
|
|
"""Run OpenSSH while promptly terminating it when cancellation is requested."""
|
|
if cancellation_check():
|
|
raise VastSshCancelledError("Vast SSH command was cancelled.")
|
|
process = subprocess.Popen(
|
|
self.command(remote_argv),
|
|
stdin=input_handle if input_handle is not None else subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
started_at = time.monotonic()
|
|
payload = input_payload
|
|
while True:
|
|
if cancellation_check():
|
|
process.terminate()
|
|
try:
|
|
process.wait(timeout=5.0)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.wait()
|
|
raise VastSshCancelledError("Vast SSH command was cancelled.")
|
|
remaining = timeout_seconds - (time.monotonic() - started_at)
|
|
if remaining <= 0:
|
|
process.kill()
|
|
stdout, stderr = process.communicate()
|
|
raise subprocess.TimeoutExpired(
|
|
self.command(remote_argv),
|
|
timeout_seconds,
|
|
output=stdout,
|
|
stderr=stderr,
|
|
)
|
|
try:
|
|
stdout, stderr = process.communicate(
|
|
input=payload,
|
|
timeout=min(0.25, remaining),
|
|
)
|
|
return subprocess.CompletedProcess(
|
|
args=self.command(remote_argv),
|
|
returncode=process.returncode,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
)
|
|
except subprocess.TimeoutExpired:
|
|
payload = None
|
|
|
|
def popen(self, remote_argv: Sequence[str]) -> subprocess.Popen[bytes]:
|
|
"""Start one streaming direct SSH command."""
|
|
try:
|
|
return subprocess.Popen(
|
|
self.command(remote_argv),
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
except FileNotFoundError as exc:
|
|
raise VastSshError(
|
|
f"The SSH client {self.ssh_binary!r} is not installed."
|
|
) from exc
|
|
|
|
|
|
@dataclass
|
|
class VastSshVolumeBackend:
|
|
"""Store content-addressed files directly on a Vast instance filesystem."""
|
|
|
|
remote_volume_epoch_scoped = True
|
|
|
|
runner: VastSshRunner
|
|
storage_root: PurePosixPath = DEFAULT_VAST_STORAGE_ROOT
|
|
_exists_cache: dict[str, bool] = field(default_factory=dict)
|
|
_materializer_path: PurePosixPath | None = field(init=False, default=None)
|
|
_r2_materializer_path: PurePosixPath | None = field(init=False, default=None)
|
|
_materializer_lock: threading.Lock = field(
|
|
init=False,
|
|
default_factory=threading.Lock,
|
|
repr=False,
|
|
)
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Require an absolute non-root storage location."""
|
|
if not self.storage_root.is_absolute() or self.storage_root == PurePosixPath("/"):
|
|
raise ValueError("Vast storage_root must be an absolute non-root path.")
|
|
self.runner.run(("mkdir", "-p", str(self.storage_root)))
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Return whether one regular file exists under the storage root."""
|
|
normalized = _validated_storage_path(remote_path)
|
|
cached = self._exists_cache.get(normalized)
|
|
if cached is not None:
|
|
return cached
|
|
result = self.runner.run(
|
|
("test", "-f", str(self.storage_root / normalized)),
|
|
check=False,
|
|
)
|
|
exists = result.returncode == 0
|
|
self._exists_cache[normalized] = exists
|
|
return exists
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Stream one local file through SSH and atomically publish it remotely."""
|
|
resolved_path = local_path.expanduser().resolve()
|
|
if not resolved_path.is_file():
|
|
raise FileNotFoundError(f"Asset not found: {resolved_path}")
|
|
size_bytes = resolved_path.stat().st_size
|
|
normalized = _validated_storage_path(remote_path)
|
|
target = str(self.storage_root / normalized)
|
|
self.runner.run(
|
|
(
|
|
"python",
|
|
"-c",
|
|
_ATOMIC_STDIN_WRITER,
|
|
str(size_bytes),
|
|
target,
|
|
),
|
|
input_file=resolved_path,
|
|
timeout_seconds=max(60.0, size_bytes / (4 * 1024 * 1024)),
|
|
)
|
|
self._exists_cache[normalized] = True
|
|
|
|
def put_file_cancellable(
|
|
self,
|
|
local_path: Path,
|
|
remote_path: str,
|
|
*,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> None:
|
|
"""Stream one file while terminating OpenSSH promptly on cancellation."""
|
|
resolved_path = local_path.expanduser().resolve()
|
|
if not resolved_path.is_file():
|
|
raise FileNotFoundError(f"Asset not found: {resolved_path}")
|
|
size_bytes = resolved_path.stat().st_size
|
|
normalized = _validated_storage_path(remote_path)
|
|
try:
|
|
self.runner.run(
|
|
(
|
|
"python",
|
|
"-c",
|
|
_ATOMIC_STDIN_WRITER,
|
|
str(size_bytes),
|
|
str(self.storage_root / normalized),
|
|
),
|
|
input_file=resolved_path,
|
|
timeout_seconds=max(60.0, size_bytes / (4 * 1024 * 1024)),
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
except VastSshCancelledError as exc:
|
|
raise InterruptedError("Vast file upload was cancelled.") from exc
|
|
self._exists_cache[normalized] = True
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Atomically upload bytes under the storage root."""
|
|
self._put_payload(payload, remote_path)
|
|
|
|
def _put_payload(self, payload: bytes, remote_path: str) -> None:
|
|
"""Stream bytes to an argv-safe remote Python atomic writer."""
|
|
normalized = _validated_storage_path(remote_path)
|
|
target = str(self.storage_root / normalized)
|
|
self.runner.run(
|
|
("python", "-c", _ATOMIC_STDIN_WRITER, str(len(payload)), target),
|
|
input_payload=payload,
|
|
timeout_seconds=max(60.0, len(payload) / (4 * 1024 * 1024)),
|
|
)
|
|
self._exists_cache[normalized] = True
|
|
|
|
def materialize_huggingface_file(
|
|
self,
|
|
source: HuggingFaceAssetSource,
|
|
remote_path: str,
|
|
*,
|
|
token: str | None,
|
|
) -> bool:
|
|
"""Ask the Vast worker to fetch and verify one immutable Hugging Face file."""
|
|
return self.materialize_huggingface_file_cancellable(
|
|
source,
|
|
remote_path,
|
|
token=token,
|
|
cancellation_check=lambda: False,
|
|
)
|
|
|
|
def materialize_huggingface_file_cancellable(
|
|
self,
|
|
source: HuggingFaceAssetSource,
|
|
remote_path: str,
|
|
*,
|
|
token: str | None,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> bool:
|
|
"""Run the current materializer bundle and allow prompt cancellation."""
|
|
normalized = _validated_storage_path(remote_path)
|
|
request_payload = json.dumps(
|
|
{
|
|
"remote_path": normalized,
|
|
"source": source.to_dict(),
|
|
"storage_root": str(self.storage_root),
|
|
"token": token,
|
|
},
|
|
sort_keys=True,
|
|
).encode("utf-8")
|
|
timeout_seconds = max(900.0, source.size_bytes / (2 * 1024 * 1024))
|
|
try:
|
|
materializer_path = self._ensure_materializer_bundle(
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
self.runner.run(
|
|
("python", str(materializer_path)),
|
|
input_payload=request_payload,
|
|
timeout_seconds=timeout_seconds,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
except VastSshCancelledError as exc:
|
|
raise InterruptedError("Vast Hugging Face download was cancelled.") from exc
|
|
except VastSshError as exc:
|
|
logger.warning(
|
|
"Vast Hugging Face materialization failed for %s; falling back to SSH upload: %s",
|
|
source.display_reference,
|
|
exc,
|
|
)
|
|
return False
|
|
self._exists_cache[normalized] = True
|
|
return True
|
|
|
|
def materialize_r2_file(
|
|
self,
|
|
request: R2DownloadRequest,
|
|
remote_path: str,
|
|
*,
|
|
cancellation_check: Callable[[], bool] | None = None,
|
|
) -> None:
|
|
"""Download and verify one signed R2 object directly on the Vast host."""
|
|
normalized = _validated_storage_path(remote_path)
|
|
active_cancellation_check = cancellation_check or (lambda: False)
|
|
try:
|
|
self._run_r2_materializer(
|
|
{
|
|
"operation": "download",
|
|
"storage_root": str(self.storage_root),
|
|
"remote_path": normalized,
|
|
"download": request.to_dict(),
|
|
},
|
|
size_bytes=request.size_bytes,
|
|
cancellation_check=active_cancellation_check,
|
|
)
|
|
except VastSshCancelledError as exc:
|
|
raise InterruptedError("Vast R2 download was cancelled.") from exc
|
|
self._exists_cache[normalized] = True
|
|
|
|
def preflight_r2_access(
|
|
self,
|
|
request: R2WorkerPreflightRequest,
|
|
*,
|
|
cancellation_check: Callable[[], bool] | None = None,
|
|
) -> None:
|
|
"""Test one read-only signed R2 request directly on the Vast host."""
|
|
active_cancellation_check = cancellation_check or (lambda: False)
|
|
try:
|
|
self._run_r2_materializer(
|
|
{
|
|
"operation": "preflight",
|
|
"preflight": request.to_dict(),
|
|
},
|
|
size_bytes=0,
|
|
cancellation_check=active_cancellation_check,
|
|
)
|
|
except VastSshCancelledError as exc:
|
|
raise InterruptedError("Vast R2 preflight was cancelled.") from exc
|
|
|
|
def upload_r2_file(
|
|
self,
|
|
plan: R2UploadPlan,
|
|
remote_path: str,
|
|
) -> R2UploadResult:
|
|
"""Upload one Vast-host file through a controller-issued R2 plan."""
|
|
return self.upload_r2_file_cancellable(
|
|
plan,
|
|
remote_path,
|
|
cancellation_check=lambda: False,
|
|
)
|
|
|
|
def upload_r2_file_cancellable(
|
|
self,
|
|
plan: R2UploadPlan,
|
|
remote_path: str,
|
|
*,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> R2UploadResult:
|
|
"""Upload one Vast-host file while yielding to foreground workflows."""
|
|
normalized = _validated_storage_path(remote_path)
|
|
try:
|
|
result = self._run_r2_materializer(
|
|
{
|
|
"operation": "upload",
|
|
"storage_root": str(self.storage_root),
|
|
"remote_path": normalized,
|
|
"upload": plan.to_dict(),
|
|
},
|
|
size_bytes=plan.size_bytes,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
except VastSshCancelledError as exc:
|
|
raise InterruptedError("Vast R2 upload was cancelled.") from exc
|
|
return R2UploadResult.from_dict(result)
|
|
|
|
def _run_r2_materializer(
|
|
self,
|
|
request: Mapping[str, Any],
|
|
*,
|
|
size_bytes: int,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> Mapping[str, object]:
|
|
"""Run the current R2 materializer with its signed URLs only on stdin."""
|
|
materializer_path = self._ensure_r2_materializer_bundle(
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
result = self.runner.run(
|
|
("python", str(materializer_path)),
|
|
input_payload=json.dumps(request, sort_keys=True).encode("utf-8"),
|
|
timeout_seconds=max(900.0, size_bytes / (2 * 1024 * 1024)),
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
try:
|
|
payload = json.loads(result.stdout_text)
|
|
except json.JSONDecodeError as exc:
|
|
raise VastSshError("Vast R2 materializer returned invalid JSON.") from exc
|
|
if not isinstance(payload, Mapping):
|
|
raise VastSshError("Vast R2 materializer returned a non-object result.")
|
|
return payload
|
|
|
|
def _ensure_materializer_bundle(
|
|
self,
|
|
*,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> PurePosixPath:
|
|
"""Upload a tiny current-source zipapp independent of the worker image version."""
|
|
with self._materializer_lock:
|
|
if self._materializer_path is None:
|
|
self._materializer_path = self._ensure_runtime_bundle(
|
|
"huggingface-materializer",
|
|
_huggingface_materializer_bundle(),
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
return self._materializer_path
|
|
|
|
def _ensure_r2_materializer_bundle(
|
|
self,
|
|
*,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> PurePosixPath:
|
|
"""Upload the current R2 transfer implementation once per backend."""
|
|
with self._materializer_lock:
|
|
if self._r2_materializer_path is None:
|
|
self._r2_materializer_path = self._ensure_runtime_bundle(
|
|
"r2-materializer",
|
|
_r2_materializer_bundle(),
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
return self._r2_materializer_path
|
|
|
|
def _ensure_runtime_bundle(
|
|
self,
|
|
bundle_name: str,
|
|
payload: bytes,
|
|
*,
|
|
cancellation_check: Callable[[], bool],
|
|
) -> PurePosixPath:
|
|
"""Publish one immutable current-source runtime zipapp under storage."""
|
|
digest = hashlib.sha256(payload).hexdigest()
|
|
normalized = _validated_storage_path(
|
|
f"runtime-tools/{bundle_name}-{digest}.pyz"
|
|
)
|
|
target = self.storage_root / normalized
|
|
if not self.exists(normalized):
|
|
self.runner.run(
|
|
("python", "-c", _ATOMIC_STDIN_WRITER, str(len(payload)), str(target)),
|
|
input_payload=payload,
|
|
timeout_seconds=60.0,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
self._exists_cache[normalized] = True
|
|
return target
|
|
|
|
|
|
def _huggingface_materializer_bundle() -> bytes:
|
|
"""Build an executable zip containing the current verified-download implementation."""
|
|
repo_root = Path(__file__).resolve().parent
|
|
members = {
|
|
"__main__.py": (
|
|
b"from remote.huggingface_materializer import main\n"
|
|
b"raise SystemExit(main())\n"
|
|
),
|
|
"remote/__init__.py": b"",
|
|
"remote/huggingface_materializer.py": (
|
|
repo_root / "remote" / "huggingface_materializer.py"
|
|
).read_bytes(),
|
|
"huggingface_assets.py": (repo_root / "huggingface_assets.py").read_bytes(),
|
|
}
|
|
output = BytesIO()
|
|
with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive:
|
|
for member_name, member_payload in members.items():
|
|
member = zipfile.ZipInfo(member_name, date_time=(1980, 1, 1, 0, 0, 0))
|
|
member.compress_type = zipfile.ZIP_DEFLATED
|
|
member.external_attr = 0o644 << 16
|
|
archive.writestr(member, member_payload)
|
|
return output.getvalue()
|
|
|
|
|
|
def _r2_materializer_bundle() -> bytes:
|
|
"""Build an executable zip containing the current signed R2 materializer."""
|
|
repo_root = Path(__file__).resolve().parent
|
|
members = {
|
|
"__main__.py": (
|
|
b"from remote.r2_materializer import main\n"
|
|
b"raise SystemExit(main())\n"
|
|
),
|
|
"remote/__init__.py": b"",
|
|
"remote/r2_materializer.py": (
|
|
repo_root / "remote" / "r2_materializer.py"
|
|
).read_bytes(),
|
|
}
|
|
output = BytesIO()
|
|
with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive:
|
|
for member_name, member_payload in members.items():
|
|
member = zipfile.ZipInfo(member_name, date_time=(1980, 1, 1, 0, 0, 0))
|
|
member.compress_type = zipfile.ZIP_DEFLATED
|
|
member.external_attr = 0o644 << 16
|
|
archive.writestr(member, member_payload)
|
|
return output.getvalue()
|
|
|
|
|
|
def vast_connection_from_lease(
|
|
*,
|
|
ssh_host: str | None,
|
|
ssh_port: int | None,
|
|
user_directory: Path,
|
|
identity_file: Path | None = None,
|
|
) -> VastSshConnection:
|
|
"""Build one direct connection from a ready lease's safe public fields."""
|
|
if not ssh_host or ssh_port is None:
|
|
raise VastSshError("Vast lease does not expose an SSH endpoint.")
|
|
return VastSshConnection(
|
|
host=ssh_host,
|
|
port=ssh_port,
|
|
known_hosts_path=(
|
|
user_directory.expanduser().resolve()
|
|
/ "comfyui-modal"
|
|
/ "vast-known-hosts"
|
|
),
|
|
identity_file=(identity_file.expanduser().resolve() if identity_file else None),
|
|
)
|
|
|
|
|
|
def _validated_storage_path(remote_path: str) -> str:
|
|
"""Return one traversal-safe path relative to Vast storage."""
|
|
path = PurePosixPath(str(remote_path).lstrip("/"))
|
|
if (
|
|
not path.parts
|
|
or any(part in {"", ".", ".."} for part in path.parts)
|
|
or "\x00" in str(path)
|
|
):
|
|
raise ValueError(f"Unsafe Vast storage path {remote_path!r}.")
|
|
return path.as_posix()
|
|
|
|
|
|
__all__ = [
|
|
"DEFAULT_VAST_STORAGE_ROOT",
|
|
"VastSshCommandResult",
|
|
"VastSshCancelledError",
|
|
"VastSshConnection",
|
|
"VastSshError",
|
|
"VastSshRunner",
|
|
"VastSshVolumeBackend",
|
|
"vast_connection_from_lease",
|
|
]
|