Add --availability-model to baseten train job update

cmd/command.train.go

@@ -421,17 +421,25 @@ var commandTrainJob = Command{
 		},
 		{
 			Name:    "update",
-			Summary: "Update a training job's queue priority",
-			Description: "Update a queued training job's priority. Higher values are dequeued first.\n\n" +
-				"Only jobs in the PENDING state can have their priority changed.",
+			Summary: "Update a queued training job's priority or capacity",
+			Description: "Update a queued training job's queue priority, its availability model, or both. " +
+				"Higher priorities are dequeued first, and the availability model selects the capacity " +
+				"the job is admitted under, so a job waiting on full dedicated capacity can be moved to " +
+				"spot without resubmitting.\n\n" +
+				"Only jobs in the PENDING state can be updated. Pass at least one of --priority or " +
+				"--availability-model.",
 			Flags: TrainJobUpdateFlags{},
 			Output: &CommandOutput[managementapi.UpdateTrainingJobResponse]{
-				TextDescription: "On success, prints the job's ID and new priority to stderr; no stdout output.",
+				TextDescription: "On success, prints the job's ID and what changed to stderr; no stdout output.",
 				Examples: []CommandExample{
 					{
 						Description: "Raise a queued job's priority.",
 						Command:     "baseten train job update --job-id p7qr9qv --priority 10",
 					},
+					{
+						Description: "Move a queued job onto spot capacity so it can dequeue sooner.",
+						Command:     "baseten train job update --job-id p7qr9qv --availability-model spot",
+					},
 				},
 				JQExample: CommandExample{
 					Description: "Print the job's new priority.",
@@ -742,7 +750,8 @@ type TrainJobUpdateFlags struct {
 	CommandFlags
 	TrainJobRefFlags
 
-	Priority int `flag:"priority" desc:"New queue priority. Higher values are dequeued first." required:"true"`
+	Priority          int    `flag:"priority" desc:"New queue priority. Higher values are dequeued first."`
+	AvailabilityModel string `flag:"availability-model" desc:"New capacity guarantee. 'dedicated' runs on on-demand capacity that is not preempted; 'spot' runs on interruptible capacity that may be preempted, and checkpointing your own progress is up to you." enum:"dedicated,spot"`
 }
 
 // TrainJobDownloadFlags configures `baseten train job download`.

internal/cmd/command.train.go

@@ -727,16 +727,40 @@ func commandTrainJobRecreate(ctx *CommandContext, flags *cmd.TrainJobRecreateFla
 }
 
 func commandTrainJobUpdate(ctx *CommandContext, flags *cmd.TrainJobUpdateFlags) error {
+	// Validate before building the client: a malformed invocation should report
+	// what is wrong with it rather than an auth failure hit on the way there.
+	//
+	// --priority 0 is a real request, so distinguish "given" from the zero value
+	// rather than inferring it from the value itself.
+	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")
+	}
+
 	cl, err := ctx.NewManagementClient()
 	if err != nil {
 		return err
 	}
+
 	ref, err := resolveTrainJob(ctx, cl.API(), flags.JobID)
 	if err != nil {
 		return err
 	}
-	resp, err := cl.API().PatchTrainingProjectsJobs(ctx, ref.ProjectID, ref.JobID,
-		managementapi.UpdateTrainingJobRequest{Priority: &flags.Priority})
+
+	var req managementapi.UpdateTrainingJobRequest
+	var changed []string
+	if setPriority {
+		req.Priority = &flags.Priority
+		changed = append(changed, fmt.Sprintf("priority to %d", flags.Priority))
+	}
+	if setAvailability {
+		model := managementapi.V1AvailabilityModel(flags.AvailabilityModel)
+		req.AvailabilityModel = &model
+		changed = append(changed, fmt.Sprintf("availability model to %s", flags.AvailabilityModel))
+	}
+
+	resp, err := cl.API().PatchTrainingProjectsJobs(ctx, ref.ProjectID, ref.JobID, req)
 	if err != nil {
 		return fmt.Errorf("update training job %s: %w", flags.JobID, err)
 	}
@@ -745,7 +769,7 @@ func commandTrainJobUpdate(ctx *CommandContext, flags *cmd.TrainJobUpdateFlags)
 		ctx.OutputJSON(resp)
 		return nil
 	}
-	ctx.Logf("Set training job %s priority to %d.\n", resp.TrainingJob.Id, flags.Priority)
+	ctx.Logf("Set training job %s %s.\n", resp.TrainingJob.Id, strings.Join(changed, " and "))
 	return nil
 }
 

internal/cmd/command.train_test.go

@@ -220,9 +220,92 @@ func Test_Train_Job_Update(t *testing.T) {
 	h.Require.NoError(h.Execute("train", "job", "update", "--job-id", "job-1", "--priority", "10"))
 	body := m.FindCall("PATCH", trainJobPath).BodyJSON(t)
 	h.Require.Equal(float64(10), body["priority"])
+	h.Require.NotContains(body, "availability_model")
 	h.Require.Contains(h.Stderr.String(), "priority to 10")
 }
 
+func Test_Train_Job_Update_AvailabilityModel(t *testing.T) {
+	h := NewCommandHarness(t)
+	m := h.MockManagementAPI()
+	mockTrainJobSearch(m, trainJobFixture("job-1", "TRAINING_JOB_PENDING"))
+	m.SetRoute("PATCH", trainJobPath, 200, map[string]any{
+		"training_job": trainJobFixture("job-1", "TRAINING_JOB_PENDING"),
+	})
+
+	h.Require.NoError(h.Execute("train", "job", "update", "--job-id", "job-1",
+		"--availability-model", "spot"))
+	body := m.FindCall("PATCH", trainJobPath).BodyJSON(t)
+	h.Require.Equal("spot", body["availability_model"])
+	// Priority is omitted entirely rather than sent as 0, which would be a real change.
+	h.Require.NotContains(body, "priority")
+	h.Require.Contains(h.Stderr.String(), "availability model to spot")
+}
+
+func Test_Train_Job_Update_PriorityAndAvailabilityModel(t *testing.T) {
+	h := NewCommandHarness(t)
+	m := h.MockManagementAPI()
+	mockTrainJobSearch(m, trainJobFixture("job-1", "TRAINING_JOB_PENDING"))
+	m.SetRoute("PATCH", trainJobPath, 200, map[string]any{
+		"training_job": trainJobFixture("job-1", "TRAINING_JOB_PENDING"),
+	})
+
+	h.Require.NoError(h.Execute("train", "job", "update", "--job-id", "job-1",
+		"--priority", "7", "--availability-model", "dedicated"))
+	body := m.FindCall("PATCH", trainJobPath).BodyJSON(t)
+	h.Require.Equal(float64(7), body["priority"])
+	h.Require.Equal("dedicated", body["availability_model"])
+	h.Require.Contains(h.Stderr.String(), "priority to 7")
+	h.Require.Contains(h.Stderr.String(), "availability model to dedicated")
+}
+
+func Test_Train_Job_Update_ZeroPriorityIsSent(t *testing.T) {
+	h := NewCommandHarness(t)
+	m := h.MockManagementAPI()
+	mockTrainJobSearch(m, trainJobFixture("job-1", "TRAINING_JOB_PENDING"))
+	m.SetRoute("PATCH", trainJobPath, 200, map[string]any{
+		"training_job": trainJobFixture("job-1", "TRAINING_JOB_PENDING"),
+	})
+
+	// 0 is a valid priority, so an explicit --priority 0 must reach the API
+	// instead of being mistaken for an unset flag.
+	h.Require.NoError(h.Execute("train", "job", "update", "--job-id", "job-1", "--priority", "0"))
+	body := m.FindCall("PATCH", trainJobPath).BodyJSON(t)
+	h.Require.Equal(float64(0), body["priority"])
+}
+
+func Test_Train_Job_Update_NoFieldsIsUsageError(t *testing.T) {
+	h := NewCommandHarness(t)
+	m := h.MockManagementAPI()
+	mockTrainJobSearch(m, trainJobFixture("job-1", "TRAINING_JOB_PENDING"))
+
+	err := h.Execute("train", "job", "update", "--job-id", "job-1")
+	h.Require.Error(err)
+	h.Require.Contains(err.Error(), "at least one of --priority or --availability-model")
+	h.Require.Nil(m.FindCall("PATCH", trainJobPath))
+}
+
+func Test_Train_Job_Update_NoFieldsReportsUsageErrorBeforeClientSetup(t *testing.T) {
+	h := NewCommandHarness(t)
+	// A malformed remote makes building the management client fail outright. A
+	// malformed invocation should still report what is wrong with the invocation
+	// rather than the config error hit on the way there.
+	t.Setenv("BASETEN_REMOTE_URL", "://bogus")
+
+	err := h.Execute("train", "job", "update", "--job-id", "job-1")
+	h.Require.Error(err)
+	h.Require.Contains(err.Error(), "at least one of --priority or --availability-model")
+	h.Require.NotContains(err.Error(), "invalid remote URL")
+}
+
+func Test_Train_Job_Update_RejectsUnknownAvailabilityModel(t *testing.T) {
+	h := NewCommandHarness(t)
+	h.MockManagementAPI()
+
+	err := h.Execute("train", "job", "update", "--job-id", "job-1",
+		"--availability-model", "bogus")
+	h.Require.Error(err)
+}
+
 func Test_Train_Job_Logs(t *testing.T) {
 	h := NewCommandHarness(t)
 	m := h.MockManagementAPI()