diff --git a/acceptance/bundle/resource_deps/job_tasks/output.txt b/acceptance/bundle/resource_deps/job_tasks/output.txt index f09bb852e99..9a3817ea2d5 100644 --- a/acceptance/bundle/resource_deps/job_tasks/output.txt +++ b/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 diff --git a/acceptance/bundle/telemetry/deploy-ai-runtime-task/databricks.yml b/acceptance/bundle/telemetry/deploy-ai-runtime-task/databricks.yml new file mode 100644 index 00000000000..f2159fcce50 --- /dev/null +++ b/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 diff --git a/acceptance/bundle/telemetry/deploy-ai-runtime-task/out.test.toml b/acceptance/bundle/telemetry/deploy-ai-runtime-task/out.test.toml new file mode 100644 index 00000000000..ae800809893 --- /dev/null +++ b/acceptance/bundle/telemetry/deploy-ai-runtime-task/out.test.toml @@ -0,0 +1,3 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] +EnvMatrix.DMS = [""] diff --git a/acceptance/bundle/telemetry/deploy-ai-runtime-task/output.txt b/acceptance/bundle/telemetry/deploy-ai-runtime-task/output.txt new file mode 100644 index 00000000000..8e8bcbd9e6d --- /dev/null +++ b/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 + } +] diff --git a/acceptance/bundle/telemetry/deploy-ai-runtime-task/script b/acceptance/bundle/telemetry/deploy-ai-runtime-task/script new file mode 100644 index 00000000000..e16013fc9c7 --- /dev/null +++ b/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" diff --git a/acceptance/bundle/telemetry/deploy-compute-type/output.txt b/acceptance/bundle/telemetry/deploy-compute-type/output.txt index fc47cfbd506..a2a7a8e8e3b 100644 --- a/acceptance/bundle/telemetry/deploy-compute-type/output.txt +++ b/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 diff --git a/acceptance/bundle/telemetry/deploy-experimental/output.txt b/acceptance/bundle/telemetry/deploy-experimental/output.txt index 4ce786a044b..0ff55600907 100644 --- a/acceptance/bundle/telemetry/deploy-experimental/output.txt +++ b/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 diff --git a/bundle/metrics/metrics.go b/bundle/metrics/metrics.go index 62ae863cd58..e5a50769ee9 100644 --- a/bundle/metrics/metrics.go +++ b/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_" 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" ) diff --git a/bundle/phases/telemetry.go b/bundle/phases/telemetry.go index cbe74f467fe..e8b7617fc55 100644 --- a/bundle/phases/telemetry.go +++ b/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_" 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_" 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_ 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 diff --git a/bundle/phases/telemetry_test.go b/bundle/phases/telemetry_test.go index 40a193077a4..0cea4df5794 100644 --- a/bundle/phases/telemetry_test.go +++ b/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") + }) + } +}