nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/profiler/callbacks/llama_index_callback_handler.py

@@ -30,6 +30,7 @@ from nat.builder.context import Context
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import ServerToolUseSchema
 from nat.data_models.intermediate_step import StreamEventData
 from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.intermediate_step import UsageInfo
@@ -64,6 +65,26 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
         self._run_id_to_tool_input = {}
         self._run_id_to_timestamp = {}
 
+    @staticmethod
+    def _extract_token_usage(response: ChatResponse) -> TokenUsageBaseModel:
+        token_usage = TokenUsageBaseModel()
+        try:
+            if response and response.additional_kwargs and "usage" in response.additional_kwargs:
+                usage = response.additional_kwargs["usage"] if "usage" in response.additional_kwargs else {}
+                token_usage.prompt_tokens = usage.input_tokens if hasattr(usage, "input_tokens") else 0
+                token_usage.completion_tokens = usage.output_tokens if hasattr(usage, "output_tokens") else 0
+
+                if hasattr(usage, "input_tokens_details") and hasattr(usage.input_tokens_details, "cached_tokens"):
+                    token_usage.cached_tokens = usage.input_tokens_details.cached_tokens
+
+                if hasattr(usage, "output_tokens_details") and hasattr(usage.output_tokens_details, "reasoning_tokens"):
+                    token_usage.reasoning_tokens = usage.output_tokens_details.reasoning_tokens
+
+        except Exception as e:
+            logger.debug("Error extracting token usage: %s", e, exc_info=True)
+
+        return token_usage
+
     def on_event_start(
         self,
         event_type: CBEventType,
@@ -167,6 +188,18 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
                 except Exception as e:
                     logger.exception("Error getting model name: %s", e)
 
+                # Append usage data to NAT usage stats
+                tool_outputs_list = []
+                # Check if message.additional_kwargs as tool_outputs indicative of server side tool calling
+                if response and response.additional_kwargs and "built_in_tool_calls" in response.additional_kwargs:
+                    tools_outputs = response.additional_kwargs["built_in_tool_calls"]
+                    if isinstance(tools_outputs, list):
+                        for tool in tools_outputs:
+                            try:
+                                tool_outputs_list.append(ServerToolUseSchema(**tool.model_dump()))
+                            except Exception:
+                                pass
+
                 # Append usage data to NAT usage stats
                 with self._lock:
                     stats = IntermediateStepPayload(
@@ -176,8 +209,9 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
                         name=model_name,
                         UUID=event_id,
                         data=StreamEventData(input=self._run_id_to_llm_input.get(event_id), output=llm_text_output),
-                        metadata=TraceMetadata(chat_responses=response.message if response.message else None),
-                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel(**response.additional_kwargs)))
+                        metadata=TraceMetadata(chat_responses=response.message if response.message else None,
+                                               tool_outputs=tool_outputs_list if tool_outputs_list else []),
+                        usage_info=UsageInfo(token_usage=self._extract_token_usage(response)))
                     self.step_manager.push_intermediate_step(stats)
 
         elif event_type == CBEventType.FUNCTION_CALL and payload:

nat/profiler/callbacks/token_usage_base_model.py

@@ -24,4 +24,6 @@ class TokenUsageBaseModel(BaseModel):
 
     prompt_tokens: int = Field(default=0, description="Number of tokens in the prompt.")
     completion_tokens: int = Field(default=0, description="Number of tokens in the completion.")
+    cached_tokens: int = Field(default=0, description="Number of tokens read from cache.")
+    reasoning_tokens: int = Field(default=0, description="Number of tokens used for reasoning.")
     total_tokens: int = Field(default=0, description="Number of tokens total.")

nat/profiler/decorators/framework_wrapper.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 import functools
 import logging
+from collections.abc import AsyncIterator
 from collections.abc import Callable
 from contextlib import AbstractAsyncContextManager as AsyncContextManager
 from contextlib import asynccontextmanager
@@ -32,35 +33,55 @@ _library_instrumented = {
     "crewai": False,
     "semantic_kernel": False,
     "agno": False,
+    "adk": False,
 }
 
 callback_handler_var: ContextVar[Any | None] = ContextVar("callback_handler_var", default=None)
 
 
 def set_framework_profiler_handler(
-    workflow_llms: dict = None,
-    frameworks: list[LLMFrameworkEnum] = None,
+    workflow_llms: dict | None = None,
+    frameworks: list[LLMFrameworkEnum] | None = None,
 ) -> Callable[[Callable[..., AsyncContextManager[Any]]], Callable[..., AsyncContextManager[Any]]]:
     """
     Decorator that wraps an async context manager function to set up framework-specific profiling.
+
+    Args:
+        workflow_llms (dict | None): A dictionary of workflow LLM configurations.
+        frameworks (list[LLMFrameworkEnum] | None): A list of LLM frameworks used in the workflow functions.
+
+    Returns:
+        Callable[[Callable[..., AsyncContextManager[Any]]], Callable[..., AsyncContextManager[Any]]]:
+        A decorator that wraps the original function with profiling setup.
     """
 
     def decorator(func: Callable[..., AsyncContextManager[Any]]) -> Callable[..., AsyncContextManager[Any]]:
+        """The actual decorator that wraps the function.
+
+        Args:
+            func (Callable[..., AsyncContextManager[Any]]): The function to wrap.
+
+        Returns:
+            Callable[..., AsyncContextManager[Any]]: The wrapped function.
+        """
 
         @functools.wraps(func)
         @asynccontextmanager
         async def wrapper(workflow_config, builder):
 
-            if LLMFrameworkEnum.LANGCHAIN in frameworks and not _library_instrumented["langchain"]:
-                from langchain_core.tracers.context import register_configure_hook
-
+            if LLMFrameworkEnum.LANGCHAIN in frameworks:
+                # Always set a fresh handler in the current context so callbacks
+                # route to the active run. Only register the hook once globally.
                 from nat.profiler.callbacks.langchain_callback_handler import LangchainProfilerHandler
 
                 handler = LangchainProfilerHandler()
                 callback_handler_var.set(handler)
-                register_configure_hook(callback_handler_var, inheritable=True)
-                _library_instrumented["langchain"] = True
-                logger.debug("LangChain/LangGraph callback handler registered")
+
+                if not _library_instrumented["langchain"]:
+                    from langchain_core.tracers.context import register_configure_hook
+                    register_configure_hook(callback_handler_var, inheritable=True)
+                    _library_instrumented["langchain"] = True
+                    logger.debug("LangChain/LangGraph callback hook registered")
 
             if LLMFrameworkEnum.LLAMA_INDEX in frameworks:
                 from llama_index.core import Settings
@@ -96,6 +117,20 @@ def set_framework_profiler_handler(
                 _library_instrumented["agno"] = True
                 logger.info("Agno callback handler registered")
 
+            if LLMFrameworkEnum.ADK in frameworks and not _library_instrumented["adk"]:
+                try:
+                    from nat.plugins.adk.adk_callback_handler import ADKProfilerHandler
+                except ImportError as e:
+                    logger.warning(
+                        "ADK profiler not available. " +
+                        "Install NAT with ADK extras: pip install \"nvidia-nat[adk]\". Error: %s",
+                        e)
+                else:
+                    handler = ADKProfilerHandler()
+                    handler.instrument()
+                    _library_instrumented["adk"] = True
+                    logger.debug("ADK callback handler registered")
+
             # IMPORTANT: actually call the wrapped function as an async context manager
             async with func(workflow_config, builder) as result:
                 yield result
@@ -114,11 +149,28 @@ def chain_wrapped_build_fn(
     Convert an original build function into an async context manager that
     wraps it with a single call to set_framework_profiler_handler, passing
     all frameworks at once.
+
+    Args:
+        original_build_fn (Callable[..., AsyncContextManager]): The original build function to wrap.
+        workflow_llms (dict): A dictionary of workflow LLM configurations.
+        function_frameworks (list[LLMFrameworkEnum]): A list of LLM frameworks used in the workflow functions.
+
+    Returns:
+        Callable[..., AsyncContextManager]: The wrapped build function.
     """
 
     # Define a base async context manager that simply calls the original build function.
     @asynccontextmanager
-    async def base_fn(*args, **kwargs):
+    async def base_fn(*args, **kwargs) -> AsyncIterator[Any]:
+        """Base async context manager that calls the original build function.
+
+        Args:
+            *args: Positional arguments to pass to the original build function.
+            **kwargs: Keyword arguments to pass to the original build function.
+
+        Yields:
+            The result of the original build function.
+        """
         async with original_build_fn(*args, **kwargs) as w:
             yield w
 

nat/profiler/decorators/function_tracking.py

@@ -18,7 +18,9 @@ import inspect
 import uuid
 from collections.abc import Callable
 from typing import Any
+from typing import TypeVar
 from typing import cast
+from typing import overload
 
 from pydantic import BaseModel
 
@@ -38,10 +40,10 @@ def _serialize_data(obj: Any) -> Any:
 
     if isinstance(obj, dict):
         return {str(k): _serialize_data(v) for k, v in obj.items()}
-    if isinstance(obj, (list, tuple, set)):
+    if isinstance(obj, list | tuple | set):
         return [_serialize_data(item) for item in obj]
 
-    if isinstance(obj, (str, int, float, bool, type(None))):
+    if isinstance(obj, str | int | float | bool | type(None)):
         return obj
 
     # Fallback
@@ -77,7 +79,24 @@ def push_intermediate_step(step_manager: IntermediateStepManager,
     step_manager.push_intermediate_step(payload)
 
 
-def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
+# Type variable for overloads
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+# Overloads for different function types
+@overload
+def track_function(func: F, *, metadata: dict[str, Any] | None = None) -> F:
+    """Overload for when a function is passed directly."""
+    ...
+
+
+@overload
+def track_function(*, metadata: dict[str, Any] | None = None) -> Callable[[F], F]:
+    """Overload for decorator factory usage (when called with parentheses)."""
+    ...
+
+
+def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None) -> Any:
     """
     Decorator that can wrap any type of function (sync, async, generator,
     async generator) and executes "tracking logic" around it.
@@ -256,6 +275,19 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
     return sync_wrapper
 
 
+# Overloads for track_unregistered_function
+@overload
+def track_unregistered_function(func: F, *, name: str | None = None, metadata: dict[str, Any] | None = None) -> F:
+    """Overload for when a function is passed directly."""
+    ...
+
+
+@overload
+def track_unregistered_function(*, name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[F], F]:
+    """Overload for decorator factory usage (when called with parentheses)."""
+    ...
+
+
 def track_unregistered_function(func: Callable[..., Any] | None = None,
                                 *,
                                 name: str | None = None,

nat/profiler/forecasting/models/linear_model.py

@@ -36,7 +36,7 @@ class LinearModel(ForecastingBaseModel):
         except ImportError:
             logger.error(
                 "scikit-learn is not installed. Please install scikit-learn to use the LinearModel "
-                "profiling model or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                "profiling model or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
             raise
 

nat/profiler/forecasting/models/random_forest_regressor.py

@@ -36,7 +36,7 @@ class RandomForestModel(ForecastingBaseModel):
         except ImportError:
             logger.error(
                 "scikit-learn is not installed. Please install scikit-learn to use the RandomForest "
-                "profiling model or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                "profiling model or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
             raise
 

nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py

@@ -304,7 +304,7 @@ def save_gantt_chart(all_nodes: list[CallNode], output_path: str) -> None:
         import matplotlib.pyplot as plt
     except ImportError:
         logger.error("matplotlib is not installed. Please install matplotlib to use generate plots for the profiler "
-                     "or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                     "or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
         raise
 

nat/profiler/inference_optimization/experimental/prefix_span_analysis.py

@@ -212,7 +212,7 @@ def run_prefixspan(sequences_map: dict[int, list[PrefixCallNode]],
         from prefixspan import PrefixSpan
     except ImportError:
         logger.error("prefixspan is not installed. Please install prefixspan to run the prefix analysis in the "
-                     "profiler or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
+                     "profiler or install \"nvidia-nat[profiler]\" to install all necessary profiling packages.")
 
         raise
 

nat/profiler/parameter_optimization/optimizable_utils.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import get_args
+from typing import get_origin
+
+from pydantic import BaseModel
+
+from nat.data_models.optimizable import SearchSpace
+
+logger = logging.getLogger(__name__)
+
+
+def walk_optimizables(obj: BaseModel, path: str = "") -> dict[str, SearchSpace]:
+    """
+    Recursively build ``{flattened.path: SearchSpace}`` for every optimizable
+    field inside *obj*.
+
+    * Honors ``optimizable_params`` on any model that mixes in
+      ``OptimizableMixin`` – only listed fields are kept.
+    * If a model contains optimizable fields **but** omits
+      ``optimizable_params``, we emit a warning and skip them.
+    """
+    spaces: dict[str, SearchSpace] = {}
+
+    allowed_params_raw = getattr(obj, "optimizable_params", None)
+    allowed_params = set(allowed_params_raw) if allowed_params_raw is not None else None
+    overrides = getattr(obj, "search_space", {}) or {}
+    has_optimizable_flag = False
+
+    for name, fld in obj.model_fields.items():
+        full = f"{path}.{name}" if path else name
+        extra = fld.json_schema_extra or {}
+
+        is_field_optimizable = extra.get("optimizable", False) or name in overrides
+        has_optimizable_flag = has_optimizable_flag or is_field_optimizable
+
+        # honour allow-list
+        if allowed_params is not None and name not in allowed_params:
+            continue
+
+        # 1. plain optimizable field or override from config
+        if is_field_optimizable:
+            space = overrides.get(name, extra.get("search_space"))
+            if space is None:
+                logger.error(
+                    "Field %s is marked optimizable but no search space was provided.",
+                    full,
+                )
+                raise ValueError(f"Field {full} is marked optimizable but no search space was provided")
+            spaces[full] = space
+
+        value = getattr(obj, name, None)
+
+        # 2. nested BaseModel
+        if isinstance(value, BaseModel):
+            spaces.update(walk_optimizables(value, full))
+
+        # 3. dict[str, BaseModel] container
+        elif isinstance(value, dict):
+            for key, subval in value.items():
+                if isinstance(subval, BaseModel):
+                    spaces.update(walk_optimizables(subval, f"{full}.{key}"))
+
+        # 4. static-type fallback for class-level annotations
+        elif isinstance(obj, type):
+            ann = fld.annotation
+            if get_origin(ann) in (dict, dict):
+                _, val_t = get_args(ann) or (None, None)
+                if isinstance(val_t, type) and issubclass(val_t, BaseModel):
+                    if allowed_params is None or name in allowed_params:
+                        spaces[f"{full}.*"] = SearchSpace(low=None, high=None)  # sentinel
+
+    if allowed_params is None and has_optimizable_flag:
+        logger.warning(
+            "Model %s contains optimizable fields but no `optimizable_params` "
+            "were defined; these fields will be ignored.",
+            obj.__class__.__name__,
+        )
+    return spaces

nat/profiler/parameter_optimization/optimizer_runtime.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from pydantic import BaseModel
+
+from nat.data_models.optimizer import OptimizerRunConfig
+from nat.experimental.decorators.experimental_warning_decorator import experimental
+from nat.profiler.parameter_optimization.optimizable_utils import walk_optimizables
+from nat.profiler.parameter_optimization.parameter_optimizer import optimize_parameters
+from nat.profiler.parameter_optimization.prompt_optimizer import optimize_prompts
+from nat.runtime.loader import load_config
+
+logger = logging.getLogger(__name__)
+
+
+@experimental(feature_name="Optimizer")
+async def optimize_config(opt_run_config: OptimizerRunConfig):
+    """Entry-point called by the CLI or runtime."""
+    # ---------------- 1. load / normalise ---------------- #
+    if not isinstance(opt_run_config.config_file, BaseModel):
+        from nat.data_models.config import Config  # guarded import
+        base_cfg: Config = load_config(config_file=opt_run_config.config_file)
+    else:
+        base_cfg = opt_run_config.config_file  # already validated
+
+    # ---------------- 2. discover search space ----------- #
+    full_space = walk_optimizables(base_cfg)
+    if not full_space:
+        logger.warning("No optimizable parameters found in the configuration. "
+                       "Skipping optimization.")
+        return base_cfg
+
+    # ---------------- 3. numeric / enum tuning ----------- #
+    tuned_cfg = base_cfg
+    if base_cfg.optimizer.numeric.enabled:
+        tuned_cfg = optimize_parameters(
+            base_cfg=base_cfg,
+            full_space=full_space,
+            optimizer_config=base_cfg.optimizer,
+            opt_run_config=opt_run_config,
+        )
+
+    # ---------------- 4. prompt optimization ------------- #
+    if base_cfg.optimizer.prompt.enabled:
+        await optimize_prompts(
+            base_cfg=tuned_cfg,
+            full_space=full_space,
+            optimizer_config=base_cfg.optimizer,
+            opt_run_config=opt_run_config,
+        )
+
+    logger.info("All optimization phases complete.")
+    return tuned_cfg

nat/profiler/parameter_optimization/parameter_optimizer.py

@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from collections.abc import Mapping as Dict
+
+import optuna
+import yaml
+
+from nat.data_models.config import Config
+from nat.data_models.optimizable import SearchSpace
+from nat.data_models.optimizer import OptimizerConfig
+from nat.data_models.optimizer import OptimizerRunConfig
+from nat.data_models.optimizer import SamplerType
+from nat.eval.evaluate import EvaluationRun
+from nat.eval.evaluate import EvaluationRunConfig
+from nat.experimental.decorators.experimental_warning_decorator import experimental
+from nat.profiler.parameter_optimization.parameter_selection import pick_trial
+from nat.profiler.parameter_optimization.update_helpers import apply_suggestions
+
+logger = logging.getLogger(__name__)
+
+
+@experimental(feature_name="Optimizer")
+def optimize_parameters(
+    *,
+    base_cfg: Config,
+    full_space: Dict[str, SearchSpace],
+    optimizer_config: OptimizerConfig,
+    opt_run_config: OptimizerRunConfig,
+) -> Config:
+    """Tune all *non-prompt* hyper-parameters and persist the best config."""
+    space = {k: v for k, v in full_space.items() if not v.is_prompt}
+
+    # Ensure output_path is not None
+    if optimizer_config.output_path is None:
+        raise ValueError("optimizer_config.output_path cannot be None")
+    out_dir = optimizer_config.output_path
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Ensure eval_metrics is not None
+    if optimizer_config.eval_metrics is None:
+        raise ValueError("optimizer_config.eval_metrics cannot be None")
+
+    metric_cfg = optimizer_config.eval_metrics
+    directions = [v.direction for v in metric_cfg.values()]
+    eval_metrics = [v.evaluator_name for v in metric_cfg.values()]
+    weights = [v.weight for v in metric_cfg.values()]
+
+    # Create appropriate sampler based on configuration
+    sampler_type = optimizer_config.numeric.sampler
+
+    if sampler_type == SamplerType.GRID:
+        # For grid search, convert the existing space to value sequences
+        grid_search_space = {param_name: search_space.to_grid_values() for param_name, search_space in space.items()}
+        sampler = optuna.samplers.GridSampler(grid_search_space)
+        logger.info("Using Grid sampler for numeric optimization")
+    else:
+        # None or BAYESIAN: let Optuna choose defaults
+        sampler = None
+        logger.info(
+            "Using Optuna default sampler types: TPESampler for single-objective, NSGAIISampler for multi-objective")
+
+    study = optuna.create_study(directions=directions, sampler=sampler)
+
+    # Create output directory for intermediate files
+    out_dir = optimizer_config.output_path
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    async def _run_eval(runner: EvaluationRun):
+        return await runner.run_and_evaluate()
+
+    def _objective(trial: optuna.Trial):
+        reps = max(1, getattr(optimizer_config, "reps_per_param_set", 1))
+
+        # build trial config
+        suggestions = {p: spec.suggest(trial, p) for p, spec in space.items()}
+        cfg_trial = apply_suggestions(base_cfg, suggestions)
+
+        async def _single_eval(trial_idx: int) -> list[float]:  # noqa: ARG001
+            eval_cfg = EvaluationRunConfig(
+                config_file=cfg_trial,
+                dataset=opt_run_config.dataset,
+                result_json_path=opt_run_config.result_json_path,
+                endpoint=opt_run_config.endpoint,
+                endpoint_timeout=opt_run_config.endpoint_timeout,
+            )
+            scores = await _run_eval(EvaluationRun(config=eval_cfg))
+            values = []
+            for metric_name in eval_metrics:
+                metric = next(r[1] for r in scores.evaluation_results if r[0] == metric_name)
+                values.append(metric.average_score)
+
+            return values
+
+        # Create tasks for all evaluations
+        async def _run_all_evals():
+            tasks = [_single_eval(i) for i in range(reps)]
+            return await asyncio.gather(*tasks)
+
+        # Calculate padding width based on total number of trials
+        trial_id_width = len(str(max(0, optimizer_config.numeric.n_trials - 1)))
+        trial_id_padded = f"{trial.number:0{trial_id_width}d}"
+        with (out_dir / f"config_numeric_trial_{trial_id_padded}.yml").open("w") as fh:
+            yaml.dump(cfg_trial.model_dump(), fh)
+
+        all_scores = asyncio.run(_run_all_evals())
+        # Persist raw per‑repetition scores so they appear in `trials_dataframe`.
+        trial.set_user_attr("rep_scores", all_scores)
+        return [sum(run[i] for run in all_scores) / reps for i in range(len(eval_metrics))]
+
+    logger.info("Starting numeric / enum parameter optimization...")
+    study.optimize(_objective, n_trials=optimizer_config.numeric.n_trials)
+    logger.info("Numeric optimization finished")
+
+    best_params = pick_trial(
+        study=study,
+        mode=optimizer_config.multi_objective_combination_mode,
+        weights=weights,
+    ).params
+    tuned_cfg = apply_suggestions(base_cfg, best_params)
+
+    # Save final results (out_dir already created and defined above)
+    with (out_dir / "optimized_config.yml").open("w") as fh:
+        yaml.dump(tuned_cfg.model_dump(mode='json'), fh)
+    with (out_dir / "trials_dataframe_params.csv").open("w") as fh:
+        # Export full trials DataFrame (values, params, timings, etc.).
+        df = study.trials_dataframe()
+
+        # Rename values_X columns to actual metric names
+        metric_names = list(metric_cfg.keys())
+        rename_mapping = {}
+        for i, metric_name in enumerate(metric_names):
+            old_col = f"values_{i}"
+            if old_col in df.columns:
+                rename_mapping[old_col] = f"values_{metric_name}"
+        if rename_mapping:
+            df = df.rename(columns=rename_mapping)
+
+        # Normalise rep_scores column naming for convenience.
+        if "user_attrs_rep_scores" in df.columns and "rep_scores" not in df.columns:
+            df = df.rename(columns={"user_attrs_rep_scores": "rep_scores"})
+        elif "user_attrs" in df.columns and "rep_scores" not in df.columns:
+            # Some Optuna versions return a dict in a single user_attrs column.
+            df["rep_scores"] = df["user_attrs"].apply(lambda d: d.get("rep_scores") if isinstance(d, dict) else None)
+            df = df.drop(columns=["user_attrs"])
+
+        # Get Pareto optimal trial numbers from Optuna study
+        pareto_trials = study.best_trials
+        pareto_trial_numbers = {trial.number for trial in pareto_trials}
+        # Add boolean column indicating if trial is Pareto optimal
+        df["pareto_optimal"] = df["number"].isin(pareto_trial_numbers)
+
+        df.to_csv(fh, index=False)
+
+    # Generate Pareto front visualizations
+    try:
+        from nat.profiler.parameter_optimization.pareto_visualizer import create_pareto_visualization
+        logger.info("Generating Pareto front visualizations...")
+        create_pareto_visualization(
+            data_source=study,
+            metric_names=eval_metrics,
+            directions=directions,
+            output_dir=out_dir / "plots",
+            title_prefix="Parameter Optimization",
+            show_plots=False  # Don't show plots in automated runs
+        )
+        logger.info("Pareto visualizations saved to: %s", out_dir / "plots")
+    except ImportError as ie:
+        logger.warning("Could not import visualization dependencies: %s. "
+                       "Have you installed nvidia-nat-profiling?",
+                       ie)
+    except Exception as e:
+        logger.warning("Failed to generate visualizations: %s", e)
+
+    return tuned_cfg