nvidia-nat 1.4.0a20251030__py3-none-any.whl → 1.4.0a20251031__py3-none-any.whl

nat/profiler/parameter_optimization/parameter_optimizer.py

@@ -111,7 +111,10 @@ def optimize_parameters(
             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:
+        # 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())
@@ -132,7 +135,7 @@ def optimize_parameters(
 
     # 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)
+        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()
@@ -154,6 +157,13 @@ def optimize_parameters(
             # 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

nat/profiler/parameter_optimization/pareto_visualizer.py

@@ -21,6 +21,8 @@ import matplotlib.pyplot as plt
 import numpy as np
 import optuna
 import pandas as pd
+from matplotlib.lines import Line2D
+from matplotlib.patches import Patch
 
 logger = logging.getLogger(__name__)
 
@@ -79,6 +81,19 @@ class ParetoVisualizer:
                        linewidths=1.5,
                        marker='*')
 
+            # Add trial number labels to Pareto optimal points
+            for idx in range(len(pareto_trials_df)):
+                trial_number = pareto_trials_df.iloc[idx]['number'] \
+                    if 'number' in pareto_trials_df.columns else pareto_trials_df.index[idx]
+                ax.annotate(f'{int(trial_number)}',
+                            xy=(pareto_x[idx], pareto_y[idx]),
+                            xytext=(8, 8),
+                            textcoords='offset points',
+                            fontsize=9,
+                            fontweight='bold',
+                            color='darkred',
+                            bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='red', alpha=0.9))
+
             # Draw Pareto front line (only for 2D)
             if len(pareto_x) > 1:
                 # Sort points for line drawing based on first objective
@@ -182,6 +197,18 @@ class ParetoVisualizer:
                     trial_values = [normalized_values[j][idx] for j in range(n_metrics)]
                     ax.plot(x_positions, trial_values, 'r-', alpha=0.8, linewidth=3)
 
+                    # Add trial number label at the rightmost point
+                    trial_number = trials_df.iloc[idx]['number'] if 'number' in trials_df.columns else idx
+                    # Position label slightly to the right and above the last point
+                    ax.annotate(f'{int(trial_number)}',
+                                xy=(x_positions[-1], trial_values[-1]),
+                                xytext=(5, 5),
+                                textcoords='offset points',
+                                fontsize=9,
+                                fontweight='bold',
+                                color='darkred',
+                                bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='red', alpha=0.8))
+
         # Customize plot
         ax.set_xticks(x_positions)
         ax.set_xticklabels([f"{name}\n({direction})" for name, direction in zip(self.metric_names, self.directions)])
@@ -191,10 +218,10 @@ class ParetoVisualizer:
         ax.grid(True, alpha=0.3)
 
         # Add legend
-        from matplotlib.lines import Line2D
         legend_elements = [
             Line2D([0], [0], color='blue', alpha=0.3, linewidth=2, label='All Trials'),
-            Line2D([0], [0], color='red', alpha=0.8, linewidth=3, label='Pareto Optimal')
+            Line2D([0], [0], color='red', alpha=0.8, linewidth=3, label='Pareto Optimal'),
+            Patch(facecolor='white', edgecolor='red', label='[n]: trial number')
         ]
         ax.legend(handles=legend_elements, loc='best')
 
@@ -262,11 +289,49 @@ class ParetoVisualizer:
                                    linewidths=1,
                                    marker='*')
 
+                        # Add trial number labels to Pareto optimal points
+                        for idx in range(len(pareto_trials_df)):
+                            trial_number = pareto_trials_df.iloc[idx]['number'] \
+                                if 'number' in pareto_trials_df.columns else pareto_trials_df.index[idx]
+                            ax.annotate(f'{int(trial_number)}',
+                                        xy=(pareto_x[idx], pareto_y[idx]),
+                                        xytext=(6, 6),
+                                        textcoords='offset points',
+                                        fontsize=8,
+                                        fontweight='bold',
+                                        color='darkred',
+                                        bbox=dict(boxstyle='round,pad=0.2',
+                                                  facecolor='white',
+                                                  edgecolor='red',
+                                                  alpha=0.8))
+
                     ax.set_xlabel(f"{self.metric_names[j]} ({self.directions[j]})")
                     ax.set_ylabel(f"{self.metric_names[i]} ({self.directions[i]})")
 
                 ax.grid(True, alpha=0.3)
 
+        # Add legend to the figure
+        legend_elements = [
+            Line2D([0], [0],
+                   marker='o',
+                   color='w',
+                   markerfacecolor='lightblue',
+                   markeredgecolor='navy',
+                   markersize=8,
+                   alpha=0.6,
+                   label='All Trials'),
+            Line2D([0], [0],
+                   marker='*',
+                   color='w',
+                   markerfacecolor='red',
+                   markeredgecolor='darkred',
+                   markersize=10,
+                   alpha=0.9,
+                   label='Pareto Optimal'),
+            Patch(facecolor='white', edgecolor='red', label='[n]: trial number')
+        ]
+        fig.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(0.98, 0.98), framealpha=0.9, fontsize=10)
+
         plt.tight_layout()
 
         if save_path:

{nvidia_nat-1.4.0a20251030.dist-info → nvidia_nat-1.4.0a20251031.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat
-Version: 1.4.0a20251030
+Version: 1.4.0a20251031
 Summary: NVIDIA NeMo Agent toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation

{nvidia_nat-1.4.0a20251030.dist-info → nvidia_nat-1.4.0a20251031.dist-info}/RECORD

@@ -371,9 +371,9 @@ nat/profiler/inference_optimization/experimental/prefix_span_analysis.py,sha256=
 nat/profiler/parameter_optimization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nat/profiler/parameter_optimization/optimizable_utils.py,sha256=93Pl8A14Zq_f3XsxSH-yFnEJ6B7W5hp7doPnPoLlRB4,3714
 nat/profiler/parameter_optimization/optimizer_runtime.py,sha256=rXmCOq81o7ZorQOUYociVjuO3NO9CIjFBbwql2u_4H4,2715
-nat/profiler/parameter_optimization/parameter_optimizer.py,sha256=vxUvso4RnSwoUF5rJkJJGaIOJojebrvYcA79WA0ZP7c,7719
+nat/profiler/parameter_optimization/parameter_optimizer.py,sha256=xKqKLOJaOi3EVZjT95zwKwHdX3AyYmMfxDxyhwtHal4,8256
 nat/profiler/parameter_optimization/parameter_selection.py,sha256=pfnNQIx1evNICgChsOJXIFQHoL1R_kmh_vNDsVMC9kg,3982
-nat/profiler/parameter_optimization/pareto_visualizer.py,sha256=QclLZmmsWINIAh4n0XAKmnIZOqGHTMr-iggZS0kxj-Y,17055
+nat/profiler/parameter_optimization/pareto_visualizer.py,sha256=mE4Az9G_hw4CsAP8jgNl9H5L9KQjq82bxmIPEr63nEg,20537
 nat/profiler/parameter_optimization/prompt_optimizer.py,sha256=_AmdeB1jRamd93qR5UqRy5LweYR3bjnD7zoLxzXYE0k,17658
 nat/profiler/parameter_optimization/update_helpers.py,sha256=NxWhrGVchbjws85pPd-jS-C14_l70QvVSvEfENndVcY,2339
 nat/registry_handlers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -475,10 +475,10 @@ nat/utils/reactive/base/observer_base.py,sha256=6BiQfx26EMumotJ3KoVcdmFBYR_fnAss
 nat/utils/reactive/base/subject_base.py,sha256=UQOxlkZTIeeyYmG5qLtDpNf_63Y7p-doEeUA08_R8ME,2521
 nat/utils/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nat/utils/settings/global_settings.py,sha256=9JaO6pxKT_Pjw6rxJRsRlFCXdVKCl_xUKU2QHZQWWNM,7294
-nvidia_nat-1.4.0a20251030.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
-nvidia_nat-1.4.0a20251030.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-nvidia_nat-1.4.0a20251030.dist-info/METADATA,sha256=EWbdokoCV8bvnCxqEIw-AssdqH85mzpC4KkEy8gC_3s,10248
-nvidia_nat-1.4.0a20251030.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat-1.4.0a20251030.dist-info/entry_points.txt,sha256=4jCqjyETMpyoWbCBf4GalZU8I_wbstpzwQNezdAVbbo,698
-nvidia_nat-1.4.0a20251030.dist-info/top_level.txt,sha256=lgJWLkigiVZuZ_O1nxVnD_ziYBwgpE2OStdaCduMEGc,8
-nvidia_nat-1.4.0a20251030.dist-info/RECORD,,
+nvidia_nat-1.4.0a20251031.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat-1.4.0a20251031.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat-1.4.0a20251031.dist-info/METADATA,sha256=qPR0zmgu2THcW57-42J5cZOLVmw059YSyZV5ngsk-Gw,10248
+nvidia_nat-1.4.0a20251031.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat-1.4.0a20251031.dist-info/entry_points.txt,sha256=4jCqjyETMpyoWbCBf4GalZU8I_wbstpzwQNezdAVbbo,698
+nvidia_nat-1.4.0a20251031.dist-info/top_level.txt,sha256=lgJWLkigiVZuZ_O1nxVnD_ziYBwgpE2OStdaCduMEGc,8
+nvidia_nat-1.4.0a20251031.dist-info/RECORD,,