Add per-task-type deploy telemetry for bundle jobs

acceptance/bundle/resource_deps/job_tasks/output.txt

@@ -16,6 +16,7 @@ dms_undeclared_service_principal false
 experimental.use_legacy_run_as false
 has_classic_interactive_compute false
 has_classic_job_compute false
+has_python_wheel_task true
 has_serverless_compute true
 local.cache.attempt true
 local.cache.miss true

acceptance/bundle/telemetry/deploy-ai-runtime-task/databricks.yml

@@ -0,0 +1,25 @@
+bundle:
+  name: test-bundle
+
+resources:
+  jobs:
+    job:
+      name: "[${bundle.target}] AI Runtime Job"
+      # Present so the ai_runtime_task_scheduled telemetry bool is exercised.
+      schedule:
+        quartz_cron_expression: "0 0 12 * * ?"
+        timezone_id: UTC
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "4"
+      tasks:
+        - task_key: train
+          environment_key: default
+          ai_runtime_task:
+            experiment: my-experiment
+            deployments:
+              - command_path: /Workspace/Shared/command.sh
+                compute:
+                  accelerator_type: GPU_1xA10
+                  accelerator_count: 1

acceptance/bundle/telemetry/deploy-ai-runtime-task/output.txt

@@ -0,0 +1,22 @@
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files...
+Created jobs.job
+Files: 2 uploaded, 0 deleted
+Resources: 1 created, 0 changed, 0 deleted, 0 unchanged
+
+>>> cat [OUTPUT_DIR]/out.requests.txt
+[
+  {
+    "key": "ai_runtime_task_multitask",
+    "value": false
+  },
+  {
+    "key": "ai_runtime_task_scheduled",
+    "value": true
+  },
+  {
+    "key": "has_ai_runtime_task",
+    "value": true
+  }
+]

acceptance/bundle/telemetry/deploy-ai-runtime-task/script

@@ -0,0 +1,7 @@
+trace $CLI bundle deploy
+
+# Only the task-type / ai_runtime_task telemetry keys, so this golden does not
+# churn when unrelated bool_values change.
+trace cat "$OUT_REQUESTS" | jq 'select(has("path") and .path == "/telemetry-ext") | .body.protoLogs[] | fromjson | .entry.databricks_cli_log.bundle_deploy_event.experimental.bool_values | map(select((.key | test("_task$")) or (.key | startswith("ai_runtime_")))) | sort_by(.key)'
+
+rm "$OUT_REQUESTS"

acceptance/bundle/telemetry/deploy-compute-type/output.txt

@@ -28,8 +28,11 @@ has_classic_interactive_compute false
 has_classic_interactive_compute true
 has_classic_job_compute false
 has_classic_job_compute true
+has_notebook_task true
 has_serverless_compute true
 has_serverless_compute true
+has_spark_python_task true
+has_spark_python_task true
 local.cache.attempt true
 local.cache.attempt true
 local.cache.hit true

acceptance/bundle/telemetry/deploy-experimental/output.txt

@@ -19,6 +19,7 @@ experimental.use_legacy_run_as true
 has_classic_interactive_compute true
 has_classic_job_compute false
 has_serverless_compute false
+has_spark_python_task true
 local.cache.attempt true
 local.cache.miss true
 permissions_section_set false

bundle/metrics/metrics.go

@@ -109,4 +109,13 @@ const (
 	DMSUndeclaredOtherUser        = "dms_undeclared_other_user"
 	DMSUndeclaredServicePrincipal = "dms_undeclared_service_principal"
 	DMSUndeclaredGroup            = "dms_undeclared_group"
+
+	// Task-type usage is recorded generically at deploy as one "has_<task_type>" key
+	// per task type present in the bundle (e.g. has_notebook_task, has_ai_runtime_task);
+	// see collectTaskTypes in bundle/phases. The two keys below are extra
+	// ai_runtime_task-specific dimensions, emitted only when the bundle declares an
+	// ai_runtime_task, so a false value means "has an ai_runtime_task, but it is not
+	// scheduled / not multi-task". GPU type and count are intentionally not recorded here.
+	AiRuntimeTaskScheduled = "ai_runtime_task_scheduled"
+	AiRuntimeTaskMultitask = "ai_runtime_task_multitask"
 )

bundle/phases/telemetry.go

@@ -4,17 +4,21 @@ import (
 	"cmp"
 	"context"
 	"math"
+	"reflect"
 	"slices"
+	"strings"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/engine"
+	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/bundle/libraries"
 	"github.com/databricks/cli/bundle/metrics"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
 )
 
 func getExecutionTimes(b *bundle.Bundle) []protos.IntMapEntry {
@@ -108,6 +112,89 @@ func uploadFileSizeHistogram(files []sizer) []int64 {
 	return hist
 }
 
+// addTaskTypeKeys records a "has_<task_type>" key in out for each task-type field
+// set on the task. A Jobs task carries its type as one of several optional pointer
+// fields (notebook_task, spark_python_task, ai_runtime_task, ...); we reflect over
+// the fields whose JSON name ends in "_task" so every task type is captured,
+// including ones added later, without changing this code.
+func addTaskTypeKeys(task jobs.Task, out map[string]bool) {
+	v := reflect.ValueOf(task)
+	t := v.Type()
+	for i := range t.NumField() {
+		field := v.Field(i)
+		if field.Kind() != reflect.Pointer || field.IsNil() {
+			continue
+		}
+		name, _, _ := strings.Cut(t.Field(i).Tag.Get("json"), ",")
+		if strings.HasSuffix(name, "_task") {
+			out["has_"+name] = true
+		}
+	}
+}
+
+// collectTaskTypes returns the sorted "has_<task_type>" telemetry keys for every
+// task type used across the bundle's jobs. A for_each_task is unwrapped so the
+// wrapped task's type is recorded too. Sorted for deterministic telemetry output.
+func collectTaskTypes(jobs map[string]*resources.Job) []string {
+	seen := map[string]bool{}
+	for _, job := range jobs {
+		if job == nil {
+			continue
+		}
+		for _, task := range job.Tasks {
+			addTaskTypeKeys(task, seen)
+			if task.ForEachTask != nil {
+				addTaskTypeKeys(task.ForEachTask.Task, seen)
+			}
+		}
+	}
+	keys := make([]string, 0, len(seen))
+	for k := range seen {
+		keys = append(keys, k)
+	}
+	slices.Sort(keys)
+	return keys
+}
+
+// aiRuntimeTaskMetrics computes the ai_runtime_task-specific deploy dimensions.
+// present is true when any job declares an ai_runtime_task (including one nested in
+// a for_each_task); scheduled and multitask describe those jobs and are meaningful
+// only when present is true. Task-type presence itself (has_ai_runtime_task and the
+// other has_*_task keys) is recorded generically by collectTaskTypes.
+//
+// code_source_path is deliberately not inspected: it is rewritten to its uploaded
+// remote path before this runs, so it would always read as remote. GPU type and
+// count are likewise not recorded here.
+func aiRuntimeTaskMetrics(jobs map[string]*resources.Job) (present, scheduled, multitask bool) {
+	for _, job := range jobs {
+		if job == nil {
+			continue
+		}
+		jobHasAiRuntimeTask := false
+		for _, task := range job.Tasks {
+			rt := task.AiRuntimeTask
+			if rt == nil && task.ForEachTask != nil {
+				rt = task.ForEachTask.Task.AiRuntimeTask
+			}
+			if rt != nil {
+				jobHasAiRuntimeTask = true
+				break
+			}
+		}
+		if !jobHasAiRuntimeTask {
+			continue
+		}
+		present = true
+		if job.Schedule != nil || job.Trigger != nil || job.Continuous != nil {
+			scheduled = true
+		}
+		if len(job.Tasks) > 1 {
+			multitask = true
+		}
+	}
+	return present, scheduled, multitask
+}
+
 // LogDeployTelemetry logs a telemetry event for a bundle deploy command.
 func LogDeployTelemetry(ctx context.Context, b *bundle.Bundle, errMsg string) {
 	errMsg = telemetry.ScrubErrorMessage(errMsg)
@@ -202,6 +289,16 @@ func LogDeployTelemetry(ctx context.Context, b *bundle.Bundle, errMsg string) {
 		}
 	}
 
+	// Record which task types the bundle uses (has_<task_type> per type present),
+	// plus the ai_runtime_task-specific scheduling / multi-task dimensions.
+	for _, key := range collectTaskTypes(b.Config.Resources.Jobs) {
+		b.Metrics.SetBoolValue(key, true)
+	}
+	if airPresent, airScheduled, airMultitask := aiRuntimeTaskMetrics(b.Config.Resources.Jobs); airPresent {
+		b.Metrics.SetBoolValue(metrics.AiRuntimeTaskScheduled, airScheduled)
+		b.Metrics.SetBoolValue(metrics.AiRuntimeTaskMultitask, airMultitask)
+	}
+
 	// Record whether the deprecated terraform engine was explicitly opted into,
 	// separately per source. Only emitted when true; absence means "not opted in
 	// via this source". An invalid env value has already failed the command

bundle/phases/telemetry_test.go

@@ -3,6 +3,8 @@ package phases
 import (
 	"testing"
 
+	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -75,3 +77,100 @@ func TestUploadFileSizeHistogramUnknownSizeOmitsHistogram(t *testing.T) {
 func TestUploadFileSizeHistogramEmpty(t *testing.T) {
 	assert.Nil(t, uploadFileSizeHistogram(nil))
 }
+
+func TestCollectTaskTypes(t *testing.T) {
+	jobs := map[string]*resources.Job{
+		"nil": nil,
+		"a": {JobSettings: jobs.JobSettings{Tasks: []jobs.Task{
+			{TaskKey: "nb", NotebookTask: &jobs.NotebookTask{NotebookPath: "/nb"}},
+			{TaskKey: "train", AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp"}},
+		}}},
+		"b": {JobSettings: jobs.JobSettings{Tasks: []jobs.Task{
+			// for_each wrapping a spark_python_task: both the wrapper and the
+			// nested task type should be recorded.
+			{TaskKey: "fan", ForEachTask: &jobs.ForEachTask{Task: jobs.Task{
+				SparkPythonTask: &jobs.SparkPythonTask{PythonFile: "main.py"},
+			}}},
+		}}},
+	}
+	assert.Equal(t, []string{
+		"has_ai_runtime_task",
+		"has_for_each_task",
+		"has_notebook_task",
+		"has_spark_python_task",
+	}, collectTaskTypes(jobs))
+
+	assert.Empty(t, collectTaskTypes(nil))
+}
+
+// aiRuntimeJob builds a job with a single ai_runtime_task, plus optionally an extra
+// no-op task to make it multi-task.
+func aiRuntimeJob(extraTask bool) *resources.Job {
+	tasks := []jobs.Task{{
+		TaskKey:       "train",
+		AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp"},
+	}}
+	if extraTask {
+		tasks = append(tasks, jobs.Task{TaskKey: "prep"})
+	}
+	return &resources.Job{JobSettings: jobs.JobSettings{Tasks: tasks}}
+}
+
+func TestAiRuntimeTaskMetrics(t *testing.T) {
+	scheduledJob := aiRuntimeJob(false)
+	scheduledJob.Schedule = &jobs.CronSchedule{QuartzCronExpression: "0 0 * * * ?"}
+
+	forEachJob := &resources.Job{JobSettings: jobs.JobSettings{Tasks: []jobs.Task{{
+		TaskKey:     "fanout",
+		ForEachTask: &jobs.ForEachTask{Task: jobs.Task{AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp"}}},
+	}}}}
+
+	tests := []struct {
+		name                          string
+		jobs                          map[string]*resources.Job
+		present, scheduled, multitask bool
+	}{
+		{
+			name: "no jobs",
+			jobs: nil,
+		},
+		{
+			name: "job without ai_runtime_task",
+			jobs: map[string]*resources.Job{"j": {JobSettings: jobs.JobSettings{Tasks: []jobs.Task{{TaskKey: "notebook"}}}}},
+		},
+		{
+			name:    "single task",
+			jobs:    map[string]*resources.Job{"j": aiRuntimeJob(false)},
+			present: true,
+		},
+		{
+			name:    "multi-task",
+			jobs:    map[string]*resources.Job{"j": aiRuntimeJob(true)},
+			present: true, multitask: true,
+		},
+		{
+			name:    "scheduled single task",
+			jobs:    map[string]*resources.Job{"j": scheduledJob},
+			present: true, scheduled: true,
+		},
+		{
+			name:    "for_each nested ai_runtime_task",
+			jobs:    map[string]*resources.Job{"j": forEachJob},
+			present: true,
+		},
+		{
+			name:    "nil job is skipped",
+			jobs:    map[string]*resources.Job{"j": nil, "k": aiRuntimeJob(false)},
+			present: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			present, scheduled, multitask := aiRuntimeTaskMetrics(tc.jobs)
+			assert.Equal(t, tc.present, present, "present")
+			assert.Equal(t, tc.scheduled, scheduled, "scheduled")
+			assert.Equal(t, tc.multitask, multitask, "multitask")
+		})
+	}
+}