Use environment variables to set thresholds in static yaml configurations

user_tools/pyproject.toml

@@ -27,7 +27,9 @@ dependencies = [
     "fastcore==1.7.10",
     "fire>=0.5.0",
     "pandas==1.4.3",
-    "pyYAML>=6.0",
+    "pyYAML>=6.0,<=7.0",
+    # This is used to resolve env-variable in yaml files. It requires netween 5.0 and 6.0
+    "pyaml-env==1.2.1",
     "tabulate==0.8.10",
     "importlib-resources==5.10.2",
     "requests==2.31.0",

user_tools/src/spark_rapids_pytools/common/prop_manager.py

@@ -21,6 +21,7 @@
 from typing import Any, Callable, Union
 
 import yaml
+from pyaml_env import parse_config
 
 from spark_rapids_tools import get_elem_from_dict, get_elem_non_safe
 
@@ -106,7 +107,7 @@ def __open_yaml_file(self):
         try:
             with open(self.prop_arg, 'r', encoding='utf-8') as yaml_file:
                 try:
-                    self.props = yaml.safe_load(yaml_file)
+                    self.props = parse_config(self.prop_arg)#yaml.safe_load(yaml_file)
                 except yaml.YAMLError as e:
                     raise RuntimeError('Incorrect format of Yaml File') from e
         except OSError as err:

user_tools/src/spark_rapids_pytools/resources/qualification-conf.yaml

@@ -249,7 +249,7 @@ local:
         columnWidth: 14
       totalCoreSecCol: 'Total Core Seconds'
       # This is total core seconds of an 8-core machine running for 24 hours
-      totalCoreSecThreshold: 691200
+      totalCoreSecThreshold: !ENV ${RAPIDS_USER_TOOLS_CORE_SECONDS_THRESHOLD:691200}
     speedupCategories:
       speedupColumnName: 'Estimated GPU Speedup'
       categoryColumnName: 'Estimated GPU Speedup Category'
@@ -293,7 +293,7 @@ local:
           columns:
             - 'stageId'
             - 'SQL Nodes(IDs)'
-        spillThresholdBytes: 10737418240
+        spillThresholdBytes: !ENV ${RAPIDS_USER_TOOLS_SPILL_BYTES_THRESHOLD:10737418240}
         allowedExecs:
           - 'Aggregate'
           - 'Join'

user_tools/src/spark_rapids_tools/utils/propmanager.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
 from typing import Union, Any, TypeVar, ClassVar, Type, Tuple, Optional
 
 import yaml
+from pyaml_env import parse_config
 from pydantic import BaseModel, ConfigDict, model_validator, ValidationError
 
 from spark_rapids_tools.exceptions import JsonLoadException, YamlLoadException, InvalidPropertiesSchema
@@ -45,7 +46,7 @@ def load_yaml(file_path: Union[str, CspPathT]) -> Any:
         file_path = CspPath(file_path)
     with file_path.open_input_stream() as fis:
         try:
-            return yaml.safe_load(fis)
+            return parse_config(data=fis.readall())
         except yaml.YAMLError as e:
             raise YamlLoadException('Incorrect format of Yaml File') from e