Leonid Pershin
91d2ce0fab
- Added new environment variables in `env.example` and `gpu-rent.vars.example` for fine-tuning LLM settings, including CUDA build options and model version pinning. - Updated `llm.md` documentation to include detailed descriptions of new configuration options and usage cases for LLM setups. - Enhanced the `provision.py` script to forward new environment variables during remote installations, improving the installation process for LLM components. - Modified the `install_llamacpp.sh` script to support conditional CUDA builds and asset URL overrides, ensuring better compatibility with various environments. - Improved logging in the installation scripts to provide clearer feedback during the setup process.
510 lines
16 KiB
Python
510 lines
16 KiB
Python
"""Wait for SSH and run remote commands."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
import shlex
|
||
import socket
|
||
import subprocess
|
||
import time
|
||
|
||
import paramiko
|
||
|
||
from gpu_rent.config import Config
|
||
from gpu_rent.errors import CloudError
|
||
|
||
|
||
def split_ssh_stream(buf: str) -> tuple[list[tuple[str, str]], str]:
|
||
"""Split SSH/PTY stdout into (kind, text) events.
|
||
|
||
kind is ``line`` (ended with \\n or \\r\\n) or ``progress`` (ended with lone \\r).
|
||
Returns (events, remainder).
|
||
"""
|
||
events: list[tuple[str, str]] = []
|
||
i = 0
|
||
start = 0
|
||
n = len(buf)
|
||
while i < n:
|
||
ch = buf[i]
|
||
if ch == "\r":
|
||
if i + 1 < n and buf[i + 1] == "\n":
|
||
events.append(("line", buf[start:i]))
|
||
i += 2
|
||
start = i
|
||
continue
|
||
events.append(("progress", buf[start:i]))
|
||
i += 1
|
||
start = i
|
||
continue
|
||
if ch == "\n":
|
||
events.append(("line", buf[start:i]))
|
||
i += 1
|
||
start = i
|
||
continue
|
||
i += 1
|
||
return events, buf[start:]
|
||
|
||
|
||
def feed_ssh_log(
|
||
log: Callable[[str], None] | None,
|
||
kind: str,
|
||
text: str,
|
||
) -> None:
|
||
"""Dispatch a stream event to the CLI log callback."""
|
||
if not log or not text:
|
||
# Empty progress is a no-op; empty line still logs blank via caller if needed.
|
||
if log and kind == "line" and text == "":
|
||
log("")
|
||
return
|
||
if kind == "progress":
|
||
log("\r" + text)
|
||
else:
|
||
log(text)
|
||
|
||
|
||
def _stream_pty_output(
|
||
stdout,
|
||
*,
|
||
log: Callable[[str], None] | None,
|
||
chunks: list[str],
|
||
) -> None:
|
||
"""Read PTY stdout in chunks; honor \\r progress and \\n lines."""
|
||
buf = ""
|
||
channel = getattr(stdout, "channel", None)
|
||
if channel is None:
|
||
while True:
|
||
piece = stdout.read(4096)
|
||
if not piece:
|
||
break
|
||
text = piece.decode("utf-8", errors="replace") if isinstance(piece, bytes) else piece
|
||
chunks.append(text)
|
||
buf += text
|
||
events, buf = split_ssh_stream(buf)
|
||
for kind, part in events:
|
||
feed_ssh_log(log, kind, part)
|
||
if buf:
|
||
feed_ssh_log(log, "line", buf.rstrip("\r"))
|
||
return
|
||
|
||
while True:
|
||
if channel.recv_ready():
|
||
data = channel.recv(4096)
|
||
if not data:
|
||
break
|
||
text = data.decode("utf-8", errors="replace")
|
||
chunks.append(text)
|
||
buf += text
|
||
events, buf = split_ssh_stream(buf)
|
||
for kind, part in events:
|
||
feed_ssh_log(log, kind, part)
|
||
continue
|
||
if channel.exit_status_ready():
|
||
while channel.recv_ready():
|
||
data = channel.recv(4096)
|
||
if not data:
|
||
break
|
||
text = data.decode("utf-8", errors="replace")
|
||
chunks.append(text)
|
||
buf += text
|
||
events, buf = split_ssh_stream(buf)
|
||
for kind, part in events:
|
||
feed_ssh_log(log, kind, part)
|
||
break
|
||
time.sleep(0.05)
|
||
if buf:
|
||
feed_ssh_log(log, "line", buf.rstrip("\r"))
|
||
|
||
|
||
def wait_tcp(host: str, port: int, timeout: float = 300.0) -> None:
|
||
deadline = time.time() + timeout
|
||
last = None
|
||
while time.time() < deadline:
|
||
try:
|
||
with socket.create_connection((host, port), timeout=8):
|
||
return
|
||
except OSError as exc:
|
||
last = exc
|
||
time.sleep(4)
|
||
raise CloudError(f"TCP {host}:{port} не открылся за {int(timeout)} с ({last})")
|
||
|
||
|
||
def probe_ssh(cfg: Config, host: str, attempts: int = 3) -> str:
|
||
"""Quick SSH check. Returns 'ok' | 'auth' | 'down'."""
|
||
import logging
|
||
|
||
logging.getLogger("paramiko").setLevel(logging.CRITICAL)
|
||
logging.getLogger("paramiko.transport").setLevel(logging.CRITICAL)
|
||
try:
|
||
with socket.create_connection((host, 22), timeout=8):
|
||
pass
|
||
except OSError:
|
||
return "down"
|
||
|
||
key = str(cfg.ssh_private_key_path)
|
||
auth_seen = False
|
||
for _ in range(max(1, attempts)):
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
try:
|
||
client.connect(
|
||
hostname=host,
|
||
username=cfg.ssh_user,
|
||
key_filename=key,
|
||
timeout=12,
|
||
banner_timeout=20,
|
||
auth_timeout=15,
|
||
allow_agent=False,
|
||
look_for_keys=False,
|
||
)
|
||
client.close()
|
||
return "ok"
|
||
except Exception as exc:
|
||
name = type(exc).__name__
|
||
if "Authentication" in name or "Authentication" in str(exc):
|
||
auth_seen = True
|
||
try:
|
||
client.close()
|
||
except Exception:
|
||
pass
|
||
time.sleep(2)
|
||
return "auth" if auth_seen else "down"
|
||
|
||
|
||
def wait_ssh(
|
||
cfg: Config,
|
||
host: str,
|
||
timeout: float = 900.0,
|
||
log: Callable[[str], None] | None = None,
|
||
) -> None:
|
||
"""Wait until sshd accepts our key. Paramiko banner noise is muted.
|
||
|
||
AuthenticationException is normal while cloud-init injects keys: keep
|
||
retrying. Only give up early after AUTH_GIVE_UP seconds of *continuous*
|
||
auth rejection (sshd up, key still wrong) so recreate-with-user_data can run.
|
||
"""
|
||
import logging
|
||
|
||
logging.getLogger("paramiko").setLevel(logging.CRITICAL)
|
||
logging.getLogger("paramiko.transport").setLevel(logging.CRITICAL)
|
||
|
||
emit = log if log is not None else print
|
||
|
||
# sshd may answer before authorized_keys is ready (boot-from-volume / user_data).
|
||
auth_give_up = 180.0
|
||
|
||
wait_tcp(host, 22, timeout=min(timeout, 300))
|
||
deadline = time.time() + timeout
|
||
key = str(cfg.ssh_private_key_path)
|
||
last = None
|
||
attempt = 0
|
||
auth_streak_started: float | None = None
|
||
while time.time() < deadline:
|
||
attempt += 1
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
try:
|
||
client.connect(
|
||
hostname=host,
|
||
username=cfg.ssh_user,
|
||
key_filename=key,
|
||
timeout=20,
|
||
banner_timeout=60,
|
||
auth_timeout=30,
|
||
allow_agent=False,
|
||
look_for_keys=False,
|
||
)
|
||
client.close()
|
||
return
|
||
except Exception as exc:
|
||
last = exc
|
||
name = type(exc).__name__
|
||
is_auth = "Authentication" in name or "Authentication" in str(exc)
|
||
if is_auth:
|
||
if auth_streak_started is None:
|
||
auth_streak_started = time.time()
|
||
else:
|
||
auth_streak_started = None
|
||
try:
|
||
client.close()
|
||
except Exception:
|
||
pass
|
||
if attempt == 1 or attempt % 6 == 0:
|
||
emit(f"жду SSH {cfg.ssh_user}@{host}… ({name})")
|
||
if (
|
||
is_auth
|
||
and auth_streak_started is not None
|
||
and (time.time() - auth_streak_started) >= auth_give_up
|
||
):
|
||
raise CloudError(
|
||
f"SSH {cfg.ssh_user}@{host}: ключ отклонён {int(auth_give_up)}с подряд. "
|
||
"Nova keypair / user_data не попал в authorized_keys. "
|
||
"up пересоздаст compute с user_data (Base64), диски оставит."
|
||
) from exc
|
||
time.sleep(5)
|
||
raise CloudError(
|
||
f"SSH {cfg.ssh_user}@{host} не принял ключ за {int(timeout)} с ({last}). "
|
||
"Часто cloud-init ещё поднимает sshd на GPU-образе. "
|
||
"Если SG узкий (VPN/WARP): в .env GPU_RENT_SSH_CIDR=0.0.0.0/0 и снова up."
|
||
)
|
||
|
||
|
||
def ssh_argv(cfg: Config, host: str, remote: list[str] | None = None) -> list[str]:
|
||
cmd = [
|
||
"ssh",
|
||
"-i",
|
||
str(cfg.ssh_private_key_path),
|
||
"-o",
|
||
"StrictHostKeyChecking=accept-new",
|
||
"-o",
|
||
"IdentitiesOnly=yes",
|
||
f"{cfg.ssh_user}@{host}",
|
||
]
|
||
if remote:
|
||
cmd.append(" ".join(remote))
|
||
return cmd
|
||
|
||
|
||
def interactive_ssh(cfg: Config, host: str) -> int:
|
||
argv = ssh_argv(cfg, host)
|
||
try:
|
||
return subprocess.call(argv)
|
||
except FileNotFoundError:
|
||
raise CloudError(
|
||
"Нет клиента ssh в PATH. Windows: установи OpenSSH Client "
|
||
f"или зайди так: ssh -i {cfg.ssh_private_key_path} {cfg.ssh_user}@{host}"
|
||
) from None
|
||
|
||
|
||
def run_ssh(cfg: Config, host: str, command: str, timeout: int = 60, check: bool = True) -> str:
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
try:
|
||
client.connect(
|
||
hostname=host,
|
||
username=cfg.ssh_user,
|
||
key_filename=str(cfg.ssh_private_key_path),
|
||
timeout=15,
|
||
)
|
||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
|
||
out = stdout.read().decode("utf-8", errors="replace")
|
||
err = stderr.read().decode("utf-8", errors="replace")
|
||
code = stdout.channel.recv_exit_status()
|
||
if check and code != 0:
|
||
raise CloudError(f"SSH `{command}` exit {code}: {err or out}")
|
||
return out
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def run_script_sudo(
|
||
cfg: Config,
|
||
host: str,
|
||
script: str,
|
||
*,
|
||
remote_path: str,
|
||
timeout: int = 1800,
|
||
env: dict[str, str] | None = None,
|
||
log: Callable[[str], None] | None = None,
|
||
) -> str:
|
||
"""Upload a script and run it with passwordless sudo, streaming stdout."""
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
chunks: list[str] = []
|
||
try:
|
||
client.connect(
|
||
hostname=host,
|
||
username=cfg.ssh_user,
|
||
key_filename=str(cfg.ssh_private_key_path),
|
||
timeout=20,
|
||
)
|
||
sftp = client.open_sftp()
|
||
with sftp.file(remote_path, "w") as fh:
|
||
fh.write(script)
|
||
sftp.chmod(remote_path, 0o755)
|
||
sftp.close()
|
||
env_s = " ".join(
|
||
f"{key}={shlex.quote(str(value))}" for key, value in (env or {}).items()
|
||
)
|
||
prefix = f"sudo -n env {env_s} " if env_s else "sudo -n "
|
||
command = f"{prefix}bash {remote_path}"
|
||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
|
||
_stream_pty_output(stdout, log=log, chunks=chunks)
|
||
code = stdout.channel.recv_exit_status()
|
||
err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
|
||
# PTY usually merges stderr; if anything remains, don't drop it past timed log.
|
||
if log and err.strip():
|
||
for line in err.strip().splitlines():
|
||
log(line)
|
||
out = "".join(chunks)
|
||
if code != 0:
|
||
raise CloudError(f"remote script exit {code}: {err or out[-2000:]}")
|
||
return out
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def run_python(
|
||
cfg: Config,
|
||
host: str,
|
||
script: str,
|
||
*,
|
||
remote_path: str,
|
||
timeout: int = 1800,
|
||
log: Callable[[str], None] | None = None,
|
||
) -> str:
|
||
"""Upload a Python script and run it as SSH user (not root)."""
|
||
client = _connect(cfg, host)
|
||
chunks: list[str] = []
|
||
try:
|
||
sftp = client.open_sftp()
|
||
with sftp.file(remote_path, "w") as fh:
|
||
fh.write(script)
|
||
sftp.chmod(remote_path, 0o755)
|
||
sftp.close()
|
||
command = f"python3 {shlex.quote(remote_path)}"
|
||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
|
||
_stream_pty_output(stdout, log=log, chunks=chunks)
|
||
code = stdout.channel.recv_exit_status()
|
||
err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
|
||
if log and err.strip():
|
||
for line in err.strip().splitlines():
|
||
log(line)
|
||
out = "".join(chunks)
|
||
if code != 0:
|
||
raise CloudError(f"remote python exit {code}: {err or out[-2000:]}")
|
||
return out
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def _connect(cfg: Config, host: str) -> paramiko.SSHClient:
|
||
client = paramiko.SSHClient()
|
||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||
client.connect(
|
||
hostname=host,
|
||
username=cfg.ssh_user,
|
||
key_filename=str(cfg.ssh_private_key_path),
|
||
timeout=20,
|
||
)
|
||
return client
|
||
|
||
|
||
def open_ssh(cfg: Config, host: str) -> paramiko.SSHClient:
|
||
"""Public alias for a connected SSH client (caller must close)."""
|
||
return _connect(cfg, host)
|
||
|
||
|
||
def put_text(cfg: Config, host: str, remote_path: str, text: str, mode: int = 0o644) -> None:
|
||
client = _connect(cfg, host)
|
||
try:
|
||
put_text_on(client, remote_path, text, mode=mode)
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def put_text_on(
|
||
client: paramiko.SSHClient, remote_path: str, text: str, mode: int = 0o644
|
||
) -> None:
|
||
sftp = client.open_sftp()
|
||
try:
|
||
_sftp_mkdirs(sftp, str(Path(remote_path).parent).replace("\\", "/"))
|
||
with sftp.file(remote_path, "w") as fh:
|
||
fh.write(text)
|
||
sftp.chmod(remote_path, mode)
|
||
finally:
|
||
sftp.close()
|
||
|
||
|
||
def put_file(cfg: Config, host: str, local: Path, remote_path: str) -> None:
|
||
client = _connect(cfg, host)
|
||
try:
|
||
put_file_on(client, local, remote_path)
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def put_file_on(client: paramiko.SSHClient, local: Path, remote_path: str) -> None:
|
||
sftp = client.open_sftp()
|
||
try:
|
||
_sftp_mkdirs(sftp, str(Path(remote_path).parent).replace("\\", "/"))
|
||
sftp.put(str(local), remote_path)
|
||
finally:
|
||
sftp.close()
|
||
|
||
|
||
def get_file(cfg: Config, host: str, remote_path: str, local: Path) -> None:
|
||
local.parent.mkdir(parents=True, exist_ok=True)
|
||
client = _connect(cfg, host)
|
||
try:
|
||
get_file_on(client, remote_path, local)
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def get_file_on(client: paramiko.SSHClient, remote_path: str, local: Path) -> None:
|
||
local.parent.mkdir(parents=True, exist_ok=True)
|
||
sftp = client.open_sftp()
|
||
try:
|
||
sftp.get(remote_path, str(local))
|
||
finally:
|
||
sftp.close()
|
||
|
||
|
||
def remote_exists(cfg: Config, host: str, remote_path: str) -> bool:
|
||
client = _connect(cfg, host)
|
||
try:
|
||
sftp = client.open_sftp()
|
||
try:
|
||
sftp.stat(remote_path)
|
||
return True
|
||
except FileNotFoundError:
|
||
return False
|
||
finally:
|
||
sftp.close()
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def remote_sha256(cfg: Config, host: str, remote_path: str) -> str | None:
|
||
client = _connect(cfg, host)
|
||
try:
|
||
return remote_sha256_on(client, remote_path)
|
||
finally:
|
||
client.close()
|
||
|
||
|
||
def remote_sha256_on(client: paramiko.SSHClient, remote_path: str) -> str | None:
|
||
_stdin, stdout, stderr = client.exec_command(
|
||
f"sha256sum {shlex.quote(remote_path)} 2>/dev/null | awk '{{print $1}}'",
|
||
timeout=60,
|
||
)
|
||
del stderr
|
||
out = stdout.read().decode("utf-8", errors="replace").strip()
|
||
return out or None
|
||
|
||
|
||
def run_ssh_on(
|
||
client: paramiko.SSHClient, command: str, timeout: int = 60, check: bool = True
|
||
) -> str:
|
||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
|
||
out = stdout.read().decode("utf-8", errors="replace")
|
||
err = stderr.read().decode("utf-8", errors="replace")
|
||
code = stdout.channel.recv_exit_status()
|
||
if check and code != 0:
|
||
raise CloudError(f"SSH `{command}` exit {code}: {err or out}")
|
||
return out
|
||
|
||
|
||
def _sftp_mkdirs(sftp, remote_dir: str) -> None:
|
||
if not remote_dir or remote_dir == "/":
|
||
return
|
||
parts = [p for p in remote_dir.split("/") if p]
|
||
cur = ""
|
||
for part in parts:
|
||
cur += "/" + part
|
||
try:
|
||
sftp.stat(cur)
|
||
except FileNotFoundError:
|
||
sftp.mkdir(cur)
|