Add truss train exec to run a local directory as a Baseten training job - #2629
Add truss train exec to run a local directory as a Baseten training job#2629claude[bot] wants to merge 8 commits into
Conversation
`truss train exec -- <command>` tars the directory the command is invoked from, ships it to a Baseten training job, and runs the command there with an on-demand SSH session available. The builder lives in truss/cli/train/exec.py, mirroring the existing workstation builder, and the CLI command sits next to `truss train workstation` in train_commands.py. Everything after `--` is passed through verbatim as the job's single start command. Unlike workstation, the job runs from the invocation directory (source_dir=Path.cwd()) and gets no cache and no checkpointing. CPU-only is the default; --accelerator opts into a GPU. --with-uv makes uv available in the job image so the command can invoke `uv run` itself.
|
|
rich_click renders errors and warnings as colorized, width-wrapped panels and turns color on under GITHUB_ACTIONS, so `--gpu-count requires --accelerator` reaches result.output with each `--flag` token wrapped in ANSI codes. Two assertions matched the plain substring and failed on every CI matrix entry while passing locally, where rich emits no color. Compare against output with ANSI codes and panel borders stripped and whitespace collapsed, so the assertions hold regardless of terminal width or color support.
The comment claimed on-demand "avoids imposing" a session timeout, but the builder passes no timeout and InteractiveSession.timeout_minutes defaults to 480, so a timeout is populated and sent regardless. State the actual reason on-demand is fine here -- the timeout applies once the job ends -- and that the model default still applies.
CI failure on
|
| attempt 1 | attempt 2 (re-run) | |
|---|---|---|
| job | 98394050633 | 98394459884 |
Install dependencies |
failure, exit 138 | failure, exit 138 |
Test bindings |
skipped | skipped |
| runner | GitHub Actions 1000807291 | GitHub Actions 1000807305 |
Why it is not this PR's:
- The PR's entire diff against
mainis three Python files:truss/cli/train/exec.py,truss/cli/train_commands.py,truss/tests/cli/train/test_exec.py. It touches nothing underbaseten-performance-client/and nothing under.github/— nopackage.json, noyarn.lock, no.rs/.js/.ts, no workflow file. - The head commit
57e0407is comment-only: 2 insertions, 1 deletion, all of them comment lines inexec.py. No executable line changed. - This exact job — same name, same
macos-latesthost, same node 18 — passed on the immediately previous headb9ffaba15 minutes earlier, with every step green includingInstall dependenciesandTest bindings. Test bindings on x86_64-apple-darwin - node@20— the same target and the sameyarn installstep, in the same workflow run — passed. So did the other 15 binding-test jobs.
A Python comment cannot affect yarn install in baseten-performance-client/node_bindings.
Why the re-run reproduced it identically. Setup node uses actions/setup-node@v4 with cache: yarn, keyed on baseten-performance-client/node_bindings/yarn.lock plus runner OS and node version — all inputs this PR leaves untouched. Re-running failed jobs within the same workflow run restores the same cache archive, so the retry is not an independent sample; it replays the same restored state. The node@20 job has a different cache key (the key includes the node version), which is consistent with it passing. The job also sets architecture: x64 on macos-latest, which is arm64, so an x64 Node 18 runs under Rosetta — SIGBUS is a characteristic crash signature there.
What I did: re-ran the single failed job once (attempt 2). It failed the same way, in the same step, with the same exit code. I have not re-run again.
Suggested next step (not taken here): this needs a fresh workflow run rather than another re-run of run 33034240537 — a new push to the branch will do it, or clear the setup-node/yarn cache entry for macos-latest + node 18 so the install step is not fed the same restored archive. I did not push an empty commit, and I changed no code, since there is no causal link to this diff.
For the record, main at 7d4daa1 does not run this workflow at all (it is path-filtered to baseten-performance-client/), so there is no base-branch comparison for this job; main does carry 5 unrelated pre-existing failures (4 × all-tests / truss-integration-tests, plus prebuild).
Generated by Claude Code
Nothing validated SecretReference names before push, so a typo surfaced only once the job failed to start. Look the names up via the existing (previously uncalled) BasetenApi.get_all_secrets and warn, pointing at the settings page, which is the only way to create a workspace secret. Warn-only and fail-open on purpose. GET v1/secrets takes no team scoping while training projects are team-scoped, so a name absent from the listing is not proof it is absent for the job; any error or an unreadable payload skips the check silently. The lookup is skipped entirely when no --secret was passed. The response shape is not pinned down by any test or doc in the repo, so the parser accepts the plausible shapes and returns None -- meaning "don't check" -- rather than guessing and warning about a secret that is really there. Secret names reach rich escaped, since a stray bracket would be read as markup.
- exec now exits non-zero when a tailed job ends in a failure state, so `truss train exec -- pytest` is not green in CI regardless of outcome. TrainingPollerMixin gains a public `failed` property rather than callers reaching for `_current_status`. Uses sys.exit, not click's Exit, which subclasses RuntimeError and would be caught by common_options' error handler. - The uv install no longer pipes curl into sh: a pipeline reports the last command's status, so a failed download fed an empty script to a shell that exited 0 and the job died later with an opaque `uv: not found`. - `--secret KEY=` is rejected instead of yielding SecretReference(name=""). An empty --env value stays legal. - The secret lookup's except now wraps only the API call, so a parser bug surfaces instead of looking like an unreachable API. - The launch line escapes interpolated text; a directory named `myproj[v2]` was reported as `myproj`, and `[/cyan]` in a path raised MarkupError. - The missing-secret warning no longer asserts absence, since GET v1/secrets takes no team parameter while training projects are team-scoped. - `--cpu-count` uses IntRange(min=1); Compute(cpu_count=-4) constructed fine. - uv detection requires a uv.lock or a [tool.uv] section, not just any pyproject.toml, which poetry, hatch, PDM and setuptools all ship. - uv detection and the new client-side --workspace-root validation both look at the directory that actually gets archived, so a bad root fails before the training project is created rather than leaving a stray empty one.
The motivating use case is a client that runs for hours, where blocking the terminal is the wrong default. This also lines exec up with `truss train push` and `truss train workstation`, where --tail is opt-in. The non-zero exit on a failed job is unchanged; it applies when --tail is passed, which is the narrower CI case that motivated it. Kept as the paired --tail/--no-tail form so --no-tail still parses and --help shows the default.
rcano-baseten
left a comment
rcano-baseten
left a comment
There was a problem hiding this comment.
left some thoughts
|
|
||
| if accelerator: | ||
| accelerator = accelerator.upper() | ||
| gpu_count = gpu_count or 1 |
There was a problem hiding this comment.
this means gpu count will default to 1 - should that be the case when a user uses compute type?
There was a problem hiding this comment.
Agreed — the CLI computed every one of those before calling, so the defaults were never exercised and just duplicated the click declarations. build_exec_project is now keyword-only with no defaults at all. Optional survives only where None is a value rather than "unspecified": accelerator (CPU-only), base_image (derive it), project (plain image, no setup steps), workspace_root (archive the invocation directory).
DEFAULT_CPU_COUNT / DEFAULT_MEMORY now have exactly one home as the click defaults, still read from Compute.model_fields so they can't drift from the model. Keyword-only seemed worth it at twelve parameters. Tests use a _build(**overrides) helper that supplies a complete baseline, so each test overrides only what it exercises, and two new tests assert a partial or positional call raises TypeError.
Generated by Claude Code
There was a problem hiding this comment.
Apologies — my previous reply on this thread was misplaced; it belongs to your build_exec_project signature comment and is now posted there. This --gpu-count question is still open and is being answered separately.
Generated by Claude Code
| def warn_about_missing_secrets( | ||
| api: BasetenApi, environment_variables: Mapping[str, Union[str, SecretReference]] | ||
| ) -> None: | ||
| """Warn when a `--secret` names a secret that isn't in the workspace listing. |
There was a problem hiding this comment.
I believe we should hard error here. If we can't see the secrets at all, we know the job will fail
There was a problem hiding this comment.
is there a team-scoped secrets api at all? we shoudl use that - i think otherwise, this set of secrets is a superset of the team secrets - lmk if that's not the case - and so we should fail closed if we can't find the secret
There was a problem hiding this comment.
Changed to a hard error (click.UsageError, no push) when a named secret is absent from a listing we successfully read.
On your question: no, there is no team-scoped secrets API in the client — only GET v1/secrets (get_all_secrets, which takes no arguments at all) and POST v1/secrets, neither with a team_id, and both were uncalled before this PR. Compare upsert_training_project, which does branch to v1/teams/{team_id}/training_projects. So there is nothing team-scoped for me to switch to. Whether the server offers one, and whether the workspace listing is genuinely a superset of a team's secrets, isn't knowable from this repo — your reading is plausible and matches the naming, but the code can't confirm it, so someone with backend context should. The error text says the listing isn't team-scoped and asks the user to report a mismatch, and I'll swap it for a team-scoped route the moment one exists.
One deliberate split I'd flag so you can overrule it: I kept fail-open when the listing call itself fails or returns a payload we can't parse. That's an API error rather than evidence the secret is missing, so hard-erroring there would break the command on a blip or a permissions quirk. Absent-from-a-readable-listing is the case that now fails closed.
Generated by Claude Code
exec.py becomes a package with one job per module: builder.py assembles the TrainingProject, secrets.py owns --env/--secret, uv.py holds everything that knows what uv is, project.py defines the Project protocol and get_project_type. __init__.py is a thin re-export, matching deploy_checkpoints/ next door. A Project protocol now supplies the base image a project type wants and the setup steps it contributes, with UvProject as the first implementation and get_project_type detecting it. No build() -- nothing is built client-side. build_start_commands no longer knows about uv. Behavior is unchanged: --with-uv still selects UvProject directly and only affects the image, the command still runs last and verbatim, and no uv sync is injected. A --secret naming a secret absent from a listing we successfully read is now a hard error instead of a warning. A failed or unparseable listing still proceeds: that is an API problem, not evidence the secret is missing. build_exec_project is keyword-only with no defaults, so a caller cannot build a partially-specified project and the CLI stays the single home for defaults. Also dropped rich escaping from the secret error: click.UsageError messages are not markup-parsed, so escaping leaked a literal backslash into the name shown to the user.
`value` was bound as str from _parse_key_value_flag, then rebound from the entries list as Union[str, SecretReference], so mypy rejected the assignment. It only surfaced in the whole-repo run, where truss_train is in the file set and SecretReference resolves precisely; checking the file alone passes either way.
3 participants
Requested by Raymond Cano · Slack thread
Linear: LPS-1193
🚀 What
Before: running a local client remotely meant doing it by hand — write a
config.pydescribing a training project, point it at a base image and compute, wire up a start command, thentruss train push config.py. So people just ran the thing on their laptop instead, holding thousands of connections open from a machine that sleeps, and only assembled a real training job when a run mattered enough to justify the setup.After:
truss train exec -- python my_script.pytars the directory you are standing in, ships it to a Baseten training job, and runs that command there. No config file, no template to copy; the flags cover the parts you would otherwise have written out, and SSH into the job is available on demand.One sentence:
truss train execlaunches the current directory as a Baseten training job running a command you supply, with no persistent storage and no checkpointing.The motivating case is running a Loops client remotely instead of on a laptop, but nothing about the command is Loops-specific.
💻 How
truss/cli/train/exec/is a package with one job per module —builder.pyassembles theTrainingProject,secrets.pyowns--env/--secret,uv.pyholds everything that knows what uv is,project.pydefines theProjectprotocol andget_project_type, and__init__.pyis a thin re-export (matchingdeploy_checkpoints/next door). A project type supplies the base image it wants and the setup steps it contributes, sobuild_start_commandsno longer knows about uv and adding pip or poetry means a sibling module plus a branch inget_project_type. The CLI command sits next totruss train workstationintruss/cli/train_commands.pyand stays thin;build_exec_projectis keyword-only with no defaults, so a caller cannot build a partially-specified project.START_COMMANDis collected withnargs=-1+click.UNPROCESSEDunderignore_unknown_options,shlex.joined, and becomes the job's single start command. The push goes throughtruss_train.public_api.pushwithsource_dir=Path.cwd()— the one real divergence from workstation, which deliberately pushes a temp dir. Team resolution,--remote,--tailand the printed connect hint follow workstation.Flags:
--accelerator(choice, default none → CPU-only),--gpu-count(1–8, default 1, requires--accelerator),--cpu-count/--memory(defaulting to theComputemodel's own defaults),--project-name,--image,--workspace-root/--exclude-dir/--external-dir,--env KEY=VALUE,--secret KEY=SECRET_NAME,--with-uv,--remote,--team,--tail/--no-tail.--envand--secretfillRuntime.environment_variables— plain strings andSecretReference(name=...)respectively. This matters because a training job gets no Baseten credential automatically;SecretReferencewas config-file-only until now, with no CLI flag on anytruss traincommand, which would have leftexecunable to run anything that calls the Baseten API. Both split on the first=only, so values containing=survive; a key set by both is an error.--secretnames are also checked against the workspace before the push, warning (never blocking) when one is missing.--with-uvmakes uv available in the job image and nothing else — the user writesuv runthemselves, e.g.truss train exec --with-uv -- uv run python my_script.py. On the CPU path it selects the official uv image; on the GPU path (and for any user-supplied--image, since we cannot know whether it ships uv) it prepends a skip-if-present install step,{ command -v uv || curl -LsSf https://astral.sh/uv/install.sh | sh ; }, so it is a no-op on an image that already has uv. Without the flag: plainpython:3.12-slimon the CPU path, nothing injected, command verbatim. If the flag is absent but the directory has auv.lockorpyproject.toml, the command warns before pushing that uv will not be in the image — it does not fail and does not auto-enable, because auto-detection would make the same command behave differently depending on directory contents.The upload already excludes
.venv(packagedtruss/util/.truss_ignore), so the local virtualenv is never shipped and the job builds its own environment.Base image, by path:
--with-uvpython:3.12-slim--with-uvghcr.io/astral-sh/uv:0.12.6-python3.12-trixie-slim--with-uv--with-uv--image X, no--with-uvX--image X+--with-uvXThe GPU image is not a single tag:
default_base_imagedelegates toworkstation.default_base_image, which returnsnvidia/cuda:13.0.3-devel-ubuntu24.04forB300andGB300andnvidia/cuda:12.8.1-devel-ubuntu24.04for every other accelerator. GPU jobs therefore inherit whateverworkstationpicks, by design, rather than pinning their own tag here.The two CPU tags are pinned and were verified to exist in their registries. uv stopped publishing
bookworm-slimvariants for current versions, so the pinned uv tag is thetrixie-slimone.Reviewer decision points
InteractiveSession(trigger=ON_DEMAND, session_provider=SSH), whereworkstationusesON_STARTUP(workstation.py:94-97). To be precise about what is sent: the builder sets no timeout, butInteractiveSession.timeout_minutescarries a model default of 480 (definitions.py:164), so the built object still populates and sendstimeout_minutes=480, along withauth_provider=MICROSOFT. Per the platform owner the timeout applies once the job ends rather than governing a running job, so underON_DEMANDit is not a concern for a long-running job — but a value is populated either way, not omitted.truss-train/truss_train/definitions.pyis untouched; widening that field's type to make omission expressible was considered and rejected..../auth_codes(truss train isession) and PATCH of the session's trigger/timeout — and the SSH cert path (sign_ssh_certificate) does not check session state client-side. So whetherssh training-job-JOB_ID-0.ssh.baseten.coworks immediately for an ON_DEMAND job, or needs the session requested first, is not knowable here. The hint points attruss train isession --job-id JOB_IDbefore thesshline, which is accurate either way; someone with backend context should confirm the wording.--is required in practice.ignore_unknown_optionsmeans a mistyped flag before--is silently folded into the start command rather than rejected. That is inherent to variadic pass-through; the help and docstring both show the--form.Runtime.start_commands, or only the first? I could not establish this from the repo. Every producer emits exactly one entry —workstation.py:79,81, the--entrypointoverride atdeployment.py:386(which collapses the list to[entrypoint]), all fivetruss-train/tests/import/fixtures — andRuntime.start_commandsis a bareList[str](definitions.py:172) passed straight to the backend with no client-side interpretation. There is no schema, contract, or doc describing multi-entry semantics. Where the repo itself needed two sequenced steps it chained them with&&inside one/bin/sh -centry (truss/templates/train/config.py:18-20), so the uv-install path does the same. If multiple entries do execute, that path could be split into two entries for cleaner log separation.--project-name, not--project-id.workstationcalls the same thing--project-ideven though it is used as the project name (workstation.py:106doesTrainingProject(name=project_id)), and--project-idelsewhere in the CLI means an actual ID.--project-nameis named for what it is; say the word if consistency withworkstationis preferred. It defaults to the current directory's name so repeated runs from the same checkout group into one project.--secretnames are looked up viaBasetenApi.get_all_secretsbefore the push. A name absent from a listing we successfully read fails the command (click.UsageError, no push), since the job would fail to start anyway. A failed or unparseable listing still proceeds — that is an API problem, not evidence the secret is missing. Because this can now block a valid run, the caveat matters more: there is no team-scoped secrets endpoint in the client (onlyGET/POST v1/secrets, neither takingteam_id, versusv1/teams/{team_id}/…for training projects), so whether the workspace listing is genuinely a superset of a team's secrets is not knowable from this repo and wants confirmation from someone with backend context. The error text says so and asks the user to report a mismatch.--tailis opt-in, matchingpushandworkstation— the motivating use case is a client that runs for hours, where blocking the terminal is the wrong default. When it is passed,execchecks the poller's terminal status and exits non-zero onTRAINING_JOB_FAILED/TRAINING_JOB_DEPLOY_FAILED, sotruss train exec --tail -- pytestis not green in CI regardless of outcome;truss train push --taildoes not do this, and that divergence remains and is deliberate.--workspace-rootis validated client-side.truss_train's own containment check runs insidepush, afterupsert_training_project, so a bad value leaves a stray empty project behind;execchecks before pushing. That ordering insidepushis pre-existing and unchanged — onlyexec's flag is covered._handle_post_create_logic,workstationandexec; extracting a shared helper would mean editing two existing commands, so it belongs in its own change rather than here.--with-uvcannot point at a package that lives below the invocation root.--workspace-rootis the natural place to build that later.🔬 Testing
New
truss/tests/cli/train/test_exec.py— 107 tests, builder and CLI. Builder: CPU defaults and image, GPU path, every supported accelerator, on-demand SSH fields, cache/checkpointing left off, workspaceNoneunless a directory flag is passed,--env/--secretparsing (values containing=, missing=, duplicate keys,SecretReferencerather than a literal), and the exactstart_commandsfor each uv combination — no flag, flag on the uv image, flag on the GPU image, flag on a custom image — asserting the user's command is last and unmodified in every case. CLI:--pass-through (including a case where the command reuses our own--memory/--tailflag names, proving they are not consumed by the parser), missing start command,--gpu-countwithout--accelerator, accelerator normalization,source_direqual to the cwd, project-name default and override,--imagewinning over both uv paths, the uv warning firing only for a real uv project, tail opt-in and the non-zero exit on a failed job, markup escaping in the launch line,--cpu-countbounds, and--workspace-rootrejection. The three fixes most worth guarding were mutation-checked: reverting the escape, the exit check, or thefailedguard each makes a specific test fail.CLI message assertions compare against ANSI-stripped, whitespace-collapsed output: rich_click renders errors and warnings as colorized, width-wrapped panels and turns color on under GITHUB_ACTIONS, which splits
--flagtokens with escape codes, so a plain substring match passes locally and fails on every CI runner.uv run ruff check,uv run ruff format, anduv run mypyare clean on all three changed files. Full non-integration suite: 1772 passed, 26 skipped (uv run pytest --durations=0 -m 'not integration'), run both normally and withGITHUB_ACTIONS=trueto match CI's color handling, withtest_workstation.py,test_cli.pyandtest_loops_cli.pygreen. The generated uv-install shell string was verified to parse as valid POSIXsh(sh -n) with the command's own quoting intact.pre-commit's hygiene hooks pass. Its three local hooks (
ruff,ruff-format,mypy) could not run in my sandbox — they shell out to bareuv run, which tries to fetch the CPython pinned in.python-versionfrom GitHub releases and gets a 403 through the environment's proxy. I ran those same three tools directly against the system interpreter instead, clean; CI will run the hooks for real.🤖 Generated with Claude Code
https://claude.ai/code/session_012vALt4HhhRV2VtW4Yxv53X