Add package data for GPU rent and update CLI documentation

- Added package data configuration for the 'gpu_rent' package in pyproject.toml. - Updated README.md to include usage instructions for Windows and Unix launchers. - Enhanced CLI documentation in cli.md to reflect new commands and their functionalities. - Revised setup.md to clarify installation steps and environment setup. - Improved error handling and command descriptions in the CLI implementation. - Added new functions for model version handling and flavor resolution in the codebase. - Updated state management to include additional properties for better tracking.
2026-08-21 03:06:51 +03:00
parent 167d07a733
commit 615cf81493
34 changed files with 2733 additions and 117 deletions
@@ -0,0 +1,274 @@
+"""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 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 wait_ssh(cfg: Config, host: str, timeout: float = 420.0) -> None:
+    wait_tcp(host, 22, timeout=min(timeout, 240))
+    deadline = time.time() + timeout
+    key = str(cfg.ssh_private_key_path)
+    last = None
+    while time.time() < deadline:
+        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=12,
+                auth_timeout=12,
+            )
+            client.close()
+            return
+        except Exception as exc:
+            last = exc
+            try:
+                client.close()
+            except Exception:
+                pass
+            time.sleep(5)
+    raise CloudError(f"SSH {cfg.ssh_user}@{host} не принял ключ ({last})")
+
+
+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}={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)
+        while True:
+            line = stdout.readline()
+            if not line:
+                break
+            chunks.append(line)
+            if log:
+                log(line.rstrip("\n\r"))
+        code = stdout.channel.recv_exit_status()
+        err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
+        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)
+        while True:
+            line = stdout.readline()
+            if not line:
+                break
+            chunks.append(line)
+            if log:
+                log(line.rstrip("\n\r"))
+        code = stdout.channel.recv_exit_status()
+        err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
+        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 put_text(cfg: Config, host: str, remote_path: str, text: str, mode: int = 0o644) -> None:
+    client = _connect(cfg, host)
+    try:
+        sftp = client.open_sftp()
+        _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)
+        sftp.close()
+    finally:
+        client.close()
+
+
+def put_file(cfg: Config, host: str, local: Path, remote_path: str) -> None:
+    client = _connect(cfg, host)
+    try:
+        sftp = client.open_sftp()
+        _sftp_mkdirs(sftp, str(Path(remote_path).parent).replace("\\", "/"))
+        sftp.put(str(local), remote_path)
+        sftp.close()
+    finally:
+        client.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:
+        sftp = client.open_sftp()
+        sftp.get(remote_path, str(local))
+        sftp.close()
+    finally:
+        client.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:
+    out = run_ssh(
+        cfg,
+        host,
+        f"sha256sum {shlex.quote(remote_path)} 2>/dev/null | awk '{{print $1}}'",
+        check=False,
+    ).strip()
+    return out or None
+
+
+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)