nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

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,153 @@
+# 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.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()]
+
+    study = optuna.create_study(directions=directions)
+
+    # 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)
+
+        with (out_dir / f"config_numeric_trial_{trial._trial_id}.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(), fh)
+    with (out_dir / "trials_dataframe_params.csv").open("w") as fh:
+        # Export full trials DataFrame (values, params, timings, etc.).
+        df = study.trials_dataframe()
+        # 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"])
+        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

nat/profiler/parameter_optimization/parameter_selection.py

@@ -0,0 +1,107 @@
+# 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.
+
+from collections.abc import Sequence
+
+import numpy as np
+import optuna
+from optuna._hypervolume import compute_hypervolume
+from optuna.study import Study
+from optuna.study import StudyDirection
+
+
+# ---------- helper ----------
+def _to_minimisation_matrix(
+    trials: Sequence[optuna.trial.FrozenTrial],
+    directions: Sequence[StudyDirection],
+) -> np.ndarray:
+    """Return array (n_trials × n_objectives) where **all** objectives are ‘smaller-is-better’."""
+    vals = np.asarray([t.values for t in trials], dtype=float)
+    for j, d in enumerate(directions):
+        if d == StudyDirection.MAXIMIZE:
+            vals[:, j] *= -1.0  # flip sign
+    return vals
+
+
+# ---------- public API ----------
+def pick_trial(
+    study: Study,
+    mode: str = "harmonic",
+    *,
+    weights: Sequence[float] | None = None,
+    ref_point: Sequence[float] | None = None,
+    eps: float = 1e-12,
+) -> optuna.trial.FrozenTrial:
+    """
+    Collapse Optuna’s Pareto front (`study.best_trials`) to a single “best compromise”.
+
+    Parameters
+    ----------
+    study      : completed **multi-objective** Optuna study
+    mode       : {"harmonic", "sum", "chebyshev", "hypervolume"}
+    weights    : per-objective weights (used only for "sum")
+    ref_point  : reference point for hyper-volume (defaults to ones after normalisation)
+    eps        : tiny value to avoid division by zero
+
+    Returns
+    -------
+    optuna.trial.FrozenTrial
+    """
+
+    # ---- 1. Pareto front ----
+    front = study.best_trials
+    if not front:
+        raise ValueError("`study.best_trials` is empty – no Pareto-optimal trials found.")
+
+    # ---- 2. Convert & normalise objectives ----
+    vals = _to_minimisation_matrix(front, study.directions)  # smaller is better
+    span = np.ptp(vals, axis=0)
+    norm = (vals - vals.min(axis=0)) / (span + eps)  # 0 = best, 1 = worst
+
+    # ---- 3. Scalarise according to chosen mode ----
+    mode = mode.lower()
+
+    if mode == "harmonic":
+        hmean = norm.shape[1] / (1.0 / (norm + eps)).sum(axis=1)
+        best_idx = hmean.argmin()  # lower = better
+
+    elif mode == "sum":
+        w = np.ones(norm.shape[1]) if weights is None else np.asarray(weights, float)
+        if w.size != norm.shape[1]:
+            raise ValueError("`weights` length must equal number of objectives.")
+        score = norm @ w
+        best_idx = score.argmin()
+
+    elif mode == "chebyshev":
+        score = norm.max(axis=1)  # worst dimension
+        best_idx = score.argmin()
+
+    elif mode == "hypervolume":
+        # Hyper-volume assumes points are *below* the reference point (minimisation space).
+        if len(front) == 0:
+            raise ValueError("Pareto front is empty - no trials to select from")
+        elif len(front) == 1:
+            best_idx = 0
+        else:
+            rp = np.ones(norm.shape[1]) if ref_point is None else np.asarray(ref_point, float)
+            base_hv = compute_hypervolume(norm, rp)
+            contrib = np.array([base_hv - compute_hypervolume(np.delete(norm, i, 0), rp) for i in range(len(front))])
+            best_idx = contrib.argmax()  # bigger contribution wins
+
+    else:
+        raise ValueError(f"Unknown mode '{mode}'. Choose from "
+                         "'harmonic', 'sum', 'chebyshev', 'hypervolume'.")
+
+    return front[best_idx]