feat(cli): add JSON output to truss train workstation

[brollb]

🚀 What

Adds -o/--output-format json to truss train workstation, so scripts can read the provisioned node SSH addresses programmatically instead of scraping the rich-formatted text block.

$ truss train workstation --node-count 2 -o json
{
  "job_id": "abc1234",
  "project": {"id": "proj123", "name": "workstation-H100"},
  "accelerator": "H100",
  "gpu_count": 8,
  "node_count": 2,
  "orchestrator": "slurm",
  "nodes": [
    {"rank": 0, "hostname": "training-job-abc1234-0.ssh.baseten.co", "is_leader": true},
    {"rank": 1, "hostname": "training-job-abc1234-1.ssh.baseten.co", "is_leader": false}
  ],
  "job": { ... raw push response ... }
}

Text output is unchanged apart from the remote-segment fix noted below.

💻 How

Follows the @json_command pattern already used by truss push (truss/cli/cli.py:748) rather than the read-only commands' cli-table|csv|json branch, because this command mutates state and emits output it doesn't own (truss_train/deployment.py push progress, --tail log streaming). The flag is spelled -o/--output-format to match the rest of the train group; json_command keys off the output_format destination, so the two compose.

Three details worth review:

• suppress_error_print. RestAPIClient._handle_error prints 4xx messages with a bare print() to stdout, which console_to_stderr() cannot redirect. Both upsert_training_project and create_training_job go through that client, so JSON mode sets suppress_error_print = True and lets the error surface as an exception for json_command to render — same as truss push does.

• JSON is printed before --tail, with an explicit flush. truss push --watch --output json prints its payload after watch mode returns. That ordering doesn't work here: the addresses are the point of the command and tailing may never end. Printing first means the payload can sit in a block-buffered pipe indefinitely, hence flush=True.

• Hostnames now include the remote segment for non-default remotes. Construction moved into workstation_ssh_hostnames() so text and JSON render from one list. The old inline string always omitted the remote, so a workstation launched with --remote dev printed an address that resolve_remote() can only resolve if the user has exactly one remote or has stamped dev as their default — otherwise it errors out. The segment is still omitted for the default baseten remote to keep the common case short.

The payload also carries the raw push response under job, so consumers aren't blocked on us adding fields.

Note the nodes are not reachable when the JSON is emitted — the job is queued, not running. Callers still need to poll status before connecting. Happy to add a --wait-for-running flag in a follow-up if scripts end up reimplementing that poll.

🔬 Testing

uv run pytest truss/tests/cli -q → 591 passed.

New tests in truss/tests/cli/train/test_workstation.py:

• workstation_ssh_hostnames() unit tests: default remote omitted, non-default remote included, one hostname per node in rank order.
• Single-node and multi-node JSON payloads (ranks, is_leader, orchestrator null for single-node, raw job passthrough).
• stdout parses as one JSON document with progress output confirmed on stderr.
• suppress_error_print set in JSON mode.
• A failing push yields {"error": {"message": ...}} on stdout with exit 1.
• Text mode output unchanged by default.

🤖 Generated with Claude Code

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

Brian Broll seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}