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

[rcano-baseten]

🚀 What

After truss train push, users currently have to know the SSH hostname format (training-job-<job_id>-<node>.ssh.baseten.co) and manually poll truss train view / truss train logs to know when SSH is reachable.

This PR adds two things:

1. Per-node SSH lines in push output. _print_training_job_success_message now appends ssh training-job-<id>-<n>.ssh.baseten.co for each node (labels node 0 (leader) for multi-node), plus a hint about truss ssh setup.
2. truss ssh --training-job-id <id> [--node-id 0]. Waits for the job to be running, then execs ssh so you land directly in the shell. No -A.

💻 How

• truss/cli/train_commands.py: small _format_ssh_commands(job_id, node_count) helper, and added an SSH section to the existing success message. Format mirrors the one already used in truss train workstation.
• truss/cli/ssh_commands.py: converted the ssh Click group to invoke_without_command=True so it accepts options directly. Flow:
  1. Pre-flight is_setup_complete() (already in truss/cli/ssh.py); errors with a truss ssh setup hint if missing.
  2. Resolve project_id via the existing get_most_recent_job(remote_provider, None, job_id) — same helper truss train logs uses.
  3. Read node_count from the job's instance_type and reject out-of-range --node-id.
  4. Reuse TrainingPollerMixin.before_polling() for the wait loop — same primitive truss train logs --tail uses. Spinner messages already handle "Waiting for GPU capacity (job is queued)…" / "Waiting for job to deploy…".
  5. Once TRAINING_JOB_RUNNING, os.execvp("ssh", ["ssh", "training-job-<id>-<n>.ssh.baseten.co"]).

Out of scope

--deployment-id for inference models — needs model_id upfront, no clean reverse lookup today. Will be a follow-up.

🔬 Testing

• Added 7 unit tests in truss/tests/cli/test_ssh.py:
  • _format_ssh_commands for single-node and multi-node.
  • truss ssh --training-job-id: correct hostname, default node-id is 0, errors when SSH setup is incomplete, errors on out-of-range --node-id, errors when job is not running.
• uv run ruff check clean on touched files.
• The 5 pre-existing TestSetupSSHConfig failures on main (env can't find python3.x via shutil.which) are unrelated.

End-to-end manual verification (to do against a real remote):

• truss train push <config> prints the SSH lines (1 line for single-node, 2+ for multi-node).
• truss ssh --training-job-id <id> against a QUEUED job shows the spinner and updates as state changes.
• Once RUNNING, lands in the shell on the right node.
• --node-id 9 (out of range) gives a clean error.
• Without truss ssh setup, gives the setup hint.

🤖 Generated with Claude Code

[CLAassistant]

All committers have signed the CLA.

[cretz]

I think this should work for inference too, though in general we recommend users use their own SSH clients configured how they want, with the args they want, etc. Is this meant to be our newly suggested way to SSH into training/inference containers, or is it just a helper?