feat(sdk): ssh for jobs

python/lightning_sdk/cli/job/__init__.py

@@ -10,11 +10,13 @@ def register_commands(group: click.Group) -> None:
     from lightning_sdk.cli.job.list import list_jobs
     from lightning_sdk.cli.job.logs import logs_job
     from lightning_sdk.cli.job.run import run_job
+    from lightning_sdk.cli.job.ssh import ssh_job
     from lightning_sdk.cli.job.stop import stop_job
 
     group.add_command(run_job, name="run")
     group.add_command(list_jobs, name="list")
     group.add_command(inspect_job, name="inspect")
     group.add_command(logs_job, name="logs")
+    group.add_command(ssh_job, name="ssh")
     group.add_command(stop_job, name="stop")
     group.add_command(delete_job, name="delete")

python/lightning_sdk/cli/job/ssh.py

@@ -0,0 +1,78 @@
+"""Job SSH command."""
+
+import subprocess
+from typing import Optional
+
+import rich_click as click
+
+from lightning_sdk.cli.utils.logging import LightningCommand
+from lightning_sdk.cli.utils.resource_resolution import resolve_job, resolve_teamspace
+from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
+from lightning_sdk.status import Status
+
+_SSH_HOST = "ssh.lightning.ai"
+_JOB_ID_PREFIX = "job_"
+
+
+def _ssh_user_for_job_id(job_id: str) -> str:
+    """Map a stored job id (``job_<ulid>``) to the SSH gateway user (``j_<ulid>``)."""
+    suffix = job_id[len(_JOB_ID_PREFIX) :] if job_id.startswith(_JOB_ID_PREFIX) else job_id
+    return f"j_{suffix}"
+
+
+@click.command("ssh", cls=LightningCommand)
+@click.argument("name")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the name of the teamspace the job lives in. "
+        "Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace). "
+        "If not specified, uses the configured default teamspace."
+    ),
+)
+def ssh_job(
+    name: str,
+    teamspace: Optional[str] = None,
+) -> None:
+    """SSH into a running job.
+
+    Example:
+        lightning job ssh my-job
+    """
+    ssh_impl(name=name, teamspace=teamspace)
+
+
+def ssh_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+) -> None:
+    if not name:
+        raise click.UsageError("Missing job name. Pass NAME.")
+
+    resolved_teamspace = resolve_teamspace(teamspace)
+    job = resolve_job(name, resolved_teamspace)
+
+    if job.status != Status.Running:
+        raise click.ClickException(
+            f"Job '{job.name}' is {job.status}, not Running. SSH is only available while the job is running."
+        )
+
+    job_id = job.id
+    if not job_id:
+        raise click.ClickException(f"Could not resolve id for job '{job.name}'.")
+
+    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())
+
+    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

python/tests/cli/job/test_ssh.py

@@ -0,0 +1,111 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import rich_click as click
+from click.testing import CliRunner
+
+from lightning_sdk.cli.job.ssh import _ssh_user_for_job_id, ssh_impl, ssh_job
+from lightning_sdk.status import Status
+from tests.cli.help import assert_help_contains, command_text, mock_command_logging
+
+
+@pytest.mark.parametrize(
+    ("job_id", "expected"),
+    [
+        ("job_01jj4hvvjj4zx1t1esm5az3zt7", "j_01jj4hvvjj4zx1t1esm5az3zt7"),
+        ("01jj4hvvjj4zx1t1esm5az3zt7", "j_01jj4hvvjj4zx1t1esm5az3zt7"),
+    ],
+)
+def test_ssh_user_for_job_id(job_id: str, expected: str) -> None:
+    assert _ssh_user_for_job_id(job_id) == expected
+
+
+@mock_command_logging
+def test_job_ssh_help() -> None:
+    result_text = command_text("lightning job ssh --help")
+
+    assert "Usage: lightning job ssh [OPTIONS] NAME" in result_text
+    assert "SSH into a running job." in result_text
+    assert "--teamspace" in result_text
+    assert "--rank" not in result_text
+
+
+@mock_command_logging
+def test_jobs_ssh_help() -> None:
+    assert_help_contains("lightning jobs ssh --help", "Usage: lightning jobs ssh", "SSH into a running job.")
+
+
+def test_ssh_resolves_before_downloading_keys() -> None:
+    """SSH lookup failures occur before key downloads."""
+    configure = MagicMock()
+    with patch(
+        "lightning_sdk.cli.job.ssh.resolve_teamspace",
+        return_value=MagicMock(name="coding-model-training"),
+    ), patch(
+        "lightning_sdk.cli.job.ssh.resolve_job",
+        side_effect=click.UsageError("Could not resolve job 'missing'."),
+    ), patch(
+        "lightning_sdk.cli.job.ssh.configure_ssh_internal",
+        configure,
+    ), pytest.raises(click.UsageError, match="Could not resolve job"):
+        ssh_impl(name="missing", teamspace=None)
+
+    configure.assert_not_called()
+
+
+def test_ssh_rejects_non_running_job() -> None:
+    job = MagicMock()
+    job.name = "train"
+    job.status = Status.Completed
+    job.id = "job_01abc"
+
+    with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch(
+        "lightning_sdk.cli.job.ssh.resolve_job", return_value=job
+    ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal") as configure, pytest.raises(
+        click.ClickException, match="not Running"
+    ):
+        ssh_impl(name="train", teamspace=None)
+
+    configure.assert_not_called()
+
+
+def test_ssh_runs_against_job_gateway_user() -> None:
+    teamspace = MagicMock()
+    job = MagicMock()
+    job.name = "train"
+    job.status = Status.Running
+    job.id = "job_01jj4hvvjj4zx1t1esm5az3zt7"
+
+    with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=teamspace), patch(
+        "lightning_sdk.cli.job.ssh.resolve_job", return_value=job
+    ) as resolve_job, patch(
+        "lightning_sdk.cli.job.ssh.configure_ssh_internal", return_value="/tmp/lightning_rsa"
+    ), patch("lightning_sdk.cli.job.ssh.subprocess.run") as run:
+        result = CliRunner().invoke(ssh_job, ["train", "--teamspace", "org/teamspace"])
+
+    assert result.exit_code == 0, result.output
+    resolve_job.assert_called_once_with("train", teamspace)
+    run.assert_called_once_with(["ssh", "-i", "/tmp/lightning_rsa", "j_01jj4hvvjj4zx1t1esm5az3zt7@ssh.lightning.ai"])
+
+
+def test_ssh_retries_with_fresh_keys_on_failure() -> None:
+    job = MagicMock()
+    job.name = "train"
+    job.status = Status.Running
+    job.id = "job_01abc"
+
+    configure = MagicMock(side_effect=["/tmp/old_key", "/tmp/new_key"])
+    run = MagicMock(side_effect=[OSError("ssh missing"), None])
+
+    with patch("lightning_sdk.cli.job.ssh.resolve_teamspace", return_value=MagicMock()), patch(
+        "lightning_sdk.cli.job.ssh.resolve_job", return_value=job
+    ), patch("lightning_sdk.cli.job.ssh.configure_ssh_internal", configure), patch(
+        "lightning_sdk.cli.job.ssh.subprocess.run", run
+    ):
+        ssh_impl(name="train", teamspace=None)
+
+    assert configure.call_count == 2
+    configure.assert_any_call(force_download=True)
+    assert run.call_args_list[0].args[0][2] == "/tmp/old_key"
+    assert run.call_args_list[1].args[0][2] == "/tmp/new_key"
+    assert run.call_args_list[1].args[0][-1] == "j_01abc@ssh.lightning.ai"