Skip to content

Commit

Permalink
Fix issue in installing UCX on UC enabled workspace (#3501)
Browse files Browse the repository at this point in the history
<!-- REMOVE IRRELEVANT COMMENTS BEFORE CREATING A PULL REQUEST -->
## Changes
<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->

This PR updates the UCX policy that is created when installing UCX. It
replaces the policy definition for spark_version from fixed to allowlist
with a default value.
When UC is enabled on a workspaces, the cluster definition take the
value as single_user and user_isolation instead of the
Legacy_Single_User and Legacy_Table_ACL and this is occurring due to the
policy which is overriding this. Changing the sparkversion from fixed to
allowlist solved this issue. This could be a API issue which needs to be
raised separately, but for now it addresses the problem by doing this
simple change

It also updates the job definition to use the default value when not
passed by setting the apply_policy_default_values to true

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

Resolves #3420 

### Functionality

- [ ] modified existing command: `databricks labs ucx ...`

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [ ] updated unit tests
- [ ] updated integration tests
- [ ] static installation test
  • Loading branch information
HariGS-DB authored Jan 15, 2025
1 parent 1c543c4 commit 3e96d4f
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 12 deletions.
2 changes: 1 addition & 1 deletion src/databricks/labs/ucx/installer/policy.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ def _definition(self, conf: dict, instance_profile: str | None, instance_pool_id
latest_lts_dbr = self._ws.clusters.select_spark_version(latest=True, long_term_support=True)
node_type_id = self._ws.clusters.select_node_type(local_disk=True, min_memory_gb=32, min_cores=4)
policy_definition = {
"spark_version": self._policy_config(latest_lts_dbr),
"spark_version": {"type": "allowlist", "values": [latest_lts_dbr], "defaultValue": latest_lts_dbr},
"node_type_id": self._policy_config(node_type_id),
}
for key, value in conf.items():
Expand Down
3 changes: 3 additions & 0 deletions src/databricks/labs/ucx/installer/workflows.py
Original file line number Diff line number Diff line change
Expand Up @@ -911,6 +911,7 @@ def _job_clusters(self, names: set[str]):
custom_tags={"ResourceClass": "SingleNode"},
num_workers=0,
policy_id=self._config.policy_id,
apply_policy_default_values=True,
),
)
)
Expand All @@ -923,6 +924,7 @@ def _job_clusters(self, names: set[str]):
spark_conf=self._job_cluster_spark_conf("tacl"),
num_workers=1, # ShowPermissionsCommand needs a worker
policy_id=self._config.policy_id,
apply_policy_default_values=True,
),
)
)
Expand All @@ -938,6 +940,7 @@ def _job_clusters(self, names: set[str]):
max_workers=self._config.max_workers,
min_workers=self._config.min_workers,
),
apply_policy_default_values=True,
),
)
)
Expand Down
26 changes: 24 additions & 2 deletions tests/integration/install/test_installation.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@
from databricks.labs.blueprint.tui import MockPrompts
from databricks.labs.blueprint.wheels import ProductInfo
from databricks.sdk import AccountClient, WorkspaceClient
from databricks.sdk.service import compute
from databricks.labs.lsql.backends import StatementExecutionBackend
from databricks.sdk.errors import (
AlreadyExists,
InvalidParameterValue,
NotFound,
)
from databricks.sdk.retries import retried
from databricks.sdk.service import compute

from databricks.labs.ucx.__about__ import __version__
from databricks.labs.ucx.config import WorkspaceConfig
Expand Down Expand Up @@ -117,7 +117,7 @@ def test_job_cluster_policy(ws, installation_ctx) -> None:
assert cluster_policy.name == f"Unity Catalog Migration ({installation_ctx.inventory_database}) ({user_name})"

spark_version = ws.clusters.select_spark_version(latest=True, long_term_support=True)
assert policy_definition["spark_version"]["value"] == spark_version
assert policy_definition["spark_version"]["values"][0] == spark_version
assert policy_definition["node_type_id"]["value"] == ws.clusters.select_node_type(local_disk=True, min_memory_gb=32)
if ws.config.is_azure:
assert (
Expand All @@ -128,6 +128,28 @@ def test_job_cluster_policy(ws, installation_ctx) -> None:
assert policy_definition["aws_attributes.availability"]["value"] == compute.AwsAvailability.ON_DEMAND.value


@retried(on=[NotFound, InvalidParameterValue], timeout=timedelta(minutes=3))
def test_job_cluster_on_uc_enabled_workpace(ws, installation_ctx) -> None:
# Set the override_cluster to empty so that the installation creates new job clusters
installation_ctx.workspace_installation.config.override_clusters = ""

installation_ctx.workspace_installation.run()
job_id = installation_ctx.install_state.jobs["assessment"]
job_clusters = installation_ctx.workspace_client.jobs.get(job_id).settings.job_clusters
for cluster in job_clusters:
if cluster.job_cluster_key == "main":
assert cluster.new_cluster.data_security_mode == compute.DataSecurityMode.LEGACY_SINGLE_USER_STANDARD
if cluster.job_cluster_key == "tacl":
assert cluster.new_cluster.data_security_mode == compute.DataSecurityMode.LEGACY_TABLE_ACL
job_id = installation_ctx.install_state.jobs["migrate-tables"]
job_clusters = installation_ctx.workspace_client.jobs.get(job_id).settings.job_clusters
for cluster in job_clusters:
if cluster.job_cluster_key == "main":
assert cluster.new_cluster.data_security_mode == compute.DataSecurityMode.LEGACY_SINGLE_USER_STANDARD
if cluster.job_cluster_key == "user_isolation":
assert cluster.new_cluster.data_security_mode == compute.DataSecurityMode.USER_ISOLATION


@retried(on=[NotFound, InvalidParameterValue], timeout=timedelta(minutes=5))
def test_running_real_remove_backup_groups_job(ws: WorkspaceClient, installation_ctx: MockInstallationContext) -> None:
ws_group_a, _ = installation_ctx.make_ucx_group(wait_for_provisioning=True)
Expand Down
18 changes: 9 additions & 9 deletions tests/unit/installer/test_policy.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ def test_cluster_policy_definition_azure_hms():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, _, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionURL": {"type": "fixed", "value": "url"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionUserName": {"type": "fixed", "value": "user1"},
Expand Down Expand Up @@ -103,7 +103,7 @@ def test_cluster_policy_definition_aws_glue():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, instance_profile, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"spark_conf.spark.databricks.hive.metastore.glueCatalog.enabled": {"type": "fixed", "value": "true"},
"aws_attributes.instance_profile_arn": {"type": "fixed", "value": "role_arn_1"},
Expand All @@ -128,7 +128,7 @@ def test_cluster_policy_definition_gcp():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, instance_profile, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionURL": {"type": "fixed", "value": "url"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionUserName": {"type": "fixed", "value": "user1"},
Expand Down Expand Up @@ -230,7 +230,7 @@ def test_cluster_policy_definition_azure_hms_warehouse():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, _, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionURL": {"type": "fixed", "value": "url"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionUserName": {"type": "fixed", "value": "user1"},
Expand Down Expand Up @@ -282,7 +282,7 @@ def test_cluster_policy_definition_aws_glue_warehouse():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, instance_profile, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"spark_conf.spark.databricks.hive.metastore.glueCatalog.enabled": {"type": "fixed", "value": "true"},
"aws_attributes.instance_profile_arn": {"type": "fixed", "value": "role_arn_1"},
Expand Down Expand Up @@ -338,7 +338,7 @@ def test_cluster_policy_definition_gcp_hms_warehouse():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, _, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionURL": {"type": "fixed", "value": "url"},
"spark_conf.spark.hadoop.javax.jdo.option.ConnectionUserName": {"type": "fixed", "value": "user1"},
Expand Down Expand Up @@ -379,7 +379,7 @@ def test_cluster_policy_definition_empty_config():
policy_installer = ClusterPolicyInstaller(MockInstallation(), ws, prompts)
policy_id, _, _, _ = policy_installer.create('ucx')
policy_definition_actual = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
"aws_attributes.zone_id": {"type": "fixed", "value": "auto"},
Expand Down Expand Up @@ -409,7 +409,7 @@ def test_cluster_policy_instance_pool():
assert instance_pool_id == "instance_pool_1"

policy_expected = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"instance_pool_id": {"type": "fixed", "value": "instance_pool_1"},
}
# test the instance pool is added to the cluster policy
Expand All @@ -422,7 +422,7 @@ def test_cluster_policy_instance_pool():
# test the instance pool is not found
ws.instance_pools.get.side_effect = NotFound()
policy_expected = {
"spark_version": {"type": "fixed", "value": "14.2.x-scala2.12"},
"spark_version": {"type": "allowlist", "values": ["14.2.x-scala2.12"], "defaultValue": "14.2.x-scala2.12"},
"node_type_id": {"type": "fixed", "value": "Standard_F4s"},
"aws_attributes.availability": {"type": "fixed", "value": "ON_DEMAND"},
"aws_attributes.zone_id": {"type": "fixed", "value": "auto"},
Expand Down

0 comments on commit 3e96d4f

Please sign in to comment.