From 606abf211dd95ffbcda4ef27fd94e72f6a89ee65 Mon Sep 17 00:00:00 2001 From: padawan Date: Tue, 11 Aug 2026 12:47:17 +0100 Subject: [PATCH] bundle: avoid num_workers with policy defaults --- .../databricks.yml | 51 ++++++++++ .../out.test.toml | 2 + .../policy-default-cluster-sizing/output.txt | 97 +++++++++++++++++++ .../sample_notebook.py | 1 + .../jobs/policy-default-cluster-sizing/script | 17 ++++ .../policy-default-cluster-sizing/test.toml | 6 ++ .../mutator/resourcemutator/cluster_fixups.go | 4 + .../resourcemutator/cluster_fixups_test.go | 69 +++++++++++++ bundle/direct/dresources/resources.yml | 5 + 9 files changed, 252 insertions(+) create mode 100644 acceptance/bundle/resources/jobs/policy-default-cluster-sizing/databricks.yml create mode 100644 acceptance/bundle/resources/jobs/policy-default-cluster-sizing/out.test.toml create mode 100644 acceptance/bundle/resources/jobs/policy-default-cluster-sizing/output.txt create mode 100644 acceptance/bundle/resources/jobs/policy-default-cluster-sizing/sample_notebook.py create mode 100644 acceptance/bundle/resources/jobs/policy-default-cluster-sizing/script create mode 100644 acceptance/bundle/resources/jobs/policy-default-cluster-sizing/test.toml create mode 100644 bundle/config/mutator/resourcemutator/cluster_fixups_test.go diff --git a/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/databricks.yml b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/databricks.yml new file mode 100644 index 00000000000..acfadd05119 --- /dev/null +++ b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/databricks.yml @@ -0,0 +1,51 @@ +bundle: + name: test-bundle + +resources: + jobs: + sample_job: + name: sample_job + + tasks: + - task_key: task_policy_defaults + notebook_task: + notebook_path: sample_notebook.py + source: WORKSPACE + new_cluster: + policy_id: policy-id + apply_policy_default_values: true + + job_clusters: + - job_cluster_key: policy_defaults + new_cluster: + policy_id: policy-id + apply_policy_default_values: true + + - job_cluster_key: policy_defaults_autoscale + new_cluster: + policy_id: policy-id + apply_policy_default_values: true + autoscale: + min_workers: 1 + max_workers: 4 + + - job_cluster_key: policy_defaults_num_workers_zero + new_cluster: + policy_id: policy-id + apply_policy_default_values: true + num_workers: 0 + spark_conf: + spark.databricks.cluster.profile: singleNode + spark.master: local[*] + custom_tags: + ResourceClass: SingleNode + + - job_cluster_key: policy_defaults_disabled + new_cluster: + policy_id: policy-id + apply_policy_default_values: false + spark_conf: + spark.databricks.cluster.profile: singleNode + spark.master: local[*] + custom_tags: + ResourceClass: SingleNode diff --git a/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/out.test.toml b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/output.txt b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/output.txt new file mode 100644 index 00000000000..fdfa4464ce5 --- /dev/null +++ b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/output.txt @@ -0,0 +1,97 @@ + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] bundle plan +Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> print_requests.py //jobs +{ + "method": "POST", + "path": "/api/2.2/jobs/create", + "body": { + "deployment": { + "kind": "BUNDLE", + "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/test-bundle/default/state/metadata.json" + }, + "edit_mode": "UI_LOCKED", + "format": "MULTI_TASK", + "job_clusters": [ + { + "job_cluster_key": "policy_defaults", + "new_cluster": { + "apply_policy_default_values": true, + "policy_id": "policy-id" + } + }, + { + "job_cluster_key": "policy_defaults_autoscale", + "new_cluster": { + "apply_policy_default_values": true, + "autoscale": { + "max_workers": 4, + "min_workers": 1 + }, + "policy_id": "policy-id" + } + }, + { + "job_cluster_key": "policy_defaults_num_workers_zero", + "new_cluster": { + "apply_policy_default_values": true, + "custom_tags": { + "ResourceClass": "SingleNode" + }, + "num_workers": 0, + "policy_id": "policy-id", + "spark_conf": { + "spark.databricks.cluster.profile": "singleNode", + "spark.master": "local[*]" + } + } + }, + { + "job_cluster_key": "policy_defaults_disabled", + "new_cluster": { + "apply_policy_default_values": false, + "custom_tags": { + "ResourceClass": "SingleNode" + }, + "num_workers": 0, + "policy_id": "policy-id", + "spark_conf": { + "spark.databricks.cluster.profile": "singleNode", + "spark.master": "local[*]" + } + } + } + ], + "max_concurrent_runs": 1, + "name": "sample_job", + "queue": { + "enabled": true + }, + "tasks": [ + { + "new_cluster": { + "apply_policy_default_values": true, + "policy_id": "policy-id" + }, + "notebook_task": { + "notebook_path": "/Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files/sample_notebook", + "source": "WORKSPACE" + }, + "task_key": "task_policy_defaults" + } + ] + } +} diff --git a/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/sample_notebook.py b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/sample_notebook.py new file mode 100644 index 00000000000..1645e04b1de --- /dev/null +++ b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/sample_notebook.py @@ -0,0 +1 @@ +# Databricks notebook source diff --git a/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/script b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/script new file mode 100644 index 00000000000..8980c3fe7b8 --- /dev/null +++ b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/script @@ -0,0 +1,17 @@ +trace $CLI bundle deploy + +job_id="$(read_id.py sample_job)" +edit_resource.py jobs "$job_id" <<'EOF' +for task in r["tasks"]: + if task["task_key"] == "task_policy_defaults": + task["new_cluster"]["autoscale"] = {"min_workers": 1, "max_workers": 4} + +for cluster in r["job_clusters"]: + if cluster["job_cluster_key"] == "policy_defaults": + cluster["new_cluster"]["autoscale"] = {"min_workers": 1, "max_workers": 4} +EOF + +trace $CLI bundle plan +trace $CLI bundle deploy +trace print_requests.py //jobs +rm out.requests.txt diff --git a/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/test.toml b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/test.toml new file mode 100644 index 00000000000..2874b03b338 --- /dev/null +++ b/acceptance/bundle/resources/jobs/policy-default-cluster-sizing/test.toml @@ -0,0 +1,6 @@ +RecordRequests = true +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +Ignore = [ + ".databricks", +] diff --git a/bundle/config/mutator/resourcemutator/cluster_fixups.go b/bundle/config/mutator/resourcemutator/cluster_fixups.go index 893cd248aa4..79c26ca0020 100644 --- a/bundle/config/mutator/resourcemutator/cluster_fixups.go +++ b/bundle/config/mutator/resourcemutator/cluster_fixups.go @@ -106,6 +106,10 @@ func initializeNumWorkers(c *compute.ClusterSpec) { if c.Autoscale != nil { return } + // Policy defaults can supply autoscaling, which conflicts with num_workers. + if c.ApplyPolicyDefaultValues { + return + } if c.NumWorkers != 0 { return } diff --git a/bundle/config/mutator/resourcemutator/cluster_fixups_test.go b/bundle/config/mutator/resourcemutator/cluster_fixups_test.go new file mode 100644 index 00000000000..cae89cf52fa --- /dev/null +++ b/bundle/config/mutator/resourcemutator/cluster_fixups_test.go @@ -0,0 +1,69 @@ +package resourcemutator + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/stretchr/testify/assert" +) + +func TestInitializeNumWorkers(t *testing.T) { + tests := []struct { + name string + in compute.ClusterSpec + want []string + }{ + { + name: "omitted", + want: []string{"NumWorkers"}, + }, + { + name: "omitted with policy defaults disabled", + in: compute.ClusterSpec{ + ForceSendFields: []string{"ApplyPolicyDefaultValues"}, + }, + want: []string{"ApplyPolicyDefaultValues", "NumWorkers"}, + }, + { + name: "omitted with policy defaults", + in: compute.ClusterSpec{ + ApplyPolicyDefaultValues: true, + }, + }, + { + name: "explicit zero with policy defaults", + in: compute.ClusterSpec{ + ApplyPolicyDefaultValues: true, + ForceSendFields: []string{"NumWorkers"}, + }, + want: []string{"NumWorkers"}, + }, + { + name: "explicit non-zero with policy defaults", + in: compute.ClusterSpec{ + ApplyPolicyDefaultValues: true, + NumWorkers: 2, + }, + }, + { + name: "autoscale with policy defaults", + in: compute.ClusterSpec{ + ApplyPolicyDefaultValues: true, + Autoscale: &compute.AutoScale{ + MinWorkers: 1, + MaxWorkers: 4, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + before := tt.in + initializeNumWorkers(&tt.in) + assert.Equal(t, before.NumWorkers, tt.in.NumWorkers) + assert.Equal(t, before.Autoscale, tt.in.Autoscale) + assert.Equal(t, tt.want, tt.in.ForceSendFields) + }) + } +} diff --git a/bundle/direct/dresources/resources.yml b/bundle/direct/dresources/resources.yml index 06502dd2017..0fe2ba65c32 100644 --- a/bundle/direct/dresources/resources.yml +++ b/bundle/direct/dresources/resources.yml @@ -138,6 +138,11 @@ resources: - field: tasks[*].for_each_task.task.new_cluster.data_security_mode - field: job_clusters[*].new_cluster.data_security_mode + # A cluster policy can supply autoscaling when apply_policy_default_values is enabled. + - field: tasks[*].new_cluster.autoscale + - field: tasks[*].for_each_task.task.new_cluster.autoscale + - field: job_clusters[*].new_cluster.autoscale + job_runs: # Every jobs.RunNow field is listed, so nothing the request carries is drift. # TestJobRunIgnoresEveryRequestField keeps the list in step with the SDK.