fix(cli): create SSH keys locally and register only the public key

python/lightning_sdk/api/ssh_key_api.py

@@ -0,0 +1,36 @@
+import json
+
+from lightning_sdk.api.utils import ApiException, cached_lightning_client
+from lightning_sdk.lightning_cloud.openapi import V1CreateSSHPublicKeyRequest, V1SSHPublicKey
+
+
+class SSHKeyRegistrationError(Exception):
+    """The platform refused to register an SSH public key. The message is the platform's own, for the user."""
+
+
+def _error_message(e: ApiException) -> str:
+    body = e.body.decode(errors="replace") if isinstance(e.body, bytes) else e.body
+    try:
+        message = json.loads(body or "")["message"]
+    except (ValueError, KeyError, TypeError):
+        message = ""
+    return message or f"HTTP {e.status} {e.reason}"
+
+
+class SSHKeyApi:
+    """Internal API client for the SSH public keys registered to the signed-in user."""
+
+    def create_public_key(self, name: str, public_key: str) -> V1SSHPublicKey:
+        """Register ``public_key`` under ``name``. If the account already has it, the platform returns it unchanged.
+
+        Raises:
+            SSHKeyRegistrationError: With the platform's reason, e.g. that the key is held by another account.
+        """
+        # Not retried: once retries run out, the client raises a bare Exception without the platform's message.
+        client = cached_lightning_client(retry=False)
+        try:
+            return client.s_sh_public_key_service_create_ssh_public_key(
+                V1CreateSSHPublicKeyRequest(name=name, public_key=public_key)
+            )
+        except ApiException as e:
+            raise SSHKeyRegistrationError(_error_message(e)) from None

python/lightning_sdk/cli/job/ssh.py

@@ -71,15 +71,8 @@ def ssh_impl(
 
     ssh_user = _ssh_user_for_job_id(job_id)
 
-    def _run(key_path: str) -> None:
-        command = f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}"
-        subprocess.run(command.split())
-
+    key_path = configure_ssh_internal()
     try:
-        _run(configure_ssh_internal())
-    except Exception:
-        # Redownload keys in case they are stale, then retry once.
-        try:
-            _run(configure_ssh_internal(force_download=True))
-        except Exception:
-            raise click.ClickException("Failed to establish SSH connection") from None
+        subprocess.run(f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}".split())
+    except OSError:
+        raise click.ClickException("Failed to establish SSH connection") from None

python/lightning_sdk/cli/legacy/configure.py

@@ -1,17 +1,10 @@
-import os
-import platform
-import uuid
-from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
 import click
-from rich.console import Console
 
-from lightning_sdk.api.utils import _get_cloud_url
-from lightning_sdk.cli.legacy.generate import _generate_ssh_config
+from lightning_sdk.cli.ssh.configure import OVERWRITE_HELP, configure_studio_ssh
 from lightning_sdk.cli.utils.auth import require_auth_header
 from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace
-from lightning_sdk.lightning_cloud.login import Auth
 
 
 @click.group(name="configure")
@@ -23,39 +16,9 @@ def _configure_ssh_internal(
     name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False
 ) -> None:
     """Internal function to configure SSH without Click decorators."""
-    auth = Auth()
     require_auth_header()
-    if not auth.api_key and not auth.load():
-        raise click.UsageError("An API key is required. Run `lightning login` first.")
-    console = Console()
-    ssh_dir = Path.home() / ".ssh"
-    ssh_dir.mkdir(parents=True, exist_ok=True)
-
-    key_path = ssh_dir / "lightning_rsa"
-    config_path = ssh_dir / "config"
-
-    # Check if the SSH key already exists
-    if key_path.exists() and (key_path.with_suffix(".pub")).exists() and not overwrite:
-        console.print(f"SSH key already exists at {key_path}")
-    else:
-        _download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name="lightning_rsa", overwrite=overwrite)
-        console.print(f"SSH key generated and saved to {key_path}")
-
-    # Check if the SSH config already contains the required configuration
     studio = resolve_studio(name, resolve_teamspace(teamspace))
-    config_content = _generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name)
-    if config_path.exists():
-        with config_path.open("r") as config_file:
-            # check if the host already exists in the config
-            if f"Host {studio.name}" in config_file.read():
-                console.print("SSH config already contains the required configuration.")
-                return
-
-    with config_path.open("a") as config_file:
-        config_file.write(os.linesep)
-        config_file.write(config_content)
-        config_file.write(os.linesep)
-        console.print(f"SSH config updated at {config_path}")
+    configure_studio_ssh(studio.name, f"s_{studio._studio.id}", overwrite=overwrite)
 
 
 @configure.command(name="ssh")
@@ -81,52 +44,8 @@ def _configure_ssh_internal(
     is_flag=True,
     flag_value=True,
     default=False,
-    help="Whether to overwrite the SSH key and config if they already exist.",
+    help=OVERWRITE_HELP,
 )
 def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
     """Get SSH config entry for a studio."""
     _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=overwrite)
-
-
-def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
-    """Download a file from a URL."""
-    import requests
-
-    if local_path.exists() and not overwrite:
-        raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
-
-    response = requests.get(url, stream=True)
-    response.raise_for_status()
-
-    with open(local_path, "wb") as file:
-        for chunk in response.iter_content(chunk_size=8192):
-            file.write(chunk)
-    if chmod is not None:
-        local_path.chmod(0o600)
-
-
-def _download_ssh_keys(
-    api_key: str,
-    key_id: str = "",
-    ssh_home: Union[str, Path] = "",
-    ssh_key_name: str = "lightning_rsa",
-    overwrite: bool = False,
-) -> None:
-    if not ssh_home:
-        ssh_home = Path.home() / ".ssh"
-    elif isinstance(ssh_home, str):
-        ssh_home = Path(ssh_home)
-    if not key_id:
-        key_id = str(uuid.uuid4())
-
-    path_key = ssh_home / ssh_key_name
-    path_pub = ssh_home / f"{ssh_key_name}.pub"
-
-    # todo: consider hitting the API to get the key pair directly instead of using wget
-    _download_file(
-        f"{_get_cloud_url()}/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
-        path_key,
-        overwrite=overwrite,
-        chmod=0o600,
-    )
-    _download_file(f"{_get_cloud_url()}/setup/ssh-public?t={api_key}&id={key_id}", path_pub, overwrite=overwrite)

python/lightning_sdk/cli/legacy/generate.py

@@ -3,6 +3,7 @@
 import click
 from rich.console import Console
 
+from lightning_sdk.cli.ssh.common import generate_ssh_config
 from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace
 
 
@@ -33,19 +34,6 @@ def ssh(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
     """Get SSH config entry for a studio."""
     studio = resolve_studio(name, resolve_teamspace(teamspace))
 
-    conf = _generate_ssh_config(key_path="~/.ssh/lightning_rsa", user=f"s_{studio._studio.id}", host=studio.name)
+    conf = generate_ssh_config(user=f"s_{studio._studio.id}", host=studio.name)
     # Print the SSH config
     Console().print(f"# ssh s_{studio._studio.id}@ssh.lightning.ai\n\n" + conf)
-
-
-def _generate_ssh_config(key_path: str, host: str, user: str) -> str:
-    return f"""Host {host}
-      User {user}
-      Hostname ssh.lightning.ai
-      IdentityFile {key_path}
-      IdentitiesOnly yes
-      ServerAliveInterval 15
-      ServerAliveCountMax 4
-      StrictHostKeyChecking no
-      UserKnownHostsFile=/dev/null
-    """

python/lightning_sdk/cli/mmt/ssh.py

@@ -71,15 +71,8 @@ def ssh_impl(
 
     ssh_user = _ssh_user_for_job_id(job_id)
 
-    def _run(key_path: str) -> None:
-        command = f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}"
-        subprocess.run(command.split())
-
+    key_path = configure_ssh_internal()
     try:
-        _run(configure_ssh_internal())
-    except Exception:
-        # Redownload keys in case they are stale, then retry once.
-        try:
-            _run(configure_ssh_internal(force_download=True))
-        except Exception:
-            raise click.ClickException("Failed to establish SSH connection") from None
+        subprocess.run(f"ssh -i {key_path} {ssh_user}@{_SSH_HOST}".split())
+    except OSError:
+        raise click.ClickException("Failed to establish SSH connection") from None

python/lightning_sdk/cli/ssh/common.py

@@ -1,14 +1,58 @@
 """Shared SSH command helpers."""
 
+from pathlib import Path
 
-def generate_ssh_config(key_path: str, host: str, user: str) -> str:
-    return f"""Host {host}
-      User {user}
-      Hostname ssh.lightning.ai
-      IdentityFile {key_path}
-      IdentitiesOnly yes
-      ServerAliveInterval 15
-      ServerAliveCountMax 4
-      StrictHostKeyChecking no
-      UserKnownHostsFile=/dev/null
+from lightning_sdk.cli.utils.ssh_connection import SSH_KEY_NAME
+
+SSH_HOST = "ssh.lightning.ai"
+
+
+def generate_ssh_config(host: str, user: str) -> str:
+    """The ``~/.ssh/config`` entry for ``host``, in the same form as the web UI's setup script writes."""
+    return (
+        f"Host {host}\n"
+        f"  User {user}\n"
+        f"  Hostname {SSH_HOST}\n"
+        f"  IdentityFile ~/.ssh/{SSH_KEY_NAME}\n"
+        "  IdentitiesOnly yes\n"
+        "  ServerAliveInterval 15\n"
+        "  ServerAliveCountMax 4\n"
+        "  StrictHostKeyChecking accept-new\n"
+        "  UserKnownHostsFile ~/.ssh/lightning_known_hosts\n"
+    )
+
+
+def host_uses_ssh_key(config: str, host: str) -> bool:
+    """Whether a ``Host`` or ``Match`` block naming ``host`` already has the Lightning key as an ``IdentityFile``."""
+    in_block = False
+    for line in config.splitlines():
+        fields = line.split()
+        if not fields:
+            continue
+        keyword = fields[0].lower()
+        if keyword in ("host", "match"):
+            in_block = host in fields[1:]
+        elif in_block and keyword.startswith("identityfile") and f".ssh/{SSH_KEY_NAME}" in line:
+            return True
+    return False
+
+
+def add_ssh_config(host: str, user: str) -> bool:
+    """Append the entry for ``host`` to ``~/.ssh/config``, unless one there already uses the Lightning key.
+
+    ssh tries every ``IdentityFile`` of every block that matches, so an older entry for ``host`` that names
+    only ``lightning_rsa`` offers both keys once this one is added after it.
+
+    Returns:
+        Whether the file was changed.
     """
+    config_path = Path.home() / ".ssh" / "config"
+    config = config_path.read_text() if config_path.exists() else ""
+    if host_uses_ssh_key(config, host):
+        return False
+    config_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    with config_path.open("a") as config_file:
+        if config and not config.endswith("\n"):
+            config_file.write("\n")
+        config_file.write("\n" + generate_ssh_config(host=host, user=user))
+    return True

python/lightning_sdk/cli/ssh/configure.py

@@ -1,19 +1,22 @@
 """SSH configure command."""
 
-from pathlib import Path
 from typing import Optional
 
 import rich_click as click
 from rich.console import Console
 
-from lightning_sdk.api.utils import _get_cloud_url
 from lightning_sdk.cli.resource_completion import complete_studio
-from lightning_sdk.cli.ssh.common import generate_ssh_config
+from lightning_sdk.cli.ssh.common import add_ssh_config
 from lightning_sdk.cli.utils.auth import require_auth_header
 from lightning_sdk.cli.utils.logging import LightningCommand
 from lightning_sdk.cli.utils.resource_resolution import resolve_studio, resolve_teamspace
-from lightning_sdk.cli.utils.ssh_connection import download_file
-from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.cli.utils.ssh_connection import ensure_ssh_key
+
+OVERWRITE_HELP = "Deprecated, and has no effect: the SSH key is never replaced."
+OVERWRITE_WARNING = (
+    "--overwrite no longer replaces the SSH key. To make a new one, delete {key_path} and its .pub file, "
+    "then run this again."
+)
 
 
 @click.command("configure", cls=LightningCommand)
@@ -37,51 +40,26 @@
     is_flag=True,
     flag_value=True,
     default=False,
-    help="Whether to overwrite the SSH key and config if they already exist.",
+    help=OVERWRITE_HELP,
 )
 def configure_ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
     """Get SSH config entry for a studio."""
-    import platform
-    import uuid
-
     resolved_teamspace = resolve_teamspace(teamspace)
     studio = resolve_studio(name, resolved_teamspace)
 
-    auth = Auth()
     require_auth_header()
-    if not auth.api_key and not auth.load():
-        raise click.UsageError("An API key is required. Run `lightning login` first.")
-    console = Console()
-    ssh_dir = Path.home() / ".ssh"
-    ssh_dir.mkdir(parents=True, exist_ok=True)
+    configure_studio_ssh(studio.name, f"s_{studio._studio.id}", overwrite=overwrite)
 
-    key_path = ssh_dir / "lightning_rsa"
-    config_path = ssh_dir / "config"
 
-    if not key_path.exists() or not (key_path.with_suffix(".pub")).exists() or overwrite:
-        key_id = str(uuid.uuid4())
-        download_file(
-            f"{_get_cloud_url()}/setup/ssh-gen?t={auth.api_key}&id={key_id}&machineName={platform.node()}",
-            key_path,
-            overwrite=overwrite,
-            chmod=0o600,
-        )
-        download_file(
-            f"{_get_cloud_url()}/setup/ssh-public?t={auth.api_key}&id={key_id}",
-            key_path.with_suffix(".pub"),
-            overwrite=overwrite,
-        )
-        console.print(f"SSH key generated and saved to {key_path}")
-    else:
-        console.print(f"SSH key already exists at {key_path}")
-
-    config_content = generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name)
-    if config_path.exists() and f"Host {studio.name}" in config_path.read_text():
-        console.print("SSH config already contains the required configuration.")
-        return
+def configure_studio_ssh(host: str, user: str, overwrite: bool = False) -> None:
+    """Set up this machine's SSH key and add the ``~/.ssh/config`` entry ``host`` for ``user``."""
+    console = Console()
+    key_path = ensure_ssh_key()
+    if overwrite:
+        console.print(OVERWRITE_WARNING.format(key_path=key_path), style="yellow")
+    console.print(f"Using SSH key {key_path}")
 
-    with config_path.open("a") as config_file:
-        config_file.write("\n")
-        config_file.write(config_content)
-        config_file.write("\n")
-    console.print(f"SSH config updated at {config_path}")
+    if add_ssh_config(host=host, user=user):
+        console.print(f"Added {host} to ~/.ssh/config")
+    else:
+        console.print(f"~/.ssh/config already uses this key for {host}")

python/lightning_sdk/cli/ssh/generate.py

@@ -31,5 +31,5 @@ def generate_ssh(name: Optional[str] = None, teamspace: Optional[str] = None) ->
     """Get SSH config entry for a studio."""
     resolved_teamspace = resolve_teamspace(teamspace)
     studio = resolve_studio(name, resolved_teamspace)
-    conf = generate_ssh_config(key_path="~/.ssh/lightning_rsa", user=f"s_{studio._studio.id}", host=studio.name)
+    conf = generate_ssh_config(user=f"s_{studio._studio.id}", host=studio.name)
     Console().print(f"# ssh s_{studio._studio.id}@ssh.lightning.ai\n\n" + conf)