Add --availability-model to baseten train job update by johnnt849 · Pull Request #69 · basetenlabs/baseten-cli

johnnt849 · 2026-08-31T22:23:15Z

What

Lets a queued training job's capacity guarantee be changed in place:

baseten train job update --job-id p7qr9qv --availability-model spot

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 --priority. Both can be set in one call.

Brings the command to parity with truss train update --availability-model, per the reminder on basetenlabs/truss#2634.

⚠️ Blocked on the SDK

Depends on basetenlabs/baseten-go#34, which regenerates the client against the current production spec. That PR changes UpdateTrainingJobRequest.Priority from int to *int and adds AvailabilityModel, so the previous struct literal here no longer compiles.

go.mod currently pins the SDK branch commit (bc949ae) so this builds and CI runs. Re-point it at a released version once baseten-go#34 merges, before merging this.

How

cmd/command.train.go — adds AvailabilityModel string with enum:"dedicated,spot" to TrainJobUpdateFlags, so unknown values are rejected before any request is made. Drops required:"true" from Priority since either field may now be sent alone. Command summary/description and examples updated — it was previously described as priority-only.
internal/cmd/command.train.go — sends whichever fields were given, and reports what actually changed instead of the previous hardcoded "Set training job %s priority to %d", which would have lied on an availability-only update. Passing neither field is a usage error, matching the API's own rejection of an empty body.

One subtlety worth flagging for review: whether --priority was given is read from ctx.Command.Flags().Changed("priority") rather than inferred from the value. 0 is a valid priority, so the usual "zero means unset" shortcut in this repo (see trussIntArg) would silently swallow an explicit --priority 0. There is a test pinning that.

Testing

go build ./..., go vet ./..., gofmt clean, and go test ./... green across every package.

New cases in internal/cmd/command.train_test.go:

availability-only — sends availability_model, omits priority rather than sending 0
priority + availability together
explicit --priority 0 reaches the API
neither flag → usage error, no request issued
--availability-model bogus → rejected by the enum

🤖 Generated with Claude Code

A queued training job's capacity guarantee can now be changed in place: baseten train job update --job-id p7qr9qv --availability-model spot Previously a job queued on dedicated capacity had to be resubmitted to run on spot. --priority loses `required` since either field may now be sent on its own, and passing neither is a usage error. Whether --priority was given is read from the flag set rather than inferred from its value: 0 is a real priority, so treating the zero value as "unset" would silently drop an explicit --priority 0. Requires the SDK bump in this commit: UpdateTrainingJobRequest.Priority became *int and gained AvailabilityModel, so the previous struct literal no longer compiles. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

CLAassistant · 2026-08-31T22:23:32Z

All committers have signed the CLA.

Copilot

Pull request overview

Adds support for updating a queued training job’s availability model (dedicated/spot) via baseten train job update, alongside the existing --priority, bringing the CLI in line with the Truss command surface.

Changes:

Extend train job update to accept --availability-model (enum: dedicated|spot) and update help text/examples.
Send a PATCH request with only the fields that were explicitly provided, and improve the non-JSON log message to reflect what changed.
Add tests covering availability-only updates, combined updates, explicit --priority 0, and usage/validation errors.

Reviewed changes

Copilot reviewed 4 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
`cmd/command.train.go`	Updates CLI command help text and adds the `--availability-model` flag with enum validation.
`internal/cmd/command.train.go`	Builds a PATCH request with only explicitly-set fields and logs the actual changes.
`internal/cmd/command.train_test.go`	Adds coverage for availability-model updates and edge cases like `--priority 0` and empty updates.
`go.mod`	Bumps `baseten-go` dependency to a newer pseudo-version to pick up the regenerated client changes.
`go.sum`	Updates checksums to match the `baseten-go` version bump.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

@@ -731,12 +731,32 @@ func commandTrainJobUpdate(ctx *CommandContext, flags *cmd.TrainJobUpdateFlags)
 	if err != nil {
 		return err
 	}


Addresses PR feedback: a no-fields invocation built the management client first, so `baseten train job update --job-id x` with a broken remote or profile reported that config failure instead of the usage error explaining what was actually wrong with the command. The regression test uses a malformed BASETEN_REMOTE_URL, which fails in AuthTransport at construction. An empty BASETEN_API_KEY does not work as a trigger here: the client sets DeferAuth, so a missing credential surfaces at request time rather than when the client is built. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

go get left the superseded baseten-go entries behind alongside the new ones, so CI's `go mod tidy && git diff --exit-code` gate failed on the two stale lines. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

cretz

Nothing blocking, though merge basetenlabs/baseten-go#34 first and update go.mod here to that commit hash before merging.

cretz · 2026-09-01T15:23:41Z

+	setPriority := ctx.Command.Flags().Changed("priority")
+	setAvailability := ctx.Command.Flags().Changed("availability-model")
+	if !setPriority && !setAvailability {
+		return cmd.NewErrUsagef("pass at least one of --priority or --availability-model")
+	}


Did this command work before when --priority was not passed? If so, is this backwards incompatible? (not that we mind in CLI, just want to understand)

Before --priority was required but this is backwards compatible because it is less strict than before.

cretz · 2026-09-02T12:48:32Z

 	github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager v0.1.21
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.101.0
-	github.com/basetenlabs/baseten-go v0.2.1-0.20260828192113-f028e27beb4b
+	github.com/basetenlabs/baseten-go v0.2.1-0.20260831221848-bc949ae00b53


baseten-cli main has a reference to baseten-go with what is needed now, can merge main and fix conflicts (accept "theirs" for go.mod and go.sum, and "mine" in command.train.go where I had to ref a pointer for priority to build)

main already references a baseten-go that carries the regenerated UpdateTrainingJobRequest (baseten-go#28, e30c99e), so go.mod and go.sum are taken from main and this branch no longer pins an unmerged SDK commit. The command.train.go conflict keeps this branch's version: it already builds the request field by field with pointers, which supersedes main's single-line `Priority: &flags.Priority` fix and additionally sends availability_model. main made no other change to that file. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…lability-model

Copilot AI lite review requested due to automatic review settings August 31, 2026 22:23

Copilot started reviewing on behalf of johnnt849 August 31, 2026 22:23 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Comment thread internal/cmd/command.train.go Outdated

Comment on lines 729 to 733

@@ -731,12 +731,32 @@ func commandTrainJobUpdate(ctx *CommandContext, flags *cmd.TrainJobUpdateFlags)

if err != nil {

return err

}

johnnt849 and others added 2 commits August 31, 2026 16:05

Tidy go.sum after the baseten-go bump

694ef5f

go get left the superseded baseten-go entries behind alongside the new ones, so CI's `go mod tidy && git diff --exit-code` gate failed on the two stale lines. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

cretz approved these changes Sep 1, 2026

View reviewed changes

cretz reviewed Sep 2, 2026

View reviewed changes

johnnt849 and others added 2 commits September 2, 2026 11:38

Merge remote-tracking branch 'origin/main' into claude/train-job-avai…

2f7d651

…lability-model

johnnt849 merged commit ee23e69 into main Sep 3, 2026
6 of 7 checks passed

johnnt849 deleted the claude/train-job-availability-model branch September 3, 2026 21:35

Navigation Menu

Uh oh!

Add --availability-model to baseten train job update - #69

Add --availability-model to baseten train job update#69
johnnt849 merged 5 commits into
mainfrom
claude/train-job-availability-model

johnnt849 commented Aug 31, 2026

Uh oh!

CLAassistant commented Aug 31, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

cretz left a comment

Uh oh!

cretz Sep 1, 2026

Uh oh!

johnnt849 Sep 2, 2026

Uh oh!

cretz Sep 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

johnnt849 commented Aug 31, 2026

What

⚠️ Blocked on the SDK

How

Testing

Uh oh!

CLAassistant commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

cretz left a comment

Choose a reason for hiding this comment

Uh oh!

cretz Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

johnnt849 Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

cretz Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

CLAassistant commented Aug 31, 2026 •

edited

Loading