feat(train): print SSH commands after push, add truss ssh --training-job-id

truss/cli/ssh_commands.py

@@ -1,28 +1,111 @@
-from typing import Optional
+import os
+import sys
+from typing import Optional, cast
 
 import rich_click as click
 from InquirerPy import inquirer
 
+from truss.cli import remote_cli
 from truss.cli.cli import truss_cli
 from truss.cli.ssh import (
     ensure_ssh_keypair,
     install_proxy_command_script,
+    is_setup_complete,
     setup_ssh_config,
 )
+from truss.cli.train.common import get_most_recent_job
+from truss.cli.train.poller import TrainingPollerMixin
 from truss.cli.utils import common
 from truss.cli.utils.common import check_is_interactive
 from truss.cli.utils.output import console
+from truss.remote.baseten.remote import BasetenRemote
 from truss.remote.remote_factory import RemoteFactory
 
 
-@click.group()
-def ssh():
-    """SSH access to Baseten workloads."""
+@click.group(invoke_without_command=True)
+@click.option(
+    "--training-job-id",
+    "training_job_id",
+    type=str,
+    required=False,
+    help="Training job ID to SSH into. Waits for the job to be running, then connects.",
+)
+@click.option(
+    "--node-id",
+    "node_id",
+    type=int,
+    default=0,
+    show_default=True,
+    help="Node index to connect to for multi-node jobs.",
+)
+@click.option("--remote", type=str, required=False, help="Remote to use.")
+@click.pass_context
+def ssh(
+    ctx: click.Context,
+    training_job_id: Optional[str],
+    node_id: int,
+    remote: Optional[str],
+):
+    """SSH access to Baseten workloads.
+
+    Pass --training-job-id <id> to wait for a training job to be running and
+    then connect via SSH. Use `truss ssh setup` for one-time SSH setup.
+    """
+    if ctx.invoked_subcommand is not None:
+        return
+    if not training_job_id:
+        click.echo(ctx.get_help())
+        ctx.exit(2)
+    _wait_and_exec_ssh_training_job(training_job_id, node_id, remote)
 
 
 truss_cli.add_command(ssh)
 
 
+def _wait_and_exec_ssh_training_job(
+    job_id: str, node_id: int, remote: Optional[str]
+) -> None:
+    if not is_setup_complete():
+        console.print(
+            "SSH is not set up yet. Run [cyan]truss ssh setup[/cyan] first.",
+            style="yellow",
+        )
+        sys.exit(1)
+
+    if node_id < 0:
+        raise click.UsageError("--node-id must be >= 0")
+
+    if not remote:
+        remote = remote_cli.inquire_remote_name()
+
+    remote_provider = cast(BasetenRemote, RemoteFactory.create(remote=remote))
+
+    project_id, job_id = get_most_recent_job(remote_provider, None, job_id)
+
+    job_resp = remote_provider.api.get_training_job(project_id, job_id)
+    training_job = job_resp.get("training_job", {})
+    instance_type = training_job.get("instance_type") or {}
+    node_count = instance_type.get("node_count") or instance_type.get("nodeCount") or 1
+    if node_id >= node_count:
+        raise click.UsageError(
+            f"--node-id {node_id} is out of range; job has {node_count} node"
+            f"{'s' if node_count != 1 else ''} (0..{node_count - 1})."
+        )
+
+    poller = TrainingPollerMixin(remote_provider.api, project_id, job_id)
+    poller.before_polling()
+    status = poller._current_status.status
+    if status != "TRAINING_JOB_RUNNING":
+        console.print(
+            f"Job is not running (status: {status}). Cannot SSH.", style="red"
+        )
+        sys.exit(1)
+
+    hostname = f"training-job-{job_id}-{node_id}.ssh.baseten.co"
+    console.print(f"Connecting to [cyan]{hostname}[/cyan]...")
+    os.execvp("ssh", ["ssh", hostname])
+
+
 @ssh.command(name="setup")
 @click.option(
     "--python",

truss/cli/train_commands.py

@@ -51,6 +51,15 @@ def train():
 truss_cli.add_command(train)
 
 
+def _format_ssh_commands(job_id: str, node_count: int) -> str:
+    lines = f"  [cyan]ssh training-job-{job_id}-0.ssh.baseten.co[/cyan]"
+    if node_count > 1:
+        lines += " (leader)"
+        for i in range(1, node_count):
+            lines += f"\n  [cyan]ssh training-job-{job_id}-{i}.ssh.baseten.co[/cyan]"
+    return lines
+
+
 def _print_training_job_success_message(
     job_id: str,
     project_id: str,
@@ -80,6 +89,17 @@ def _print_training_job_success_message(
         f"🌐 View job in the UI: {common.format_link(core.status_page_url(remote_provider.remote_url, project_id, job_id))}"
     )
 
+    node_count = job_object.compute.node_count if job_object else 1
+    ssh_lines = _format_ssh_commands(job_id, node_count)
+    console.print(
+        f"\n🔑 Once the job is running, SSH in with:\n"
+        f"{ssh_lines}\n"
+        f"   Or wait + connect automatically: "
+        f"[cyan]truss ssh --training-job-id {job_id}"
+        f"{' --node-id <0..%d>' % (node_count - 1) if node_count > 1 else ''}[/cyan]\n"
+        f"   First time? Run [cyan]truss ssh setup[/cyan]."
+    )
+
 
 def _handle_post_create_logic(
     job_resp: dict, remote_provider: BasetenRemote, tail: bool

truss/tests/cli/test_ssh.py

@@ -2,8 +2,9 @@
 from unittest import mock
 
 import pytest
+from click.testing import CliRunner
 
-from truss.cli import proxy_command
+from truss.cli import proxy_command, ssh_commands
 from truss.cli import ssh as ssh_mod
 from truss.cli.proxy_command import (
     WORKLOAD_MODEL,
@@ -23,6 +24,7 @@
     is_setup_complete,
     setup_ssh_config,
 )
+from truss.cli.train_commands import _format_ssh_commands
 
 
 class TestParseHostname:
@@ -427,3 +429,115 @@ def test_returns_false_without_proxy_script(self, tmp_path):
     def test_returns_false_without_key(self, tmp_path):
         (tmp_path / "proxy-command.py").touch()
         assert is_setup_complete(tmp_path) is False
+
+
+class TestFormatSSHCommands:
+    def test_single_node(self):
+        out = _format_ssh_commands("wgvj7gw", node_count=1)
+        assert out == "  [cyan]ssh training-job-wgvj7gw-0.ssh.baseten.co[/cyan]"
+        assert "(leader)" not in out
+
+    def test_multi_node(self):
+        out = _format_ssh_commands("wgvj7gw", node_count=3)
+        assert "ssh training-job-wgvj7gw-0.ssh.baseten.co[/cyan] (leader)" in out
+        assert "ssh training-job-wgvj7gw-1.ssh.baseten.co" in out
+        assert "ssh training-job-wgvj7gw-2.ssh.baseten.co" in out
+        # No -3 line.
+        assert "training-job-wgvj7gw-3" not in out
+
+
+class _FakePoller:
+    """Minimal stand-in for TrainingPollerMixin used in tests."""
+
+    next_status = "TRAINING_JOB_RUNNING"
+
+    def __init__(self, api, project_id, job_id):
+        self.api = api
+        self.project_id = project_id
+        self.job_id = job_id
+        self._current_status = mock.Mock(status=self.next_status, error_message=None)
+
+    def before_polling(self):
+        return None
+
+
+class TestSSHTrainingJobCommand:
+    def _api_with_node_count(self, node_count=2):
+        api = mock.Mock()
+        api.get_training_job.return_value = {
+            "training_job": {
+                "current_status": "TRAINING_JOB_RUNNING",
+                "instance_type": {"name": "fake", "node_count": node_count},
+            }
+        }
+        return api
+
+    def _invoke(self, *, args, api, poller_status="TRAINING_JOB_RUNNING"):
+        remote_provider = mock.Mock()
+        remote_provider.api = api
+
+        class _StatusPoller(_FakePoller):
+            next_status = poller_status
+
+        runner = CliRunner()
+        with (
+            mock.patch.object(ssh_commands, "is_setup_complete", return_value=True),
+            mock.patch.object(
+                ssh_commands.RemoteFactory, "create", return_value=remote_provider
+            ),
+            mock.patch.object(
+                ssh_commands, "get_most_recent_job", return_value=("proj_id", "wgvj7gw")
+            ),
+            mock.patch.object(ssh_commands, "TrainingPollerMixin", _StatusPoller),
+            mock.patch.object(ssh_commands.os, "execvp") as mock_exec,
+        ):
+            result = runner.invoke(ssh_commands.ssh, args)
+        return result, mock_exec
+
+    def test_exec_ssh_with_correct_hostname(self):
+        result, mock_exec = self._invoke(
+            args=["--training-job-id", "wgvj7gw", "--node-id", "1", "--remote", "dev"],
+            api=self._api_with_node_count(node_count=2),
+        )
+        assert result.exit_code == 0, result.output
+        mock_exec.assert_called_once_with(
+            "ssh", ["ssh", "training-job-wgvj7gw-1.ssh.baseten.co"]
+        )
+
+    def test_default_node_id_is_zero(self):
+        result, mock_exec = self._invoke(
+            args=["--training-job-id", "wgvj7gw", "--remote", "dev"],
+            api=self._api_with_node_count(node_count=1),
+        )
+        assert result.exit_code == 0, result.output
+        mock_exec.assert_called_once_with(
+            "ssh", ["ssh", "training-job-wgvj7gw-0.ssh.baseten.co"]
+        )
+
+    def test_exits_when_setup_incomplete(self):
+        runner = CliRunner()
+        with mock.patch.object(ssh_commands, "is_setup_complete", return_value=False):
+            result = runner.invoke(
+                ssh_commands.ssh, ["--training-job-id", "wgvj7gw", "--remote", "dev"]
+            )
+        assert result.exit_code == 1
+        assert "truss ssh setup" in result.output
+
+    def test_node_id_out_of_range(self):
+        result, mock_exec = self._invoke(
+            args=["--training-job-id", "wgvj7gw", "--node-id", "5", "--remote", "dev"],
+            api=self._api_with_node_count(node_count=2),
+        )
+        assert result.exit_code != 0
+        assert "out of range" in result.output
+        mock_exec.assert_not_called()
+
+    def test_errors_when_job_not_running(self):
+        result, mock_exec = self._invoke(
+            args=["--training-job-id", "wgvj7gw", "--remote", "dev"],
+            api=self._api_with_node_count(node_count=1),
+            poller_status="TRAINING_JOB_FAILED",
+        )
+        assert result.exit_code == 1
+        assert "Cannot SSH" in result.output
+        mock_exec.assert_not_called()