Increased memory for ucx clusters (#1366)

databrickslabs · Apr 24, 2024 · 491f792 · 491f792
1 parent 6f7696f
commit 491f792
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 10 deletions.
diff --git a/src/databricks/labs/ucx/installer/policy.py b/src/databricks/labs/ucx/installer/policy.py
@@ -9,6 +9,7 @@
 from databricks.sdk.service import compute
 from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -89,9 +90,10 @@ def _get_instance_pool_id(self) -> str | None:
 
     def _definition(self, conf: dict, instance_profile: str | None, instance_pool_id: str | None) -> str:
         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=16)
         policy_definition = {
             "spark_version": self._policy_config(latest_lts_dbr),
-            "node_type_id": self._policy_config(self._ws.clusters.select_node_type(local_disk=True)),
+            "node_type_id": self._policy_config(node_type_id),
         }
         for key, value in conf.items():
             policy_definition[f"spark_conf.{key}"] = self._policy_config(value)

diff --git a/src/databricks/labs/ucx/mixins/fixtures.py b/src/databricks/labs/ucx/mixins/fixtures.py
@@ -697,7 +697,7 @@ def create(
                 kwargs["spark_conf"] = {"spark.databricks.cluster.profile": "singleNode", "spark.master": "local[*]"}
             kwargs["custom_tags"] = {"ResourceClass": "SingleNode"}
         if "instance_pool_id" not in kwargs:
-            kwargs["node_type_id"] = ws.clusters.select_node_type(local_disk=True)
+            kwargs["node_type_id"] = ws.clusters.select_node_type(local_disk=True, min_memory_gb=16)
 
         return ws.clusters.create(
             cluster_name=cluster_name,
@@ -738,7 +738,7 @@ def create(*, instance_pool_name=None, node_type_id=None, **kwargs):
         if instance_pool_name is None:
             instance_pool_name = f"sdk-{make_random(4)}"
         if node_type_id is None:
-            node_type_id = ws.clusters.select_node_type(local_disk=True)
+            node_type_id = ws.clusters.select_node_type(local_disk=True, min_memory_gb=16)
         return ws.instance_pools.create(instance_pool_name, node_type_id, **kwargs)
 
     yield from factory("instance pool", create, lambda item: ws.instance_pools.delete(item.instance_pool_id))
@@ -761,7 +761,7 @@ def create(**kwargs):
                         description=make_random(4),
                         new_cluster=compute.ClusterSpec(
                             num_workers=1,
-                            node_type_id=ws.clusters.select_node_type(local_disk=True),
+                            node_type_id=ws.clusters.select_node_type(local_disk=True, min_memory_gb=16),
                             spark_version=ws.clusters.select_spark_version(latest=True),
                             spark_conf=task_spark_conf,
                         ),
@@ -776,7 +776,7 @@ def create(**kwargs):
                         description=make_random(4),
                         new_cluster=compute.ClusterSpec(
                             num_workers=1,
-                            node_type_id=ws.clusters.select_node_type(local_disk=True),
+                            node_type_id=ws.clusters.select_node_type(local_disk=True, min_memory_gb=16),
                             spark_version=ws.clusters.select_spark_version(latest=True),
                         ),
                         notebook_task=jobs.NotebookTask(notebook_path=make_notebook()),
@@ -817,7 +817,7 @@ def create(**kwargs) -> pipelines.CreatePipelineResponse:
         if "clusters" not in kwargs:
             kwargs["clusters"] = [
                 pipelines.PipelineCluster(
-                    node_type_id=ws.clusters.select_node_type(local_disk=True),
+                    node_type_id=ws.clusters.select_node_type(local_disk=True, min_memory_gb=16),
                     label="default",
                     num_workers=1,
                     custom_tags={

diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py
@@ -149,7 +149,7 @@ def test_job_cluster_policy(ws, installation_ctx):
 
     spark_version = ws.clusters.select_spark_version(latest=True, long_term_support=True)
     assert policy_definition["spark_version"]["value"] == spark_version
-    assert policy_definition["node_type_id"]["value"] == ws.clusters.select_node_type(local_disk=True)
+    assert policy_definition["node_type_id"]["value"] == ws.clusters.select_node_type(local_disk=True, min_memory_gb=16)
     if ws.config.is_azure:
         assert (
             policy_definition["azure_attributes.availability"]["value"]

diff --git a/tests/unit/installer/test_policy.py b/tests/unit/installer/test_policy.py
@@ -28,7 +28,7 @@ def common():
 
     w.cluster_policies.list.return_value = [policy]
     w.clusters.select_spark_version = lambda **_: "14.2.x-scala2.12"
-    w.clusters.select_node_type = lambda local_disk: "Standard_F4s"
+    w.clusters.select_node_type = lambda **_: "Standard_F4s"
     w.current_user.me = lambda: iam.User(user_name="[email protected]", groups=[iam.ComplexValue(display="admins")])
     prompts = MockPrompts(
         {
@@ -210,7 +210,7 @@ def test_update_job_policy():
                 '123',
                 new_cluster=ClusterSpec(
                     num_workers=1,
-                    node_type_id=ws.clusters.select_node_type(local_disk=True),
+                    node_type_id=ws.clusters.select_node_type(local_disk=True, min_memory_gb=16),
                     spark_version=ws.clusters.select_spark_version(latest=True),
                 ),
             )

diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py
@@ -142,7 +142,7 @@ def download(path: str) -> io.StringIO | io.BytesIO:
     workspace_client.clusters.list.return_value = mock_clusters()
     workspace_client.cluster_policies.create.return_value = CreatePolicyResponse(policy_id="foo")
     workspace_client.clusters.select_spark_version = lambda **_: "14.2.x-scala2.12"
-    workspace_client.clusters.select_node_type = lambda local_disk: "Standard_F4s"
+    workspace_client.clusters.select_node_type = lambda **_: "Standard_F4s"
     workspace_client.workspace.download = download
 
     return workspace_client