From c7a56bdd566e2421105a3b46dcffdd699340d720 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 04:37:09 +0300 Subject: [PATCH] Enhance SSH handling in session management - Introduced `probe_ssh` function for quick SSH checks, improving the handling of server bootstrap scenarios. - Updated `cmd_up` to utilize `probe_ssh`, enhancing the logic for managing server states based on SSH availability. - Improved logging for server state transitions and error handling during SSH connection attempts. - Adjusted timeout settings in `wait_ssh` for better performance and reliability in SSH key acceptance checks. --- src/gpu_rent/session.py | 61 ++++++++++++++++++----------------------- src/gpu_rent/ssh_ops.py | 52 ++++++++++++++++++++++++++++++----- 2 files changed, 71 insertions(+), 42 deletions(-) diff --git a/src/gpu_rent/session.py b/src/gpu_rent/session.py index a1b3a58..e146ec3 100644 --- a/src/gpu_rent/session.py +++ b/src/gpu_rent/session.py @@ -46,7 +46,7 @@ from gpu_rent.os_client import ( iter_volume_types, ) from gpu_rent.ssh_keys import ensure_ed25519 -from gpu_rent.ssh_ops import wait_ssh +from gpu_rent.ssh_ops import probe_ssh, wait_ssh from gpu_rent.state import SessionState, load_state, save_state, utc_now Log = Callable[[str], None] @@ -147,7 +147,8 @@ def cmd_up( save_state(state) _bind_access(conn, existing, state, cfg, log) return state - # First boot / stuck without injected key: probe SSH briefly. + + # Bootstrap не завершён: почти всегда VM без authorized_keys. fip = state.floating_ip if not fip: try: @@ -160,43 +161,33 @@ def cmd_up( 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 не создаём") + + outcome = probe_ssh(cfg, fip, attempts=2) if fip else "down" + if outcome == "ok": + log("сервер ACTIVE, SSH ок — продолжаем bootstrap") state.phase = "ready_cloud" save_state(state) _bind_access(conn, existing, state, cfg, log) return state + + log( + f"ACTIVE без bootstrap, SSH={outcome} — " + "удаляем compute (диски оставляем), create с user_data Base64" + ) + 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 if existing is not None and status in {"EXPIRED", "SHELVED", "SHELVED_OFFLOADED"}: existing = unshelve(conn, existing, log) state.server_id = existing.id diff --git a/src/gpu_rent/ssh_ops.py b/src/gpu_rent/ssh_ops.py index 5592ee4..6a8f609 100644 --- a/src/gpu_rent/ssh_ops.py +++ b/src/gpu_rent/ssh_ops.py @@ -28,6 +28,48 @@ def wait_tcp(host: str, port: int, timeout: float = 300.0) -> None: 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) -> None: """Wait until sshd accepts our key. Paramiko banner noise is muted.""" import logging @@ -70,17 +112,13 @@ def wait_ssh(cfg: Config, host: str, timeout: float = 900.0) -> None: if attempt == 1 or attempt % 6 == 0: 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: + if auth_fails >= 3: 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" + "up пересоздаст compute с user_data (Base64), диски оставит." ) from exc - time.sleep(8) + time.sleep(5) raise CloudError( f"SSH {cfg.ssh_user}@{host} не принял ключ за {int(timeout)} с ({last}). " "Часто cloud-init ещё поднимает sshd на GPU-образе. "