chore: Centralize OmegaConf resolver registration (#1882)

Signed-off-by: ruit <ruit@nvidia.com>

Author: RayenTian

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -48,11 +48,13 @@
 from nemo_rl.environments.utils import create_env
 from nemo_rl.experience.rollouts import run_async_nemo_gym_rollout
 from nemo_rl.models.generation import configure_generation_config
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 from nemo_rl.utils.logger import get_next_experiment_dir
 
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-
 
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
@@ -118,6 +120,7 @@ def collect_trajectories(
 def main() -> None:
     """Main entry point."""
     # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

examples/run_distillation.py

@@ -22,11 +22,13 @@
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.models.generation import configure_generation_config
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 from nemo_rl.utils.logger import get_next_experiment_dir
 
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-
 
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
@@ -46,6 +48,7 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
 def main() -> None:
     """Main entry point."""
     # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

examples/run_grpo.py

@@ -23,11 +23,13 @@
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.models.generation import configure_generation_config
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 from nemo_rl.utils.logger import get_next_experiment_dir
 
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-
 
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
@@ -45,6 +47,7 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
 def main() -> None:
     """Main entry point."""
     # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

examples/run_grpo_sliding_puzzle.py

@@ -34,11 +34,13 @@
     SlidingPuzzleMetadata,
 )
 from nemo_rl.models.generation import configure_generation_config
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 from nemo_rl.utils.logger import get_next_experiment_dir
 
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-
 
 def parse_args():
     """Parse command line arguments."""
@@ -192,6 +194,7 @@ def setup_puzzle_data(
 def main():
     """Main entry point."""
     # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

examples/run_sft.py

@@ -30,12 +30,13 @@
     update_single_dataset_config,
 )
 from nemo_rl.distributed.virtual_cluster import init_ray
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 from nemo_rl.utils.logger import get_next_experiment_dir
 
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-OmegaConf.register_new_resolver("max", lambda a, b: max(a, b))
-
 
 def parse_args():
     """Parse command line arguments."""
@@ -148,6 +149,7 @@ def setup_data(tokenizer: AutoTokenizer, data_config: DataConfig):
 def main(is_vlm: bool = False):
     """Main entry point."""
     # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

examples/run_vlm_grpo.py

@@ -23,11 +23,13 @@
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
 from nemo_rl.models.generation import configure_generation_config
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 from nemo_rl.utils.logger import get_next_experiment_dir
 
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-
 
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
@@ -42,6 +44,8 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
 
 def main() -> None:
     """Main entry point."""
+    # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

nemo_rl/utils/config.py

@@ -186,3 +186,11 @@ def parse_hydra_overrides(cfg: DictConfig, overrides: list[str]) -> DictConfig:
         return cfg
     except Exception as e:
         raise OverridesError(f"Failed to parse Hydra overrides: {str(e)}") from e
+
+
+def register_omegaconf_resolvers() -> None:
+    """Register shared OmegaConf resolvers used in configs."""
+    if not OmegaConf.has_resolver("mul"):
+        OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
+    if not OmegaConf.has_resolver("max"):
+        OmegaConf.register_new_resolver("max", lambda a, b: max(a, b))

research/template_project/single_update.py

@@ -41,9 +41,11 @@
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.models.generation.vllm import VllmGeneration
 from nemo_rl.models.policy.lm_policy import Policy
-from nemo_rl.utils.config import load_config, parse_hydra_overrides
-
-OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
+from nemo_rl.utils.config import (
+    load_config,
+    parse_hydra_overrides,
+    register_omegaconf_resolvers,
+)
 
 
 def main(config: MasterConfig) -> None:
@@ -178,6 +180,7 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
 
 if __name__ == "__main__":
     # Parse arguments
+    register_omegaconf_resolvers()
     args, overrides = parse_args()
 
     if not args.config:

tests/unit/test_config_validation.py

@@ -27,16 +27,15 @@
 from nemo_rl.algorithms.rm import MasterConfig as RMMasterConfig
 from nemo_rl.algorithms.sft import MasterConfig as SFTMasterConfig
 from nemo_rl.evals.eval import MasterConfig as EvalMasterConfig
-from nemo_rl.utils.config import load_config_with_inheritance
+from nemo_rl.utils.config import (
+    load_config_with_inheritance,
+    register_omegaconf_resolvers,
+)
 
 # All tests in this module should run first
 pytestmark = pytest.mark.run_first
 
-if not OmegaConf.has_resolver("mul"):
-    OmegaConf.register_new_resolver("mul", lambda a, b: a * b)
-
-if not OmegaConf.has_resolver("max"):
-    OmegaConf.register_new_resolver("max", lambda a, b: max(a, b))
+register_omegaconf_resolvers()
 
 
 def validate_config_section(