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..e820cb047 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,8 @@ 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') + # a zero-worker (single-node) cluster has no `executors` entry at all + 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..d16bd5fef 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 @@ -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. @@ -244,7 +244,8 @@ 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') + # a zero-worker (single-node) cluster has no `executors` entry at all + 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..da2123ab6 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,13 +14,22 @@ """Test Identifying cluster from properties""" +import json + import pytest +from spark_rapids_pytools.cloud_api.sp_types import CspEnv, get_platform +from spark_rapids_pytools.common.prop_manager import JSONPropertiesContainer from spark_rapids_tools import CspPath from spark_rapids_tools.cloud import ClientCluster from spark_rapids_tools.exceptions import InvalidPropertiesSchema from .conftest import SparkRapidsToolsUT, all_cpu_cluster_props +databricks_cluster_props = [ + pytest.param(CspEnv.DATABRICKS_AWS, 'cluster/databricks/aws-cpu-00.json', id='databricks_aws'), + pytest.param(CspEnv.DATABRICKS_AZURE, 'cluster/databricks/azure-cpu-00.json', id='databricks_azure') +] + class TestClusterCSP(SparkRapidsToolsUT): # pylint: disable=too-few-public-methods """ @@ -36,3 +45,32 @@ 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 + + +class TestDatabricksClusterWorkers(SparkRapidsToolsUT): # pylint: disable=too-few-public-methods + """ + Class testing how the Databricks platforms build the worker nodes of a cluster from its + `clusters get` properties when the `executors` entry is absent, which is what the API + returns for a zero-worker (single-node) cluster and for a terminated cluster + """ + @staticmethod + def _load_cluster(csp_enum, prop_path, data_dir, num_workers): + with open(f'{data_dir}/{prop_path}', encoding='utf8') as prop_file: + props = json.load(prop_file) + props.pop('executors', None) + props['num_workers'] = num_workers + platform = get_platform(csp_enum)(ctxt_args={}) + return platform.load_cluster_by_prop(JSONPropertiesContainer(prop_arg=props, file_load=False)) + + @pytest.mark.parametrize('csp_enum,prop_path', databricks_cluster_props) + def test_zero_worker_cluster_raises_no_workers_error(self, csp_enum, prop_path, get_ut_data_dir): + # a single-node cluster has num_workers 0 and no executors entry; the platform's own + # "no worker nodes" error is the expected answer, not a TypeError from iterating None + with pytest.raises(RuntimeError, match='The cluster has no worker nodes'): + self._load_cluster(csp_enum, prop_path, get_ut_data_dir, num_workers=0) + + @pytest.mark.parametrize('csp_enum,prop_path', databricks_cluster_props) + def test_terminated_cluster_generates_workers(self, csp_enum, prop_path, get_ut_data_dir): + # a terminated multi-node cluster has no executors entry either; its workers are generated + cluster = self._load_cluster(csp_enum, prop_path, get_ut_data_dir, num_workers=2) + assert cluster.get_workers_count() == 2