From 8601399dfab9fdea94d5027993198d4c9f5766e7 Mon Sep 17 00:00:00 2001 From: Ching-Wei Kang Date: Thu, 10 Sep 2026 13:56:55 -0500 Subject: [PATCH] fix(user_tools): handle zero-worker Databricks clusters Signed-off-by: Ching-Wei Kang --- .../cloud_api/databricks_aws.py | 2 +- .../cloud_api/databricks_azure.py | 2 +- .../spark_rapids_tools_ut/test_cluster.py | 29 ++++++++++++++++++- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/user_tools/src/spark_rapids_pytools/cloud_api/databricks_aws.py b/user_tools/src/spark_rapids_pytools/cloud_api/databricks_aws.py index 8f4fc554e..980dff30f 100644 --- a/user_tools/src/spark_rapids_pytools/cloud_api/databricks_aws.py +++ b/user_tools/src/spark_rapids_pytools/cloud_api/databricks_aws.py @@ -193,7 +193,7 @@ def _set_name_from_props(self) -> None: def _init_nodes(self): # assume that only one master node master_nodes_from_conf = self.props.get_value_silent('driver') - worker_nodes_from_conf = self.props.get_value_silent('executors') + worker_nodes_from_conf = self.props.get_value_silent('executors') or [] num_workers = self.props.get_value_silent('num_workers') if num_workers is None and self.props.get_value_silent('autoscale') is not None: target_workers = self.props.get_value_silent('autoscale', 'target_workers') diff --git a/user_tools/src/spark_rapids_pytools/cloud_api/databricks_azure.py b/user_tools/src/spark_rapids_pytools/cloud_api/databricks_azure.py index 5705b2018..beb13f9b0 100644 --- a/user_tools/src/spark_rapids_pytools/cloud_api/databricks_azure.py +++ b/user_tools/src/spark_rapids_pytools/cloud_api/databricks_azure.py @@ -244,7 +244,7 @@ def _set_name_from_props(self) -> None: def _init_nodes(self): # assume that only one driver node driver_nodes_from_conf = self.props.get_value_silent('driver') - worker_nodes_from_conf = self.props.get_value_silent('executors') + worker_nodes_from_conf = self.props.get_value_silent('executors') or [] num_workers = self.props.get_value_silent('num_workers') if num_workers is None and self.props.get_value_silent('autoscale') is not None: target_workers = self.props.get_value_silent('autoscale', 'target_workers') diff --git a/user_tools/tests/spark_rapids_tools_ut/test_cluster.py b/user_tools/tests/spark_rapids_tools_ut/test_cluster.py index 01805cc05..b49a948fb 100644 --- a/user_tools/tests/spark_rapids_tools_ut/test_cluster.py +++ b/user_tools/tests/spark_rapids_tools_ut/test_cluster.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023-2025, NVIDIA CORPORATION. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,11 +14,16 @@ """Test Identifying cluster from properties""" +import json +from types import SimpleNamespace + import pytest from spark_rapids_tools import CspPath from spark_rapids_tools.cloud import ClientCluster from spark_rapids_tools.exceptions import InvalidPropertiesSchema +from spark_rapids_pytools.cloud_api.databricks_aws import DatabricksCluster, DatabricksNode +from spark_rapids_pytools.cloud_api.databricks_azure import DatabricksAzureCluster, DatabricksAzureNode from .conftest import SparkRapidsToolsUT, all_cpu_cluster_props @@ -36,3 +41,25 @@ def test_cluster_invalid_path(self, get_ut_data_dir): def test_define_cluster_type_from_schema(self, csp, prop_path, get_ut_data_dir): client_cluster = ClientCluster(CspPath(f'{get_ut_data_dir}/{prop_path}')) assert client_cluster.platform_name == csp + + @pytest.mark.parametrize( + ('cluster_type', 'node_type', 'prop_path'), + [ + (DatabricksCluster, DatabricksNode, 'cluster/databricks/aws-cpu-00.json'), + (DatabricksAzureCluster, DatabricksAzureNode, 'cluster/databricks/azure-cpu-00.json'), + ], + ) + def test_databricks_zero_worker_cluster_reports_validation_error( + self, monkeypatch, get_ut_data_dir, cluster_type, node_type, prop_path): + cluster_props = json.loads((get_ut_data_dir / prop_path).read_text(encoding='utf-8')) + cluster_props.pop('executors') + cluster_props['num_workers'] = 0 + + platform = SimpleNamespace(cli=SimpleNamespace( + get_region=lambda: 'test-region', + get_env_var=lambda _: 'test-region', + )) + monkeypatch.setattr(node_type, 'fetch_and_set_hw_info', lambda *_: None) + + with pytest.raises(RuntimeError, match='The cluster has no worker nodes'): + cluster_type(platform).set_connection(props=json.dumps(cluster_props))