Fix test_olive_evaluator.py - ClassVar for _nested_field_name and safe user_config access

Co-authored-by: shaahji <96227573+shaahji@users.noreply.github.com>

Author: Copilot

olive/data/config.py

@@ -6,7 +6,7 @@
 import re
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, ClassVar, Optional, Union
 
 from olive.common.config_utils import ConfigBase, NestedConfig, validate_lowercase
 from olive.common.import_lib import import_user_module
@@ -21,7 +21,7 @@
 
 
 class DataComponentConfig(NestedConfig):
-    _nested_field_name = "params"
+    _nested_field_name: ClassVar[str] = "params"
 
     type: str = None
     params: dict = Field(default_factory=dict)

olive/evaluator/metric.py

@@ -3,7 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import logging
-from typing import Any, Optional, Union
+from typing import Any, Callable, ClassVar, Optional, Union
 
 from olive.common.config_utils import ConfigBase, NestedConfig, validate_config
 from olive.common.pydantic_v1 import field_validator
@@ -105,13 +105,13 @@ def validate_goal(cls, v, info):
 
 
 class Metric(NestedConfig):
-    _nested_field_name = "user_config"
+    _nested_field_name: ClassVar[str] = "user_config"
 
     name: str
     type: MetricType
     backend: Optional[str] = "torch_metrics"
     sub_types: list[SubMetric]
-    user_config: ConfigBase = None
+    user_config: Optional[ConfigBase] = None
     data_config: Optional[DataConfig] = None
 
     def get_inference_settings(self, framework):

olive/evaluator/olive_evaluator.py

@@ -10,7 +10,7 @@
 from functools import partial
 from numbers import Number
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, Optional, Union
 
 import numpy as np
 import torch
@@ -121,15 +121,16 @@ def generate_metric_user_config_with_model_io(metric: Metric, model: "OliveModel
 
     @staticmethod
     def get_user_config(framework: Framework, metric: Metric):
-        assert metric.user_config, "user_config is not specified in the metric config"
-
         dataloader = None
         eval_func = None
         post_func = None
 
         # load the evaluate function
         # priority: evaluate_func > metric_func
         if metric.type == MetricType.CUSTOM:
+            if not metric.user_config:
+                raise ValueError("user_config is required for CUSTOM metric type")
+            
             evaluate_func = getattr(metric.user_config, "evaluate_func", None)
             kwargs = getattr(metric.user_config, "evaluate_func_kwargs", None) or {}
             if not evaluate_func:
@@ -217,7 +218,7 @@ def device_string_to_torch_device(device: Device):
 
     @classmethod
     def io_bind_enabled(cls, metric: Metric, inference_settings: dict) -> bool:
-        if metric.user_config.io_bind:
+        if metric.user_config and metric.user_config.io_bind:
             return True
 
         return inference_settings and inference_settings.get("io_bind")
@@ -307,7 +308,7 @@ def _evaluate_custom(
         execution_providers=None,
     ) -> MetricResult:
         raw_res = None
-        if metric.user_config.evaluate_func:
+        if metric.user_config and metric.user_config.evaluate_func:
             raw_res = eval_func(model, device, execution_providers)
         else:
             inference_output, targets = self._inference(
@@ -645,15 +646,15 @@ def _evaluate_distributed_latency_worker(config) -> list[float]:
         batch = next(iter(dataloader))
         input_data = OliveEvaluator.extract_input_data(batch)
         input_feed = format_data(input_data, io_config)
-        kv_cache_ortvalues = {} if metric.user_config.shared_kv_buffer else None
+        kv_cache_ortvalues = {} if (metric.user_config and getattr(metric.user_config, 'shared_kv_buffer', None)) else None
 
         io_bind = OnnxEvaluator.io_bind_enabled(metric, model.inference_settings)
         if io_bind:
             io_bind_op = prepare_io_bindings(
                 session,
                 input_feed,
                 Device.GPU,
-                shared_kv_buffer=metric.user_config.shared_kv_buffer,
+                shared_kv_buffer=getattr(metric.user_config, 'shared_kv_buffer', None) if metric.user_config else None,
                 kv_cache_ortvalues=kv_cache_ortvalues,
             )
         latencies = []
@@ -1123,7 +1124,7 @@ def evaluate(
 
 
 class OliveEvaluatorConfig(NestedConfig):
-    _nested_field_name = "type_args"
+    _nested_field_name: ClassVar[str] = "type_args"
 
     name: Optional[str] = None
     type: Optional[str] = None

olive/model/config/hf_config.py

@@ -4,7 +4,7 @@
 # --------------------------------------------------------------------------
 import logging
 from copy import deepcopy
-from typing import Any, Optional, Union
+from typing import Any, ClassVar, Optional, Union
 
 import torch
 import transformers
@@ -22,7 +22,7 @@ class HfLoadKwargs(NestedConfig):
     Refer to https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.from_pretrained
     """
 
-    _nested_field_name = "extra_args"
+    _nested_field_name: ClassVar[str] = "extra_args"
 
     torch_dtype: str = Field(
         None,

olive/passes/pytorch/train_utils.py

@@ -5,7 +5,7 @@
 import dataclasses
 import logging
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union
 
 import torch
 import transformers
@@ -33,7 +33,7 @@
 class BaseHFTrainingArguments(NestedConfig):
     """Training arguments for transformers.Trainer."""
 
-    _nested_field_name = "extra_args"
+    _nested_field_name: ClassVar[str] = "extra_args"
 
     gradient_checkpointing: bool = Field(True, description="Use gradient checkpointing. Recommended.")
     report_to: Union[str, list[str]] = Field(

olive/workflows/run/config.py

@@ -4,7 +4,7 @@
 # --------------------------------------------------------------------------
 import shutil
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, ClassVar, Optional, Union
 
 from olive.cache import CacheConfig
 from olive.common.config_utils import NestedConfig, validate_config
@@ -109,7 +109,7 @@ class RunConfig(NestedConfig):
     evaluators, engine, passes, and auto optimizer.
     """
 
-    _nested_field_name = "engine"
+    _nested_field_name: ClassVar[str] = "engine"
 
     workflow_id: str = Field(
         DEFAULT_WORKFLOW_ID, description="Workflow ID. If not provided, use the default ID 'default_workflow'."

test/evaluator/test_olive_evaluator.py

@@ -379,7 +379,12 @@ def test_evaluator_get_inference_session(self, metric_inference_settings, model_
         """
         metric = get_latency_metric(LatencySubType.AVG)
         if metric_inference_settings:
-            metric.user_config.inference_settings = {"onnx": metric_inference_settings.copy()}
+            # Initialize user_config if needed to set inference_settings
+            if metric.user_config is None:
+                from olive.common.config_utils import ConfigBase
+                metric.user_config = ConfigBase()
+            # Use object.__setattr__ to set dynamic attributes on ConfigBase
+            object.__setattr__(metric.user_config, 'inference_settings', {"onnx": metric_inference_settings.copy()})
         model = get_onnx_model()
         model.inference_settings = model_inference_settings.copy() if model_inference_settings else None
         inference_settings = OnnxEvaluator.get_inference_settings(metric, model)