diff --git a/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py b/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py index 78fe0361c..500644666 100644 --- a/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py +++ b/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py @@ -16,6 +16,7 @@ """Implementation specific to Dataproc""" from collections import defaultdict +import json from dataclasses import dataclass, field from typing import Any, List, Union, Optional @@ -25,7 +26,8 @@ from spark_rapids_pytools.cloud_api.sp_types import PlatformBase, CMDDriverBase, \ ClusterBase, ClusterNode, SysInfo, GpuHWInfo, SparkNodeType, ClusterState, GpuDevice, \ NodeHWInfo, ClusterGetAccessor -from spark_rapids_pytools.common.prop_manager import JSONPropertiesContainer, is_valid_gpu_device +from spark_rapids_tools.utils.propmanager import AbstractPropContainer +from spark_rapids_pytools.common.prop_manager import is_valid_gpu_device from spark_rapids_pytools.common.sys_storage import FSUtil from spark_rapids_pytools.common.utilities import Utils from spark_rapids_pytools.pricing.dataproc_pricing import DataprocPriceProvider @@ -114,10 +116,9 @@ def create_saving_estimator(self, source_cost: float = None): raw_pricing_config = self.configs.get_value_silent('pricing') if raw_pricing_config: - pricing_config = JSONPropertiesContainer(prop_arg=raw_pricing_config, - file_load=False) + pricing_config = AbstractPropContainer(props=raw_pricing_config) else: - pricing_config: JSONPropertiesContainer = None + pricing_config: AbstractPropContainer = None pricing_provider = DataprocPriceProvider(region=self.cli.get_region(), pricing_configs={'gcloud': pricing_config}) saving_estimator = DataprocSavingsEstimator(price_provider=pricing_provider, @@ -319,7 +320,9 @@ def extract_gpu_name(gpu_description: str) -> str: return gpu_name.upper() processed_instance_descriptions = {} - raw_instances_descriptions = JSONPropertiesContainer(prop_arg=instance_descriptions, file_load=False) + if isinstance(instance_descriptions, str): + instance_descriptions = json.loads(instance_descriptions) + raw_instances_descriptions = AbstractPropContainer(props=instance_descriptions) for instance in raw_instances_descriptions.props: instance_content = {} instance_content['VCpuCount'] = int(instance.get('guestCpus', -1)) @@ -505,7 +508,7 @@ def _init_nodes(self): for worker_node in worker_nodes_from_conf: worker_props = { 'name': worker_node, - 'props': JSONPropertiesContainer(prop_arg=raw_worker_prop, file_load=False), + 'props': AbstractPropContainer(props=raw_worker_prop), # set the node zone based on the wrapper defined zone 'zone': self.zone } @@ -516,7 +519,7 @@ def _init_nodes(self): raw_master_props = self.props.get_value('config', 'masterConfig') master_props = { 'name': master_nodes_from_conf[0], - 'props': JSONPropertiesContainer(prop_arg=raw_master_props, file_load=False), + 'props': AbstractPropContainer(props=raw_master_props), # set the node zone based on the wrapper defined zone 'zone': self.zone } @@ -527,7 +530,7 @@ def _init_nodes(self): SparkNodeType.MASTER: master_node } - def _set_zone_from_props(self, prop_container: JSONPropertiesContainer): + def _set_zone_from_props(self, prop_container: AbstractPropContainer): """ Extracts the 'zoneUri' from the properties container and updates the environment variable dictionary. """ diff --git a/user_tools/src/spark_rapids_pytools/common/cluster_inference.py b/user_tools/src/spark_rapids_pytools/common/cluster_inference.py index 8f5f2a2d7..9a398e332 100644 --- a/user_tools/src/spark_rapids_pytools/common/cluster_inference.py +++ b/user_tools/src/spark_rapids_pytools/common/cluster_inference.py @@ -14,15 +14,16 @@ """This module provides functionality for cluster inference""" +import json from dataclasses import dataclass, field from enum import Enum -from typing import Optional from logging import Logger +from typing import Optional import pandas as pd from spark_rapids_pytools.cloud_api.sp_types import PlatformBase, ClusterBase -from spark_rapids_pytools.common.prop_manager import JSONPropertiesContainer +from spark_rapids_tools.utils.propmanager import AbstractPropContainer from spark_rapids_pytools.common.utilities import ToolLogging from spark_rapids_tools import CspEnv @@ -145,7 +146,9 @@ def infer_cluster(self, cluster_info_df: pd.DataFrame) -> Optional[ClusterBase]: cluster_conf = self.platform.generate_cluster_configuration(cluster_template_args) if cluster_conf is None: return None - cluster_props_new = JSONPropertiesContainer(cluster_conf, file_load=False) + if isinstance(cluster_conf, str): + cluster_conf = json.loads(cluster_conf) + cluster_props_new = AbstractPropContainer(props=cluster_conf) return self.platform.load_cluster_by_prop(cluster_props_new, is_inferred=True) except Exception as e: # pylint: disable=broad-except self.logger.error('Error while inferring cluster: %s', str(e)) diff --git a/user_tools/tests/spark_rapids_tools_ut/test_json_boundaries.py b/user_tools/tests/spark_rapids_tools_ut/test_json_boundaries.py new file mode 100644 index 000000000..59d3cc8b2 --- /dev/null +++ b/user_tools/tests/spark_rapids_tools_ut/test_json_boundaries.py @@ -0,0 +1,51 @@ +import json + +import pandas as pd +import pytest + +from spark_rapids_pytools.cloud_api.dataproc import DataprocCMDDriver +from spark_rapids_pytools.common.cluster_inference import ClusterInference +from spark_rapids_tools import CspEnv + + +@pytest.mark.parametrize('instance_descriptions', [ + json.dumps([{'name': 'n2-standard-4', 'guestCpus': 4, 'memoryMb': 16384}]), + [{'name': 'n2-standard-4', 'guestCpus': 4, 'memoryMb': 16384}], +]) +def test_dataproc_process_instance_description_decodes_json(instance_descriptions): + driver = object.__new__(DataprocCMDDriver) + + assert driver._process_instance_description(instance_descriptions) == { + 'n2-standard-4': {'VCpuCount': 4, 'MemoryInMB': 16384} + } + + +@pytest.mark.parametrize('cluster_conf', [ + json.dumps({'NUM_DRIVER_NODES': 1, 'NUM_WORKER_NODES': 2}), + {'NUM_DRIVER_NODES': 1, 'NUM_WORKER_NODES': 2}, +]) +def test_cluster_inference_decodes_rendered_json(cluster_conf): + captured = {} + expected_cluster = object() + + class Platform: + def get_platform_name(self): + return CspEnv.ONPREM + + def generate_cluster_configuration(self, cluster_template_args): + return cluster_conf + + def load_cluster_by_prop(self, prop_container, is_inferred=False): + captured['props'] = prop_container.props + captured['is_inferred'] = is_inferred + return expected_cluster + + cluster_info = pd.DataFrame([{ + 'App ID': 'app-1', + 'Num Worker Nodes': 2, + 'Cores Per Executor': 4, + 'Num Executors Per Node': 2, + }]) + + assert ClusterInference(platform=Platform()).infer_cluster(cluster_info) is expected_cluster + assert captured == {'props': {'NUM_DRIVER_NODES': 1, 'NUM_WORKER_NODES': 2}, 'is_inferred': True}