diff --git a/truss/cli/ssh_commands.py b/truss/cli/ssh_commands.py index 479e76762..91f051fd6 100644 --- a/truss/cli/ssh_commands.py +++ b/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 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", diff --git a/truss/cli/train_commands.py b/truss/cli/train_commands.py index 5fa7ac3c6..224443796 100644 --- a/truss/cli/train_commands.py +++ b/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 diff --git a/truss/tests/cli/test_ssh.py b/truss/tests/cli/test_ssh.py index b3b68401c..2f5d1aa4b 100644 --- a/truss/tests/cli/test_ssh.py +++ b/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()