Implement SSH user data handling and adjust volume size

- Added functions to generate SSH user data in cloud-config format and encode it in Base64 for server creation. - Updated the `create_gpu_server` function to include public key handling for SSH access. - Adjusted the default boot volume size from 40GB to 30GB in `os_client.py`. - Enhanced the `cmd_up` function to manage SSH key injection and server state more effectively. - Improved logging for SSH connection attempts and error handling in `wait_ssh` to provide clearer feedback on authentication issues.
2026-08-21 04:35:02 +03:00
parent d6d247875b
commit 1f7237f7db
6 changed files with 155 additions and 22 deletions
@@ -12,9 +12,8 @@ GPU_RENT_AZ=ru-7a

 SSH_PRIVATE_KEY_PATH=
 SSH_USER=ubuntu
-# Optional: force SSH security-group CIDR (default = your public IP /32).
-# If WARP/VPN and SSH hangs: GPU_RENT_SSH_CIDR=0.0.0.0/0
-# GPU_RENT_SSH_CIDR=0.0.0.0/0
+# Spike / WARP: open SSH. Tighten to your /32 later.
+GPU_RENT_SSH_CIDR=0.0.0.0/0

 BOOT_VOLUME_ID=
 DATA_VOLUME_ID=
@@ -335,6 +335,30 @@ def _tag_server(conn, server, spot: bool) -> None:
            pass


+def _ssh_user_data(public_key: str) -> str:
+    """Minimal cloud-config (Selectel docs: ssh_authorized_keys)."""
+    key = public_key.strip()
+    return (
+        "#cloud-config\n"
+        "ssh_authorized_keys:\n"
+        f"  - {key}\n"
+        "users:\n"
+        "  - default\n"
+        "  - name: ubuntu\n"
+        "    lock_passwd: true\n"
+        "    shell: /bin/bash\n"
+        "    sudo: ['ALL=(ALL) NOPASSWD:ALL']\n"
+        "    ssh_authorized_keys:\n"
+        f"      - {key}\n"
+    )
+
+
+def _user_data_b64(public_key: str) -> str:
+    import base64
+
+    return base64.b64encode(_ssh_user_data(public_key).encode("utf-8")).decode("ascii")
+
+
 def create_gpu_server(
    conn,
    *,
@@ -345,6 +369,7 @@ def create_gpu_server(
    data_volume_id: str,
    az: str,
    spot: bool,
+    public_key: str,
    log: Callable[[str], None],
 ) -> Any:
    bdm = [
@@ -366,7 +391,7 @@ def create_gpu_server(
    tags = [RESOURCE_TAG]
    if spot:
        tags.append(PREEMPTIBLE_TAG)
-    kwargs: dict[str, Any] = {
+    base: dict[str, Any] = {
        "name": SERVER_NAME,
        "flavor_id": flavor_id,
        "networks": [{"uuid": net_id}],
@@ -374,16 +399,40 @@ def create_gpu_server(
        "availability_zone": az,
        "block_device_mapping_v2": bdm,
        "security_groups": [{"name": sg_name}],
-        "tags": tags,
    }
-    try:
-        server = conn.compute.create_server(**kwargs)
-    except Exception as exc:
-        kwargs.pop("tags", None)
+    ud_b64 = _user_data_b64(public_key)
+    ud_plain = _ssh_user_data(public_key)
+
+    # Selectel: OpenStack API expects Base64 user_data; config_drive often unsupported.
+    attempts: list[tuple[str, dict[str, Any]]] = [
+        ("tags+user_data_b64", {**base, "tags": tags, "user_data": ud_b64}),
+        ("user_data_b64", {**base, "user_data": ud_b64}),
+        ("user_data_plain", {**base, "user_data": ud_plain}),
+        ("key_name_only", {**base, "tags": tags}),
+        ("minimal", dict(base)),
+    ]
+    server = None
+    last_exc: BaseException | None = None
+    used = ""
+    for label, kwargs in attempts:
        try:
            server = conn.compute.create_server(**kwargs)
-        except Exception as exc2:
-            raise _wrap(exc2, "create server") from exc
+            used = label
+            break
+        except Exception as exc:
+            last_exc = exc
+            log(f"create {label}: {exc}")
+    if server is None:
+        raise _wrap(last_exc or RuntimeError("create failed"), "create server")
+
+    if "user_data" not in used:
+        log(
+            f"create OK ({used}) без user_data — ключ может не попасть на boot-from-volume; "
+            "смотри docs Selectel user-data (Base64)"
+        )
+    else:
+        log(f"create OK ({used})")
+
    log("ждём Nova ACTIVE (GPU create может занять несколько минут)")
    server = wait_server(conn, server, "ACTIVE")
    _tag_server(conn, server, spot)
@@ -21,7 +21,7 @@ SG_NAME = "gpu-rent"
 BOOT_VOLUME_NAME = "gpu-rent-boot"
 DATA_VOLUME_NAME = "gpu-rent-data"
 SUBNET_CIDR = "192.168.77.0/24"
-BOOT_VOLUME_SIZE_GB = 40
+BOOT_VOLUME_SIZE_GB = 30


 def connect(cfg: Config):
@@ -18,6 +18,7 @@ from gpu_rent.cloud import (
    pick_existing_server,
    server_status,
    unshelve,
+    wait_volume,
 )
 from gpu_rent.bootstrap import run_bootstrap
 from gpu_rent.provision import provision_vm
@@ -140,12 +141,63 @@ def cmd_up(
            state.server_id = existing.id
            state.server_name = getattr(existing, "name", None)
            if status == "ACTIVE":
-                log("сервер уже ACTIVE — второй GPU не создаём")
-                state.phase = "ready_cloud"
-                save_state(state)
-                _bind_access(conn, existing, state, cfg, log)
-                return state
-            if status in {"EXPIRED", "SHELVED", "SHELVED_OFFLOADED"}:
+                if state.bootstrapped and state.floating_ip:
+                    log("сервер уже ACTIVE — второй GPU не создаём")
+                    state.phase = "ready_cloud"
+                    save_state(state)
+                    _bind_access(conn, existing, state, cfg, log)
+                    return state
+                # First boot / stuck without injected key: probe SSH briefly.
+                fip = state.floating_ip
+                if not fip:
+                    try:
+                        fip, fip_id = ensure_floating_ip(
+                            conn, existing, state.floating_ip_id, state.floating_ip, log
+                        )
+                        state.floating_ip = fip
+                        if fip_id:
+                            state.floating_ip_id = fip_id
+                        save_state(state)
+                    except Exception as exc:
+                        log(f"FIP: {exc}")
+                if fip:
+                    try:
+                        wait_ssh(cfg, fip, timeout=90)
+                        log("сервер уже ACTIVE — второй GPU не создаём")
+                        state.phase = "ready_cloud"
+                        save_state(state)
+                        _bind_access(conn, existing, state, cfg, log)
+                        return state
+                    except CloudError as exc:
+                        msg = str(exc)
+                        if "ключ отклонён" in msg or "Authentication" in msg:
+                            log(
+                                "SSH ключ не на VM (boot-from-volume) — "
+                                "удаляем compute, диски оставляем, create с config_drive"
+                            )
+                            delete_server(conn, existing, log)
+                            state.server_id = None
+                            state.server_name = None
+                            state.bootstrapped = False
+                            state.phase = "idle"
+                            save_state(state)
+                            for vid in (state.boot_volume_id, state.data_volume_id):
+                                if not vid:
+                                    continue
+                                try:
+                                    wait_volume(conn, conn.block_storage.get_volume(vid), "available")
+                                except Exception as vol_exc:
+                                    log(f"wait volume {vid}: {vol_exc}")
+                            existing = None
+                        else:
+                            raise
+                if existing is not None:
+                    log("сервер уже ACTIVE — второй GPU не создаём")
+                    state.phase = "ready_cloud"
+                    save_state(state)
+                    _bind_access(conn, existing, state, cfg, log)
+                    return state
+            if existing is not None and status in {"EXPIRED", "SHELVED", "SHELVED_OFFLOADED"}:
                existing = unshelve(conn, existing, log)
                state.server_id = existing.id
                state.phase = "ready_cloud"
@@ -153,7 +205,8 @@ def cmd_up(
                save_state(state)
                _bind_access(conn, existing, state, cfg, log)
                return state
-            raise CloudError(f"сервер gpu-rent в статусе {status} — разбери в панели")
+            if existing is not None:
+                raise CloudError(f"сервер gpu-rent в статусе {status} — разбери в панели")

        if state.server_id:
            log(f"в state был server {state.server_id}, в облаке нет — создаём заново")
@@ -220,7 +273,12 @@ def cmd_up(
        net, _subnet = ensure_network(conn, log)
        cidr = guess_operator_cidr()
        if cidr == "0.0.0.0/0":
-            log("не удалось узнать твой IP — SG откроет SSH с 0.0.0.0/0")
+            import os
+
+            if (os.environ.get("GPU_RENT_SSH_CIDR") or "").strip():
+                log("SG SSH: GPU_RENT_SSH_CIDR=0.0.0.0/0")
+            else:
+                log("не удалось узнать твой IP — SG откроет SSH с 0.0.0.0/0")
        sg = ensure_security_group(conn, cidr, log)

        boot = ensure_boot_volume(
@@ -263,6 +321,7 @@ def cmd_up(
            data_volume_id=data.id,
            az=cfg.gpu_rent_az,
            spot=spot,
+            public_key=public_key,
            log=log,
        )
        state.server_id = server.id
@@ -40,6 +40,7 @@ def wait_ssh(cfg: Config, host: str, timeout: float = 900.0) -> None:
    key = str(cfg.ssh_private_key_path)
    last = None
    attempt = 0
+    auth_fails = 0
    while time.time() < deadline:
        attempt += 1
        client = paramiko.SSHClient()
@@ -59,13 +60,26 @@ def wait_ssh(cfg: Config, host: str, timeout: float = 900.0) -> None:
            return
        except Exception as exc:
            last = exc
+            name = type(exc).__name__
+            if "Authentication" in name or "Authentication" in str(exc):
+                auth_fails += 1
            try:
                client.close()
            except Exception:
                pass
            if attempt == 1 or attempt % 6 == 0:
-                # Progress without paramiko traceback spam.
-                print(f"жду SSH {cfg.ssh_user}@{host}… ({type(exc).__name__})", flush=True)
+                print(f"жду SSH {cfg.ssh_user}@{host}… ({name})", flush=True)
+            # Key never injected (boot-from-volume): don't burn full timeout.
+            if auth_fails >= 8:
+                raise CloudError(
+                    f"SSH {cfg.ssh_user}@{host}: ключ отклонён (AuthenticationException). "
+                    "Nova keypair не попал в authorized_keys при boot-from-volume. "
+                    "Останови up (Ctrl+C), затем снова: gpu-rent up --yes "
+                    "(create теперь с config_drive + user_data). "
+                    "Или в консоли панели (root): "
+                    f"добавь содержимое {cfg.ssh_private_key_path}.pub в "
+                    f"/home/{cfg.ssh_user}/.ssh/authorized_keys"
+                ) from exc
            time.sleep(8)
    raise CloudError(
        f"SSH {cfg.ssh_user}@{host} не принял ключ за {int(timeout)} с ({last}). "
@@ -0,0 +1,12 @@
+import base64
+
+from gpu_rent.cloud import _ssh_user_data, _user_data_b64
+
+
+def test_user_data_b64_is_selectel_cloud_config():
+    key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFakeKeyMaterialHere gpu-rent"
+    plain = _ssh_user_data(key)
+    assert plain.startswith("#cloud-config")
+    assert key in plain
+    raw = base64.b64decode(_user_data_b64(key))
+    assert raw.decode("utf-8") == plain