From 1f7237f7db525f983fe4ca80fc4bf41f077790ba Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 04:35:02 +0300 Subject: [PATCH] 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. --- env.example | 5 ++- src/gpu_rent/cloud.py | 65 ++++++++++++++++++++++++++++----- src/gpu_rent/os_client.py | 2 +- src/gpu_rent/session.py | 75 ++++++++++++++++++++++++++++++++++----- src/gpu_rent/ssh_ops.py | 18 ++++++++-- tests/test_user_data.py | 12 +++++++ 6 files changed, 155 insertions(+), 22 deletions(-) create mode 100644 tests/test_user_data.py diff --git a/env.example b/env.example index 650b789..fbb75af 100644 --- a/env.example +++ b/env.example @@ -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= diff --git a/src/gpu_rent/cloud.py b/src/gpu_rent/cloud.py index f226fd6..7e97542 100644 --- a/src/gpu_rent/cloud.py +++ b/src/gpu_rent/cloud.py @@ -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) diff --git a/src/gpu_rent/os_client.py b/src/gpu_rent/os_client.py index 769974b..78b7f29 100644 --- a/src/gpu_rent/os_client.py +++ b/src/gpu_rent/os_client.py @@ -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): diff --git a/src/gpu_rent/session.py b/src/gpu_rent/session.py index ed38c1b..a1b3a58 100644 --- a/src/gpu_rent/session.py +++ b/src/gpu_rent/session.py @@ -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 diff --git a/src/gpu_rent/ssh_ops.py b/src/gpu_rent/ssh_ops.py index f179f5d..5592ee4 100644 --- a/src/gpu_rent/ssh_ops.py +++ b/src/gpu_rent/ssh_ops.py @@ -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}). " diff --git a/tests/test_user_data.py b/tests/test_user_data.py new file mode 100644 index 0000000..0de8df3 --- /dev/null +++ b/tests/test_user_data.py @@ -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