[air] Add compute.priority_class for AI Runtime reservations (#6533)

Write
https://github.com/databricks-eng/universe/pull/2557225?timeline_per_page=5
to Go.

Tests: Unit tests.

## Changes
<!-- Brief summary of your changes that is easy to understand -->

## Why
<!-- Why are these changes needed? Provide the context that the reviewer
might be missing.
For example, were there any decisions behind the change that are not
reflected in the code itself? -->

## Tests
<!-- How have you tested the changes? -->

<!-- If your PR needs to be included in the release notes for next
release,
add a changelog fragment: create .nextchanges/<section>/<name>.md with a
one-line description (e.g. .nextchanges/cli/quickstart.md). See
.nextchanges/README.md. -->

Co-authored-by: Isaac <no-reply@databricks.com>

Author: vinchenzo-db

acceptance/experimental/air/config-help/output.txt

@@ -66,6 +66,7 @@ config.compute
     num_accelerators         Total number of GPUs to allocate.
     accelerator_type         Which accelerator to run on, e.g. GPU_1xA10.
     provisioned_capacity_id  Pre-provisioned AIR capacity reservation id.
+    priority_class           Scheduling priority within the reservation: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest).
 
 Use "-h config.compute.<field>" for details on a field.
 
@@ -79,7 +80,14 @@ config.mlflow_artifact_location
 
 >>> [CLI] experimental air run -h config.compute.provisioned_capacity_id
 config.compute.provisioned_capacity_id
-  Pre-provisioned AIR capacity reservation id. Must be 1-255 characters.
+  Pre-provisioned AIR capacity reservation id. Must be 1-255 characters. Contact your Databricks account team to provision capacity.
+
+  Type:     string
+  Required: no
+
+>>> [CLI] experimental air run -h config.compute.priority_class
+config.compute.priority_class
+  Scheduling priority within the reservation: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires provisioned_capacity_id.
 
   Type:     string
   Required: no
@@ -127,7 +135,7 @@ config.compute.num_accelerators
 >>> [CLI] experimental air run -h config.compute.acclerator_type
 Error: unknown config field "config.compute.acclerator_type"; did you mean "accelerator_type"?
 
-fields under "config.compute" are: accelerator_type, num_accelerators, provisioned_capacity_id
+fields under "config.compute" are: accelerator_type, num_accelerators, priority_class, provisioned_capacity_id
 
 === free-form map keys are not schema fields
 >>> [CLI] experimental air run -h config.parameters.learning_rate

acceptance/experimental/air/config-help/script

@@ -15,6 +15,7 @@ trace $CLI experimental air run -h config.compute
 title "new submission fields are documented"
 trace $CLI experimental air run -h config.mlflow_artifact_location
 trace $CLI experimental air run -h config.compute.provisioned_capacity_id
+trace $CLI experimental air run -h config.compute.priority_class
 trace $CLI experimental air run -h config.environment.dependencies
 
 title "leaf field"

experimental/air/cmd/compute.go

@@ -37,6 +37,40 @@ func parseGPUType(value string) (gpuType, error) {
 	return "", fmt.Errorf("invalid GPU type %q: %s", value, validGPUTypesHint())
 }
 
+// priorityClass is the AI Runtime scheduling priority a run is given within its
+// reservation. The values match the AiRuntimeTask.PriorityClass proto enum and
+// AICM's PriorityClassValidator (the server-side contract).
+type priorityClass string
+
+const (
+	priorityClassBestEffort priorityClass = "BEST_EFFORT"
+	priorityClassNormal     priorityClass = "NORMAL"
+	priorityClassCritical   priorityClass = "CRITICAL"
+)
+
+// priorityClasses lists every valid value, low to high. Used for validation
+// error messages.
+var priorityClasses = []priorityClass{priorityClassBestEffort, priorityClassNormal, priorityClassCritical}
+
+func validPriorityClassesHint() string {
+	names := make([]string, len(priorityClasses))
+	for i, p := range priorityClasses {
+		names[i] = string(p)
+	}
+	return "valid values are: " + strings.Join(names, ", ")
+}
+
+// parsePriorityClass normalizes a YAML priority_class to the enum's contract
+// value. Unlike accelerator_type, the match is case-insensitive: the input is
+// trimmed and upper-cased, so `critical` and `CRITICAL` both resolve.
+func parsePriorityClass(value string) (priorityClass, error) {
+	switch p := priorityClass(strings.ToUpper(strings.TrimSpace(value))); p {
+	case priorityClassBestEffort, priorityClassNormal, priorityClassCritical:
+		return p, nil
+	}
+	return "", fmt.Errorf("invalid priority_class %q: %s", value, validPriorityClassesHint())
+}
+
 // gpusPerNode returns the per-node GPU count, which is the partition count from
 // the name (GPU_1xH100 -> 1, GPU_8xH100 -> 8). num_accelerators must be a
 // round multiple of this since accelerators are allocated in whole nodes.
@@ -57,7 +91,8 @@ func gpusPerNode(g gpuType) (int, error) {
 type computeConfig struct {
 	NumAccelerators       int     `yaml:"num_accelerators" help:"Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types."`
 	AcceleratorType       string  `yaml:"accelerator_type" help:"Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively."`
-	ProvisionedCapacityID *string `yaml:"provisioned_capacity_id" help:"Pre-provisioned AIR capacity reservation id. Must be 1-255 characters."`
+	ProvisionedCapacityID *string `yaml:"provisioned_capacity_id" help:"Pre-provisioned AIR capacity reservation id. Must be 1-255 characters. Contact your Databricks account team to provision capacity."`
+	PriorityClass         *string `yaml:"priority_class" help:"Scheduling priority within the reservation: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires provisioned_capacity_id."`
 }
 
 // validate checks the compute block against the backend's constraints.
@@ -90,5 +125,18 @@ func (c *computeConfig) validate() error {
 		*c.ProvisionedCapacityID = v
 	}
 
+	if c.PriorityClass != nil {
+		p, err := parsePriorityClass(*c.PriorityClass)
+		if err != nil {
+			return fmt.Errorf("compute.priority_class: %w", err)
+		}
+		// A priority class only ranks pending work within a reservation, so it
+		// requires one.
+		if c.ProvisionedCapacityID == nil {
+			return errors.New("compute.priority_class requires compute.provisioned_capacity_id — priority applies only to a pre-provisioned capacity reservation")
+		}
+		*c.PriorityClass = string(p)
+	}
+
 	return nil
 }

experimental/air/cmd/compute_test.go

@@ -59,6 +59,30 @@ func TestGPUsPerNode(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestParsePriorityClass(t *testing.T) {
+	// Case-insensitive: mixed/lower case and surrounding whitespace normalize to
+	// the enum's contract value.
+	for _, in := range []string{"CRITICAL", "critical", "Critical", "  critical  "} {
+		got, err := parsePriorityClass(in)
+		require.NoError(t, err)
+		assert.Equal(t, priorityClassCritical, got)
+	}
+
+	for _, in := range []string{"urgent", "", "high"} {
+		_, err := parsePriorityClass(in)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "valid values are")
+	}
+}
+
+func TestComputeConfigValidateNormalizesPriorityClass(t *testing.T) {
+	// A lower-case value is upper-cased in place so the submitted payload carries
+	// the enum's contract value.
+	cfg := computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new("cap"), PriorityClass: new("critical")}
+	require.NoError(t, cfg.validate())
+	assert.Equal(t, "CRITICAL", *cfg.PriorityClass)
+}
+
 func TestComputeConfigValidate(t *testing.T) {
 	tests := []struct {
 		name    string
@@ -75,6 +99,9 @@ func TestComputeConfigValidate(t *testing.T) {
 		{"legacy type rejected", computeConfig{NumAccelerators: 8, AcceleratorType: "h100_80gb"}, "accelerator_type"},
 		{"non-positive count", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_1xH100"}, "must be positive"},
 		{"count not a multiple", computeConfig{NumAccelerators: 4, AcceleratorType: "GPU_8xH100"}, "multiple of 8"},
+		{"priority class", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new("cap"), PriorityClass: new("critical")}, ""},
+		{"priority class requires reservation", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PriorityClass: new("NORMAL")}, "requires compute.provisioned_capacity_id"},
+		{"invalid priority class", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new("cap"), PriorityClass: new("urgent")}, "invalid priority_class"},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

experimental/air/cmd/runconfig_test.go

@@ -489,7 +489,7 @@ func TestResolveConfigField_Containers(t *testing.T) {
 	compute, err := resolveConfigField("config.compute")
 	require.NoError(t, err)
 	assert.Equal(t, "object", compute.typeName)
-	require.Len(t, compute.children, 3)
+	require.Len(t, compute.children, 4)
 }
 
 func TestResolveConfigField_Errors(t *testing.T) {

experimental/air/cmd/runsubmit.go

@@ -114,8 +114,12 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri
 	}
 }
 
-func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.SubmitRun, provisionedCapacityID string) (int64, error) {
-	if provisionedCapacityID == "" {
+func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.SubmitRun, provisionedCapacityID, priorityClass string) (int64, error) {
+	// Neither reservation field is modeled by the SDK's AiRuntimeTask, so a run
+	// that sets either has to go through the raw /api/2.2 body. priority_class only
+	// ever appears alongside a reservation (validation enforces it), but route on
+	// both so it can never be silently dropped.
+	if provisionedCapacityID == "" && priorityClass == "" {
 		wait, err := w.Jobs.Submit(ctx, payload)
 		if err != nil {
 			return 0, err
@@ -133,7 +137,7 @@ func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.
 	if err := decoder.Decode(&body); err != nil {
 		return 0, fmt.Errorf("failed to decode AIR submit payload: %w", err)
 	}
-	if err := injectProvisionedCapacityID(body, provisionedCapacityID); err != nil {
+	if err := injectReservationFields(body, provisionedCapacityID, priorityClass); err != nil {
 		return 0, err
 	}
 
@@ -149,33 +153,52 @@ func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.
 	return response.RunId, nil
 }
 
-func injectProvisionedCapacityID(body map[string]any, provisionedCapacityID string) error {
+// injectReservationFields sets the reservation-only fields the SDK does not
+// model onto the decoded submit body: priority_class rides directly on the
+// ai_runtime_task, while provisioned_capacity_id rides on the deployment's
+// compute spec. Each is set only when non-empty.
+func injectReservationFields(body map[string]any, provisionedCapacityID, priorityClass string) error {
+	aiRuntimeTask, err := aiRuntimeTaskFromSubmitBody(body)
+	if err != nil {
+		return err
+	}
+	if priorityClass != "" {
+		aiRuntimeTask["priority_class"] = priorityClass
+	}
+	if provisionedCapacityID != "" {
+		deployments, ok := aiRuntimeTask["deployments"].([]any)
+		if !ok || len(deployments) != 1 {
+			return errors.New("AIR submit payload must contain exactly one deployment")
+		}
+		deployment, ok := deployments[0].(map[string]any)
+		if !ok {
+			return errors.New("AIR submit payload deployment has an invalid shape")
+		}
+		computeSpec, ok := deployment["compute"].(map[string]any)
+		if !ok {
+			return errors.New("AIR submit payload is missing deployment compute")
+		}
+		computeSpec["provisioned_capacity_id"] = provisionedCapacityID
+	}
+	return nil
+}
+
+// aiRuntimeTaskFromSubmitBody navigates a decoded runs/submit body to its single
+// ai_runtime_task map, erroring if the payload isn't the expected single-task shape.
+func aiRuntimeTaskFromSubmitBody(body map[string]any) (map[string]any, error) {
 	tasks, ok := body["tasks"].([]any)
 	if !ok || len(tasks) != 1 {
-		return errors.New("AIR submit payload must contain exactly one task")
+		return nil, errors.New("AIR submit payload must contain exactly one task")
 	}
 	task, ok := tasks[0].(map[string]any)
 	if !ok {
-		return errors.New("AIR submit payload task has an invalid shape")
+		return nil, errors.New("AIR submit payload task has an invalid shape")
 	}
 	aiRuntimeTask, ok := task["ai_runtime_task"].(map[string]any)
 	if !ok {
-		return errors.New("AIR submit payload is missing ai_runtime_task")
-	}
-	deployments, ok := aiRuntimeTask["deployments"].([]any)
-	if !ok || len(deployments) != 1 {
-		return errors.New("AIR submit payload must contain exactly one deployment")
-	}
-	deployment, ok := deployments[0].(map[string]any)
-	if !ok {
-		return errors.New("AIR submit payload deployment has an invalid shape")
+		return nil, errors.New("AIR submit payload is missing ai_runtime_task")
 	}
-	computeSpec, ok := deployment["compute"].(map[string]any)
-	if !ok {
-		return errors.New("AIR submit payload is missing deployment compute")
-	}
-	computeSpec["provisioned_capacity_id"] = provisionedCapacityID
-	return nil
+	return aiRuntimeTask, nil
 }
 
 // submitToken resolves the idempotency token: the --idempotency-key flag wins,
@@ -313,8 +336,12 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 	if cfg.Compute.ProvisionedCapacityID != nil {
 		provisionedCapacityID = *cfg.Compute.ProvisionedCapacityID
 	}
+	priorityClass := ""
+	if cfg.Compute.PriorityClass != nil {
+		priorityClass = *cfg.Compute.PriorityClass
+	}
 	// Submit returns as soon as the run is created; we don't wait for it to finish.
-	runID, err := submitRun(ctx, w, payload, provisionedCapacityID)
+	runID, err := submitRun(ctx, w, payload, provisionedCapacityID, priorityClass)
 	if err != nil {
 		return 0, "", err
 	}

experimental/air/cmd/runsubmit_test.go

@@ -109,11 +109,42 @@ func TestSubmitRunInjectsProvisionedCapacityID(t *testing.T) {
 		Compute:        &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1},
 	}, "/command.sh", "4", "", snapshotResult{}, nil)
 
-	runID, err := submitRun(t.Context(), w, payload, "capacity-1")
+	runID, err := submitRun(t.Context(), w, payload, "capacity-1", "")
 	require.NoError(t, err)
 	assert.Equal(t, int64(42), runID)
 }
 
+func TestSubmitRunInjectsPriorityClass(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
+		var body map[string]any
+		require.NoError(t, json.Unmarshal(req.Body, &body))
+		tasks := body["tasks"].([]any)
+		task := tasks[0].(map[string]any)
+		airTask := task["ai_runtime_task"].(map[string]any)
+		// priority_class rides directly on the ai_runtime_task, next to
+		// provisioned_capacity_id on the deployment compute.
+		assert.Equal(t, "CRITICAL", airTask["priority_class"])
+		deployment := airTask["deployments"].([]any)[0].(map[string]any)
+		compute := deployment["compute"].(map[string]any)
+		assert.Equal(t, "capacity-1", compute["provisioned_capacity_id"])
+		return jobs.SubmitRunResponse{RunId: 7}
+	})
+
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token", WorkspaceID: "123"})
+	require.NoError(t, err)
+	payload := buildSubmitPayload(&runConfig{
+		ExperimentName: "exp",
+		Command:        new("x"),
+		Compute:        &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1},
+	}, "/command.sh", "4", "", snapshotResult{}, nil)
+
+	runID, err := submitRun(t.Context(), w, payload, "capacity-1", "CRITICAL")
+	require.NoError(t, err)
+	assert.Equal(t, int64(7), runID)
+}
+
 func TestBuildSubmitPayloadDefaultRetries(t *testing.T) {
 	// max_retries unset defaults to 3 (matching the Python native path), so both
 	// retry fields are sent.

experimental/air/cmd/validateconfig.go

@@ -72,6 +72,11 @@ func validateConfigRequest(cfg *runConfig, commandPath string) map[string]any {
 		"experiment":  cfg.ExperimentName,
 		"deployments": []any{map[string]any{"command_path": commandPath, "compute": compute}},
 	}
+	if cfg.Compute != nil {
+		// priority_class rides on the ai_runtime_task (task-level), not the deployment
+		// compute where provisioned_capacity_id lives.
+		putOpt(task, "priority_class", cfg.Compute.PriorityClass)
+	}
 	putOpt(task, "mlflow_run", cfg.MLflowRunName)
 	putOpt(task, "mlflow_experiment_directory", cfg.MLflowExperimentDirectory)
 	putOpt(task, "mlflow_artifact_location", cfg.MLflowArtifactLocation)

experimental/air/cmd/validateconfig_test.go

@@ -110,6 +110,23 @@ func TestValidateConfigRequestShape(t *testing.T) {
 	assert.Equal(t, map[string]any{"HF_HOME": "/tmp/hf"}, runOptions["env_variables"])
 }
 
+func TestValidateConfigRequestCarriesPriorityClass(t *testing.T) {
+	var gotReq map[string]any
+	srv := validateServer(t, http.StatusOK, `{}`, &gotReq)
+
+	cfg := baseRunConfig()
+	cfg.Compute.ProvisionedCapacityID = new("cap-8xh100-res")
+	cfg.Compute.PriorityClass = new("CRITICAL")
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh")
+	require.NoError(t, err)
+
+	task := gotReq["task"].(map[string]any)
+	// priority_class is task-level; provisioned_capacity_id stays on the compute spec.
+	assert.Equal(t, "CRITICAL", task["priority_class"])
+	compute := task["deployments"].([]any)[0].(map[string]any)["compute"].(map[string]any)
+	assert.Equal(t, "cap-8xh100-res", compute["provisioned_capacity_id"])
+}
+
 func TestValidateConfigRequestOmitsUnsetOptions(t *testing.T) {
 	// A minimal config carries no run_options and only the fields it set, so the
 	// server never validates values the user didn't provide.