Add --availability-model to truss train update - #2634
Open
johnnt849 wants to merge 3 commits into
Open
Conversation
A queued training job's capacity guarantee can now be changed in place:
truss train update --job-id <id> --availability-model spot
Previously a dedicated job blocked on full capacity had to be resubmitted to
run on spot. It can now be switched while it sits in the queue, alongside the
existing --priority update.
Follows the same three layers as the priority update: the option on the
`train update` command, the enum-typed parameter in cli/train/core.py, and the
PATCH body field in remote/baseten/api.py. Choices are derived from the
AvailabilityModel enum so they cannot drift from it, and a choice option rather
than a bare --spot flag (as `truss train push` has) since an update also needs
to switch a job back to dedicated.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
TestUpdateTrainingJob in truss-train/tests/test_recreate.py asserts the exact kwargs passed to the API, so adding availability_model to the call broke test_update_training_job_success. Fix the assertion, and move the two availability-model cases from truss/tests/cli/train/test_train_cli_core.py into TestUpdateTrainingJob beside their siblings — the no-fields case there was an exact duplicate of one already in that class. All five update_training_job call assertions now live in one file, so a signature change can't pass one suite while breaking the other. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
cretz
reviewed
Aug 31, 2026
Contributor
There was a problem hiding this comment.
Setting reminder to also add this to https://github.com/basetenlabs/baseten-cli (unless you want to)
Contributor
Author
There was a problem hiding this comment.
I'll add a PR for that in a bit!
test_recreate.py is about job recreation, and it lives in truss-train/tests/ even though it exercises truss.cli.train.core. TestUpdateTrainingJob moves to truss/tests/cli/train/test_job_update.py, which names what it tests and sits in the tree for the package under test. All five update assertions stay together in the new file, so a change to the call signature still can't pass one suite while breaking another. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🚀 What
Lets a queued training job's capacity guarantee be changed in place:
Previously a job queued on dedicated capacity had to be cancelled and resubmitted to run on spot. It can now be switched while it sits in the queue, alongside the existing
--priorityupdate. Priority and availability model can be set in one call.The backing
PATCH /v1/training_projects/{project_id}/jobs/{job_id}endpoint acceptsavailability_model(dedicated/spot); only queued jobs can be updated, and the server rejects anything else.💻 How
Follows the same three layers as the existing
--priorityupdate:truss/cli/train_commands.py— the--availability-modeloption ontrain update. Choices are derived from theAvailabilityModelenum rather than hardcoded, so they cannot drift from it. The "at least one field" guard now covers both fields.truss/cli/train/core.py—update_training_jobtakesavailability_model: Optional[AvailabilityModel]and unwraps it to its wire value at the API boundary.truss/remote/baseten/api.py— addsavailability_modelto the PATCH body only when set, matching howpriorityis handled. Typed asstrhere rather than the enum, sincetruss/remote/sits belowtruss_trainin the layering.A choice option rather than a bare
--spotflag (astruss train pushhas), because an update also needs to switch a job back todedicated.🔬 Testing
56 tests pass;
ruff check/ruff formatclean; mypy clean via pre-commit.truss/tests/remote/baseten/test_api.py— availability-only and priority+availability request bodies.truss/tests/cli/train/test_train_cli_core.py— enum → wire-value conversion, both fields forwarded together, and the no-fieldsValueErrorraised before any API call is made.Verified
truss train update --helprenders the new option.🤖 Generated with Claude Code