Feature/optimizer module and registry#1614
Feature/optimizer module and registry#1614afourniernv wants to merge 4 commits intoNVIDIA:developfrom
Conversation
- Add packages/nvidia_nat_optimizer with pyproject.toml, README, src/nat/optimizer, tests - Move parameter_optimization modules and tests from core to optimizer package - Core CLI optimize: resilient import; on missing optimizer, suggest nvidia-nat-optimizer or nvidia-nat - Meta-package default and most extras include nvidia-nat-optimizer; add uv.sources path - Update LangChain prompt_optimizer register and examples (DEVELOPER_NOTES, email_phishing test) - Docs: installation (default includes optimizer, extras), optimizer (prerequisites) Signed-off-by: afourniernv <afournier@nvidia.com>
- evolutionary_base: shared evaluation loop and template method; subclasses implement fitness and persistence - ga_prompt_optimizer: GA fitness, selection, oracle feedback, checkpointing - Oracle feedback moves out of core config into model_extra and ga_config - Replace prompt_optimizer with ga_prompt_optimizer; update runtime, langchain, tests Signed-off-by: afourniernv <afournier@nvidia.com>
- Add OPTIMIZER to ComponentEnum, OptimizerStrategyBaseConfig (TypedBaseModel) - Extend NumericOptimizationConfig and PromptGAOptimizationConfig with name= - Add register_optimizer, get_optimizer, RegisteredOptimizerInfo in type_registry - Add register_optimizer decorator in register_workflow - Add nat.optimizer.register with GA and parameter optimizer registration - Update optimizer_runtime to dispatch via registry instead of direct imports - Update test_optimizer_runtime_extra for registry-based flow
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
|
||
|
|
||
| class PromptGAOptimizationConfig(BaseModel): | ||
| class PromptGAOptimizationConfig(OptimizerStrategyBaseConfig, name="ga"): |
There was a problem hiding this comment.
I feel like this should just be PromptOptimizationConfig with the genetic algorithm being a registered variant.
| class PromptGAOptimizationConfig(OptimizerStrategyBaseConfig, name="ga"): | ||
| """ | ||
| Configuration for prompt optimization using a Genetic Algorithm. | ||
| Oracle feedback and other implementation-specific options are not part of this |
There was a problem hiding this comment.
The GA-specific registered variant of the PromptOptimization config can also have oracle feedback. Passing it in as model extra limits the amount of type checking and validation we can do, as well as the discoverability of available fields.
There was a problem hiding this comment.
Theses tests will need to change once we pull oracle feedback and associated config items back out of model extra
| from nat.data_models.optimizer import PromptGAOptimizationConfig | ||
|
|
||
|
|
||
| class GAOptimizerConfig(BaseOptimizerConfig): |
There was a problem hiding this comment.
Why do we have PromptGAOptimizationConfig and also this?
| from nat.optimizer.update_helpers import apply_suggestions | ||
|
|
||
|
|
||
| class BaseEvolutionaryPromptOptimizer(ABC): |
There was a problem hiding this comment.
Seems like we need one more level of interface above this. Something like BasePromptOptimizer which defines methods that all prompt optimizers must implement. And then the GAPromptOptimizer is a subclass of that. BaseEvolutionaryPromptOptimizer seems unecessary unless we expect to have many evolutionary algorithms AND other algorithms we want to support.
There was a problem hiding this comment.
Another suggestion/note on the directory structure. Perhaps the dilenation would be clearer if we have two directories under packages/nvidia_nat_optimizer/src/nat/optimizer. One for parameters and one for prompts
|
|
||
|
|
||
| @experimental(feature_name="Optimizer") | ||
| def optimize_parameters( |
There was a problem hiding this comment.
I feel like we also need an ABC for parameter optimization so that developers can also register new ways to do parameter optimization but still enforce a standard interface.
2 participants
WIP DRAFT PR
Description
Extracts the optimizer into a standalone
nvidia-nat-optimizerpackage, refactors the GA as an evolutionary-base implementation, and introduces a pluggable registry for optimizer strategies—enabling third-party optimizers without modifying core code.Package extraction
packages/nvidia_nat_optimizerpackage with its ownpyproject.tomland minimal deps. Installable asnvidia-nat-optimizer(standalone) or included by default vianvidia-nat.nat_coreprofiler intonat.optimizer. Corenat optimizeCLI imports optimizer resiliently and suggests installation when missing.Evolutionary base and GA refactor
evolutionary_base.pywith a shared evaluation loop and template method; subclasses implement fitness computation and persistence.prompt_optimizerwithga_prompt_optimizeras the first evolutionary implementation. Oracle feedback moves out of the shared config intomodel_extraand GA-specific config.Pluggable optimizer registry
OptimizerStrategyBaseConfigbase and registry intype_registry(register_optimizer,get_optimizer).NumericOptimizationConfigandPromptGAOptimizationConfigextend it withname="numeric"andname="ga".register_optimizerdecorator inregister_workflow.py(mirrorsregister_evaluator). Optimizer runtime dispatches viaget_optimizer(type(config))instead of direct imports.nat.optimizer.registerregisters built-in strategies when the package loads. Existing YAML configs remain valid; strategy selection uses the inferred config type.Note on diff size
~5k additions are mostly: uv.lock for the new package (~3.6k lines) and one new file
ga_prompt_optimizer.py(~526 lines). Moved files (e.g.oracle_feedback.py,update_helpers.py) show as "renamed without changes" so blame is preserved.